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_sp 和 fuse_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 进行配置。
代码位置。
- Pass:
vllm/compilation/passes/fusion/allreduce_rms_fusion.py - FlashInfer all-reduce:
vllm/distributed/device_communicators/flashinfer_all_reduce.py - 基准测试:
benchmarks/kernels/benchmark_fused_collective.py
Attention + 量化 (fuse_attn_quant)¶
信息
fuse_attn_quant 目前在任何优化级别中默认均未启用,必须显式设置。它要求模型图全貌可见(Inductor 分区或 splitting_ops=[])。
它融合了什么。 在 Attention 计算之后直接融合 Attention 输出量化,消除了 Attention 输出的全精度内存往返。此融合支持标准 Attention 和 MLAAttention(由 DeepSeek-V2/V3/R1 模型使用)。涵盖的模式
Attention → FP8 static quant:
TRITON_ATTN: CUDA, ROCmFLASHINFER: 带有 FlashInfer 的 CUDA sm100+ROCM_ATTN: ROCmROCM_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 后端尚不支持融合输出量化。
代码位置。
- Pass (Attention):
vllm/compilation/passes/fusion/attn_quant_fusion.py - Pass (MLAAttention):
vllm/compilation/passes/fusion/mla_attn_quant_fusion.py - Attention 后端:
vllm/v1/attention/backends/
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 → RMSNorm→ReduceScatter → RMSNorm → AllGather - 中间块:
AllReduce → fused_add_RMSNorm→ReduceScatter → 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-scatter→fused_matmul_reduce_scatterall-gather → GEMM→all_gather_matmul- 两种模式的 FP8 缩放变体
支持的硬件:具有对称内存 (torch.distributed._symmetric_memory) 支持的 NVIDIA CUDA。
在 B200 上,不支持模式匹配 fp8 FlashInfer 缩放 MM,因此必须禁用它( #27893)
代码位置。
- Pass:
vllm/compilation/passes/fusion/collective_fusion.py - 序列并行传递:
vllm/compilation/passes/fusion/sequence_parallelism.py
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 上测试。
代码位置。
- Pass:
vllm/compilation/passes/fusion/qk_norm_rope_fusion.py - CUDA 内核:
csrc/ops.h(fused_qk_norm_rope)
RMSNorm + 量化 (fuse_norm_quant)¶
警告
在 NVIDIA 上,Inductor 生成的融合内核实际上比我们的自定义 CUDA 内核更快。因此,此融合仅在 rms_norm 或 quant_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
代码位置。
- Pass:
vllm/compilation/passes/fusion/rms_quant_fusion.py - ROCm AITER 传递:
vllm/compilation/passes/fusion/rocm_aiter_fusion.py - CUDA/HIP 内核:
csrc/layernorm_quant_kernels.cu
SiLU+Mul + 量化 (fuse_act_quant)¶
警告
与 fuse_norm_quant 相同:在 NVIDIA 上,Inductor 生成的融合内核比我们的自定义算子更快。此融合仅在 silu_and_mul 或 quant_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
代码位置。
- Pass:
vllm/compilation/passes/fusion/act_quant_fusion.py - ROCm AITER 传递:
vllm/compilation/passes/fusion/rocm_aiter_fusion.py - CUDA/HIP 内核:
csrc/quantization/ - 融合 SiLU+Mul+BlockQuant 内核:
csrc/quantization/fused_kernels/fused_silu_mul_block_quant.cu
RMSNorm + 填充 (fuse_act_padding)¶
信息
仅限 ROCm/AITER。针对 GPT-OSS 模型。
它融合了什么。 将残差相加 + RMSNorm 与随后的填充操作融合,将隐藏维度填充到下游 AITER Triton GEMM 内核所需的倍数。
要求:启用 AITER RMSNorm 的 AMD ROCm。当隐藏大小为 2880 且 未 启用 AITER Triton GEMM 时,在优化级别 O1 及以上默认启用。
代码位置。
参见¶
- 优化级别 — 设置融合默认值的高级预设。
- vLLM 中的 torch.compile — Inductor 传递管道的工作原理。
- Attention 后端 — attention 特定的内核选择。