NVIDIA Blackwell GPU上FLUX.2模型的FP4推理优化技术

用户576110558132
4 阅读9分钟

在2025年，某机构与Black Forest Labs合作优化了FLUX.1和FLUX.2文生图模型，在Blackwell GeForce RTX 50系列和数据中心GPU上实现了FP4和FP8量化，显著降低了消费级和企业级部署的内存需求和延迟。

FLUX.2 [dev] 集成了先进的推理优化技术，包括采用微块缩放的NVFP4量化、时间步嵌入感知缓存 (TeaCache)、CUDA Graphs、torch.compile 以及通过TensorRT-LLM visual_gen 实现的序列并行多GPU支持。与H200 GPU相比，该优化实现了高达10.2倍的加速，并在B200、GB200、B300和GB300架构上实现了近乎线性的扩展。该优化流程保持了高输出保真度，NVFP4量化和TeaCache分别提供了约2倍的加速效果，同时保持了视觉质量，而文本编码器的FP8量化为大规模多GPU部署带来了更多优势。所有优化均在 NVIDIA/TensorRT-LLM/visual_gen 代码库中开源。

作为潜在扩散模型的自然延伸，FLUX.1 Kontext [dev] 证明了上下文学习不仅适用于大语言模型，对视觉生成模型同样可行。为了使这种体验更广泛可用，某机构与BFL合作，利用低精度量化实现了近乎实时的编辑体验。

FLUX.2 是一个巨大的飞跃，为公众提供了多图像参考功能，其质量可媲美最佳的企业级模型。然而，由于FLUX.2 [dev] 需要大量计算资源，BFL、Comfy和某机构合作实现了一项重大突破：将FLUX.2 [dev] 的内存需求降低了40%以上，并通过ComfyUI实现了本地部署。这项使用FP8精度的优化使得FLUX.2 [dev] 成为图像生成领域最受欢迎的模型之一。

随着FLUX.2 [dev] 成为开源权重模型的黄金标准，某机构团队与BFL合作，现在分享了性能的下一飞跃：在功能最强大的数据中心NVIDIA Blackwell GPU（包括某机构DGX B200和某机构DGX B300）上，为FLUX.2 [dev] 实现4位加速。

本文介绍了团队用于在这些某机构数据中心架构上加速FLUX.2 [dev] 的各种推理优化技术，包括代码片段和入门步骤。这些优化的综合效果显著降低了延迟，实现了在数据中心GPU上的高效部署。

BF16与NVFP4在FLUX.2 [dev] 上的视觉对比

在深入了解细节之前，先来查看FLUX.2 [dev] 在默认BF16精度下与使用NVFP4（图1和图2）时在输出质量上惊人的相似性。

图1的提示词为：“一只猫在一张舒适的沙发上安详地打盹。沙发位于一棵从月球表面生长的高树顶上。地球悬挂在远方，是黑暗太空中一颗充满活力的蓝绿色宝石。一艘光滑的宇宙飞船盘旋在附近，为场景投下柔和的光线，整个数字艺术作品呈现出梦幻般的质感。”

图1. 使用BF16精度（左）和NVFP4量化（右）的FLUX.2 [dev] 生成的图像

图2的提示词为：“一幅油画，描绘了一对穿着晚礼服的夫妇在倾盆大雨中回家，没有带伞。”在这个案例中，差异更难发现。最明显的是BF16图像中男士的微笑以及NVFP4图像背景中的多把雨伞。除此之外，两幅图像的前景和背景中的大部分细节都得到了保留。

图2. 使用BF16精度（左）和NVFP4量化（右）的FLUX.2 [dev] 生成的图像

优化 FLUX.2 [dev]

FLUX.2 [dev] 模型由三个关键组件构成：文本嵌入模型（具体为Mistral Small 3）、扩散变换器模型和一个自编码器。某机构团队将以下优化技术应用于开源的diffusers实现，使用了TensorRT-LLM/feat/visual_gen分支中的原型运行时：

  • NVFP4量化
  • 时间步嵌入感知缓存 (TeaCache)
  • CUDA Graphs
  • Torch compile
  • 多GPU支持

NVFP4量化

NVFP4通过引入两级微块缩放策略，推进了微缩放数据格式的概念。这种方法旨在最小化精度损失，并具有两种不同的机制：每张量缩放和每块缩放。

每张量缩放是一个以FP32精度存储的值，它调整整个张量分布，可以静态或动态计算。相比之下，每块缩放通过将张量划分为16个元素块来实时动态计算。

为了最大灵活性，用户可以选择将特定层保留为更高精度并应用动态量化，如下面使用FLUX.2 [dev] 的示例所示：

exclude_pattern = r"^(?!.*(embedder|norm_out|proj_out|to_add_out|to_added_qkv|stream)).*"

NVFP4计算的应用使用以下语句：

from visual_gen.layers import apply_visual_gen_linear
apply_visual_gen_linear(
    model,
    load_parameters=True,
    quantize_weights=True,
    exclude_pattern=exclude_pattern,
)

TeaCache

TeaCache技术用于加速推理过程。TeaCache通过使用扩散过程中生成的先前潜变量有条件地跳过扩散步骤。为了量化这种效果，进行了测试：在包含20个提示词和50步推理过程的场景中，TeaCache平均跳过了16步，这导致推理延迟降低了约30%。

为了确定TeaCache超参数的最佳配置，采用了网格搜索方法。该配置在计算速度和生成质量之间取得了最佳平衡。

