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 and distilled 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 and FLUX 3 Image Klein Base. 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 and Task Scheduling 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 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 or VC-Attention FP8, and sparse attention with Sol-Attention. VC-Attention FP8 is a TensorRT LLM variant of the method described in the VC-Attention paper.
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.
docker run --rm -it --ipc=host --gpus all -p 8000:8000 \
--name tensorrt_llm \
-e TLLM_FMHA_LIBS=+prims_ts \
-e HF_TOKEN=<your_hf_token> \
-v ~/.cache:/root/.cache:rw \
nvcr.io/nvidia/tensorrt-llm/release:feat-flux3-image-klein-1.4.0a0.dev202610280436
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.
Get started with TensorRT LLM VisualGen #
With 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.
These examples do not expand prompts. If you use BFL’s prompt upsampler, run it separately before inference. Reported latencies exclude prompt upsampling.
Offline inference
Inside the container, run the FLUX 3 example:
cd /app/tensorrt_llm/examples/visual_gen
python models/flux3.py \
--model black-forest-labs/flux-3-image-klein-base \
--visual_gen_args configs/flux3-images-1gpu.yaml \
--prompt "A small robot tending a rooftop garden at sunrise" \
--height 2048 --width 2048 --seed 0 \
--output_path flux3.png
For 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.
Online serving endpoint
From the same container directory, launch the OpenAI-compatible /v1/images/generations endpoint:
cd /app/tensorrt_llm/examples/visual_gen
trtllm-serve black-forest-labs/flux-3-image-klein-base \
--host 0.0.0.0 --port 8000 \
--visual_gen_args configs/flux3-images-1gpu.yaml
Once the server is ready, open a second terminal, enter the running container, and run the client:
docker exec -it tensorrt_llm bash
python examples/visual_gen/serve/sync_image_gen.py \
--base-url http://localhost:8000/v1 \
--model black-forest-labs/flux-3-image-klein-base \
--size 2048x2048 \
--prompt "A small robot tending a rooftop garden at sunrise" \
--output flux3-served.png
For NVFP4 serving, add pipeline_config: {variant: nvfp4} to the YAML before starting the server; use fp8r for FP8.
Optimized single-GPU deployment
For 2048×2048 inference on one B200, save the following as configs/ flux3-nvfp4-vc-fp8-1gpu.yaml, which combines NVFP4 weights, torch.compile, and VC-Attention FP8.
pipeline_config:
variant: nvfp4
torch_compile_config:
enable: true
compilation_config:
resolutions: [[2048, 2048]]
num_frames: [1]
modes: [t2i, i2i]
attention_config:
backend: TRTLLM
quant_attention_config:
qk_dtype: fp8
v_dtype: fp8
vc_attention: true
vc_repair_budget: 0.005
For Sol-Attention and Cache-DiT, copy this file to configs/ flux3-nvfp4-sol-cache-1gpu.yaml, replace attention_config, and add cache_config:
pipeline_config:
variant: nvfp4
torch_compile_config:
enable: true
compilation_config:
resolutions: [[2048, 2048]]
num_frames: [1]
modes: [t2i, i2i]
attention_config:
backend: TRTLLM
sparse_attention_config:
algorithm: sol_attn
tau: 1.0
disabled_until_timestep: 0.96
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
The 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.
Set --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.
Conclusion and next steps #
TensorRT 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.
Get started:
- Run the end-to-end FLUX 3 example
- Explore the TensorRT LLM VisualGen runtime and optimization utilities
- File issues or share results on the TensorRT LLM GitHub
Acknowledgements #
Nikita Korobov, Harrison Zhang, Pamela Peng, Kanghwan Jang, Yuhang He, and Charlotte Mallo