跳到内容

Fusion torch.compile 过程

vLLM 在编译时应用一组内核/算子融合(通过自定义 torch.compile Inductor 传递),旨在将优化与模型定义分离,并避免破坏模型代码中的层抽象。这些融合由 PassConfig 中的字段控制,并在适当的 优化级别 下自动启用。

快速参考

下表将每个融合映射到其控制标志/配置旋钮、它所融合的操作、默认启用的级别以及预期的加速效果。Fullgraph(全图)列表示该融合是否需要可见整个模型图(通过 Inductor 分区或 splitting_ops=[]),最后一列表示该融合是在所有 num_tokens 下激活,还是仅在低端或高端激活。

信息

加速效果在很大程度上取决于具体的模型、批处理大小和硬件。如果手动调整性能,请务必针对您的具体用例在启用和禁用融合的情况下进行基准测试,以验证其影响。

融合 PassConfig 标志 融合操作 默认启用级别 端到端加速 Fullgraph num_tokens
AllReduce + RMSNorm fuse_allreduce_rms All-reduce → RMSNorm (+residual_add) (→ 量化) O2 (Hopper/Blackwell + TP > 1) 5-20%
Attention + 量化 fuse_attn_quant Attention 输出 → FP8/NVFP4 量化 默认关闭 3-7% 必填 总是
MLA Attention + 量化 fuse_attn_quant MLA Attention 输出 → FP8/NVFP4 量化 默认关闭 待定 必填 总是
RoPE + KV-Cache 更新 fuse_rope_kvcache 旋转位置编码 → KV cache 写入 O2 (仅限 ROCm/AITER) 2-4%
QK Norm + RoPE enable_qk_norm_rope_fusion Q/K RMSNorm → 旋转位置编码 默认关闭 2-3%
序列并行 enable_sp AllReduce → ReduceScatter + AllGather 默认关闭 AsyncTP 前置条件 必填
AsyncTP GEMM + 集体通信 fuse_gemm_comms GEMM → reduce-scatter / all-gather → GEMM 默认关闭 7-10% 必填
RMSNorm + 量化 fuse_norm_quant RMSNorm (+residual add) → FP8/FP4 量化 O1 (有条件) 1-4% 总是
SiLU+Mul + 量化 fuse_act_quant SiLU+Mul 激活 → FP8/FP4 量化 O1 (有条件) 1-4% 总是
RMSNorm + 填充 fuse_act_padding Residual add + RMSNorm → 填充 O1 (仅限 ROCm/AITER) 待定 总是

支持矩阵

下表列出了每个平台对每种融合所支持的量化方案。 表示该融合在该平台上不可用。最新工作进展及正在进行的工作可在跟踪 issue 中查看: #36066

融合 SM100 (Blackwell) SM90 (Hopper) SM89 (Ada) SM80 (Ampere) ROCm
fuse_allreduce_rms FP16/BF16, FP8 static, NVFP4 FP16/BF16, FP8 static
fuse_attn_quant* FP8 static*, NVFP4* FP8 static* FP8 static* FP8 static*
fuse_attn_quant (MLA)* FP8 static*, NVFP4* FP8 static* FP8 static* FP8 static(untested)*
fuse_rope_kvcache FP16/BF16
enable_qk_norm_rope_fusion FP16/BF16 FP16/BF16 FP16/BF16† FP16/BF16†
enable_sp FP16/BF16, FP8 static† FP16/BF16, FP8 static FP16/BF16† FP16/BF16†
fuse_gemm_comms FP16/BF16, FP8 static† FP16/BF16, FP8 static FP16/BF16† FP16/BF16†
fuse_norm_quant FP8 static, FP8 per-token, FP8 per-group FP8 static, FP8 per-token, FP8 per-group FP8 static, FP8 per-token, FP8 per-group FP8 static, FP8 per-token, FP8 per-group
fuse_act_quant FP8 static, NVFP4 FP8 static, FP8 per-group (128/64) FP8 static, FP8 per-group (128/64) FP8 per-group
fuse_act_padding FP16/BF16

* fuse_attn_quant 的支持取决于所使用的 attention 后端;并非所有后端都支持融合量化输出。有关每个后端的详细信息,请参见 fuse_attn_quant 部分

enable_spfuse_gemm_comms 目前仅针对 SM90 自动配置;其他架构的支持需要显式设置 PassConfig.sp_min_token_num。SM100 支持还需要设置 VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel

启用 / 禁用融合

融合通过 PassConfig 公开,该配置嵌套在 CompilationConfig 内部。

from vllm import LLM
from vllm.config import CompilationConfig, PassConfig

llm = LLM(
    model="...",
    optimization_level=2, # Default optimization level
    compilation_config=CompilationConfig(
        pass_config=PassConfig(
            fuse_norm_quant=True,
            fuse_act_quant=True,
            fuse_allreduce_rms=False,  # disable a specific fusion
        )
    ),
)

融合还可以通过任何 vllm ... 命令的命令行标志启用。

# Enable O2 defaults, but turn off allreduce fusion
vllm serve meta-llama/Llama-3.1-8B-Instruct -O2 -cc.pass_config.fuse_allreduce_rms=False

