A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix minimax_qk_norm_fusion (#44983)
Jee Jee Li committed
Jun 9, 2026 at 21:43 UTC
dc10e467a9859f220ba9d36c31c50fb22bc38373
1 file changed
+41
-9
vllm/model_executor/layers/minimax_rms_norm/rms_norm_tp.py
+41
-9
@@ -12,10 +12,13 @@ from vllm.distributed.parallel_state import (
12
get_tensor_model_parallel_world_size,
13
get_tp_group,
14
)
15
+from vllm.logger import init_logger
16
from vllm.model_executor.custom_op import CustomOp
17
from vllm.platforms import current_platform
18
from vllm.utils.torch_utils import direct_register_custom_op
19
20
+logger = init_logger(__name__)
21
+
22
# Max number of tokens supported by the Lamport fused allreduce+RMSNorm kernel.
23
# Larger batches fall back to the eager allreduce + RMSNorm path.
24
MINIMAX_QK_NORM_MAX_TOKEN_NUM = 2048
@@ -23,6 +26,20 @@ MINIMAX_QK_NORM_MAX_TOKEN_NUM = 2048
26
_MINIMAX_FUSED_AR_RMS_QK = getattr(torch.ops._C, "minimax_allreduce_rms_qk", None)
27
28
29
+def _all_reduce_variance(var: torch.Tensor) -> torch.Tensor:
30
+ """All-reduce a per-token variance tensor across the TP group.
31
+
32
+ Variance is accumulated in fp32 for numerical stability. The FlashInfer
33
+ fused all-reduce caches a single global workspace keyed to the model's
34
+ 16-bit activation dtype (``use_fp32_lamport=False``); routing an fp32
35
+ reduction through it would read against a mismatched workspace and corrupt
36
+ the result. FlashInfer's fast-path only triggers for 2D inputs, so reducing
37
+ a flattened (1D) view keeps these fp32 reductions on custom all-reduce /
38
+ pynccl, both of which handle fp32 correctly.
39
+ """
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(
45
qkv: torch.Tensor,
@@ -42,7 +59,7 @@ def _minimax_qk_norm_fallback(
59
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)
45
- qk_var = tensor_model_parallel_all_reduce(qk_var) / tp_world
62
+ qk_var = _all_reduce_variance(qk_var) / tp_world
63
q_var, k_var = qk_var.chunk(2, dim=-1)
64
q = q * torch.rsqrt(q_var + eps) * q_weight
65
k = k * torch.rsqrt(k_var + eps) * k_weight
@@ -143,12 +160,27 @@ class MiniMaxText01RMSNormTP(CustomOp):
160
get_allreduce_workspace,
161
)
162
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
- )
163
+ # The Lamport workspace exchanges CUDA IPC handles and enables peer
164
+ # access between GPUs. This requires P2P (IPC peer access) to be
165
+ # available; on topologies where it is not (e.g. consumer PCIe cards
166
+ # with P2P disabled in the driver), allocation raises. Fall back to
167
+ # the eager allreduce + RMSNorm path instead of failing model load.
168
+ try:
169
+ self.workspace = get_allreduce_workspace(
170
+ rank=self.tp_rank,
171
+ world_size=self.tp_world,
172
+ max_tokens=MINIMAX_QK_NORM_MAX_TOKEN_NUM,
173
+ process_group=get_tp_group().cpu_group,
174
+ )
175
+ except Exception as e:
176
+ logger.warning_once(
177
+ "Failed to initialize MiniMax fused allreduce+RMSNorm "
178
+ "Lamport workspace: %s. This is expected on GPUs without "
179
+ "P2P (IPC peer access) support. Falling back to the eager "
180
+ "allreduce + RMSNorm path.",
181
+ e,
182
+ )
183
+ self.workspace = None
184
185
@staticmethod
186
def weight_loader(
@@ -174,7 +206,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
206
x = x.to(torch.float32)
207
variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32)
208
if self.tp_world > 1:
177
- variance = tensor_model_parallel_all_reduce(variance) / self.tp_world
209
+ variance = _all_reduce_variance(variance) / self.tp_world
210
x = x * torch.rsqrt(variance + self.variance_epsilon)
211
x = (x * self.weight).to(orig_dtype)
212
return x
@@ -201,7 +233,7 @@ class MiniMaxText01RMSNormTP(CustomOp):
233
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)
204
- qk_var = tensor_model_parallel_all_reduce(qk_var) / q_norm.tp_world
236
+ qk_var = _all_reduce_variance(qk_var) / q_norm.tp_world
237
q_var, k_var = qk_var.chunk(2, dim=-1)
238
q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
239
k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight