返回文章 apply CMS 文章

在 NVIDIA Blackwell GPU 上使用 TensorRT LLM 加速 FLUX 3 Image Klein 图像生成

用 TensorRT LLM VisualGen 的量化、注意力、缓存与多 GPU 优化，把 FLUX 3 Image Klein 的推理延迟压到原来的三分之一左右。

NVIDIA

FLUX 3 Image KleinTensorRT LLMNVIDIA BlackwellNVFP4

成长分 / 100 75 综合收获、行动、留存与影响

在 NVIDIA Blackwell GPU 上使用 TensorRT LLM 加速 FLUX 3 Image Klein 图像生成

为什么值得读了解 FLUX 3 Image Klein 在 Blackwell 上的完整推理优化路径，从 NVFP4 量化到注意力、缓存和多 GPU 并行逐层叠加。

获得可直接复用的 YAML 配置与部署说明，覆盖 Base 与 distilled 两种 checkpoint 的不同调优策略。

关键洞察

在单张 B200 上，组合优化后 FLUX 3 Image Klein Base 的 2048×2048 T2I 提速 3.46 倍、I2I 提速 4.45 倍（对比未编译 BF16 推理）。
NVFP4 相比 BF16 在 Blackwell 上带来最高 1.62 倍 T2I 和 1.58 倍 I2I 提速，是后续优化的起点。
去噪阶段占总生成时间的 84%–97.8%，因此 DiT 加速是优化重点；VANILLA BF16 注意力约占 DiT GPU kernel 时间的 60%。

转成行动

深入阅读

正文与原文对照

原文保真覆盖：全文原文字符：24419

在设计工具、图像编辑器和视觉生产工作流中部署扩散模型时，开发者必须考虑推理延迟。生成一张图像需要对具有数十亿参数的模型进行多次顺序前向传递，而更高的分辨率会增加注意力开销。

FLUX 3 Image Klein 是 Black Forest Labs（BFL）推出的图像编辑与生成模型，在 NVIDIA GB200 NVL72 上训练，并针对 NVIDIA GPU 进行了优化。它提供 Base 和 distilled 检查点，以满足不同的延迟和保真度需求。

本教程展示如何在 NVIDIA TensorRT LLM VisualGen 中配置 FLUX 3 Image Klein，以利用融合内核、多 GPU 并行、FP8/NVFP4 支持、缓存和注意力优化。它比较了各项优化下的延迟和生成输出的变化。在单块 NVIDIA B200 GPU 上报告的 2048×2048 测试中，与未编译的 BF16 推理相比，组合优化使 FLUX 3 Image Klein Base 的文本到图像（T2I）生成速度最高提升 3.46 倍，图像到图像（I2I）生成速度最高提升 4.45 倍。本教程包含部署说明和 YAML 配置，用于将这些优化应用于你的工作负载。

使用 TensorRT LLM VisualGen 进行 FLUX 3 推理

FLUX 3 Image Klein 结合了用于文本条件化的 Qwen3-VL-4B-Instruct、用于细化图像潜变量的扩散 Transformer（DiT），以及用于将潜变量解码为图像并编码 I2I 参考输入的自动编码器。Klein 独立处理文本、含噪目标图像和可选参考图像流，然后通过联合注意力将它们组合起来。

TensorRT LLM VisualGen 是 TensorRT LLM 中的专用运行时，发布时支持 BF16、FP8 和 NVFP4 变体。在 NVIDIA Blackwell GPU 上，NVFP4 的 T2I 生成速度比 BF16 最高提升 1.62 倍，I2I 生成速度最高提升 1.58 倍。图 2 比较了生成延迟，图 1 比较了示例输出。以下各节以 NVFP4 为起点进行进一步优化。

图 1 比较了 2048×2048 下的输出。概括的 T2I 提示描述了一个厨房静物；概括的 I2I 提示要求修改一个太空场景，包括一颗红色气态巨行星、一条银龙和一块雕刻铭牌。

图 1. T2I 和 I2I 的 BF16 与 NVFP4 输出

FLUX 3 各模型和精度变体的生成延迟比较。 — 图 2. FLUX 3 Image Klein 生成延迟

以 NVFP4 为起点，VisualGen 优化转向去噪，去噪在 B200 上主导了整体生成时间。因此，加速 FLUX 3 的 DiT 是以下各节的重点。

