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 4, 15:00 UTC
[Bugfix] Fix k_norm weight sharding in MiniMaxM2Attention when total_num_kv_heads < tp_size (#38191)
Signed-off-by: wxsIcey <1790571317@qq.com> Signed-off-by: Icey <1790571317@qq.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Xiaoshuang Wang committed
Apr 27, 2026 at 20:57 UTC
f8ac0c7cf0e3d4ac8894346005bdffe3bd7bd378
2 files changed
+44
-12
vllm/model_executor/layers/mamba/linear_attn.py
+28
-9
@@ -3,6 +3,7 @@
3
4
import math
5
from collections.abc import Callable
6
+from functools import partial
7
8
import torch
9
import torch.nn.functional as F
@@ -33,28 +34,46 @@ from vllm.v1.attention.backend import AttentionMetadata
34
from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
35
36
37
+@CustomOp.register("minimax_text01_rmsnorm_tp")
38
class MiniMaxText01RMSNormTP(CustomOp):
37
- name = "MiniMaxText01RMSNormTP"
38
-
39
- def __init__(self, hidden_size: int, eps: float = 1e-6) -> None:
39
+ def __init__(
40
+ self,
41
+ hidden_size: int,
42
+ eps: float = 1e-6,
43
+ *,
44
+ weight_shard_world_size: int | None = None,
45
+ weight_shard_rank: int | None = None,
46
+ ) -> None:
47
super().__init__()
48
self.tp_world = get_tensor_model_parallel_world_size()
49
self.tp_rank = get_tensor_model_parallel_rank()
43
- self.weight = nn.Parameter(torch.ones(int(hidden_size / self.tp_world)))
50
+ self.weight_shard_world = weight_shard_world_size or self.tp_world
51
+ self.weight_shard_rank = (
52
+ self.tp_rank if weight_shard_rank is None else weight_shard_rank
53
+ )
54
45
- self.weight.weight_loader = self.weight_loader
55
+ self.weight = nn.Parameter(torch.ones(hidden_size // self.weight_shard_world))
56
+ self.weight.weight_loader = partial(
57
+ self.weight_loader,
58
+ shard_world_size=self.weight_shard_world,
59
+ shard_rank=self.weight_shard_rank,
60
+ )
61
self.variance_epsilon = eps
62
63
@staticmethod
64
def weight_loader(
65
param: nn.Parameter,
66
loaded_weight: torch.Tensor,
67
+ shard_world_size: int | None = None,
68
+ shard_rank: int | None = None,
69
) -> None:
53
- tp_world = get_tensor_model_parallel_world_size()
54
- tp_rank = get_tensor_model_parallel_rank()
70
+ if shard_world_size is None:
71
+ shard_world_size = get_tensor_model_parallel_world_size()
72
+ if shard_rank is None:
73
+ shard_rank = get_tensor_model_parallel_rank()
74
56
- shard_size = loaded_weight.shape[0] // tp_world
57
- shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size)
75
+ shard_size = loaded_weight.shape[0] // shard_world_size
76
+ shard = slice(shard_rank * shard_size, (shard_rank + 1) * shard_size)
77
param.data.copy_(loaded_weight[shard])
78
79
def _forward(
vllm/model_executor/models/minimax_m2.py
+16
-3
@@ -35,6 +35,7 @@ from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, ModelConfig, VllmConfig
36
from vllm.distributed import (
37
get_pp_group,
38
+ get_tensor_model_parallel_rank,
39
get_tensor_model_parallel_world_size,
40
)
41
from vllm.model_executor.layers.attention import Attention
@@ -220,9 +221,21 @@ class MiniMaxM2Attention(nn.Module):
221
self.q_norm = MiniMaxText01RMSNormTP(
222
self.head_dim * self.total_num_heads, eps=rms_norm_eps
223
)
223
- self.k_norm = MiniMaxText01RMSNormTP(
224
- self.head_dim * self.total_num_kv_heads, eps=rms_norm_eps
225
- )
224
+ if self.total_num_kv_heads >= tp_size:
225
+ self.k_norm = MiniMaxText01RMSNormTP(
226
+ self.head_dim * self.total_num_kv_heads, eps=rms_norm_eps
227
+ )
228
+ else:
229
+ # KV heads are replicated across TP ranks; shard k_norm weight by
230
+ # total_num_kv_heads rather than tp_size to avoid incorrect sharding.
231
+ num_kv_head_replicas = tp_size // self.total_num_kv_heads
232
+ self.k_norm = MiniMaxText01RMSNormTP(
233
+ self.head_dim * self.total_num_kv_heads,
234
+ eps=rms_norm_eps,
235
+ weight_shard_world_size=self.total_num_kv_heads,
236
+ weight_shard_rank=get_tensor_model_parallel_rank()
237
+ // num_kv_head_replicas,
238
+ )
239
240
def forward(
241
self,