[Model]Fix MiniMaxM2ForCausalLM perf regression (#45935)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Jun 22, 2026 at 00:28 UTC
745bba5ea8fa17dd6ae3751daf43c3d1bd8522df
2 files changed
+215
-16
tests/kernels/core/test_minimax_reduce_rms.py
+59
-3
@@ -10,8 +10,12 @@ 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.minimax_rms_norm import MiniMaxText01RMSNormTP
13
+from vllm.model_executor.layers.minimax_rms_norm import (
14
+ MiniMaxText01RMSNormTP,
15
+ rms_norm_tp,
16
+)
17
from vllm.platforms import current_platform
18
+from vllm.triton_utils import HAS_TRITON
19
from vllm.utils.network_utils import get_open_port
20
from vllm.utils.torch_utils import set_random_seed
21
@@ -54,8 +58,19 @@ def _worker_forward_qk(
58
torch.manual_seed(seed + 1000 + local_rank)
59
qkv = torch.randn(num_tokens, hq + hk + hk, dtype=dtype, device="cuda")
60
57
- q_ref, k_ref, v_ref = qkv.clone().split([hq, hk, hk], dim=-1)
58
- ref_q, ref_k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q_ref, k_ref)
61
+ # Reference: eager all-reduce path. ``forward_qk`` no longer all-reduces
62
+ # the variance (it is the tp==1 / already-reduced building block), so the
63
+ # multi-rank reference must use the eager path that performs the global
64
+ # variance all-reduce, matching the fused kernel below.
65
+ ref_q, ref_k = rms_norm_tp._minimax_qk_norm_tp_eager(
66
+ qkv.clone(),
67
+ q_norm.weight,
68
+ k_norm.weight,
69
+ hq,
70
+ hk,
71
+ world_size,
72
+ eps,
73
+ )
74
75
# Set up Lamport workspace.
76
from vllm.distributed.parallel_state import get_tp_group
@@ -150,3 +165,44 @@ def test_minimax_reduce_rms_qk(
165
nprocs=world_size,
166
join=True,
167
)
168
+
169
+
170
+@pytest.mark.skipif(
171
+ not current_platform.is_cuda() or not HAS_TRITON,
172
+ reason="CUDA and Triton required",
173
+)
174
+@pytest.mark.parametrize("num_tokens", [1, 7, 128, 333, 2049])
175
+@pytest.mark.parametrize("hidden_dims", [(3072, 512), (768, 256), (3000, 500)])
176
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
177
+@pytest.mark.parametrize("tp_world", [1, 4, 8])
178
+@pytest.mark.parametrize("eps", [1e-6])
179
+@pytest.mark.parametrize("seed", [42])
180
+def test_minimax_qk_norm_triton_fallback(
181
+ monkeypatch, num_tokens, hidden_dims, dtype, tp_world, eps, seed
182
+):
183
+ """Single-GPU check: Triton fallback kernels vs the pure-torch reference.
184
+
185
+ The all-reduce is a TP communication barrier, so it is monkeypatched to
186
+ identity here; both the Triton path and the reference see the same
187
+ (patched) reduction. This validates the kernel math and the folded
188
+ ``/ tp_world`` scaling without needing multiple ranks -- ``hidden_dims``
189
+ are the per-rank q/k segment widths.
190
+ """
191
+ monkeypatch.setattr(rms_norm_tp, "_all_reduce_variance", lambda v: v)
192
+
193
+ q_size, kv_size = hidden_dims
194
+ device = "cuda"
195
+ torch.manual_seed(seed)
196
+ qkv = torch.randn(num_tokens, q_size + 2 * kv_size, dtype=dtype, device=device)
197
+ q_weight = torch.randn(q_size, dtype=dtype, device=device)
198
+ k_weight = torch.randn(kv_size, dtype=dtype, device=device)
199
+
200
+ q_triton, k_triton = rms_norm_tp._minimax_qk_norm_tp_fallback(
201
+ qkv, q_weight, k_weight, q_size, kv_size, 0, tp_world, eps
202
+ )
203
+ q_ref, k_ref = rms_norm_tp._minimax_qk_norm_tp_eager(
204
+ qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
205
+ )
206
+
207
+ torch.testing.assert_close(q_triton, q_ref, atol=3e-2, rtol=3e-2)
208
+ torch.testing.assert_close(k_triton, k_ref, atol=3e-2, rtol=3e-2)
vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py
+156
-13
@@ -14,7 +14,7 @@ from vllm.distributed.parallel_state import (
14
)
15
from vllm.logger import init_logger
16
from vllm.model_executor.custom_op import CustomOp
17
-from vllm.platforms import current_platform
17
+from vllm.triton_utils import HAS_TRITON, tl, triton
18
from vllm.utils.torch_utils import direct_register_custom_op
19
20
logger = init_logger(__name__)
@@ -40,32 +40,171 @@ def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
40
return tensor_model_parallel_all_reduce(var.flatten()).view_as(var)
41
42
43
-@torch.compile(backend=current_platform.simple_compile_backend, dynamic=True)
44
-def _minimax_qk_norm_fallback(
43
+@triton.jit
44
+def _minimax_qk_var_kernel(
45
+ qkv_ptr, # [num_tokens, hidden], 16-bit activations
46
+ var_ptr, # [num_tokens, 2], fp32
47
+ row_stride, # element stride between tokens in qkv
48
+ q_size: tl.constexpr, # constant per deployment -> loops unroll, mask elides
49
+ kv_size: tl.constexpr,
50
+ BLOCK: tl.constexpr,
51
+):
52
+ """TP-pre stage: per-token mean-of-squares for the q and k segments.
53
+
54
+ Accumulates in fp32 while reading the 16-bit qkv in place, so no fp32
55
+ copy of q/k is materialized. ``var[:, 0]`` is the q variance and
56
+ ``var[:, 1]`` the k variance; both are the local-shard means, ready for
57
+ the all-reduce that follows.
58
+ """
59
+ token = tl.program_id(0)
60
+ base = qkv_ptr + token * row_stride
61
+
62
+ q_acc = 0.0
63
+ for off in range(0, q_size, BLOCK):
64
+ idx = off + tl.arange(0, BLOCK)
65
+ mask = idx < q_size
66
+ x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
67
+ q_acc += tl.sum(x * x, axis=0)
68
+
69
+ k_acc = 0.0
70
+ for off in range(0, kv_size, BLOCK):
71
+ idx = off + tl.arange(0, BLOCK)
72
+ mask = idx < kv_size
73
+ x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
74
+ k_acc += tl.sum(x * x, axis=0)
75
+
76
+ tl.store(var_ptr + token * 2 + 0, q_acc / q_size)
77
+ tl.store(var_ptr + token * 2 + 1, k_acc / kv_size)
78
+
79
+
80
+@triton.jit
81
+def _minimax_rms_apply_kernel(
82
+ qkv_ptr, # [num_tokens, hidden]
83
+ var_ptr, # [num_tokens, 2], fp32, all-reduced sum of per-shard means
84
+ q_w_ptr, # [q_size], q per-channel weight
85
+ k_w_ptr, # [kv_size], k per-channel weight
86
+ q_out_ptr, # [num_tokens, q_size], contiguous
87
+ k_out_ptr, # [num_tokens, kv_size], contiguous
88
+ row_stride, # element stride between tokens in qkv
89
+ q_size: tl.constexpr, # constant per deployment -> loops unroll, mask elides
90
+ kv_size: tl.constexpr,
91
+ tp_world: tl.constexpr, # folds the post-all-reduce /tp_world into rsqrt
92
+ eps: tl.constexpr,
93
+ BLOCK: tl.constexpr,
94
+):
95
+ """TP-post stage: ``x * rsqrt(var / tp_world + eps) * weight``.
96
+
97
+ A single program normalizes both the q and k segments of one token, so q
98
+ and k share one launch instead of two. The all-reduce yields the sum of
99
+ per-shard means, so the ``/ tp_world`` that recovers the global
100
+ mean-of-squares is folded into the ``rsqrt`` here rather than run as a
101
+ separate elementwise pass over the ``[num_tokens, 2]`` variance tensor.
102
+ """
103
+ token = tl.program_id(0)
104
+ base = qkv_ptr + token * row_stride
105
+
106
+ q_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 0) / tp_world + eps)
107
+ q_out_row = q_out_ptr + token * q_size
108
+ for off in range(0, q_size, BLOCK):
109
+ idx = off + tl.arange(0, BLOCK)
110
+ mask = idx < q_size
111
+ x = tl.load(base + idx, mask=mask, other=0.0).to(tl.float32)
112
+ w = tl.load(q_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
113
+ y = x * q_inv * w
114
+ tl.store(q_out_row + idx, y.to(q_out_ptr.dtype.element_ty), mask=mask)
115
+
116
+ k_inv = tl.rsqrt(tl.load(var_ptr + token * 2 + 1) / tp_world + eps)
117
+ k_out_row = k_out_ptr + token * kv_size
118
+ for off in range(0, kv_size, BLOCK):
119
+ idx = off + tl.arange(0, BLOCK)
120
+ mask = idx < kv_size
121
+ x = tl.load(base + q_size + idx, mask=mask, other=0.0).to(tl.float32)
122
+ w = tl.load(k_w_ptr + idx, mask=mask, other=0.0).to(tl.float32)
123
+ y = x * k_inv * w
124
+ tl.store(k_out_row + idx, y.to(k_out_ptr.dtype.element_ty), mask=mask)
125
+
126
+
127
+def _minimax_qk_norm_tp_eager(
128
qkv: torch.Tensor,
129
q_weight: torch.Tensor,
130
k_weight: torch.Tensor,
131
q_size: int,
132
kv_size: int,
50
- tp_rank: int,
133
tp_world: int,
134
eps: float,
135
) -> tuple[torch.Tensor, torch.Tensor]:
136
+ """Pure-torch reference path used when Triton is unavailable."""
137
q, k, _ = qkv.split([q_size, kv_size, kv_size], dim=-1)
138
orig_dtype = q.dtype
139
q = q.to(torch.float32)
140
k = k.to(torch.float32)
141
q_var = q.pow(2).mean(dim=-1, keepdim=True)
142
k_var = k.pow(2).mean(dim=-1, keepdim=True)
60
- if tp_world > 1:
61
- qk_var = torch.cat([q_var, k_var], dim=-1)
62
- qk_var = _all_reduce_variance(qk_var) / tp_world
63
- q_var, k_var = qk_var.chunk(2, dim=-1)
143
+
144
+ qk_var = torch.cat([q_var, k_var], dim=-1)
145
+ qk_var = _all_reduce_variance(qk_var) / tp_world
146
+ q_var, k_var = qk_var.chunk(2, dim=-1)
147
q = q * torch.rsqrt(q_var + eps) * q_weight
148
k = k * torch.rsqrt(k_var + eps) * k_weight
149
return q.to(orig_dtype), k.to(orig_dtype)
150
151
152
+def _minimax_qk_norm_tp_fallback(
153
+ qkv: torch.Tensor,
154
+ q_weight: torch.Tensor,
155
+ k_weight: torch.Tensor,
156
+ q_size: int,
157
+ kv_size: int,
158
+ tp_rank: int,
159
+ tp_world: int,
160
+ eps: float,
161
+) -> tuple[torch.Tensor, torch.Tensor]:
162
+ """All-reduce + QK RMSNorm without the Lamport fused kernel.
163
+
164
+ The all-reduce is a TP communication barrier and cannot live inside a
165
+ single kernel, so the eager-torch path is split into two Triton kernels
166
+ around it: a variance reduction before the all-reduce and a normalize
167
+ after. Compared to the ``torch.compile`` path this avoids materializing
168
+ fp32 copies of q/k and the ``cat``/``chunk`` temporaries.
169
+ """
170
+ if not HAS_TRITON:
171
+ return _minimax_qk_norm_tp_eager(
172
+ qkv, q_weight, k_weight, q_size, kv_size, tp_world, eps
173
+ )
174
+
175
+ num_tokens = qkv.shape[0]
176
+ row_stride = qkv.stride(0)
177
+ BLOCK = 1024
178
+ grid = (num_tokens,)
179
+
180
+ qk_var = torch.empty(num_tokens, 2, dtype=torch.float32, device=qkv.device)
181
+ _minimax_qk_var_kernel[grid](
182
+ qkv, qk_var, row_stride, q_size=q_size, kv_size=kv_size, BLOCK=BLOCK
183
+ )
184
+
185
+ # All-reduce sums the per-shard means; the /tp_world that turns this back
186
+ # into the global mean is folded into the apply kernel's rsqrt below.
187
+ qk_var = _all_reduce_variance(qk_var)
188
+
189
+ q_out = torch.empty(num_tokens, q_size, dtype=qkv.dtype, device=qkv.device)
190
+ k_out = torch.empty(num_tokens, kv_size, dtype=qkv.dtype, device=qkv.device)
191
+ _minimax_rms_apply_kernel[grid](
192
+ qkv,
193
+ qk_var,
194
+ q_weight,
195
+ k_weight,
196
+ q_out,
197
+ k_out,
198
+ row_stride,
199
+ q_size=q_size,
200
+ kv_size=kv_size,
201
+ tp_world=tp_world,
202
+ eps=eps,
203
+ BLOCK=BLOCK,
204
+ )
205
+ return q_out, k_out
206
+
207
+
208
def _minimax_qk_norm_fusion(
209
qkv: torch.Tensor,
210
q_weight: torch.Tensor,
@@ -96,7 +235,7 @@ def _minimax_qk_norm_fusion(
235
tp_world,
236
eps,
237
)
99
- return _minimax_qk_norm_fallback(
238
+ return _minimax_qk_norm_tp_fallback(
239
qkv, q_weight, k_weight, q_size, kv_size, tp_rank, tp_world, eps
240
)
241
@@ -231,10 +370,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
370
k = k.to(torch.float32)
371
q_var = q.pow(2).mean(dim=-1, keepdim=True)
372
k_var = k.pow(2).mean(dim=-1, keepdim=True)
234
- if q_norm.tp_world > 1:
235
- qk_var = torch.cat([q_var, k_var], dim=-1)
236
- qk_var = _all_reduce_variance(qk_var) / q_norm.tp_world
237
- q_var, k_var = qk_var.chunk(2, dim=-1)
373
+
374
q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
375
k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight
376
q = q.to(orig_dtype)
@@ -250,7 +386,14 @@ class MiniMaxText01RMSNormTP(CustomOp):
386
kv_size: int,
387
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
388
assert qkv.ndim == 2
389
+
390
assert q_norm.variance_epsilon == k_norm.variance_epsilon
391
+ # Case 0： tp_size=1
392
+ if get_tensor_model_parallel_world_size() == 1:
393
+ q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
394
+ q, k = MiniMaxText01RMSNormTP.forward_qk(q_norm, k_norm, q, k)
395
+ return q, k, v
396
+ # Case : tp_size>1
397
q, k = torch.ops.vllm.minimax_qk_norm_fusion(
398
qkv,
399
q_norm.weight,