表 1 展示了在单块 NVIDIA B200 GPU 上使用 NVFP4、VANILLA 注意力和 torch.compile 时，2048×2048 下的流水线延迟。结果是在两次预热请求后对 10 次请求取平均值。Base 使用 50 个去噪步骤，无分类器引导（CFG）缩放为 3.5；distilled Klein 使用 16 个步骤，不使用 CFG。

表 1. NVIDIA B200 上的 FLUX 3 Image Klein 流水线延迟

任务	模型	文本 + 准备	去噪	解码 + 后处理	总计	去噪占比
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%

在 TensorRT LLM VisualGen 中优化 FLUX 3

TensorRT LLM VisualGen 通过优化的通用矩阵乘法（GEMM）内核和 epilogue 融合、低精度和稀疏注意力、缓存以及多 GPU 执行来加速 NVFP4 DiT 推理。除非另有说明，以下基准测试使用 torch.compile 和种子 0。Base 使用 50 个去噪步骤和 3.5 的 CFG 缩放；distilled Klein 使用 16 个去噪步骤。

NVFP4 推理

VisualGen 支持 BFL 为 FLUX 3 Image Klein 和 FLUX 3 Image Klein Base 提供的预量化 NVFP4 DiT 检查点。要使用这些检查点，请在 pipeline_config 中将 variant 设置为 nvfp4：

pipeline_config:
  variant: nvfp4

GEMM 内核优化

TensorRT LLM VisualGen 使用 FlashInfer 的 PrimsTS FP8/NVFP4 GEMM 内核，这些内核基于 CUTLASS Primitives 和 Task Scheduling 构建，以加速量化推理。它们针对 FLUX 3 的张量形状进行了调优并默认启用，将优化的 GEMM 与融合 epilogue 相结合，包括带 SwiGLU 的 GEMM，以及带查询/键归一化和旋转位置嵌入（RoPE）的查询、键和值（QKV）投影。更多详情请参见 FlashInfer PR #5758。

发布时，固定版本的 NVFP4 PrimsTS 内核在 B200 上的 T2I 生成速度比自动调优的 TensorRT LLM NVFP4 GEMM 实现快 1.02–1.10 倍。

延迟比较显示使用 PrimsTS GEMM 生成更快。 — 图 3. 使用 PrimsTS 和 TensorRT LLM GEMM 的 T2I 延迟

量化注意力与稀疏注意力

注意力是高分辨率图像生成的主要瓶颈。在 B200 上报告的 2048×2048 NVFP4 T2I 测试中，VANILLA BF16 注意力在 Base 和蒸馏版 Klein 上均约占 DiT GPU 内核时间的 60%。VisualGen 支持两种降低该开销的方法：使用 cuDNN FP8 或 VC-Attention FP8 的量化注意力，以及使用 Sol-Attention 的稀疏注意力。VC-Attention FP8 是 VC-Attention 论文中所述方法的 TensorRT LLM 变体。

学习感知图像块相似度（LPIPS）衡量生成图像之间的感知差异。在以下比较中，分数越低表示与对应 VANILLA 基线输出的相似度越高。

与 VANILLA BF16 注意力相比，cuDNN FP8 为 Base 带来 1.08 倍的生成加速，为 16 步蒸馏版 Klein 带来 1.09 倍的加速。对应的 LPIPS 分数分别为 0.072 和 0.161。VC-Attention FP8 为 Base 带来 1.18 倍的加速，为蒸馏版 Klein 带来 1.16 倍的加速，对应的 LPIPS 分数分别为 0.086 和 0.151。

图 4 比较了在启用 CUDA Graphs 的情况下，FLUX 3 Image Klein Base NVFP4 在 2048×2048 分辨率下的输出。

这些测量结果表明，在 B200 上以 2048×2048 分辨率进行量化注意力时应使用 VC-Attention FP8。测试配置使用 ρ=0.005 的修复预算，以选择量化误差最高的 token 比例进行残差校正。要使用量化注意力，请将 attention_config 替换为以下配置之一：

对于 cuDNN FP8：

attention_config:
  backend: CUDNN
  quant_attention_config:
    qk_dtype: fp8
    v_dtype: fp8

对于 VC-Attention FP8：

attention_config:
    backend: TRTLLM
    quant_attention_config:
        qk_dtype: fp8
        v_dtype: fp8
        vc_attention: true
        vc_repair_budget: 0.005

