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 10, 21:30 UTC
Optimize QKNorm for MiniMax-M2/M2.1 (#31493)
Signed-off-by: xuebi <xuebi@minimaxi.com> Co-authored-by: xuebi <xuebi@minimaxi.com>
Roger Young committed
Dec 30, 2025 at 00:30 UTC
5bc664110f120adf789560b2ecfdb19a1cc72ee0
2 files changed
+25
-2
vllm/model_executor/layers/mamba/linear_attn.py
+22
@@ -79,6 +79,28 @@ class MiniMaxText01RMSNormTP(CustomOp):
79
assert residual is None, "RMSNorm does not support residual connection."
80
return self._forward(x)
81
82
+ @staticmethod
83
+ def forward_qk(
84
+ q_norm: "MiniMaxText01RMSNormTP",
85
+ k_norm: "MiniMaxText01RMSNormTP",
86
+ q: torch.Tensor,
87
+ k: torch.Tensor,
88
+ ) -> tuple[torch.Tensor, torch.Tensor]:
89
+ orig_dtype = q.dtype
90
+ q = q.to(torch.float32)
91
+ k = k.to(torch.float32)
92
+ q_var = q.pow(2).mean(dim=-1, keepdim=True)
93
+ k_var = k.pow(2).mean(dim=-1, keepdim=True)
94
+ if q_norm.tp_world > 1:
95
+ qk_var = torch.cat([q_var, k_var], dim=-1)
96
+ qk_var = tensor_model_parallel_all_reduce(qk_var) / q_norm.tp_world
97
+ q_var, k_var = qk_var.chunk(2, dim=-1)
98
+ q = q * torch.rsqrt(q_var + q_norm.variance_epsilon) * q_norm.weight
99
+ k = k * torch.rsqrt(k_var + k_norm.variance_epsilon) * k_norm.weight
100
+ q = q.to(orig_dtype)
101
+ k = k.to(orig_dtype)
102
+ return q, k
103
+
104
105
class MiniMaxText01LinearKernel:
106
@staticmethod
vllm/model_executor/models/minimax_m2.py
+3
-2
@@ -234,8 +234,9 @@ class MiniMaxM2Attention(nn.Module):
234
) -> torch.Tensor:
235
qkv, _ = self.qkv_proj(hidden_states)
236
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
237
- q = self.q_norm(q)
238
- k = self.k_norm(k)
237
+ q, k = MiniMaxText01RMSNormTP.forward_qk(
238
+ self.q_norm, self.k_norm, q.contiguous(), k.contiguous()
239
+ )
240
q, k = self.rotary_emb(positions, q, k)
241
attn_output = self.attn(q, k, v)
242
output, _ = self.o_proj(attn_output)