Speed Up Image Generation with FLUX 3 Image Klein on NVIDIA Blackwell GPUs NVIDIA reported that FLUX 3 Image Klein, Black Forest Labs' image generation and editing model, runs up to 3.46x faster for text-to-image and 4.45x faster for image-to-image generation at 2048×2048 on a single NVIDIA B200 GPU when configured in TensorRT LLM VisualGen with combined optimizations, compared with uncompiled BF16 inference. NVIDIA said NVFP4 alone delivered up to 1.62x faster T2I and 1.58x faster I2I than BF16 on Blackwell GPUs, and that denoising dominates runtime, accounting for 96.8% of Base T2I latency (14.24 s of 14.70 s) and 97.8% of Base I2I latency (39.40 s of 40.29 s). The tutorial covers fused GEMM kernels, low-precision and sparse attention, caching, and multi-GPU parallelism, with deployment instructions and YAML configurations for applying the optimizations. When deploying diffusion models in design tools, image editors, and visual production workflows, developers must account for inference latency. Generating an image requires multiple sequential passes through a model with billions of parameters, and higher resolutions increase attention costs. FLUX 3 Image Klein is an image editing and generation model from Black Forest Labs BFL trained on NVIDIA GB200 NVL72 and optimized for NVIDIA GPUs. It offers Base https://huggingface.co/black-forest-labs/flux-3-image-klein-base and distilled https://huggingface.co/black-forest-labs/flux-3-image-klein checkpoints for different latency and fidelity requirements. This tutorial shows how to configure FLUX 3 Image Klein in NVIDIA TensorRT LLM VisualGen to take advantage of fused kernels, multi-GPU parallelism, FP8/NVFP4 support, caching, and attention optimizations. It compares latency and changes in generated outputs across optimizations. In the reported 2048×2048 tests on a single NVIDIA B200 GPU, combined optimizations delivered up to 3.46x faster text-to-image T2I and 4.45x faster image-to-image I2I generation for FLUX 3 Image Klein Base compared with uncompiled BF16 inference. The tutorial includes deployment instructions and YAML configurations for applying these optimizations to your workloads. FLUX 3 inference with TensorRT LLM VisualGen FLUX 3 Image Klein combines Qwen3-VL-4B-Instruct for text conditioning, a diffusion transformer DiT to refine image latents, and an autoencoder to decode latents into images and encode I2I reference inputs. Klein processes text, noisy target-image, and optional reference-image streams independently before combining them through joint attention. TensorRT LLM VisualGen, a dedicated runtime within TensorRT LLM, supported BF16, FP8, and NVFP4 variants at launch. On an NVIDIA Blackwell GPU, NVFP4 delivered up to 1.62x faster T2I and 1.58x faster I2I generation than BF16. Figure 2 compares generation latency, while Figure 1 compares example outputs. The following sections use NVFP4 as the starting point for further optimization. Figure 1 compares outputs at 2048×2048. The summarized T2I prompt describes a kitchen still life; the summarized I2I prompt requests changes to a space scene, including a red gas giant, a silver dragon, and an engraved plaque. Figure 1. BF16 and NVFP4 outputs for T2I and I2I With NVFP4 as the starting point, VisualGen optimizations turned to denoising, which dominates overall generation time on B200. Accelerating FLUX 3’s DiT is therefore the focus of the following sections. Table 1 shows pipeline latency at 2048×2048 on one NVIDIA B200 GPU using NVFP4, VANILLA attention, and torch.compile. Results are averaged over 10 requests after two warmup requests. Base uses 50 denoising steps with a classifier-free guidance CFG scale of 3.5; the distilled Klein uses 16 steps without CFG. | Task | Model | Text + preparation | Denoising | Decode + postprocess | Total | Denoising share | |---|---|---|---|---|---|---| | T2I | Base | 47 ms | 14.24 s | 416 ms | 14.70 s | 96.8% | | T2I | Distilled | 24 ms | 2.28 s | 410 ms | 2.71 s | 84.0% | | I2I | Base | 468 ms | 39.40 s | 417 ms | 40.29 s | 97.8% | | I2I | Distilled | 407 ms | 6.65 s | 418 ms | 7.48 s | 89.0% | Table 1. FLUX 3 Image Klein pipeline latency on NVIDIA B200 Optimizing FLUX 3 in TensorRT LLM VisualGen TensorRT LLM VisualGen accelerates NVFP4 DiT inference with optimized general matrix multiplication GEMM kernels and epilogue fusion, low-precision and sparse attention, caching, and multi-GPU execution. Unless otherwise noted, the following benchmarks use torch.compile and seed 0. Base uses 50 denoising steps and a CFG scale of 3.5; the distilled Klein uses 16 denoising steps. NVFP4 inference VisualGen supports BFL’s prequantized NVFP4 DiT checkpoints for FLUX 3 Image Klein https://huggingface.co/black-forest-labs/flux-3-image-klein/tree/main/variants/nvfp4 and FLUX 3 Image Klein Base https://huggingface.co/black-forest-labs/flux-3-image-klein-base/tree/main/variants/nvfp4 . To use these checkpoints, set variant to nvfp4 in pipeline config : pipeline config: variant: nvfp4 GEMM kernel optimization TensorRT LLM VisualGen uses FlashInfer’s PrimsTS FP8/NVFP4 GEMM kernels, built with CUTLASS Primitives https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/primitives.html and Task Scheduling https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/ts.html to accelerate quantized inference. Tuned for FLUX 3 tensor shapes and enabled by default, they combine optimized GEMMs with fused epilogues, including GEMM with SwiGLU and query, key, and value QKV projections with query/key normalization and rotary position embeddings RoPE . See FlashInfer PR 5758 https://github.com/flashinfer-ai/flashinfer/pull/5758 for more details. At launch, the pinned NVFP4 PrimsTS kernels delivered 1.02–1.10x faster T2I generation on B200 than the autotuned TensorRT LLM NVFP4 GEMM implementation. Quantized and sparse attention Attention is a major bottleneck in high-resolution image generation. In the reported 2048×2048 NVFP4 T2I tests on B200, VANILLA BF16 attention accounted for approximately 60% of DiT GPU kernel time for both Base and distilled Klein. VisualGen supports two approaches to reducing this cost: quantized attention with cuDNN FP8 https://github.com/NVIDIA/cudnn-frontend/blob/v1.27.0/docs/operations/Attention.md or VC-Attention https://arxiv.org/abs/2609.15810 FP8, and sparse attention with Sol-Attention https://arxiv.org/abs/2607.24027 . VC-Attention FP8 is a TensorRT LLM variant of the method described in the VC-Attention paper https://arxiv.org/abs/2609.15810 . Learned Perceptual Image Patch Similarity LPIPS measures perceptual differences between generated images. In the following comparisons, lower scores indicate greater similarity to the corresponding VANILLA baseline output. Compared with VANILLA BF16 attention, cuDNN FP8 delivered generation speedups of 1.08x for Base and 1.09x for the 16-step distilled Klein. The corresponding LPIPS scores were 0.072 and 0.161. VC-Attention FP8 delivered speedups of 1.18x for Base and 1.16x for distilled Klein, with corresponding LPIPS scores of 0.086 and 0.151. Figure 4 compares FLUX 3 Image Klein Base NVFP4 outputs at 2048×2048 with CUDA Graphs enabled. These measurements suggest using VC-Attention FP8 for quantized attention at 2048×2048 on B200. The tested configuration uses a repair budget of ρ=0.005 to select the fraction of tokens with the highest quantization error for residual correction. To use quantized attention, replace attention config with one of the following configurations: For cuDNN FP8: attention config: backend: CUDNN quant attention config: qk dtype: fp8 v dtype: fp8 For VC-Attention FP8: attention config: backend: TRTLLM quant attention config: qk dtype: fp8 v dtype: fp8 vc attention: true vc repair budget: 0.005 Figure 4. BF16, cuDNN FP8, and VC-Attention FP8 example outputs Sol-Attention reduces DiT computation with sparsity controlled by τ. The reported tests use the PrimsTS-based general block-sparse attention kernel in TensorRT LLM. Configure Sol-Attention as follows: attention config: backend: TRTLLM sparse attention config: algorithm: sol attn tau: 1.0 disabled until timestep: 0.96 For NVFP4 FLUX 3 Image Klein Base at 2048×2048 on B200, Sol-Attention with τ=1.0 delivered 1.24x faster denoising plus autoencoder decoding than compiled dense VANILLA with 16 initial dense steps, or 1.30x with 12. These measurements suggest using 16 initial dense steps to prioritize fidelity or 12 to prioritize speed. Higher τ reduced output quality with little additional speedup. Figure 5 compares NVFP4 FLUX 3 Image Klein Base outputs at 2048×2048 with CUDA Graphs enabled and τ=1.0. Figure 5. Dense attention and Sol-Attention outputs with 16 or 12 initial dense steps Step caching Cache-DiT reduces FLUX 3 Image Klein inference latency by reusing intermediate transformer outputs across denoising steps without retraining or modifying model weights. A residual-difference threshold controls when cached outputs are reused, so you can adjust the balance between inference speed and output quality. Configure Cache-DiT with your desired parameters: cache config: cache backend: cache dit Fn compute blocks: 1 Bn compute blocks: 0 max warmup steps: 4 max continuous cached steps: 3 residual diff threshold: 0.24 At a threshold of 0.24, Cache-DiT nearly halved NVFP4 Base T2I latency at both tested resolutions, with smaller gains for the distilled Klein. At 1024×1024, speedups were 1.97x for Base and 1.26x for distilled Klein, with corresponding LPIPS scores of 0.038 and 0.036 against their respective uncached outputs. At 2048×2048, speedups were 1.94x for Base and 1.20x for distilled Klein Figure 6 . Figure 6. Base and distilled Klein outputs with and without Cache-DiT Multi-GPU inference The VisualGen runtime in TensorRT LLM scales FLUX 3 Image Klein using classifier-free guidance CFG and Ulysses sequence parallelism. Base assigns positive/negative branches to separate GPU groups and partitions each branch’s sequence with Ulysses; distilled runs without CFG and uses Ulysses alone. VisualGen handles sequence splitting, distributed attention, and output gathering. For eight-GPU NVFP4 Base inference, use the following configuration: parallel config: cfg size: 2 ulysses size: 4 For the distilled model, set cfg size: 1 , ulysses size: 8 , and use the flux-3-image-klein checkpoint. At 2048×2048, scaling from one to eight GPUs on an NVIDIA DGX B200 reduced median T2I latency from 14.90 to 3.00 seconds for Base and 2.73 to 1.08 seconds for distilled. Denoising speedups were larger, at 6.01x for Base and 4.44x for distilled Klein, but roughly half of eight-GPU distilled generation time remained outside denoising, limiting end-to-end gains. At smaller resolutions, CUDA Graphs complement parallelism by reducing kernel-launch overhead as work per GPU shrinks. In a separate 1024×1024 NVFP4 experiment on eight B200s with VANILLA attention, adding CUDA Graphs to compilation reduced Base latency from 1.823 to 0.809 seconds 2.25x, CFG across two GPU groups with Ulysses across four GPUs per group and distilled latency from 0.791 to 0.366 seconds 2.16x, Ulysses across eight GPUs . Enable CUDA Graphs with the following configuration: cuda graph config: enable: true FLUX 3 end-to-end performance At 2048×2048, Figure 8 compares end-to-end generation speedups across successive TensorRT LLM VisualGen optimization stages against uncompiled BF16 VANILLA inference. On a single B200, compilation, NVFP4 with PrimsTS GEMMs, Sol-Attention, and Cache-DiT delivered a 3.46x generation speedup for T2I, reducing latency from 22.36 to 6.46 seconds, and a 4.45x speedup for I2I, reducing latency from 54.64 to 12.28 seconds. The 16-step distilled Klein achieved up to 2.34x end-to-end speedups. On a single NVIDIA RTX PRO 6000 Blackwell Server Edition GPU, compilation, NVFP4 with fused CuTe DSL kernels, and Cache-DiT delivered 1.73-3.25x generation speedups, extending the gains across both platforms. Figures 9 and 10 compare example outputs across optimization stages. The examples in Figures 9 and 10 were generated on a single B200 GPU. Figure 9. FLUX 3 Image Klein Base T2I outputs across optimization stages The editing prompt requests a weathered appearance for the canoe while preserving its shape, position, and surrounding scene. Figure 10. FLUX 3 Image Klein Base I2I outputs across optimization stages To apply these optimizations to your own workloads, follow the runtime setup, inference commands, and YAML configurations in the following sections for offline generation or online serving. Prerequisites Start the prebuilt TensorRT LLM container with the following command. To build it from the feat/flux3-image-klein branch, follow the FLUX 3 deployment guide https://github.com/NVIDIA/TensorRT-LLM/blob/feat/flux3-image-klein/docs/source/deployment-guide/deployment-guide-for-flux3-image-klein-on-trtllm.md . docker run --rm -it --ipc=host --gpus all -p 8000:8000 \ --name tensorrt llm \ -e TLLM FMHA LIBS=+prims ts \ -e HF TOKEN=