茶场景；Dense 基线 · 1.000x。 — 茶场景 · Dense 基线 · 1.000x · LPIPS 0.000

茶场景；cuDNN FP8 · 1.08x。 — 茶场景 · cuDNN FP8 · 1.08x · LPIPS 0.072

茶场景；VC-Attn FP8 ρ=0.005 · 1.18x。 — 茶场景 · VC-Attn FP8 ρ=0.005 · 1.18x · LPIPS 0.086

猫场景；Dense 基线 · 1.000x。 — 猫场景 · Dense 基线 · 1.000x · LPIPS 0.000

猫场景；cuDNN FP8 · 1.08x。 — 猫场景 · cuDNN FP8 · 1.08x · LPIPS 0.066

猫场景；VC-Attn FP8 ρ=0.005 · 1.18x。 — 猫场景 · VC-Attn FP8 ρ=0.005 · 1.18x · LPIPS 0.049

图 4. BF16、cuDNN FP8 和 VC-Attention FP8 示例输出

Sol-Attention 通过由 τ 控制的稀疏性减少 DiT 计算。报告的测试使用 TensorRT LLM 中基于 PrimsTS 的通用块稀疏注意力内核。按如下方式配置 Sol-Attention：

attention_config:
  backend: TRTLLM
  sparse_attention_config:
    algorithm: sol_attn
    tau: 1.0
    disabled_until_timestep: 0.96

对于 B200 上 2048×2048 分辨率的 NVFP4 FLUX 3 Image Klein Base，τ=1.0 的 Sol-Attention 在 16 个初始密集步下，去噪加自编码器解码的速度比编译后的密集 VANILLA 快 1.24 倍，在 12 个初始密集步下快 1.30 倍。

这些测量结果表明，应使用 16 个初始密集步以优先保证保真度，或使用 12 个以优先保证速度。更高的 τ 会降低输出质量，而额外加速很小。

图 5 比较了在启用 CUDA Graphs 且 τ=1.0 的情况下，NVFP4 FLUX 3 Image Klein Base 在 2048×2048 分辨率下的输出。

茶场景；Sol-Attn τ=1.0 · 16 个密集步 · 1.235x。 — 茶场景 · Sol-Attn τ=1.0 · 16 个密集步 · 1.235x · LPIPS 0.075

茶场景；Sol-Attn τ=1.0 · 12 个密集步 · 1.297x。 — 茶场景 · Sol-Attn τ=1.0 · 12 个密集步 · 1.297x · LPIPS 0.151

猫场景；Sol-Attn τ=1.0 · 16 个密集步 · 1.235x。 — 猫场景 · Sol-Attn τ=1.0 · 16 个密集步 · 1.235x · LPIPS 0.116

猫场景；Sol-Attn τ=1.0 · 12 个密集步 · 1.297x。 — 猫场景 · Sol-Attn τ=1.0 · 12 个密集步 · 1.297x · LPIPS 0.140

图 5. 使用 16 或 12 个初始密集步的 Dense 注意力和 Sol-Attention 输出

步缓存

Cache-DiT 通过在去噪步之间复用中间 transformer 输出，在不重新训练或修改模型权重的情况下降低 FLUX 3 Image Klein 推理延迟。残差差异阈值控制何时复用缓存输出，因此可以调整推理速度与输出质量之间的平衡。使用所需参数配置 Cache-DiT：

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

在阈值为 0.24 时，Cache-DiT 在两个测试分辨率下几乎将 NVFP4 Base T2I 延迟减半，对蒸馏版 Klein 的增益较小。在 1024×1024 下，Base 的加速为 1.97 倍，蒸馏版 Klein 为 1.26 倍，相对于各自未缓存输出的对应 LPIPS 分数分别为 0.038 和 0.036。在 2048×2048 下，Base 的加速为 1.94 倍，蒸馏版 Klein 为 1.20 倍（图 6）。

日落时的灯塔；FLUX 3 Image Klein Base · 1x。 — Base 未使用 Cache-DiT · 1x · LPIPS 0.000

日落时的灯塔；+ Cache-DiT · 1.94x。 — Base 使用 Cache-DiT · 1.94x · LPIPS 0.033

日落时的灯塔；FLUX 3 Image Klein 蒸馏版 · 1x。 — 蒸馏版 Klein 未使用 Cache-DiT · 1x · LPIPS 0