dit_configs = {
        ...
   "teacache": {
       "enable_teacache": True,
       "use_ret_steps": True,
       "teacache_thresh": 0.05,
       "ret_steps": 10,
       "cutoff_steps": 50,
   },
        ...
}

缓存机制的缩放因子通过经验确定，并使用三次多项式进行近似。这个多项式是通过一个校准数据集拟合的，该数据集包含文生图和多参考图生成示例。

图3展示了这种经验方法，绘制了原始校准数据点以及拟合出的三次多项式曲线（红色所示），该曲线模拟了调制输入差值与模型输出差值之间的关系。

图3. 调制输入差值与模型预测输出差值之间的相关性

CUDA Graphs

TensorRT-LLM visual_gen 提供了一个现成的包装器来支持CUDA Graphs捕获。只需导入包装器并替换前向函数：

from visual_gen.utils.cudagraph import cudagraph_wrapper
model.forward = cudagraph_wrapper(model.forward)

Torch compile

在所有团队的实验中，除了基线运行外，torch.compile 都被启用，因为FLUX.2 [dev] 默认未启用它。

model = torch.compile(model)

多GPU支持

使用TensorRT-LLM visual_gen 启用多GPU涉及四个步骤：

  1. 修改模型的 forward 函数，插入处理GPU间通信的代码
  2. 将模型中的注意力实现替换为 ditAttnProcessor
  3. 在配置中选择并行算法并设置并行度
  4. 使用 torchrun 启动

以下代码片段提供了一个示例。在 model.forward 的开头插入拆分代码，将输入数据分布到多个GPU上：

from visual_gen.utils import (
    dit_sp_gather,
    dit_sp_split,
)

# ...
hidden_states = dit_sp_split(hidden_states, dim=1)
encoder_hidden_states = dit_sp_split(encoder_hidden_states, dim=1)
img_ids = dit_sp_split(img_ids, dim=1)
txt_ids = dit_sp_split(txt_ids, dim=1)

随后，在返回之前在 model.forward 的末尾插入收集代码：

output = dit_sp_gather(output, dim=1)

然后用提供的注意力处理器替换原始的注意力实现，以确保跨多个GPU的正确通信：

from visual_gen.layers import ditAttnProcessor

# ...
def attention(...):
    # ...
    x = ditAttnProcessor().visual_gen_attn(q, k, v, tensor_layout="HND")
    # ...

在配置中设置正确的并行大小。例如，要在四个GPU上使用Ulysses并行：

dit_config = {
...
    "parallel": {
          "dit_ulysses_size": 4,
    }
...
}

最后，调用 setup_configs API 来激活配置：

visual_gen.setup_configs(**dit_configs)

当使用多个GPU时，脚本必须使用 torchrun 启动。TensorRT-LLM visual_gen 将使用来自 torchrun 的排名信息，并正确处理所有通信和任务拆分。

性能分析

所有的推理优化都已包含在一个端到端的FLUX.2 [dev] 示例中——低精度内核、缓存技术和多GPU推理。

如图4所示，即使在默认BF16精度下，某机构DGX B200架构相比某机构H200也实现了1.7倍的代际提升。此外，应用分层推理优化——包括CUDA Graphs、torch.compile、NVFP4精度和TeaCache——将单B200的性能从该基线逐步提升至显著的6.3倍加速。最终，在两B200配置上的多GPU推理相比当前行业标准的H200，实现了10.2倍的性能提升。

图4. 在NVIDIA B200 GPU上FLUX.2 [dev] 的推理延迟比较

基线是未进行任何优化且未启用 torch.compile 的原始 FLUX.2 [dev]。优化系列包括启用 torch.compile、CUDA Graphs、NVFP4 和 TeaCache。基准测试中使用的扩散步数为50步。

在单GPU上，团队发现NVFP4和TeaCache在加速和输出质量之间提供了良好的权衡，各自提供了约2倍的加速。torch.compile 是一种近乎无损的加速技术，大多数开发者都很熟悉，但其收益有限。CUDA Graphs主要对多GPU推理有益，通过在某机构B200上使用多个GPU实现了增量扩展。最后，整个流程证明了文本编码器的FP8量化的鲁棒性，为大规模部署提供了额外的好处。

在多GPU上，TensorRT-LLM visual_gen 序列并行在增加更多GPU时实现了近乎线性的扩展。在NVIDIA Blackwell B200和GB200以及NVIDIA Blackwell Ultra B300和GB300 GPU上也观察到了同样的效果。针对NVIDIA Blackwell Ultra GPU的额外优化正在进行中。

图5. 在Blackwell GPU上的FLUX.2 [dev] 多GPU推理扩展

在NVIDIA Blackwell GPU上开始使用FLUX.2

FLUX.2 是图像生成领域的重大进步，成功地将高质量输出与用户友好的部署选项相结合。某机构团队与BFL合作，在功能最强大的某机构数据中心GPU上实现了FLUX.2 [dev] 的显著加速。

将新技术应用于FLUX.2 [dev] 模型，包括NVFP4量化和TeaCaching，带来了推理速度的强大代际飞跃。这些优化的综合效果显著降低了延迟，实现了在某机构数据中心GPU上的高效部署。

要开始使用这些最先进的优化构建自己的推理流程，请查看 NVIDIA/TensorRT-LLM/visual_gen GitHub 代码库中的端到端FLUX.2示例和附带代码。FINISHED

avatar
文章
阅读
粉丝