Optimizing FLUX 3 Image Klein on NVIDIA Blackwell GPUs with TensorRT LLM VisualGen

FLUX 3 Image Klein is an image generation and editing model from Black Forest Labs (BFL). The model was trained on NVIDIA GB200 NVL72 hardware and the tutorial demonstrates how to optimize its inference on NVIDIA GPUs using TensorRT LLM VisualGen. Diffusion models require many sequential denoising steps on very large parameter counts, and attention costs increase at higher resolutions, so inference latency is a key practical constraint for production tools.

BFL publishes both Base and distilled checkpoints for FLUX 3 Image Klein to address different latency and fidelity needs.

Optimization building blocks used

TensorRT LLM VisualGen combines multiple techniques to reduce inference time:

• NVFP4/FP8 quantized weight variants
• FlashInfer/PrimsTS tuned GEMM kernels with fused epilogues
• low‑precision and sparse attention (cuDNN FP8, VC‑Attention FP8, Sol‑Attention)
• Cache‑DiT to reuse intermediate transformer outputs across denoising steps
• torch.compile based compilation
• multi‑GPU parallelism using Ulysses sequence parallelism and CFG partitioning
• CUDA Graphs to reduce kernel‑launch overhead at smaller per‑GPU workloads

The experiments reported focus on NVIDIA Blackwell data center GPUs (for example B200) and NVIDIA RTX PRO 6000 Blackwell Server Edition GPUs.

Model pipeline overview

FLUX 3 Image Klein uses Qwen3‑VL‑4B‑Instruct for text conditioning, a diffusion transformer (DiT) to refine image latents, and an autoencoder to decode latents into images and to encode I2I reference inputs. Klein handles text, noisy target image, and optional reference image streams independently before combining them via joint attention.

Experimental defaults in the tutorial: Base runs with 50 denoising steps and a CFG scale of 3.5; the distilled Klein runs with 16 steps and no CFG. Benchmarks generally use torch.compile with seed 0.

Pipeline latency on NVIDIA B200 (2048×2048)

Measured on a single NVIDIA B200 with NVFP4, VANILLA attention and torch.compile (averages over 10 requests after 2 warmups):

• T2I Base: Text+preparation 47 ms, denoising 14.24 s, decode+postprocess 416 ms, total 14.70 s (denoising share 96.8%)
• T2I Distilled: 24 ms, 2.28 s, 410 ms, total 2.71 s (84.0%)
• I2I Base: 468 ms, 39.40 s, 417 ms, total 40.29 s (97.8%)
• I2I Distilled: 407 ms, 6.65 s, 418 ms, total 7.48 s (89.0%)

Denoising (the DiT) dominates total generation time, especially for the Base checkpoint.

Specific optimizations and measured effects

• GEMM kernels: PrimsTS FP8/NVFP4 GEMMs (FlashInfer) are tuned for FLUX 3 shapes and include fused epilogues (SwiGLU, QKV projections, RoPE). On B200, pinned NVFP4 PrimsTS kernels produced 1.02–1.10× faster T2I than the autotuned TensorRT LLM NVFP4 GEMMs.

• Quantized attention: cuDNN FP8 and VC‑Attention FP8 reduce attention cost. In 2048×2048 NVFP4 T2I tests on B200, cuDNN FP8 yielded 1.08× speedup (Base) and 1.09× (16‑step distilled) with LPIPS 0.072 and 0.161 respectively. VC‑Attention FP8 gave 1.18× (Base) and 1.16× (distilled) speedups with LPIPS 0.086 and 0.151. Based on these tests, VC‑Attention FP8 with a repair budget ρ=0.005 is recommended for 2048×2048 on B200.

• Sol‑Attention (sparse attention): Sol‑Attention with τ controlling sparsity was tested using the PrimsTS block‑sparse kernel. For NVFP4 Base at 2048×2048 on B200, Sol‑Attention with τ=1.0 produced 1.24× faster denoising+decode than compiled dense VANILLA with 16 initial dense steps, or 1.30× with 12 dense steps. The authors suggest using 16 initial dense steps for fidelity or 12 if speed is prioritized; larger τ reduced quality with little speed benefit.

