A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
[DOC] Add fuse_minimax_qk_norm (#39782)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Apr 18, 2026 at 15:41 UTC
bfde49e287cb5522fb0625c8e2b4e03cac20cbb2
1 file changed
+34
docs/design/fusions.md
+34
@@ -21,6 +21,7 @@ or just on the low or high end.
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
+| [MiniMax QK Norm](#minimax-qk-norm-fuse_minimax_qk_norm) | `fuse_minimax_qk_norm` | Q/K variance all-reduce → Q/K RMSNorm | Off by default | 2-3% | No | Low |
25
| [Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | Attention output → FP8/NVFP4 quant | Off by default | 3-7% | Yes | Always |
26
| [MLA Attention + Quant](#attention--quantization-fuse_attn_quant) | `fuse_attn_quant` | MLA Attention output → FP8/NVFP4 quant | Off by default | TBD | Yes | Always |
27
| [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 |
@@ -40,6 +41,7 @@ The table below lists the quantization schemes supported by each fusion on each
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_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
45
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
46
| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
47
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
@@ -54,6 +56,9 @@ The table below lists the quantization schemes supported by each fusion on each
56
fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)
57
for per-backend details.
58
59
+\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
60
+tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
61
+
62
† `enable_sp` and `fuse_gemm_comms` are only autoconfigured for SM90 today;
63
other architectures support requires setting `PassConfig.sp_min_token_num` explicitly.
64
SM100 support also requires setting `VLLM_DISABLED_KERNELS=FlashInferFP8ScaledMMLinearKernel`.
@@ -184,6 +189,35 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
189
190
- 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)
191
192
+### MiniMax QK Norm (`fuse_minimax_qk_norm`)
193
+
194
+!!! info
195
+ This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
196
+ the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
197
+ and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
198
+ optimization level.
199
+
200
+**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
201
+per-token Q/K variances before applying RMS normalization to Q and K.
202
+
203
+This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
204
+`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
205
+`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
206
+
207
+Example:
208
+
209
+```bash
210
+vllm serve MiniMaxAI/MiniMax-M2.5 \
211
+ --tensor-parallel-size 4 \
212
+ --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
213
+```
214
+
215
+**Code locations.**
216
+
217
+- Pass: [`vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py)
218
+- CUDA op: [`csrc/minimax_reduce_rms_kernel.cu`](https://github.com/vllm-project/vllm/blob/main/csrc/minimax_reduce_rms_kernel.cu) (`minimax_allreduce_rms_qk`)
219
+- Workspace helper: [`vllm/model_executor/layers/mamba/lamport_workspace.py`](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/lamport_workspace.py)
220
+
221
### Sequence Parallelism (`enable_sp`)
222
223
**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,