{"slug": "speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus", "title": "Speed Up Image Generation with FLUX 3 Image Klein on NVIDIA Blackwell GPUs", "summary": "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.", "body_md": "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.\n\nFLUX 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.\n\nThis 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.\n\n## FLUX 3 inference with TensorRT LLM VisualGen\n\nFLUX 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.\n\nTensorRT 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.\n\nFigure 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.\n\n*Figure 1. BF16 and NVFP4 outputs for T2I and I2I*\n\nWith 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.\n\nTable 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.\n\n| **Task** | **Model** | **Text + preparation** | **Denoising** | **Decode + postprocess** | **Total** | **Denoising share** | \n|---|---|---|---|---|---|---|\n| T2I | Base | 47 ms | 14.24 s | 416 ms | 14.70 s | 96.8% | \n| T2I | Distilled | 24 ms | 2.28 s | 410 ms | 2.71 s | 84.0% | \n| I2I | Base | 468 ms | 39.40 s | 417 ms | 40.29 s | 97.8% | \n| I2I | Distilled | 407 ms | 6.65 s | 418 ms | 7.48 s | 89.0% | \n\n*Table 1. FLUX 3 Image Klein pipeline latency on NVIDIA B200*\n\n## Optimizing FLUX 3 in TensorRT LLM VisualGen\n\nTensorRT 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.\n\n### NVFP4 inference\n\nVisualGen 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`:\n\n```\npipeline_config:\n  variant: nvfp4\n```\n\n### GEMM kernel optimization\n\nTensorRT 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.\n\nAt launch, the pinned NVFP4 PrimsTS kernels delivered 1.02–1.10x faster T2I generation on B200 than the autotuned TensorRT LLM NVFP4 GEMM implementation.\n\n### Quantized and sparse attention\n\nAttention 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).\n\nLearned 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.\n\nCompared 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.\n\nFigure 4 compares FLUX 3 Image Klein Base NVFP4 outputs at 2048×2048 with CUDA Graphs enabled.\n\nThese 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:\n\nFor cuDNN FP8:\n\n```\nattention_config:\n  backend: CUDNN\n  quant_attention_config:\n    qk_dtype: fp8\n    v_dtype: fp8\n```\n\nFor VC-Attention FP8:\n\n```\nattention_config:\n    backend: TRTLLM\n    quant_attention_config:\n        qk_dtype: fp8\n        v_dtype: fp8\n        vc_attention: true\n        vc_repair_budget: 0.005\n```\n\n*Figure 4. BF16, cuDNN FP8, and VC-Attention FP8 example outputs*\n\nSol-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:\n\n```\nattention_config:\n  backend: TRTLLM\n  sparse_attention_config:\n    algorithm: sol_attn\n    tau: 1.0\n    disabled_until_timestep: 0.96\n```\n\nFor 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.\n\nThese measurements suggest using 16 initial dense steps to prioritize fidelity or 12 to prioritize speed. Higher τ reduced output quality with little additional speedup.\n\nFigure 5 compares NVFP4 FLUX 3 Image Klein Base outputs at 2048×2048 with CUDA Graphs enabled and τ=1.0.\n\n*Figure 5. Dense attention and Sol-Attention outputs with 16 or 12 initial dense steps*\n\n### Step caching\n\nCache-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:\n\n```\ncache_config:\n  cache_backend: cache_dit\n  Fn_compute_blocks: 1\n  Bn_compute_blocks: 0\n  max_warmup_steps: 4\n  max_continuous_cached_steps: 3\n  residual_diff_threshold: 0.24\n```\n\nAt 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).\n\n*Figure 6. Base and distilled Klein outputs with and without Cache-DiT*\n\n### Multi-GPU inference\n\nThe 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.\n\nVisualGen handles sequence splitting, distributed attention, and output gathering. For eight-GPU NVFP4 Base inference, use the following configuration:\n\n```\nparallel_config:\n    cfg_size: 2\n    ulysses_size: 4\n```\n\nFor the distilled model, set `cfg_size: 1`, `ulysses_size: 8`, and use the `flux-3-image-klein` checkpoint.\n\nAt 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.\n\nAt 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:\n\n```\ncuda_graph_config:\n    enable: true\n```\n\n## FLUX 3 end-to-end performance\n\nAt 2048×2048, Figure 8 compares end-to-end generation speedups across successive TensorRT LLM VisualGen optimization stages against uncompiled BF16 VANILLA inference.\n\nOn 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.\n\nOn 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.\n\nThe examples in Figures 9 and 10 were generated on a single B200 GPU.\n\n*Figure 9. FLUX 3 Image Klein Base T2I outputs across optimization stages*\n\nThe editing prompt requests a weathered appearance for the canoe while preserving its shape, position, and surrounding scene.\n\n*Figure 10. FLUX 3 Image Klein Base I2I outputs across optimization stages*\n\nTo 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.\n\n## Prerequisites\n\nStart 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).\n\n```\ndocker run --rm -it --ipc=host --gpus all -p 8000:8000 \\\n      --name tensorrt_llm \\\n      -e TLLM_FMHA_LIBS=+prims_ts \\\n      -e HF_TOKEN=<your_hf_token> \\\n      -v ~/.cache:/root/.cache:rw \\\n      nvcr.io/nvidia/tensorrt-llm/release:feat-flux3-image-klein-1.4.0a0.dev202610280436\n```\n\n`TLLM_FMHA_LIBS=+prims_ts` enables the PrimsTS attention kernels VC-Attention requires. The optimizations and results in this post focus on NVIDIA Blackwell data center GPUs and NVIDIA RTX PRO 6000 Blackwell Server Edition GPUs.\n\n## Get started with TensorRT LLM VisualGen\n\nWith the runtime environment ready, you can run FLUX 3 offline or serve it through an image generation endpoint. Use the sample YAML configurations in the following sections to apply the optimization settings used for the results reported earlier.\n\nThese examples do not expand prompts. If you use BFL’s prompt upsampler, run it separately before inference. Reported latencies exclude prompt upsampling.\n\n### Offline inference\n\nInside the container, run the [FLUX 3 example](https://github.com/NVIDIA/TensorRT-LLM/blob/feat/flux3-image-klein/examples/visual_gen/models/flux3.py):\n\n```\ncd /app/tensorrt_llm/examples/visual_gen\npython models/flux3.py \\\n  --model black-forest-labs/flux-3-image-klein-base \\\n  --visual_gen_args configs/flux3-images-1gpu.yaml \\\n  --prompt \"A small robot tending a rooftop garden at sunrise\" \\\n  --height 2048 --width 2048 --seed 0 \\\n  --output_path flux3.png\n```\n\nFor NVFP4 on Blackwell, add `--variant nvfp4`; use `--variant fp8r` for row-wise FP8 or omit `--variant` for BF16. Select `black-forest-labs/flux-3-image-klein` for the distilled model. For editing, add an editing prompt and `--image /path/to/reference.png`; repeat `--image` for multiple references.\n\n### Online serving endpoint\n\nFrom the same container directory, launch the OpenAI-compatible `/v1/images/generations` endpoint:\n\n```\ncd /app/tensorrt_llm/examples/visual_gen\ntrtllm-serve black-forest-labs/flux-3-image-klein-base \\\n  --host 0.0.0.0 --port 8000 \\\n  --visual_gen_args configs/flux3-images-1gpu.yaml\n```\n\nOnce the server is ready, open a second terminal, enter the running container, and run the client:\n\n```\ndocker exec -it tensorrt_llm bash\npython examples/visual_gen/serve/sync_image_gen.py \\\n  --base-url http://localhost:8000/v1 \\\n  --model black-forest-labs/flux-3-image-klein-base \\\n  --size 2048x2048 \\\n  --prompt \"A small robot tending a rooftop garden at sunrise\" \\\n  --output flux3-served.png\n```\n\nFor NVFP4 serving, add `pipeline_config: {variant: nvfp4}` to the YAML before starting the server; use `fp8r` for FP8.\n\n### Optimized single-GPU deployment\n\nFor 2048×2048 inference on one B200, save the following as `configs/`[` flux3-nvfp4-vc-fp8-1gpu.yaml`](https://github.com/NVIDIA/TensorRT-LLM/blob/feat/flux3-image-klein/examples/visual_gen/configs/flux3-nvfp4-vc-fp8-1gpu.yaml), which combines NVFP4 weights, torch.compile, and VC-Attention FP8.\n\n```\npipeline_config:\n  variant: nvfp4\ntorch_compile_config:\n  enable: true\ncompilation_config:\n  resolutions: [[2048, 2048]]\n  num_frames: [1]\n  modes: [t2i, i2i]\nattention_config:\n  backend: TRTLLM\n  quant_attention_config:\n    qk_dtype: fp8\n    v_dtype: fp8\n    vc_attention: true\n    vc_repair_budget: 0.005\n```\n\nFor Sol-Attention and Cache-DiT, copy this file to `configs/`[` flux3-nvfp4-sol-cache-1gpu.yaml`](https://github.com/NVIDIA/TensorRT-LLM/blob/feat/flux3-image-klein/examples/visual_gen/configs/flux3-nvfp4-sol-cache-1gpu.yaml), replace `attention_config`, and add `cache_config`:\n\n```\npipeline_config:\n  variant: nvfp4\ntorch_compile_config:\n  enable: true\ncompilation_config:\n  resolutions: [[2048, 2048]]\n  num_frames: [1]\n  modes: [t2i, i2i]\nattention_config:\n  backend: TRTLLM\n  sparse_attention_config:\n    algorithm: sol_attn\n    tau: 1.0\n    disabled_until_timestep: 0.96\ncache_config:\n  cache_backend: cache_dit\n  Fn_compute_blocks: 1\n  Bn_compute_blocks: 0\n  max_warmup_steps: 4\n  max_continuous_cached_steps: 3\n  residual_diff_threshold: 0.24\n```\n\nThe timestep cutoff of 0.96 keeps the first 12 denoising steps dense for 50-step FLUX 3 Image Klein Base and the first four steps dense for 16-step distilled Klein under the packaged α=6.93 schedule. For the 16-dense-step setting recommended above, set `disabled_until_timestep: 0.94`. The cutoff specifies a noise level rather than a number of steps. Remove `cache_config` to use Sol-Attention without Cache-DiT.\n\nSet `--visual_gen_args` in the offline inference or serving commands to the path of your selected YAML file. Set the image size to 2048×2048 to match the performance examples.\n\n## Conclusion and next steps\n\nTensorRT LLM VisualGen supports FLUX 3 Image Klein inference with configurable performance optimizations. NVFP4, fused kernels, compilation, multi-GPU execution, and CUDA Graphs accelerate execution; quantized and sparse attention and caching provide further latency reductions with measurable effects on output quality. Together, they deliver up to 3.46x faster T2I and 4.45x faster I2I generation using a single GPU on an NVIDIA DGX B200, compared with uncompiled BF16 VANILLA inference.\n\n**Get started:**\n\n- Run the end-to-end [FLUX 3 example](https://github.com/NVIDIA/TensorRT-LLM/blob/feat/flux3-image-klein/examples/visual_gen/models/flux3.py)\n- Explore the TensorRT LLM VisualGen [runtime and optimization utilities](https://github.com/NVIDIA/TensorRT-LLM/tree/feat/visual_gen/tensorrt_llm/visual_gen/)\n- File issues or share results on the [TensorRT LLM GitHub](https://github.com/NVIDIA/TensorRT-LLM)\n\n## Acknowledgements\n\nNikita Korobov, Harrison Zhang, Pamela Peng, Kanghwan Jang, Yuhang He, and Charlotte Mallo", "url": "https://wpnews.pro/news/speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus", "canonical_source": "https://developer.nvidia.com/blog/speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus/", "published_at": "2026-11-06 22:20:20+00:00", "updated_at": "2026-11-06 22:21:25.894338+00:00", "lang": "en", "topics": ["generative-ai", "ai-infrastructure", "ai-tools", "ai-chips"], "entities": ["NVIDIA", "Black Forest Labs", "FLUX 3 Image Klein", "TensorRT LLM VisualGen", "NVIDIA B200", "NVIDIA GB200 NVL72", "Qwen3-VL-4B-Instruct"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus", "markdown": "https://wpnews.pro/news/speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus.md", "text": "https://wpnews.pro/news/speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus.txt", "jsonld": "https://wpnews.pro/news/speed-up-image-generation-with-flux-3-image-klein-on-nvidia-blackwell-gpus.jsonld"}}