日落时的灯塔；+ Cache-DiT · 1.20x。 — 蒸馏版 Klein 使用 Cache-DiT · 1.20x · LPIPS 0.022

图 6. 使用和不使用 Cache-DiT 的 Base 与蒸馏版 Klein 输出

多 GPU 推理

TensorRT LLM 中的 VisualGen 运行时使用无分类器引导（CFG）和 Ulysses 序列并行来扩展 FLUX 3 Image Klein。Base 将正/负分支分配到不同的 GPU 组，并使用 Ulysses 对每个分支的序列进行划分；蒸馏版不使用 CFG 运行，仅使用 Ulysses。

VisualGen 处理序列拆分、分布式注意力和输出收集。对于八 GPU NVFP4 Base 推理，使用以下配置：

parallel_config:
    cfg_size: 2
    ulysses_size: 4

对于蒸馏模型，设置 cfg_size: 1、ulysses_size: 8，并使用 flux-3-image-klein 检查点。

在 2048×2048 下，在 NVIDIA DGX B200 上从一 GPU 扩展到八 GPU，将 Base 的中位 T2I 延迟从 14.90 秒降至 3.00 秒，蒸馏版从 2.73 秒降至 1.08 秒。去噪加速更大，Base 为 6.01 倍，蒸馏版 Klein 为 4.44 倍，但八 GPU 蒸馏版生成时间中约有一半仍在去噪之外，限制了端到端增益。

多 GPU 扩展比较显示，增加 B200 GPU 可降低生成延迟。 — 图 7. FLUX 3 Image Klein 多 GPU 扩展

在较小分辨率下，随着每个 GPU 的工作量减少，CUDA Graphs 通过降低内核启动开销来补充并行性。在另一项 1024×1024 NVFP4 实验中，使用 VANILLA 注意力在八块 B200 上，将 CUDA Graphs 加入编译后，Base 延迟从 1.823 秒降至 0.809 秒（2.25 倍，CFG 跨两个 GPU 组，每组内 Ulysses 跨四个 GPU），蒸馏版延迟从 0.791 秒降至 0.366 秒（2.16 倍，Ulysses 跨八个 GPU）。使用以下配置启用 CUDA Graphs：

cuda_graph_config:
    enable: true

FLUX 3 端到端性能

在 2048×2048 下，图 8 比较了连续 TensorRT LLM VisualGen 优化阶段相对于未编译 BF16 VANILLA 推理的端到端生成加速。

在单块 B200 上，编译、使用 PrimsTS GEMM 的 NVFP4、Sol-Attention 和 Cache-DiT 为 T2I 带来了 3.46 倍的生成加速，将延迟从 22.36 秒降至 6.46 秒；为 I2I 带来了 4.45 倍的加速，将延迟从 54.64 秒降至 12.28 秒。16 步蒸馏版 Klein 实现了高达 2.34 倍的端到端加速。

在单块 NVIDIA RTX PRO 6000 Blackwell Server Edition GPU 上，编译、使用融合 CuTe DSL 内核的 NVFP4 和 Cache-DiT 带来了 1.73-3.25 倍的生成加速，将增益扩展到两个平台。图 9 和图 10 比较了各优化阶段的示例输出。

连续优化提升了 B200 和 NVIDIA RTX PRO 6000 Blackwell Server Edition GPU 上的生成速度。 — 图 8. 各优化阶段的累计单 GPU 加速

图 9 和图 10 中的示例在单块 B200 GPU 上生成。

秋日森林中的红狐；BF16，VANILLA 注意力。 — 1. BF16，VANILLA 注意力 · LPIPS 0

秋日森林中的红狐；+ torch.compile。 — 2. + torch.compile · LPIPS 0

秋日森林中的红狐；+ PrimsTS NVFP4。 — 3. + PrimsTS NVFP4 · LPIPS 0.044

秋日森林中的红狐；+ Sol-Attention。 — 4. + Sol-Attention · LPIPS 0.084

秋日森林中的红狐；+ Cache-DiT。 — 5. + Cache-DiT · LPIPS 0.095

图 9. FLUX 3 Image Klein Base T2I 在各优化阶段的输出

编辑提示要求独木舟呈现风化外观，同时保留其形状、位置和周围场景。

