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 2, 19:30 UTC
[Kernel] Porting fuse_minimax_qk_norm to manual fusion (#43410)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
May 27, 2026 at 04:16 UTC
6e503868caa46f3afa87e8d3365495464fd75fb3
12 files changed
+262
-490
docs/design/fusions.md
-34
@@ -21,7 +21,6 @@ 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 |
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 |
@@ -42,7 +41,6 @@ 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 | — | — | — |
45
-| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
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 |
@@ -58,9 +56,6 @@ 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
61
-\* `fuse_minimax_qk_norm` is a model-specific pass for `MiniMaxM2ForCausalLM`. It also requires
62
-tensor parallelism (`tp_size > 1`) and the CUDA custom op `minimax_allreduce_rms_qk`.
63
-
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`.
@@ -191,35 +186,6 @@ If these conditions are set, the fusion is enabled automatically for optimizatio
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
194
-### MiniMax QK Norm (`fuse_minimax_qk_norm`)
195
-
196
-!!! info
197
- This is a MiniMax-specific compile pass. It is currently only enabled when all of the following hold:
198
- the model architecture is `MiniMaxM2ForCausalLM`, tensor parallelism is enabled (`tp_size > 1`),
199
- and the CUDA custom op `minimax_allreduce_rms_qk` is available. It is not enabled by default at any
200
- optimization level.
201
-
202
-**What it fuses.** Fuses the MiniMax M2 Q/K normalization path that performs an all-reduce over the
203
-per-token Q/K variances before applying RMS normalization to Q and K.
204
-
205
-This pass is distinct from [`enable_qk_norm_rope_fusion`](#qk-norm--rope-enable_qk_norm_rope_fusion):
206
-`fuse_minimax_qk_norm` targets MiniMax M2's tensor-parallel all-reduce + RMSNorm sequence, while
207
-`enable_qk_norm_rope_fusion` targets the later Q/K RMSNorm + RoPE sequence used by several other models.
208
-
209
-Example:
210
-
211
-```bash
212
-vllm serve MiniMaxAI/MiniMax-M2.5 \
213
- --tensor-parallel-size 4 \
214
- --compilation-config '{"mode": 3, "pass_config": {"fuse_minimax_qk_norm": true}}'
215
-```
216
-
217
-**Code locations.**
218
-
219
-- 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)
220
-- 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`)
221
-- 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)
222
-
189
### Sequence Parallelism (`enable_sp`)
190
191
**What it fuses.** Replaces all-reduce collectives with reduce-scatter + local RMSNorm + all-gather,
tests/kernels/core/test_minimax_reduce_rms.py
+2
-2
@@ -10,7 +10,7 @@ from torch.multiprocessing import spawn
10
from tests.kernels.utils import opcheck
11
from tests.utils import ensure_current_vllm_config, init_test_distributed_environment
12
from vllm.distributed import cleanup_dist_env_and_memory
13
-from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP
13
+from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
14
from vllm.platforms import current_platform
15
from vllm.utils.network_utils import get_open_port
16
from vllm.utils.torch_utils import set_random_seed
@@ -59,7 +59,7 @@ def _worker_forward_qk(
59
60
# Set up Lamport workspace.
61
from vllm.distributed.parallel_state import get_tp_group
62
- from vllm.model_executor.layers.mamba.lamport_workspace import (
62
+ from vllm.model_executor.layers.minimax_rms_norm.lamport_workspace import (
63
get_allreduce_workspace,
64
)
65
vllm/compilation/passes/fusion/minimax_qk_norm_fusion.py
deleted
-340
@@ -1,340 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-"""
5
-Fusion pass: replace MiniMax QK allreduce + RMS norm with the Lamport
6
-fused kernel (minimax_allreduce_rms_qk) for decode-size batches.
7
-
8
-Pattern (inlined forward_qk in compiled graph):
9
- q, k, v = qkv.split([q_size, kv_size, kv_size], -1)
10
- q_fp32 = q.to(float32); k_fp32 = k.to(float32)
11
- q_var = q_fp32.pow(2).mean(-1, keepdim=True)
12
- k_var = k_fp32.pow(2).mean(-1, keepdim=True)
13
- qk_var = cat([q_var, k_var], -1)
14
- qk_var = allreduce(qk_var) / tp_world
15
- q_var, k_var = qk_var.chunk(2, -1)
16
- q_out = (q_fp32 * rsqrt(q_var + eps) * q_weight).to(orig_dtype)
17
- k_out = (k_fp32 * rsqrt(k_var + eps) * k_weight).to(orig_dtype)
18
- return q_out, k_out, v
19
-
20
-Replacement (pure, no in-place on qkv/q/k):
21
- q_out, k_out = minimax_qk_norm_fused(qkv, q_weight, k_weight, workspace, ...)
22
- v = qkv.split([q_size, kv_size, kv_size], -1)[2]
23
- return q_out, k_out, v
24
-
25
-is_applicable_for_range: only fires for compile_range.end <= max_decode_tokens
26
-so that large prefill batches fall through to the original forward_qk (= main).
27
-"""
28
-
29
-import torch
30
-import torch._inductor.pattern_matcher as pm
31
-import torch.fx as fx
32
-from torch._inductor.pattern_matcher import PatternMatcherPass
33
-
34
-from vllm.config import VllmConfig
35
-from vllm.config.utils import Range
36
-from vllm.distributed import tensor_model_parallel_all_reduce
37
-from vllm.distributed.parallel_state import (
38
- get_tensor_model_parallel_rank,
39
- get_tensor_model_parallel_world_size,
40
-)
41
-from vllm.logger import init_logger
42
-from vllm.utils.torch_utils import direct_register_custom_op
43
-
44
-from ..inductor_pass import enable_fake_mode
45
-from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
46
-
47
-logger = init_logger(__name__)
48
-
49
-MAX_TOKEN_NUM = 2048
50
-
51
-_MINIMAX_QK_NORM_FUSED_OP = None
52
-if hasattr(torch.ops._C, "minimax_allreduce_rms_qk"):
53
-
54
- def _minimax_qk_norm_fused(
55
- qkv: torch.Tensor,
56
- norm_weight_q: torch.Tensor,
57
- norm_weight_k: torch.Tensor,
58
- q_size: int,
59
- kv_size: int,
60
- rank: int,
61
- nranks: int,
62
- eps: float,
63
- max_tokens: int,
64
- ) -> tuple[torch.Tensor, torch.Tensor]:
65
- from vllm.distributed.parallel_state import get_tp_group
66
- from vllm.model_executor.layers.mamba.lamport_workspace import (
67
- get_allreduce_workspace,
68
- )
69
-
70
- workspace = get_allreduce_workspace(
71
- rank=rank,
72
- world_size=nranks,
73
- max_tokens=max_tokens,
74
- process_group=get_tp_group().cpu_group,
75
- )
76
- return torch.ops._C.minimax_allreduce_rms_qk(
77
- qkv,
78
- norm_weight_q,
79
- norm_weight_k,
80
- workspace,
81
- q_size,
82
- kv_size,
83
- rank,
84
- nranks,
85
- eps,
86
- )
87
-
88
- def _minimax_qk_norm_fused_fake(
89
- qkv: torch.Tensor,
90
- norm_weight_q: torch.Tensor,
91
- norm_weight_k: torch.Tensor,
92
- q_size: int,
93
- kv_size: int,
94
- rank: int,
95
- nranks: int,
96
- eps: float,
97
- max_tokens: int,
98
- ) -> tuple[torch.Tensor, torch.Tensor]:
99
- T = qkv.shape[0]
100
- return (
101
- torch.empty([T, q_size], dtype=qkv.dtype, device=qkv.device),
102
- torch.empty([T, kv_size], dtype=qkv.dtype, device=qkv.device),
103
- )
104
-
105
- direct_register_custom_op(
106
- op_name="minimax_qk_norm_fused",
107
- op_func=_minimax_qk_norm_fused,
108
- fake_impl=_minimax_qk_norm_fused_fake,
109
- mutates_args=[],
110
- )
111
- _MINIMAX_QK_NORM_FUSED_OP = torch.ops.vllm.minimax_qk_norm_fused.default
112
-
113
-
114
-class MiniMaxQKNormPattern:
115
- """
116
- Match the forward_qk allreduce+rms pattern and replace with Lamport kernel.
117
- """
118
-
119
- def __init__(
120
- self,
121
- q_size: int,
122
- kv_size: int,
123
- eps: float,
124
- tp_world: int,
125
- tp_rank: int,
126
- max_tokens: int,
127
- dtype: torch.dtype,
128
- device: str | None,
129
- ) -> None:
130
- self.q_size = q_size
131
- self.kv_size = kv_size
132
- self.eps = eps
133
- self.tp_world = tp_world
134
- self.tp_rank = tp_rank
135
- self.max_tokens = max_tokens
136
- self.dtype = dtype
137
- self.device = device
138
-
139
- def get_inputs(self) -> list[torch.Tensor]:
140
- T = 4
141
- qkv = torch.empty(
142
- [T, self.q_size + 2 * self.kv_size],
143
- device=self.device,
144
- dtype=self.dtype,
145
- )
146
- q_weight = torch.empty([self.q_size], device=self.device, dtype=self.dtype)
147
- k_weight = torch.empty([self.kv_size], device=self.device, dtype=self.dtype)
148
- return [qkv, q_weight, k_weight]
149
-
150
- def register(self, pm_pass: PatternMatcherPass) -> None:
151
- q_size = self.q_size
152
- kv_size = self.kv_size
153
- eps = self.eps
154
- tp_world = self.tp_world
155
- max_tokens = self.max_tokens
156
- tp_rank = self.tp_rank
157
- dtype = self.dtype
158
-
159
- def pattern(
160
- qkv: torch.Tensor,
161
- q_weight: torch.Tensor,
162
- k_weight: torch.Tensor,
163
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
164
- q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
165
- q_fp32 = q.to(torch.float32)
166
- k_fp32 = k.to(torch.float32)
167
- q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
168
- k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
169
- qk_var = torch.cat([q_var, k_var], dim=-1)
170
- qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
171
- q_var, k_var = qk_var.chunk(2, dim=-1)
172
- q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
173
- k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
174
- return q_out, k_out, v
175
-
176
- def replacement(
177
- qkv: torch.Tensor,
178
- q_weight: torch.Tensor,
179
- k_weight: torch.Tensor,
180
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
181
- assert _MINIMAX_QK_NORM_FUSED_OP is not None
182
- q_out, k_out = torch.ops.vllm.minimax_qk_norm_fused(
183
- qkv,
184
- q_weight,
185
- k_weight,
186
- q_size,
187
- kv_size,
188
- tp_rank,
189
- tp_world,
190
- eps,
191
- max_tokens,
192
- )
193
- _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
194
- return q_out, k_out, v
195
-
196
- pm.register_replacement(
197
- pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
198
- )
199
-
200
- # Second pattern: three separate split_with_sizes nodes (one per output),
201
- # each with _users=1. This occurs when the QKV projection uses a
202
- # functional GEMM kernel (e.g. cutlass_scaled_mm via auto_functionalized),
203
- # which causes inductor to generate one split per consumer.
204
- def pattern_split3(
205
- qkv: torch.Tensor,
206
- q_weight: torch.Tensor,
207
- k_weight: torch.Tensor,
208
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
209
- q = qkv.split([q_size, kv_size, kv_size], dim=-1)[0]
210
- k = qkv.split([q_size, kv_size, kv_size], dim=-1)[1]
211
- v = qkv.split([q_size, kv_size, kv_size], dim=-1)[2]
212
- q_fp32 = q.to(torch.float32)
213
- k_fp32 = k.to(torch.float32)
214
- q_var = q_fp32.pow(2).mean(dim=-1, keepdim=True)
215
- k_var = k_fp32.pow(2).mean(dim=-1, keepdim=True)
216
- qk_var = torch.cat([q_var, k_var], dim=-1)
217
- qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
218
- q_var, k_var = qk_var.chunk(2, dim=-1)
219
- q_out = (q_fp32 * torch.rsqrt(q_var + eps) * q_weight).to(dtype)
220
- k_out = (k_fp32 * torch.rsqrt(k_var + eps) * k_weight).to(dtype)
221
- return q_out, k_out, v
222
-
223
- pm.register_replacement(
224
- pattern_split3, replacement, self.get_inputs(), pm.fwd_only, pm_pass
225
- )
226
-
227
-
228
-class MiniMaxQKNormPass(VllmPatternMatcherPass):
229
- """
230
- Replace forward_qk allreduce+norm with the Lamport fused kernel.
231
- Only applied for decode-size compile ranges (small token counts).
232
- """
233
-
234
- def __init__(self, config: VllmConfig) -> None:
235
- super().__init__(config)
236
- self.disabled = True
237
-
238
- if _MINIMAX_QK_NORM_FUSED_OP is None:
239
- logger.warning_once(
240
- "minimax_allreduce_rms_qk op not found, MiniMaxQKNormPass disabled."
241
- )
242
- return
243
-
244
- tp_world = get_tensor_model_parallel_world_size()
245
- if tp_world <= 1:
246
- logger.warning_once("MiniMaxQKNormPass disabled: tp_size <= 1.")
247
- return
248
-
249
- if config.model_config is None:
250
- logger.warning_once("MiniMaxQKNormPass disabled: no model_config.")
251
- return
252
-
253
- hf_cfg = config.model_config.hf_config
254
-
255
- model_name = getattr(hf_cfg, "architectures", "")[0]
256
- if model_name != "MiniMaxM2ForCausalLM":
257
- return
258
-
259
- num_attention_heads = getattr(hf_cfg, "num_attention_heads", 0)
260
- num_key_value_heads = getattr(hf_cfg, "num_key_value_heads", 0)
261
- hidden_size = getattr(hf_cfg, "hidden_size", 0)
262
- head_dim = getattr(hf_cfg, "head_dim", 0)
263
- eps: float = getattr(hf_cfg, "rms_norm_eps", 1e-6)
264
-
265
- if (
266
- num_attention_heads != 48
267
- or num_key_value_heads != 8
268
- or hidden_size != 3072
269
- or head_dim != 128
270
- ):
271
- logger.warning_once(
272
- "MiniMaxQKNormPass disabled: cannot infer model info from hf_config."
273
- )
274
- return
275
-
276
- num_heads_per_rank = num_attention_heads // tp_world
277
- num_kv_heads_per_rank = max(1, num_key_value_heads // tp_world)
278
- q_size = num_heads_per_rank * head_dim
279
- kv_size = num_kv_heads_per_rank * head_dim
280
-
281
- self.max_token_num = min(
282
- MAX_TOKEN_NUM, config.scheduler_config.max_num_batched_tokens
283
- )
284
-
285
- tp_rank = get_tensor_model_parallel_rank()
286
- # Allocate Lamport workspace first.
287
- from vllm.distributed.parallel_state import get_tp_group
288
- from vllm.model_executor.layers.mamba.lamport_workspace import (
289
- get_allreduce_workspace,
290
- )
291
-
292
- get_allreduce_workspace(
293
- rank=tp_rank,
294
- world_size=tp_world,
295
- max_tokens=self.max_token_num,
296
- process_group=get_tp_group().cpu_group,
297
- )
298
-
299
- self.patterns: PatternMatcherPass = PatternMatcherPass(
300
- pass_name="minimax_qk_norm_pass"
301
- )
302
- self._register_patterns(q_size, kv_size, eps, tp_world, tp_rank)
303
- self.dump_patterns(config, self.patterns)
304
- self.disabled = False
305
-
306
- @enable_fake_mode
307
- def _register_patterns(
308
- self,
309
- q_size: int,
310
- kv_size: int,
311
- eps: float,
312
- tp_world: int,
313
- tp_rank: int,
314
- ) -> None:
315
- MiniMaxQKNormPattern(
316
- q_size=q_size,
317
- kv_size=kv_size,
318
- eps=eps,
319
- tp_world=tp_world,
320
- tp_rank=tp_rank,
321
- max_tokens=self.max_token_num,
322
- dtype=self.model_dtype,
323
- device=self.device,
324
- ).register(self.patterns)
325
-
326
- def is_applicable_for_range(self, compile_range: Range) -> bool:
327
- if self.disabled:
328
- return False
329
-
330
- return bool(compile_range.end <= self.max_token_num)
331
-
332
- @VllmInductorPass.time_and_log
333
- def __call__(self, graph: fx.Graph) -> None:
334
- if self.disabled:
335
- return
336
- self.matched_count = self.patterns.apply(graph)
337
- logger.debug("MiniMaxQKNormPass replaced %s patterns", self.matched_count)
338
-
339
- def uuid(self) -> str:
340
- return VllmInductorPass.hash_source(self, MiniMaxQKNormPattern)
vllm/compilation/passes/pass_manager.py
-4
@@ -44,7 +44,6 @@ if current_platform.is_cuda_alike():
44
if current_platform.is_cuda():
45
from .fusion.allreduce_rms_fusion import AllReduceFusionPass
46
from .fusion.collective_fusion import AsyncTPPass
47
- from .fusion.minimax_qk_norm_fusion import MiniMaxQKNormPass
47
48
from .inductor_pass import (
49
CustomGraphPass,
@@ -154,9 +153,6 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
153
else:
154
self.passes += [AllReduceFusionPass(config)]
155
157
- if self.pass_config.fuse_minimax_qk_norm:
158
- self.passes += [MiniMaxQKNormPass(config)]
159
-
156
if self.pass_config.fuse_norm_quant:
157
if rocm_aiter_ops.is_enabled():
158
self.passes += [
vllm/config/compilation.py
+10
-1
@@ -135,7 +135,9 @@ class PassConfig:
135
fuse_allreduce_rms: bool = None # type: ignore[assignment]
136
"""Enable flashinfer allreduce fusion."""
137
fuse_minimax_qk_norm: bool = None # type: ignore[assignment]
138
- """Enable fused allreduce+RMSNorm for MiniMax QK norm."""
138
+ """Deprecated. The MiniMax QK norm fusion is now applied automatically at
139
+ runtime (see `MiniMaxText01RMSNormTP.forward_qkv`). This flag is kept for
140
+ backward compatibility and has no effect; it will be removed in v0.23."""
141
enable_qk_norm_rope_fusion: bool = None # type: ignore[assignment]
142
"""Enable fused Q/K RMSNorm + RoPE pass."""
143
fuse_rope_kvcache_cat_mla: bool = None # type: ignore[assignment]
@@ -294,6 +296,13 @@ class PassConfig:
296
"current platform is not CUDA or ROCm. The fusion will be disabled."
297
)
298
self.fuse_rope_kvcache_cat_mla = False
299
+ if self.fuse_minimax_qk_norm is not None:
300
+ logger.warning_once(
301
+ "`fuse_minimax_qk_norm` is deprecated and has no effect; "
302
+ "the MiniMax QK norm fusion is now applied automatically at "
303
+ "runtime when its conditions are met. This flag will be "
304
+ "removed in v0.23."
305
+ )
306
307
def log_enabled_passes(self) -> None:
308
"""
vllm/config/vllm.py
-16
@@ -1828,22 +1828,6 @@ class VllmConfig:
1828
compile_range_end,
1829
)
1830
1831
- if compilation_config.pass_config.fuse_minimax_qk_norm:
1832
- from vllm.compilation.passes.fusion.minimax_qk_norm_fusion import (
1833
- MAX_TOKEN_NUM,
1834
- )
1835
-
1836
- max_token_num = min(
1837
- MAX_TOKEN_NUM, self.scheduler_config.max_num_batched_tokens
1838
- )
1839
- if compile_range_end is not None and max_token_num < compile_range_end:
1840
- computed_compile_ranges_endpoints.append(max_token_num)
1841
- else:
1842
- logger.debug(
1843
- "Max num batched tokens below MiniMax QK norm fusion threshold, "
1844
- "MiniMax QK norm fusion enabled for all num_tokens."
1845
- )
1846
-
1831
if compilation_config.compile_ranges_endpoints is not None:
1832
for x in compilation_config.compile_ranges_endpoints:
1833
assert isinstance(x, int)
vllm/model_executor/layers/mamba/linear_attn.py
+1
-89
@@ -3,7 +3,6 @@
3
4
import math
5
from collections.abc import Callable
6
-from functools import partial
6
7
import torch
8
import torch.nn.functional as F
@@ -11,13 +10,11 @@ from einops import rearrange
10
from torch import nn
11
12
from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
14
-from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
13
from vllm.distributed.parallel_state import (
14
get_tensor_model_parallel_rank,
15
get_tensor_model_parallel_world_size,
16
)
17
from vllm.forward_context import ForwardContext, get_forward_context
20
-from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.layers.lightning_attn import (
19
lightning_attention,
20
linear_decode_forward_triton,
@@ -28,6 +25,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
25
MambaStateDtypeCalculator,
26
MambaStateShapeCalculator,
27
)
28
+from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
29
from vllm.model_executor.layers.quantization import QuantizationConfig
30
from vllm.utils.torch_utils import direct_register_custom_op
31
from vllm.v1.attention.backend import AttentionMetadata
@@ -35,92 +33,6 @@ from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
33
from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
34
35
38
-@CustomOp.register("minimax_text01_rmsnorm_tp")
39
-class MiniMaxText01RMSNormTP(CustomOp):
40
- def __init__(
41
- self,
42
- hidden_size: int,
43
- eps: float = 1e-6,
44
- *,
45
- weight_shard_world_size: int | None = None,
46
- weight_shard_rank: int | None = None,
47
- ) -> None:
48
- super().__init__()
49
- self.tp_world = get_tensor_model_parallel_world_size()
50
- self.tp_rank = get_tensor_model_parallel_rank()
51
- self.weight_shard_world = weight_shard_world_size or self.tp_world
52
- self.weight_shard_rank = (
53
- self.tp_rank if weight_shard_rank is None else weight_shard_rank
54
- )
55
-
56
- self.weight = nn.Parameter(torch.ones(hidden_size // self.weight_shard_world))
57
- self.weight.weight_loader = partial(
58
- self.weight_loader,
59
- shard_world_size=self.weight_shard_world,
60
- shard_rank=self.weight_shard_rank,
61
- )
62
- self.variance_epsilon = eps
63
-
64
- @staticmethod
65
- def weight_loader(
66
- param: nn.Parameter,
67
- loaded_weight: torch.Tensor,
68
- shard_world_size: int | None = None,
69
- shard_rank: int | None = None,
70
- ) -> None:
71
- if shard_world_size is None:
72
- shard_world_size = get_tensor_model_parallel_world_size()
73
- if shard_rank is None:
74
- shard_rank = get_tensor_model_parallel_rank()
75
-
76
- shard_size = loaded_weight.shape[0] // shard_world_size
77
- shard = slice(shard_rank * shard_size, (shard_rank + 1) * shard_size)
78
- param.data.copy_(loaded_weight[shard])
79
-
80
- def _forward(
81
- self,
82
- x: torch.Tensor,
83
- ) -> torch.Tensor:
84
- orig_dtype = x.dtype
85
- x = x.to(torch.float32)
86
- variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
87
- if self.tp_world > 1:
88
- variance = tensor_model_parallel_all_reduce(variance) / self.tp_world
89
- x = x * torch.rsqrt(variance + self.variance_epsilon)
90
- x = (x * self.weight).to(orig_dtype)
91
- return x
92
-
93
- def forward(
94
- self,
95
- x: torch.Tensor,
96
- residual: torch.Tensor | None = None,
97
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
98
- assert residual is None, "RMSNorm does not support residual connection."
99
- return self._forward(x)
100
-
101
- @staticmethod
102
- def forward_qk(
103
- q_norm: "MiniMaxText01RMSNormTP",
104
- k_norm: "MiniMaxText01RMSNormTP",
105
- q: torch.Tensor,
106
- k: torch.Tensor,
107
- ) -> tuple[torch.Tensor, torch.Tensor]:
108
- orig_dtype = q.dtype
109
- q = q.to(torch.float32)
110
- k = k.to(torch.float32)
111
- q_var = q.pow(2).mean(dim=-1, keepdim=True)
112
- k_var = k.pow(2).mean(dim=-1, keepdim=True)
113
- if q_norm.tp_world > 1:
114
- qk_var = torch.cat([q_var, k_var], dim=-1)
115
- qk_var = tensor_model_parallel_all_reduce(qk_var) / q_norm.tp_world
116
- q_var, k_var = qk_var.chunk(2, dim=-1)
117
- q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
118
- k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight
119
- q = q.to(orig_dtype)
120
- k = k.to(orig_dtype)
121
- return q, k
122
-
123
-
36
def clear_linear_attention_cache_for_new_sequences(
37
kv_cache: torch.Tensor,
38
state_indices_tensor: torch.Tensor,
vllm/model_executor/layers/minimax_rms_norm/__init__.py
new
+10
@@ -0,0 +1,10 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.model_executor.layers.minimax_rms_norm.rms_norm_tp import (
5
+ MiniMaxText01RMSNormTP,
6
+)
7
+
8
+__all__ = [
9
+ "MiniMaxText01RMSNormTP",
10
+]
vllm/model_executor/layers/minimax_rms_norm/lamport_workspace.py
renamed
vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py
new
+234
@@ -0,0 +1,234 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from functools import partial
5
+
6
+import torch
7
+from torch import nn
8
+
9
+from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
10
+from vllm.distributed.parallel_state import (
11
+ get_tensor_model_parallel_rank,
12
+ get_tensor_model_parallel_world_size,
13
+ get_tp_group,
14
+)
15
+from vllm.model_executor.custom_op import CustomOp
16
+from vllm.platforms import current_platform
17
+from vllm.utils.torch_utils import direct_register_custom_op
18
+
19
+# Max number of tokens supported by the Lamport fused allreduce+RMSNorm kernel.
20
+# Larger batches fall back to the eager allreduce + RMSNorm path.
21
+MINIMAX_QK_NORM_MAX_TOKEN_NUM = 2048
22
+
23
+_MINIMAX_FUSED_AR_RMS_QK = getattr(torch.ops._C, "minimax_allreduce_rms_qk", None)
24
+
25
+
26
+@torch.compile(backend=current_platform.simple_compile_backend, dynamic=True)
27
+def _minimax_qk_norm_fallback(
28
+ qkv: torch.Tensor,
29
+ q_weight: torch.Tensor,
30
+ k_weight: torch.Tensor,
31
+ q_size: int,
32
+ kv_size: int,
33
+ tp_rank: int,
34
+ tp_world: int,
35
+ eps: float,
36
+) -> tuple[torch.Tensor, torch.Tensor]:
37
+ q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
38
+ orig_dtype = q.dtype
39
+ q = q.to(torch.float32)
40
+ k = k.to(torch.float32)
41
+ q_var = q.pow(2).mean(dim=-1, keepdim=True)
42
+ k_var = k.pow(2).mean(dim=-1, keepdim=True)
43
+ if tp_world > 1:
44
+ qk_var = torch.cat([q_var, k_var], dim=-1)
45
+ qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
46
+ q_var, k_var = qk_var.chunk(2, dim=-1)
47
+ q = q * torch.rsqrt(q_var + eps) * q_weight
48
+ k = k * torch.rsqrt(k_var + eps) * k_weight
49
+ return q.to(orig_dtype), k.to(orig_dtype)
50
+
51
+
52
+def _minimax_qk_norm_fusion(
53
+ qkv: torch.Tensor,
54
+ q_weight: torch.Tensor,
55
+ k_weight: torch.Tensor,
56
+ q_size: int,
57
+ kv_size: int,
58
+ tp_rank: int,
59
+ tp_world: int,
60
+ eps: float,
61
+ workspace: torch.Tensor | None,
62
+) -> tuple[torch.Tensor, torch.Tensor]:
63
+ assert qkv.ndim == 2
64
+ num_tokens = qkv.shape[0]
65
+ if (
66
+ workspace is not None
67
+ and tp_world > 1
68
+ and num_tokens <= MINIMAX_QK_NORM_MAX_TOKEN_NUM
69
+ and _MINIMAX_FUSED_AR_RMS_QK is not None
70
+ ):
71
+ return _MINIMAX_FUSED_AR_RMS_QK(
72
+ qkv,
73
+ q_weight,
74
+ k_weight,
75
+ workspace,
76
+ q_size,
77
+ kv_size,
78
+ tp_rank,
79
+ tp_world,
80
+ eps,
81
+ )
82
+ return _minimax_qk_norm_fallback(
83
+ qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps
84
+ )
85
+
86
+
87
+def _minimax_qk_norm_fusion_fake(
88
+ qkv: torch.Tensor,
89
+ q_weight: torch.Tensor,
90
+ k_weight: torch.Tensor,
91
+ q_size: int,
92
+ kv_size: int,
93
+ tp_rank: int,
94
+ tp_world: int,
95
+ eps: float,
96
+ workspace: torch.Tensor | None,
97
+) -> tuple[torch.Tensor, torch.Tensor]:
98
+ assert qkv.ndim == 2
99
+ num_tokens = qkv.shape[0]
100
+ return (
101
+ torch.empty([num_tokens, q_size], dtype=qkv.dtype, device=qkv.device),
102
+ torch.empty([num_tokens, kv_size], dtype=qkv.dtype, device=qkv.device),
103
+ )
104
+
105
+
106
+direct_register_custom_op(
107
+ op_name="minimax_qk_norm_fusion",
108
+ op_func=_minimax_qk_norm_fusion,
109
+ fake_impl=_minimax_qk_norm_fusion_fake,
110
+ mutates_args=[],
111
+)
112
+
113
+
114
+@CustomOp.register("minimax_text01_rmsnorm_tp")
115
+class MiniMaxText01RMSNormTP(CustomOp):
116
+ def __init__(
117
+ self,
118
+ hidden_size: int,
119
+ eps: float = 1e-6,
120
+ *,
121
+ weight_shard_world_size: int | None = None,
122
+ weight_shard_rank: int | None = None,
123
+ ) -> None:
124
+ super().__init__()
125
+ self.tp_world = get_tensor_model_parallel_world_size()
126
+ self.tp_rank = get_tensor_model_parallel_rank()
127
+ self.weight_shard_world = weight_shard_world_size or self.tp_world
128
+ self.weight_shard_rank = (
129
+ self.tp_rank if weight_shard_rank is None else weight_shard_rank
130
+ )
131
+
132
+ self.weight = nn.Parameter(torch.ones(hidden_size // self.weight_shard_world))
133
+ self.weight.weight_loader = partial(
134
+ self.weight_loader,
135
+ shard_world_size=self.weight_shard_world,
136
+ shard_rank=self.weight_shard_rank,
137
+ )
138
+ self.variance_epsilon = eps
139
+
140
+ self.workspace = None
141
+ if _MINIMAX_FUSED_AR_RMS_QK is not None and self.tp_world > 1:
142
+ from .lamport_workspace import (
143
+ get_allreduce_workspace,
144
+ )
145
+
146
+ self.workspace = get_allreduce_workspace(
147
+ rank=self.tp_rank,
148
+ world_size=self.tp_world,
149
+ max_tokens=MINIMAX_QK_NORM_MAX_TOKEN_NUM,
150
+ process_group=get_tp_group().cpu_group,
151
+ )
152
+
153
+ @staticmethod
154
+ def weight_loader(
155
+ param: nn.Parameter,
156
+ loaded_weight: torch.Tensor,
157
+ shard_world_size: int | None = None,
158
+ shard_rank: int | None = None,
159
+ ) -> None:
160
+ if shard_world_size is None:
161
+ shard_world_size = get_tensor_model_parallel_world_size()
162
+ if shard_rank is None:
163
+ shard_rank = get_tensor_model_parallel_rank()
164
+
165
+ shard_size = loaded_weight.shape[0] // shard_world_size
166
+ shard = slice(shard_rank * shard_size, (shard_rank + 1) * shard_size)
167
+ param.data.copy_(loaded_weight[shard])
168
+
169
+ def _forward(
170
+ self,
171
+ x: torch.Tensor,
172
+ ) -> torch.Tensor:
173
+ orig_dtype = x.dtype
174
+ x = x.to(torch.float32)
175
+ variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
176
+ if self.tp_world > 1:
177
+ variance = tensor_model_parallel_all_reduce(variance) / self.tp_world
178
+ x = x * torch.rsqrt(variance + self.variance_epsilon)
179
+ x = (x * self.weight).to(orig_dtype)
180
+ return x
181
+
182
+ def forward(
183
+ self,
184
+ x: torch.Tensor,
185
+ residual: torch.Tensor | None = None,
186
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
187
+ assert residual is None, "RMSNorm does not support residual connection."
188
+ return self._forward(x)
189
+
190
+ @staticmethod
191
+ def forward_qk(
192
+ q_norm: "MiniMaxText01RMSNormTP",
193
+ k_norm: "MiniMaxText01RMSNormTP",
194
+ q: torch.Tensor,
195
+ k: torch.Tensor,
196
+ ) -> tuple[torch.Tensor, torch.Tensor]:
197
+ orig_dtype = q.dtype
198
+ q = q.to(torch.float32)
199
+ k = k.to(torch.float32)
200
+ q_var = q.pow(2).mean(dim=-1, keepdim=True)
201
+ k_var = k.pow(2).mean(dim=-1, keepdim=True)
202
+ if q_norm.tp_world > 1:
203
+ qk_var = torch.cat([q_var, k_var], dim=-1)
204
+ qk_var = tensor_model_parallel_all_reduce(qk_var) / q_norm.tp_world
205
+ q_var, k_var = qk_var.chunk(2, dim=-1)
206
+ q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
207
+ k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight
208
+ q = q.to(orig_dtype)
209
+ k = k.to(orig_dtype)
210
+ return q, k
211
+
212
+ @staticmethod
213
+ def forward_qkv(
214
+ q_norm: "MiniMaxText01RMSNormTP",
215
+ k_norm: "MiniMaxText01RMSNormTP",
216
+ qkv: torch.Tensor,
217
+ q_size: int,
218
+ kv_size: int,
219
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
220
+ assert qkv.ndim == 2
221
+ assert q_norm.variance_epsilon == k_norm.variance_epsilon
222
+ q, k = torch.ops.vllm.minimax_qk_norm_fusion(
223
+ qkv,
224
+ q_norm.weight,
225
+ k_norm.weight,
226
+ q_size,
227
+ kv_size,
228
+ q_norm.tp_rank,
229
+ q_norm.tp_world,
230
+ q_norm.variance_epsilon,
231
+ q_norm.workspace,
232
+ )
233
+ _, _, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
234
+ return q, k, v
vllm/model_executor/models/bailing_moe_linear.py
+1
-1
@@ -39,7 +39,6 @@ from vllm.model_executor.layers.mamba.abstract import MambaBase
39
from vllm.model_executor.layers.mamba.linear_attn import (
40
MiniMaxText01LinearAttention,
41
MiniMaxText01LinearKernel,
42
- MiniMaxText01RMSNormTP,
42
clear_linear_attention_cache_for_new_sequences,
43
linear_attention_decode,
44
linear_attention_prefill_and_mix,
@@ -49,6 +48,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
48
MambaStateDtypeCalculator,
49
MambaStateShapeCalculator,
50
)
51
+from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
52
from vllm.model_executor.layers.mla import MLAModules, MultiHeadLatentAttentionWrapper
53
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
54
from vllm.model_executor.layers.rotary_embedding import get_rope
vllm/model_executor/models/minimax_m2.py
+4
-3
@@ -50,7 +50,7 @@ from vllm.model_executor.layers.linear import (
50
RowParallelLinear,
51
)
52
from vllm.model_executor.layers.logits_processor import LogitsProcessor
53
-from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01RMSNormTP
53
+from vllm.model_executor.layers.minimax_rms_norm import MiniMaxText01RMSNormTP
54
from vllm.model_executor.layers.quantization import QuantizationConfig
55
from vllm.model_executor.layers.rotary_embedding import get_rope
56
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -243,8 +243,9 @@ class MiniMaxM2Attention(nn.Module):
243
hidden_states: torch.Tensor,
244
) -> torch.Tensor:
245
qkv, _ = self.qkv_proj(hidden_states)
246
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
247
- q, k = MiniMaxText01RMSNormTP.forward_qk(self.q_norm, self.k_norm, q, k)
246
+ q, k, v = MiniMaxText01RMSNormTP.forward_qkv(
247
+ self.q_norm, self.k_norm, qkv, self.q_size, self.kv_size
248
+ )
249
q, k = self.rotary_emb(positions, q, k)
250
attn_output = self.attn(q, k, v)
251
output, _ = self.o_proj(attn_output)