[Bugfix] Don't transpose fused MoE quantization scales in `RoutedExperts.load_weights` (#50137)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Harry Mellor committed
Jul 31, 2026 at 16:46 UTC
94e9ef0768d0297299069502bf46813dc7479ce5
2 files changed
+77
-8
tests/kernels/moe/test_moe_weight_loading_padded.py
+49
@@ -61,6 +61,55 @@ class TestGetHiddenDim:
61
RoutedExperts._get_hidden_dim(shard_dim=0, ndim=3)
62
63
64
+class TestOrientFusedWeight:
65
+ """Unit tests for _orient_fused_weight.
66
+
67
+ E=8 experts, hidden=3072, intermediate=1024.
68
+ """
69
+
70
+ HIDDEN = 3072
71
+
72
+ def test_w13_standard_orientation_is_untouched(self):
73
+ weight = torch.randn(8, 2048, self.HIDDEN)
74
+ result = RoutedExperts._orient_fused_weight(weight, "w1", self.HIDDEN)
75
+ assert result.shape == (8, 2048, self.HIDDEN)
76
+
77
+ def test_w13_transposed_checkpoint_is_normalised(self):
78
+ # e.g. Qwen3 VL MoE stores [experts, hidden, 2 * intermediate]
79
+ weight = torch.randn(8, self.HIDDEN, 2048)
80
+ result = RoutedExperts._orient_fused_weight(weight, "w3", self.HIDDEN)
81
+ assert result.shape == (8, 2048, self.HIDDEN)
82
+
83
+ def test_w2_standard_orientation_is_untouched(self):
84
+ weight = torch.randn(8, self.HIDDEN, 1024)
85
+ result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN)
86
+ assert result.shape == (8, self.HIDDEN, 1024)
87
+
88
+ def test_w2_transposed_checkpoint_is_normalised(self):
89
+ weight = torch.randn(8, 1024, self.HIDDEN)
90
+ result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN)
91
+ assert result.shape == (8, self.HIDDEN, 1024)
92
+
93
+ def test_w13_per_channel_scale_is_untouched(self):
94
+ # A fused per-channel scale has no hidden dim, so transposing it would
95
+ # leave chunk()/TP sharding operating on the wrong axis.
96
+ scale = torch.randn(8, 2048, 1)
97
+ result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN)
98
+ assert result.shape == (8, 2048, 1)
99
+ assert result.chunk(2, dim=1)[0].shape == (8, 1024, 1)
100
+
101
+ def test_w2_per_channel_scale_is_untouched(self):
102
+ scale = torch.randn(8, self.HIDDEN, 1)
103
+ result = RoutedExperts._orient_fused_weight(scale, "w2", self.HIDDEN)
104
+ assert result.shape == (8, self.HIDDEN, 1)
105
+
106
+ def test_block_scale_is_untouched(self):
107
+ # Block scales are [experts, 2 * intermediate / block, hidden / block]
108
+ scale = torch.randn(8, 16, 24)
109
+ result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN)
110
+ assert result.shape == (8, 16, 24)
111
+
112
+
113
class TestNarrowExpertDataForPadding:
114
"""Unit tests for _narrow_expert_data_for_padding."""
115
vllm/model_executor/layers/fused_moe/routed_experts.py
+28
-8
@@ -422,6 +422,28 @@ class RoutedExperts(PluggableLayer):
422
f"for a {ndim}D tensor (expected {dim_a} or {dim_b})"
423
)
424
425
+ @staticmethod
426
+ def _orient_fused_weight(
427
+ fused_weight: torch.Tensor, shard_id: str, unpadded_hidden: int
428
+ ) -> torch.Tensor:
429
+ """Normalise a fused expert tensor from either checkpoint orientation
430
+ to (intermediate, hidden) for w1/w3 and (hidden, intermediate) for w2.
431
+
432
+ Only transposes when the hidden dim is definitively on the wrong axis,
433
+ leaving tensors that have no hidden dim alone (e.g. a fused per-channel
434
+ scale, (2 * intermediate, 1)).
435
+ """
436
+ if shard_id == "w2":
437
+ hidden_axis, intermediate_axis = -2, -1
438
+ else:
439
+ hidden_axis, intermediate_axis = -1, -2
440
+ if (
441
+ fused_weight.shape[hidden_axis] != unpadded_hidden
442
+ and fused_weight.shape[intermediate_axis] == unpadded_hidden
443
+ ):
444
+ return fused_weight.transpose(-1, -2)
445
+ return fused_weight
446
+
447
@staticmethod
448
def _narrow_expert_data_for_padding(
449
expert_data: torch.Tensor,
@@ -900,7 +922,9 @@ class RoutedExperts(PluggableLayer):
922
self, weights: Iterable[tuple[str, torch.Tensor]]
923
) -> Iterable[str]:
924
expert_mapping = self.get_expert_mapping(include_fused=True)
903
- unpadded_hidden = self.moe_config.hidden_dim_unpadded
925
+ unpadded_hidden = (
926
+ self.moe_config.hidden_dim_unpadded or self.moe_config.hidden_dim
927
+ )
928
for expert_name, loaded_weight in weights:
929
qual_name = f"{self.layer_name}.{expert_name}"
930
# Fused expert weights can be identified by their 3D tensors
@@ -919,17 +943,13 @@ class RoutedExperts(PluggableLayer):
943
# w1 and w3 share one fused tensor; use a local copy so the
944
# transpose below doesn't mutate loaded_weight across
945
# iterations (else w3 is transposed twice and wrongly chunked)
922
- fused_weight = loaded_weight
946
+ fused_weight = self._orient_fused_weight(
947
+ loaded_weight, shard_id, unpadded_hidden
948
+ )
949
if shard_id in {"w1", "w3"}:
924
- if fused_weight.shape[-1] != unpadded_hidden:
925
- # [..., hidden, intermediate] -> [..., intermediate, hidden]
926
- fused_weight = fused_weight.transpose(-1, -2)
950
# Repurpose expert_id for deconcatenating w1 and w3
951
experts_shard = fused_weight.chunk(2, dim=1)[expert_id]
952
else:
930
- if fused_weight.shape[-2] != unpadded_hidden:
931
- # [..., intermediate, hidden] -> [..., hidden, intermediate]
932
- fused_weight = fused_weight.transpose(-1, -2)
953
experts_shard = fused_weight
954
start = 0
955
else: