@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 418 lines 24.4 KB
Rendered Raw
1 # Fusion torch.compile passes
2
3 vLLM applies a set of kernel/operator fusions at compile time (via custom [`torch.compile`](torch_compile.md) Inductor passes)
4 to separate optimizations from model definitions and avoid breaking layer abstractions in model code.
5 These fusions are controlled by fields in [`PassConfig`][vllm.config.compilation.PassConfig] and are automatically enabled
6 at appropriate [optimization levels](optimization_levels.md).
7
8 ## Quick Reference
9
10 The table below maps each fusion to its controlling flag/config knob, the
11 operations it fuses, what level enables it by default, and an indicative speedup.
12 The Fullgraph column indicates whether the fusion requires the entire model graph to be
13 visible (either via Inductor partition or `splitting_ops=[]`),
14 and the last column indicates whether the fusion activates for all `num_tokens`
15 or just on the low or high end.
16
17 !!! info
18 Speedup depends heavily on the exact model, batch size, and hardware.
19 If tuning performance by hand, always benchmark your exact use-case with and without the fusion to verify the impact.
20
21 | Fusion | `PassConfig` flag | Fused operations | Default at | E2E Speedup | Fullgraph | `num_tokens` |
22 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |
23 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |
24 | [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always |
25 | [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always |
26 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |
27 | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low |
28 | [Sequence Parallelism](#sequence-parallelism-enable_sp) | `enable_sp` | AllReduce → ReduceScatter + AllGather | Off by default | Prereq for AsyncTP | Yes | High |
29 | [AsyncTP GEMM + collective](#asynctp-gemm--collective-overlap-fuse_gemm_comms) | `fuse_gemm_comms` | GEMM → reduce-scatter / all-gather → GEMM | Off by default | 7-10% | Yes | High |
30 | [RMSNorm + Quant](#rmsnorm--quantization-fuse_norm_quant) | `fuse_norm_quant` | RMSNorm (+residual add) → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
31 | [SiLU+Mul + Quant](#silumul--quantization-fuse_act_quant) | `fuse_act_quant` | SiLU+Mul activation → FP8/FP4 quant | O1 (conditional) | 1-4% | No | Always |
32 | [RMSNorm + Padding](#rmsnorm--padding-fuse_act_padding) | `fuse_act_padding` | Residual add + RMSNorm → padding | O1 (ROCm/AITER only) | TBD | No | Always |
33 | [MLA Dual RMSNorm](#mla-dual-rmsnorm-fuse_mla_dual_rms_norm) | `fuse_mla_dual_rms_norm` | Paired Q + KV RMSNorm (+ FP8 quant) → 1 kernel | O1 (ROCm/AITER only) | 1-2% | No | Always |
34
35 ## Support Matrix
36
37 The table below lists the quantization schemes supported by each fusion on each platform.
38 **—** means the fusion is not available on that platform. The latest and in-progress work is available in the tracking issue:
39 [#36066](https://github.com/vllm-project/vllm/issues/36066)
40
41 | Fusion | SM100 (Blackwell) | SM90 (Hopper) | SM89 (Ada) | SM80 (Ampere) | ROCm |
42 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |
43 | `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
44 | `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
45 | `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
46 | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
47 | `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — |
48 | `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
49 | `fuse_gemm_comms` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
50 | `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 |
51 | `fuse_act_quant` | FP8 static, NVFP4 | FP8 static, FP8 per-group (128/64) | FP8 static, FP8 per-group (128/64) | — | FP8 per-group |
52 | `fuse_act_padding` | — | — | — | — | FP16/BF16 |
53 | `fuse_mla_dual_rms_norm` | — | — | — | — | BF16 |
54
55 \* `fuse_attn_quant` support depends on the attention backend in use; not all backends support
56 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
57 for per-backend details.
58
59 † `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
60 other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
61 SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
62
63 ## Enabling / Disabling Fusions
64
65 Fusions are exposed through `PassConfig`, which is nested inside `CompilationConfig`:
66
67 ```python
68 from vllm import LLM
69 from vllm.config import CompilationConfig, PassConfig
70
71 llm = LLM(
72 model="...",
73 optimization_level=2, # Default optimization level
74 compilation_config=CompilationConfig(
75 pass_config=PassConfig(
76 fuse_norm_quant=True,
77 fuse_act_quant=True,
78 fuse_allreduce_rms=False, # disable a specific fusion
79 )
80 ),
81 )
82 ```
83
84 Fusions can also be enabled using command-line flags with any `vllm ...` command:
85
86 ```bash
87 # Enable O2 defaults, but turn off allreduce fusion
88 vllm serve meta-llama/Llama-3.1-8B-Instruct -O2 -cc.pass_config.fuse_allreduce_rms=False
89
90 # The above is equivalent to the more verbose:
91 vllm serve meta-llama/Llama-3.1-8B-Instruct -O2 --compilation-config '{"pass_config": {"fuse_allreduce_rms": false}}'
92
93 # Same syntax in other commands, e.g. vllm bench:
94 vllm bench latency --model=meta-llama/Llama-3.1-8B-Instruct -O2 -cc.pass_config.fuse_allreduce_rms=False
95 ```
96
97 Fields set explicitly by the user always take precedence over optimization-level defaults.
98
99 ## Fusion Details
100
101 ### AllReduce + RMSNorm (`fuse_allreduce_rms`)
102
103 !!! warning
104 TP+DP and TP+PP combinations are currently broken
105 ([#34458](https://github.com/vllm-project/vllm/issues/34458) and
106 [#35426](https://github.com/vllm-project/vllm/issues/35426)).
107 Only supported on NVIDIA Hopper (SM90) and Blackwell (SM100) with FlashInfer installed.
108
109 **What it fuses.** Fuses the tensor-parallel all-reduce collective with the subsequent residual add,
110 RMSNorm, and optionally a quantization step into a single FlashInfer / TRT-LLM communication kernel.
111 This fusion is only profitable for small `num_tokens`,
112 so the fusion is only performed in the lower compiled range.
113
114 Patterns covered:
115
116 - `AllReduce → RMSNorm(+residual_add)`: CUDA sm90+ with FlashInfer
117 - `AllReduce → RMSNorm(+residual_add) → FP8 static quant`: CUDA sm90+ with FlashInfer
118 - `AllReduce → RMSNorm(+residual_add) → NVFP4 dynamic quant`: CUDA sm100+ with FlashInfer
119
120 The maximum tensor size below which the fused kernel is used is hardware-dependent (64 MB for TP=2
121 on SM90/SM100) and configurable via `PassConfig.fi_allreduce_fusion_max_size_mb`.
122
123 **Code locations.**
124
125 - Pass: [`vllm/compilation/passes/fusion/allreduce_rms_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/allreduce_rms_fusion.py)
126 - FlashInfer all-reduce: [`vllm/distributed/device_communicators/flashinfer_all_reduce.py`](https://github.com/vllm-project/vllm/blob/main/vllm/distributed/device_communicators/flashinfer_all_reduce.py)
127 - Benchmark: [`benchmarks/kernels/benchmark_fused_collective.py`](https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_fused_collective.py)
128
129 ### Attention + Quantization (`fuse_attn_quant`)
130
131 !!! info
132 `fuse_attn_quant` is currently not enabled at any optimization level by default and must be set
133 explicitly. It requires the full model graph to be visible (Inductor partition or `splitting_ops=[]`).
134
135 **What it fuses.** Fuses the attention output quantization directly after the attention computation,
136 eliminating a full-precision memory round-trip of the attention output. This fusion supports both
137 standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patterns covered:
138
139 `Attention → FP8 static quant`:
140
141 - `TRITON_ATTN`: CUDA, ROCm
142 - `FLASHINFER`: CUDA sm100+ with FlashInfer installed
143 - `ROCM_ATTN`: ROCm
144 - `ROCM_AITER_UNIFIED_ATTN`: ROCm with AITER
145
146 `Attention → NVFP4 dynamic quant`:
147
148 - `FLASHINFER`: CUDA sm100+ with FlashInfer installed
149
150 `MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`
151
152 The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
153 with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
154 the kernel writes FP8 output directly), no MLA prefill or decode backend currently supports direct
155 FP8/FP4 output. The fusion writes to an intermediate buffer and quantizes in a separate step, so
156 there is no memory round-trip elimination yet.
157
158 !!! info
159 The MLA attention fusion is not expected to yield a measurable speedup yet.
160 This will improve once MLA prefill/decode kernels support direct FP8/FP4 output.
161
162 Other attention backends do not support fused output quantization yet.
163
164 **Code locations.**
165
166 - Pass (Attention): [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py)
167 - Pass (MLAAttention): [`vllm/compilation/passes/fusion/mla_attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py)
168 - Attention backends: [`vllm/v1/attention/backends/`](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/)
169
170 ### RoPE + KV-Cache Update (`fuse_rope_kvcache`)
171
172 !!! info
173 ROCm/AITER-only. Not available on NVIDIA CUDA or CPU. The fusion is only enabled for
174 `num_tokens ≤ 256` by default due to AITER fused kernel performance issues.
175 This threshold is configurable via `PassConfig.rope_kvcache_fusion_max_token_num`.
176
177 **What it fuses.** Fuses the rotary positional embedding kernel with the KV-cache scatter/write into
178 a single kernel, avoiding separate reads and writes of the key and value tensors.
179
180 Requires: AMD ROCm with AITER enabled, the `rotary_embedding` custom op active (automatic),
181 and the `kv_cache` update op visible in the graph: either by using Inductor graph partition
182 or removed from `splitting_ops`.
183 If these conditions are set, the fusion is enabled automatically for optimization level O1 and above.
184
185 **Code locations.**
186
187 - Pass: [`vllm/compilation/passes/fusion/rope_kvcache_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rope_kvcache_fusion.py)
188
189 ### Sequence Parallelism (`enable_sp`)
190
191 **What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
192 splitting the sequence dimension across TP ranks. This restructures the graph so the subsequent AsyncTP
193 pass can fuse the reduce-scatter / all-gather with the surrounding GEMMs.
194
195 Sequence Parallelism itself does not directly improve performance; it is a prerequisite for the
196 AsyncTP pass (`fuse_gemm_comms`). SP is only applied above a minimum token threshold that is
197 autoconfigured based on device capability and model `hidden_size`. Currently only active on
198 H100/SM90 for models with `hidden_size >= 8192`. The threshold is configurable via
199 `PassConfig.sp_min_token_num`.
200
201 The general transformation:
202
203 ```text
204 Input → AllReduce → RMSNorm → Output
205 becomes:
206 Input → ReduceScatter → local RMSNorm → AllGather → Output
207 ```
208
209 Patterns covered:
210
211 - First block: `AllReduce → RMSNorm` → `ReduceScatter → RMSNorm → AllGather`
212 - Middle blocks: `AllReduce → fused_add_RMSNorm` → `ReduceScatter → fused_add_RMSNorm → AllGather`
213 - Both with optional `→ FP8 static quant` suffix
214
215 Requires: `use_inductor_graph_partition=True` **or** piecewise compilation with static sizes
216 divisible by `tensor_parallel_size`.
217
218 Supported hardware: Only tested on NVIDIA CUDA, possibly works on ROCm. FP8 all-gather requires sm90+.
219
220 **Code locations.**
221
222 - Pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py)
223
224 ### AsyncTP GEMM + Collective Overlap (`fuse_gemm_comms`)
225
226 !!! info
227 Requires `enable_sp=True` (enabled automatically). This pass is a no-op if Sequence Parallelism has not been applied.
228
229 **What it fuses.** After Sequence Parallelism transforms the graph, fuses GEMM kernels with the
230 surrounding reduce-scatter (output projection) and all-gather (input projection) using
231 `torch.ops.symm_mem` symmetric-memory primitives, overlapping communication and computation.
232 This overlap is only profitable for large `num_tokens`, so the fusion (and preceding SP)
233 is only performed in the higher compiled range above `PassConfig.sp_min_token_num`.
234
235 Patterns covered:
236
237 - `GEMM → reduce-scatter` → `fused_matmul_reduce_scatter`
238 - `all-gather → GEMM` → `all_gather_matmul`
239 - FP8 scaled variants of both patterns
240
241 Supported hardware: NVIDIA CUDA with symmetric-memory (`torch.distributed._symmetric_memory`) support.
242
243 On B200, pattern-matching fp8 FlashInfer scaled MM is not supported, so it must be disabled
244 ([#27893](https://github.com/vllm-project/vllm/issues/27893))
245
246 ```shell
247 VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel ...
248 ```
249
250 **Code locations.**
251
252 - Pass: [`vllm/compilation/passes/fusion/collective_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/collective_fusion.py)
253 - Sequence parallelism pass: [`vllm/compilation/passes/fusion/sequence_parallelism.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/sequence_parallelism.py)
254
255 ### QK Norm + RoPE (`enable_qk_norm_rope_fusion`)
256
257 !!! info
258 Only applicable to models that apply per-head RMSNorm to Q and K before rotary positional
259 embedding (e.g. Qwen). Not enabled by default at any optimization level due to perf issues on H100:
260 [#34391](https://github.com/vllm-project/vllm/issues/34391)
261
262 **What it fuses.** Fuses the sequence: split QKV → reshape → Q/K RMSNorm → reshape → rotary
263 embedding into a single `fused_qk_norm_rope` CUDA kernel.
264
265 ```text
266 # Unfused:
267 q, k, v = split(qkv)
268 q_norm = rms_norm(q.view(heads))
269 k_norm = rms_norm(k.view(kv_heads))
270 q_rope, k_rope = rotary_embedding(q_norm, k_norm, ...)
271
272 # Fused:
273 fused_qk_norm_rope(qkv, ...)
274 ```
275
276 Supported hardware: CUDA (sm80+) only, tested only on sm90 and sm100.
277
278 **Code locations.**
279
280 - Pass: [`vllm/compilation/passes/fusion/qk_norm_rope_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/qk_norm_rope_fusion.py)
281 - CUDA kernel: [`csrc/ops.h`](https://github.com/vllm-project/vllm/blob/main/csrc/ops.h) (`fused_qk_norm_rope`)
282
283 ### RMSNorm + Quantization (`fuse_norm_quant`)
284
285 !!! warning
286 On NVIDIA, Inductor actually generates a faster fused kernel than our custom CUDA kernel.
287 Hence, this fusion is only enabled when either `rms_norm` or `quant_fp8` is using a custom kernel.
288
289 **What it fuses.** Combines the custom `rms_norm` / `fused_add_rms_norm`
290 operations with subsequent quantization into a single fused kernel,
291 eliminating an intermediate read/write of the full-precision activation tensor.
292 Two variants are fused:
293
294 - *Plain RMSNorm + quant*: `rms_norm(x) → quant_fp8(y)`
295 - *Fused-add RMSNorm + quant*: `fused_add_rms_norm(x, residual) → quant_fp8(y)` — also updates the residual in-place.
296
297 Note that AITER fusions are currently in a separate pass in `vllm.compilation.passes.fusion.rocm_aiter_fusion`.
298
299 Supported quantization scheme/hardware combinations:
300
301 - FP8 static per-tensor: CUDA & HIP kernel
302 - FP8 dynamic per-token: CUDA & HIP kernel, AITER
303 - FP8 dynamic per-token-group (128/64): CUDA & HIP kernel, AITER
304
305 **Code locations.**
306
307 - Pass: [`vllm/compilation/passes/fusion/rms_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rms_quant_fusion.py)
308 - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
309 - CUDA/HIP kernels: [`csrc/libtorch_stable/layernorm_quant_kernels.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/layernorm_quant_kernels.cu)
310
311 ### SiLU+Mul + Quantization (`fuse_act_quant`)
312
313 !!! warning
314 Same as `fuse_norm_quant`: on NVIDIA, Inductor generates a faster fused kernel than our custom ops.
315 This fusion is only enabled when either `silu_and_mul` or `quant_fp8` are using a custom kernel,
316 or for NVFP4-quantized models (where FP4 quant is always a custom op).
317
318 **What it fuses.** Fuses the `silu_and_mul` gate-up projection activation with subsequent quantization into a single kernel,
319 avoiding materialization of the full-precision post-activation tensor.
320
321 Note that AITER fusions are in a separate pass in `vllm.compilation.passes.fusion.rocm_aiter_fusion`.
322
323 Supported quantization scheme/hardware combinations:
324
325 - FP8 static per-tensor: CUDA & HIP kernel
326 - FP8 dynamic per-group (128/64): CUDA kernel (sm89+, not active when DeepGemm is used on sm100+)
327 - NVFP4 dynamic: CUDA sm100+ only with FlashInfer
328 - FP8 per-token-group (128): ROCm AITER only
329
330 **Code locations.**
331
332 - Pass: [`vllm/compilation/passes/fusion/act_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/act_quant_fusion.py)
333 - ROCm AITER pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py)
334 - CUDA/HIP kernels: [`csrc/quantization/`](https://github.com/vllm-project/vllm/blob/main/csrc/quantization/)
335 - Fused SiLU+Mul+BlockQuant kernel: [`csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/libtorch_stable/quantization/fused_kernels/fused_silu_mul_block_quant.cu)
336
337 ### RMSNorm + Padding (`fuse_act_padding`)
338
339 !!! info
340 ROCm/AITER-only. Targeted at GPT-OSS models.
341
342 **What it fuses.** Fuses a residual add + RMSNorm with a subsequent padding operation that pads
343 the hidden dimension to a multiple required by downstream AITER Triton GEMM kernels.
344
345 Requires: AMD ROCm with AITER RMSNorm enabled. Enabled by default in optimization level O1 and above
346 when the hidden size is 2880 and AITER Triton GEMMs *not* enabled.
347
348 **Code locations.**
349
350 - Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`RocmAiterTritonAddRMSNormPadFusionPass`)
351
352 ### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`)
353
354 !!! info
355 ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention.
356
357 !!! note
358 When the native implementation of `rms_norm` is used (the default on CUDA and
359 ROCm for now), Inductor's built-in fusion already handles merging these norms
360 automatically. This explicit pass targets the case where AITER's custom
361 `rms_norm` op is active, which Inductor cannot fuse on its own.
362
363 **What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm
364 operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER,
365 reducing kernel launch overhead from 2 launches to 1 per MLA layer.
366
367 ```text
368 # Unfused:
369 q_c, kv_lora = split(projected, [q_dim, kv_dim])
370 kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
371 q_c = rms_norm(q_c, q_weight, eps)
372 kv_c = rms_norm(kv_c, kv_weight, eps)
373
374 # Fused:
375 q_c, kv_lora = split(projected, [q_dim, kv_dim])
376 kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
377 q_normed, kv_normed = fused_mla_dual_rms_norm(
378 q_c, q_weight, kv_c, kv_weight, eps1, eps2)
379 ```
380
381 Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above
382 when AITER is available.
383
384 **FP8 attention variant (per-token quant).** With a per-token FP8 `q_b_proj`,
385 only the *q* latent is FP8-quantized while *kv* stays bf16.
386 `RocmAiterRMSNormQuantFusionPass` first folds the q side into
387 `rocm_aiter_rmsnorm_fused_dynamic_quant`, leaving kv a plain
388 `rms_norm` — breaking the symmetric pattern above. The same pass then matches
389 this asymmetric pair and lowers it to `fused_mla_dual_rms_norm_per_token_quant`.
390
391 ```text
392 # Unfused (q norm+quant fused; kv still plain rms_norm):
393 q_c, kv_lora = split(projected, [q_dim, kv_dim])
394 kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
395 q_fp8, q_scale = rocm_aiter_rmsnorm_fused_dynamic_quant(q_c, q_weight, eps, fp8)
396 kv_normed = rms_norm(kv_c, kv_weight, eps) # bf16
397
398 # Fused:
399 q_c, kv_lora = split(projected, [q_dim, kv_dim])
400 kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])
401 q_fp8, q_scale, kv_normed = fused_mla_dual_rms_norm_per_token_quant(
402 q_c, q_weight, kv_c, kv_weight, eps1, eps2)
403 ```
404
405 **Code locations.**
406
407 - Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`, `MLADualRMSPerTokenQuantPattern`)
408 - Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`, `fused_mla_dual_rms_norm_per_token_quant`)
409 - AITER kernels: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442), `fused_qk_rmsnorm_per_token_quant`
410
411 ## See Also
412
413 - [Optimization Levels](optimization_levels.md) — high-level presets that set
414 fusion defaults.
415 - [torch.compile in vLLM](torch_compile.md) — how the Inductor pass pipeline
416 works.
417 - [Attention Backends](attention_backends.md) — attention-specific kernel
418 selection.