# The above is equivalent to the more verbose:
vllm serve meta-llama/Llama-3.1-8B-Instruct -O2 --compilation-config '{"pass_config": {"fuse_allreduce_rms": false}}'

# Same syntax in other commands, e.g. vllm bench:
vllm bench latency --model=meta-llama/Llama-3.1-8B-Instruct -O2 -cc.pass_config.fuse_allreduce_rms=False

用户显式设置的字段总是优先于优化级别的默认值。

融合详情

AllReduce + RMSNorm (fuse_allreduce_rms)

警告

TP+DP 和 TP+PP 组合目前存在问题( #34458 #35426)。仅在安装了 FlashInfer 的 NVIDIA Hopper (SM90) 和 Blackwell (SM100) 上受支持。

它融合了什么。 将张量并行 All-Reduce 集体通信与随后的残差相加、RMSNorm 和可选的量化步骤融合成一个 FlashInfer / TRT-LLM 通信内核。此融合仅在 num_tokens 较小时才有利可图,因此仅在编译后的较低范围内执行。

涵盖的模式

  • AllReduce → RMSNorm(+residual_add): 带有 FlashInfer 的 CUDA sm90+
  • AllReduce → RMSNorm(+residual_add) → FP8 static quant: 带有 FlashInfer 的 CUDA sm90+
  • AllReduce → RMSNorm(+residual_add) → NVFP4 dynamic quant: 带有 FlashInfer 的 CUDA sm100+

使用融合内核的最大张量大小取决于硬件(SM90/SM100 上 TP=2 时为 64 MB),可通过 PassConfig.fi_allreduce_fusion_max_size_mb 进行配置。

代码位置。

Attention + 量化 (fuse_attn_quant)

信息

fuse_attn_quant 目前在任何优化级别中默认均未启用,必须显式设置。它要求模型图全貌可见(Inductor 分区或 splitting_ops=[])。

它融合了什么。 在 Attention 计算之后直接融合 Attention 输出量化,消除了 Attention 输出的全精度内存往返。此融合支持标准 AttentionMLAAttention(由 DeepSeek-V2/V3/R1 模型使用)。涵盖的模式

Attention → FP8 static quant:

  • TRITON_ATTN: CUDA, ROCm
  • FLASHINFER: 带有 FlashInfer 的 CUDA sm100+
  • ROCM_ATTN: ROCm
  • ROCM_AITER_UNIFIED_ATTN: 带有 AITER 的 ROCm

Attention → NVFP4 dynamic quant:

  • FLASHINFER: 带有 FlashInfer 的 CUDA sm100+

MLAAttention → FP8 static quant / MLAAttention → NVFP4 dynamic quant

MLA 融合在图级别对 unified_mla_attention_with_output 操作进行处理,适用于所有 MLA decode 和 prefill 后端组合。与标准 Attention 后端不同(内核直接写入 FP8 输出),目前没有任何 MLA prefill 或 decode 后端支持直接 FP8/FP4 输出。该融合会写入中间缓冲区并在单独的步骤中进行量化,因此尚未消除内存往返。

信息

MLA Attention 融合目前预计不会产生可测量的加速。一旦 MLA prefill/decode 内核支持直接 FP8/FP4 输出,这种情况将会改善。

其他 Attention 后端尚不支持融合输出量化。

代码位置。

RoPE + KV-Cache 更新 (fuse_rope_kvcache)

信息

仅限 ROCm/AITER。在 NVIDIA CUDA 或 CPU 上不可用。由于 AITER 融合内核性能问题,此融合仅在 num_tokens ≤ 256 时默认启用。此阈值可通过 PassConfig.rope_kvcache_fusion_max_token_num 进行配置。

它融合了什么。 将旋转位置编码内核与 KV-cache 散射/写入融合成单个内核,避免了对 key 和 value 张量的单独读取和写入。

要求:启用 AITER 的 AMD ROCm,处于活跃状态的 rotary_embedding 自定义算子(自动),以及图中的 kv_cache 更新算子可见:通过使用 Inductor 图分区或从 splitting_ops 中移除。如果设置了这些条件,融合会在优化级别 O1 及以上自动启用。

代码位置。

序列并行 (enable_sp)

它融合了什么。 用 reduce-scatter + 本地 RMSNorm + all-gather 替换 All-reduce 集体通信,在 TP 秩之间分割序列维度。这重构了图,以便后续的 AsyncTP 传递可以将 reduce-scatter / all-gather 与周围的 GEMM 进行融合。

序列并行本身并不直接提高性能;它是 AsyncTP 传递 (fuse_gemm_comms) 的先决条件。SP 仅在超过基于设备能力和模型 hidden_size 自动配置的最小 token 阈值时应用。目前仅在 hidden_size >= 8192 的模型的 H100/SM90 上激活。该阈值可通过 PassConfig.sp_min_token_num 配置。

通用转换

Input → AllReduce → RMSNorm → Output
becomes:
Input → ReduceScatter → local RMSNorm → AllGather → Output

涵盖的模式

  • 第一块: AllReduce → RMSNormReduceScatter → RMSNorm → AllGather
  • 中间块: AllReduce → fused_add_RMSNormReduceScatter → fused_add_RMSNorm → AllGather
  • 两者均带有可选的 → FP8 static quant 后缀

要求:use_inductor_graph_partition=True 静态大小可被 tensor_parallel_size 整除的分段编译。

支持的硬件:仅在 NVIDIA CUDA 上测试,可能在 ROCm 上有效。FP8 all-gather 要求 sm90+。

代码位置。

AsyncTP GEMM + 集体通信重叠 (fuse_gemm_comms)

信息

要求 enable_sp=True(自动启用)。如果未应用序列并行,则此传递为空操作。

它融合了什么。 在序列并行转换图之后,使用 torch.ops.symm_mem 对称内存原语将 GEMM 内核与周围的 reduce-scatter(输出投影)和 all-gather(输入投影)融合,实现通信与计算的重叠。这种重叠仅对大的 num_tokens 有利,因此融合(以及前面的 SP)仅在高于 PassConfig.sp_min_token_num 的较高编译范围内执行。

涵盖的模式

  • GEMM → reduce-scatterfused_matmul_reduce_scatter
  • all-gather → GEMMall_gather_matmul
  • 两种模式的 FP8 缩放变体

支持的硬件:具有对称内存 (torch.distributed._symmetric_memory) 支持的 NVIDIA CUDA。

在 B200 上,不支持模式匹配 fp8 FlashInfer 缩放 MM,因此必须禁用它( #27893

VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel ...

代码位置。

QK Norm + RoPE (enable_qk_norm_rope_fusion)

信息

仅适用于在旋转位置编码之前对 Q 和 K 应用每头 RMSNorm 的模型(例如 Qwen)。由于 H100 上的性能问题,在任何优化级别中默认均未启用: #34391

它融合了什么。 将序列:分割 QKV → 重塑 → Q/K RMSNorm → 重塑 → 旋转位置编码融合成单个 fused_qk_norm_rope CUDA 内核。

# Unfused:
q, k, v = split(qkv)
q_norm = rms_norm(q.view(heads))
k_norm = rms_norm(k.view(kv_heads))
q_rope, k_rope = rotary_embedding(q_norm, k_norm, ...)

# Fused:
fused_qk_norm_rope(qkv, ...)

支持的硬件:仅限 CUDA (sm80+),仅在 sm90 和 sm100 上测试。

代码位置。

RMSNorm + 量化 (fuse_norm_quant)

警告

在 NVIDIA 上,Inductor 生成的融合内核实际上比我们的自定义 CUDA 内核更快。因此,此融合仅在 rms_normquant_fp8 使用自定义内核时启用。

它融合了什么。 将自定义的 rms_norm / fused_add_rms_norm 操作与随后的量化结合成单个融合内核,消除了全精度激活张量的中间读/写。两种变体被融合

  • 普通 RMSNorm + 量化: rms_norm(x) → quant_fp8(y)
  • 融合加 RMSNorm + 量化: fused_add_rms_norm(x, residual) → quant_fp8(y) — 同时就地更新残差。

注意,AITER 融合目前在 vllm.compilation.passes.fusion.rocm_aiter_fusion 中的单独传递中。

支持的量化方案/硬件组合

  • FP8 static per-tensor: CUDA & HIP 内核
  • FP8 dynamic per-token: CUDA & HIP 内核, AITER
  • FP8 dynamic per-token-group (128/64): CUDA & HIP 内核, AITER

代码位置。

SiLU+Mul + 量化 (fuse_act_quant)

警告

fuse_norm_quant 相同:在 NVIDIA 上,Inductor 生成的融合内核比我们的自定义算子更快。此融合仅在 silu_and_mulquant_fp8 使用自定义内核时,或对于 NVFP4 量化模型(其中 FP4 量化始终是自定义算子)才启用。

它融合了什么。silu_and_mul 门控上投影激活与随后的量化融合成单个内核,避免了全精度激活后张量的具体化。

注意,AITER 融合在 vllm.compilation.passes.fusion.rocm_aiter_fusion 中的单独传递中。

支持的量化方案/硬件组合

  • FP8 static per-tensor: CUDA & HIP 内核
  • FP8 dynamic per-group (128/64): CUDA 内核 (sm89+,在 sm100+ 上使用 DeepGemm 时不激活)
  • NVFP4 dynamic: 仅限带有 FlashInfer 的 CUDA sm100+
  • FP8 per-token-group (128): 仅限 ROCm AITER

代码位置。

RMSNorm + 填充 (fuse_act_padding)

信息

仅限 ROCm/AITER。针对 GPT-OSS 模型。

它融合了什么。 将残差相加 + RMSNorm 与随后的填充操作融合,将隐藏维度填充到下游 AITER Triton GEMM 内核所需的倍数。

要求:启用 AITER RMSNorm 的 AMD ROCm。当隐藏大小为 2880 且 启用 AITER Triton GEMM 时,在优化级别 O1 及以上默认启用。

代码位置。

参见