• Cache‑DiT: Caching intermediate transformer outputs across denoising steps can greatly reduce latency without retraining. A residual‑difference threshold controls reuse. At threshold 0.24, Cache‑DiT nearly halved NVFP4 Base T2I latency at tested resolutions, with smaller gains for the distilled model. At 1024×1024 speedups were 1.97× (Base) and 1.26× (distilled); at 2048×2048 speedups were 1.94× (Base) and 1.20× (distilled). Corresponding LPIPS deltas against uncached outputs are reported in the tutorial.

• Multi‑GPU scaling: VisualGen uses CFG partitioning and Ulysses sequence parallelism. For eight‑GPU NVFP4 Base inference, a recommended config is cfg_size: 2, ulysses_size: 4; for distilled cfg_size: 1, ulysses_size: 8. At 2048×2048, scaling from 1→8 B200 GPUs reduced median T2I latency from 14.90 s to 3.00 s for Base and 2.73 s to 1.08 s for distilled. Denoising saw larger speedups (Base 6.01×, distilled 4.44×), but non‑denoising stages limited full end‑to‑end gains.

• CUDA Graphs: On smaller resolutions, CUDA Graphs lower kernel launch overhead. In a 1024×1024 NVFP4 test on eight B200 GPUs, enabling CUDA Graphs reduced Base latency from 1.823 s to 0.809 s (2.25×) and distilled from 0.791 s to 0.366 s (2.16×).

Cumulative end‑to‑end improvements

Combining compilation, NVFP4 with PrimsTS fused GEMMs, Sol‑Attention and Cache‑DiT on a single B200 at 2048×2048 produced:

• T2I: up to 3.46× speedup, reducing latency from 22.36 s to 6.46 s
• I2I: up to 4.45× speedup, reducing latency from 54.64 s to 12.28 s The 16‑step distilled Klein achieved up to 2.34× end‑to‑end speedups. On an NVIDIA RTX PRO 6000 Blackwell Server Edition, combinations of compilation, NVFP4 with fused CuTe DSL kernels, and Cache‑DiT delivered 1.73–3.25× speedups.

The tutorial includes LPIPS comparisons showing perceptual differences across optimization stages; cumulative LPIPS values are provided for several stepwise configurations.

How to apply these optimizations (examples)

The tutorial provides runnable examples and YAML snippets. Key steps include:

• start the prebuilt TensorRT LLM container with TLLM_FMHA_LIBS=+prims_ts and an HF token
• run an offline example (python models/flux3.py ...) and add --variant nvfp4 for NVFP4 on Blackwell or --variant fp8r for row‑wise FP8
• start an OpenAI‑compatible /v1/images/generations endpoint via trtllm-serve and run the provided client script

Sample YAML files are given, for example configs/flux3-nvfp4-vc-fp8-1gpu.yaml (NVFP4 + torch.compile + VC‑Attention FP8) and configs/flux3-nvfp4-sol-cache-1gpu.yaml (NVFP4 + Sol‑Attention + Cache‑DiT). The tutorial explains settings like disabled_until_timestep to control how many initial denoising steps remain dense.

Conclusion

TensorRT LLM VisualGen provides a set of complementary optimizations—NVFP4/FP8 quantization, fused GEMMs, sparse attention modes, caching, compilation, CUDA Graphs and multi‑GPU execution—that can substantially reduce FLUX 3 Image Klein inference latency. The documented combination of techniques achieves up to 3.46× faster T2I and 4.45× faster I2I on a single NVIDIA DGX B200 at 2048×2048 versus uncompiled BF16 VANILLA inference, with perceptual differences quantified by LPIPS and configuration recipes provided to reproduce the results.

Acknowledgements

Nikita Korobov, Harrison Zhang, Pamela Peng, Kanghwan Jang, Yuhang He, and Charlotte Mallo.