山间湖泊上风化的独木舟；BF16，VANILLA 注意力。 — 1. BF16，VANILLA 注意力 · LPIPS 0

山间湖泊上风化的独木舟；+ torch.compile。 — 2. + torch.compile · LPIPS 0

山间湖泊上风化的独木舟；+ PrimsTS NVFP4。 — 3. + PrimsTS NVFP4 · LPIPS 0.031

山间湖泊上风化的独木舟；+ Sol-Attention。 — 4. + Sol-Attention · LPIPS 0.029

山间湖泊上风化的独木舟；+ Cache-DiT。 — 5. + Cache-DiT · LPIPS 0.045

图 10. FLUX 3 Image Klein Base I2I 在各优化阶段的输出

要将这些优化应用于您自己的工作负载，请按照以下各节中的运行时设置、推理命令和 YAML 配置进行离线生成或在线服务。

前提条件

使用以下命令启动预构建的 TensorRT LLM 容器。要从 feat/flux3-image-klein 分支构建它，请遵循 FLUX 3 部署指南。

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 启用 VC-Attention 所需的 PrimsTS 注意力内核。本文中的优化和结果聚焦于 NVIDIA Blackwell 数据中心 GPU 和 NVIDIA RTX PRO 6000 Blackwell Server Edition GPU。

开始使用 TensorRT LLM VisualGen

运行时环境就绪后，您可以离线运行 FLUX 3，或通过图像生成端点提供服务。使用以下各节中的示例 YAML 配置来应用前文报告结果所用的优化设置。

这些示例不会扩展提示。如果您使用 BFL 的提示上采样器，请在推理前单独运行它。报告的延迟不包括提示上采样。

离线推理

在容器内，运行 FLUX 3 示例：

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

对于 Blackwell 上的 NVFP4，添加 --variant nvfp4；使用 --variant fp8r 表示逐行 FP8，或省略 --variant 表示 BF16。为蒸馏模型选择 black-forest-labs/flux-3-image-klein。对于编辑，添加编辑提示和 --image /path/to/reference.png；为多个参考重复 --image。

在线服务端点

从同一容器目录，启动 OpenAI 兼容的 /v1/images/generations 端点：

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

服务器就绪后，打开第二个终端，进入正在运行的容器，并运行客户端：

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

对于 NVFP4 服务，在启动服务器前向 YAML 添加 pipeline_config: {variant: nvfp4}；使用 fp8r 表示 FP8。

优化的单 GPU 部署

对于单块 B200 上的 2048×2048 推理，将以下内容保存为 configs/flux3-nvfp4-vc-fp8-1gpu.yaml，它结合了 NVFP4 权重、torch.compile 和 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

对于 Sol-Attention 和 Cache-DiT，将此文件复制到 configs/flux3-nvfp4-sol-cache-1gpu.yaml，替换 attention_config，并添加 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

在打包的 α=6.93 调度下，0.96 的时间步截止值使 50 步 FLUX 3 Image Klein Base 的前 12 个去噪步骤保持稠密，使 16 步蒸馏版 Klein 的前四个步骤保持稠密。对于上文推荐的 16 稠密步设置，设置 disabled_until_timestep: 0.94。截止值指定的是噪声水平而非步骤数。移除 cache_config 以使用不带 Cache-DiT 的 Sol-Attention。

在离线推理或服务命令中，将 --visual_gen_args 设置为您所选 YAML 文件的路径。将图像大小设置为 2048×2048 以匹配性能示例。

结论与后续步骤

TensorRT LLM VisualGen 支持 FLUX 3 Image Klein 推理，并提供可配置的性能优化。NVFP4、融合内核、编译、多 GPU 执行和 CUDA Graphs 可加速执行；量化和稀疏注意力以及缓存可进一步降低延迟，并对输出质量产生可测量的影响。与未编译的 BF16 VANILLA 推理相比，在 NVIDIA DGX B200 上使用单块 GPU，它们共同实现了高达 3.46 倍的 T2I 加速和 4.45 倍的 I2I 生成加速。

开始使用：

运行端到端 FLUX 3 示例
探索 TensorRT LLM VisualGen 运行时和优化实用工具
在 TensorRT LLM GitHub 上提交问题或分享结果

致谢

Nikita Korobov、Harrison Zhang、Pamela Peng、Kanghwan Jang、Yuhang He 和 Charlotte Mallo