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 5, 04:45 UTC
[fix][MOE] Fix MOE experts `intermediate_size` dimension not being narrowed before weight loading (#39688)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
fxmarty-amd committed
Apr 14, 2026 at 11:35 UTC
4eee77b877d9ecf493967433707eb44caf5c7558
2 files changed
+63
-15
tests/kernels/moe/test_moe_weight_loading_padded.py
+35
@@ -257,6 +257,41 @@ class TestWeightLoadingWithPaddedHiddenSize:
257
258
assert torch.equal(expert_data_full, loaded_weight)
259
260
+ def test_narrow_shard_dim(self):
261
+ """Simulate loading w2 when both hidden_size and intermediate_size
262
+ are padded.
263
+ """
264
+ padded_hidden = 3072
265
+ original_hidden = 2688
266
+ padded_intermediate = 1024
267
+ original_intermediate = 896
268
+
269
+ expert_data_full = torch.zeros(padded_hidden, padded_intermediate)
270
+ loaded_weight = torch.randn(original_hidden, original_intermediate)
271
+
272
+ shard_dim = 1
273
+ hidden_dim = FusedMoE._get_hidden_dim(shard_dim=shard_dim, ndim=2)
274
+ expert_data = FusedMoE._narrow_expert_data_for_padding(
275
+ expert_data_full,
276
+ loaded_weight,
277
+ hidden_dim=hidden_dim,
278
+ shard_dim=shard_dim,
279
+ )
280
+ expert_data.copy_(loaded_weight)
281
+
282
+ assert torch.equal(
283
+ expert_data_full[:original_hidden, :original_intermediate],
284
+ loaded_weight,
285
+ )
286
+ assert torch.equal(
287
+ expert_data_full[original_hidden:, :],
288
+ torch.zeros(padded_hidden - original_hidden, padded_intermediate),
289
+ )
290
+ assert torch.equal(
291
+ expert_data_full[:original_hidden, original_intermediate:],
292
+ torch.zeros(original_hidden, padded_intermediate - original_intermediate),
293
+ )
294
+
295
def test_bnb_shape_mismatch_raises(self):
296
"""BnB + padded hidden_size should raise via weight_loader."""
297
from unittest.mock import MagicMock
vllm/model_executor/layers/fused_moe/layer.py
+28
-15
@@ -842,7 +842,10 @@ class FusedMoE(CustomOp):
842
if shard_id == "w2":
843
hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
844
expert_data = self._narrow_expert_data_for_padding(
845
- expert_data, loaded_weight, hidden_dim=hidden_dim
845
+ expert_data,
846
+ loaded_weight,
847
+ hidden_dim=hidden_dim,
848
+ shard_dim=shard_dim,
849
)
850
expert_data.copy_(loaded_weight)
851
elif shard_id in ("w1", "w3"):
@@ -882,29 +885,33 @@ class FusedMoE(CustomOp):
885
expert_data: torch.Tensor,
886
loaded_weight: torch.Tensor,
887
hidden_dim: int,
888
+ shard_dim: int | None = None,
889
) -> torch.Tensor:
886
- """Narrow expert_data hidden dim to match loaded_weight for padded
887
- hidden_size.
890
+ """Narrow expert_data to match loaded_weight for padded dimensions.
891
892
When backends (e.g., DeepEP) round up hidden_size, weight parameters
893
are larger than checkpoint weights. Narrow the padded hidden dimension
891
- before copying.
894
+ before copying. Similarly, when padding occurs on the shard
895
+ (intermediate) dimension (e.g. for MXFP4 GEMM), narrow that dimension
896
+ as well.
897
898
Args:
899
expert_data: The (possibly padded) parameter tensor to narrow.
900
loaded_weight: The checkpoint weight tensor with original size.
901
hidden_dim: The dimension index corresponding to hidden_size.
902
Must be non-negative.
903
+ shard_dim: The dimension index corresponding to the shard
904
+ (intermediate) dimension. Defaults to `None`.
905
"""
899
- if (
900
- loaded_weight.ndim > 0
901
- and 0 <= hidden_dim < expert_data.ndim
902
- and hidden_dim < loaded_weight.ndim
903
- and expert_data.shape[hidden_dim] > loaded_weight.shape[hidden_dim]
904
- ):
905
- expert_data = expert_data.narrow(
906
- hidden_dim, 0, loaded_weight.shape[hidden_dim]
907
- )
906
+ dims = (hidden_dim,) if shard_dim is None else (hidden_dim, shard_dim)
907
+ if loaded_weight.ndim > 0:
908
+ for dim in dims:
909
+ if (
910
+ 0 <= dim < expert_data.ndim
911
+ and dim < loaded_weight.ndim
912
+ and expert_data.shape[dim] > loaded_weight.shape[dim]
913
+ ):
914
+ expert_data = expert_data.narrow(dim, 0, loaded_weight.shape[dim])
915
return expert_data
916
917
def _load_w13(
@@ -946,7 +953,10 @@ class FusedMoE(CustomOp):
953
expert_data = expert_data.narrow(shard_dim, shard_size, shard_size)
954
hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
955
expert_data = self._narrow_expert_data_for_padding(
949
- expert_data, loaded_weight, hidden_dim=hidden_dim
956
+ expert_data,
957
+ loaded_weight,
958
+ hidden_dim=hidden_dim,
959
+ shard_dim=shard_dim,
960
)
961
expert_data.copy_(loaded_weight)
962
@@ -979,7 +989,10 @@ class FusedMoE(CustomOp):
989
# w2, down_proj: Load into only logical weight of w2.
990
hidden_dim = self._get_hidden_dim(shard_dim, expert_data.ndim)
991
expert_data = self._narrow_expert_data_for_padding(
982
- expert_data, loaded_weight, hidden_dim=hidden_dim
992
+ expert_data,
993
+ loaded_weight,
994
+ hidden_dim=hidden_dim,
995
+ shard_dim=shard_dim,
996
)
997
expert_data.copy_(loaded_weight)
998