[Bugfix][MoE] Fix fused block-scale orientation (#50727)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Andreas Karatzas committed
Aug 11, 2026 at 15:05 UTC
36f4630d8926b8bd16fd6fc80b71e9defce2ad92
5 files changed
+80
-29
tests/kernels/moe/test_moe_weight_loading_padded.py
+27
-7
@@ -73,43 +73,63 @@ class TestOrientFusedWeight:
73
74
def test_w13_standard_orientation_is_untouched(self):
75
weight = torch.randn(8, 2048, self.HIDDEN)
76
- result = RoutedExperts._orient_fused_weight(weight, "w1", self.HIDDEN)
76
+ result = RoutedExperts._orient_fused_weight(weight, False)
77
assert result.shape == (8, 2048, self.HIDDEN)
78
79
def test_w13_transposed_checkpoint_is_normalised(self):
80
# e.g. Qwen3 VL MoE stores [experts, hidden, 2 * intermediate]
81
weight = torch.randn(8, self.HIDDEN, 2048)
82
- result = RoutedExperts._orient_fused_weight(weight, "w3", self.HIDDEN)
82
+ result = RoutedExperts._orient_fused_weight(weight, True)
83
assert result.shape == (8, 2048, self.HIDDEN)
84
85
def test_w2_standard_orientation_is_untouched(self):
86
weight = torch.randn(8, self.HIDDEN, 1024)
87
- result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN)
87
+ result = RoutedExperts._orient_fused_weight(weight, False)
88
assert result.shape == (8, self.HIDDEN, 1024)
89
90
def test_w2_transposed_checkpoint_is_normalised(self):
91
weight = torch.randn(8, 1024, self.HIDDEN)
92
- result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN)
92
+ result = RoutedExperts._orient_fused_weight(weight, True)
93
assert result.shape == (8, self.HIDDEN, 1024)
94
95
def test_w13_per_channel_scale_is_untouched(self):
96
# A fused per-channel scale has no hidden dim, so transposing it would
97
# leave chunk()/TP sharding operating on the wrong axis.
98
scale = torch.randn(8, 2048, 1)
99
- result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN)
99
+ result = RoutedExperts._orient_fused_weight(scale, False)
100
assert result.shape == (8, 2048, 1)
101
assert result.chunk(2, dim=1)[0].shape == (8, 1024, 1)
102
103
def test_w2_per_channel_scale_is_untouched(self):
104
scale = torch.randn(8, self.HIDDEN, 1)
105
- result = RoutedExperts._orient_fused_weight(scale, "w2", self.HIDDEN)
105
+ result = RoutedExperts._orient_fused_weight(scale, False)
106
assert result.shape == (8, self.HIDDEN, 1)
107
108
def test_block_scale_is_untouched(self):
109
# Block scales are [experts, 2 * intermediate / block, hidden / block]
110
scale = torch.randn(8, 16, 24)
111
- result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN)
111
+ result = RoutedExperts._orient_fused_weight(scale, False)
112
assert result.shape == (8, 16, 24)
113
+ assert result.data_ptr() == scale.data_ptr()
114
+
115
+ @pytest.mark.parametrize(
116
+ "checkpoint_shape",
117
+ [
118
+ # Qwen3-VL stores block scales in checkpoint weight orientation.
119
+ (8, 16, 12),
120
+ (8, 6, 16),
121
+ (8, 16, 16),
122
+ ],
123
+ )
124
+ def test_qwen3_vl_transposed_block_scale_uses_explicit_layout(
125
+ self,
126
+ checkpoint_shape: tuple[int, ...],
127
+ ):
128
+ scale = torch.arange(torch.tensor(checkpoint_shape).prod()).reshape(
129
+ checkpoint_shape
130
+ )
131
+ result = RoutedExperts._orient_fused_weight(scale, True)
132
+ torch.testing.assert_close(result, scale.transpose(-1, -2))
133
134
135
class TestNarrowExpertDataForPadding:
vllm/model_executor/layers/fused_moe/layer.py
+4
@@ -130,6 +130,7 @@ def FusedMoEFactory(
130
is_sequence_parallel: bool = False,
131
reduce_results: bool = True,
132
ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),
133
+ is_fused_checkpoint_transposed: bool = False,
134
n_shared_experts: int | None = None,
135
router_logits_dtype: torch.dtype | None = None,
136
gate: torch.nn.Module | None = None,
@@ -193,6 +194,8 @@ def FusedMoEFactory(
194
the late-AR path.
195
ckpt_names: Checkpoint parameter name tuple (gate_proj, down_proj,
196
up_proj) used for weight loading
197
+ is_fused_checkpoint_transposed: Whether fused checkpoint weights and
198
+ block scales use transposed storage.
199
n_shared_experts: Number of shared experts to fuse into the routed
200
grouped GEMM (ROCm; requires aiter FSE or the router-append path)
201
router_logits_dtype: Data type for router logits buffers
@@ -378,6 +381,7 @@ def FusedMoEFactory(
381
ckpt_gate_proj_name=ckpt_names[0],
382
ckpt_down_proj_name=ckpt_names[1],
383
ckpt_up_proj_name=ckpt_names[2],
384
+ is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
385
# Extra params that are needed by quant_methods, pass along for now
386
# Prefer getting these from other sources, e.g. moe_config or
387
# router object
vllm/model_executor/layers/fused_moe/routed_experts.py
+15
-20
@@ -61,6 +61,7 @@ class RoutedExperts(PluggableLayer):
61
ckpt_gate_proj_name: str = "gate_proj",
62
ckpt_down_proj_name: str = "down_proj",
63
ckpt_up_proj_name: str = "up_proj",
64
+ is_fused_checkpoint_transposed: bool = False,
65
#
66
# Extra params that are needed by quant_methods, pass along for now
67
# Prefer getting these from other sources, e.g. moe_config or
@@ -86,6 +87,7 @@ class RoutedExperts(PluggableLayer):
87
self.ckpt_gate_proj_name = ckpt_gate_proj_name
88
self.ckpt_down_proj_name = ckpt_down_proj_name
89
self.ckpt_up_proj_name = ckpt_up_proj_name
90
+ self.is_fused_checkpoint_transposed = is_fused_checkpoint_transposed
91
self.expert_map_manager = expert_map_manager
92
self.hidden_size = moe_config.hidden_dim
93
self.global_num_experts = moe_config.num_experts
@@ -423,23 +425,11 @@ class RoutedExperts(PluggableLayer):
425
426
@staticmethod
427
def _orient_fused_weight(
426
- fused_weight: torch.Tensor, shard_id: str, unpadded_hidden: int
428
+ fused_weight: torch.Tensor,
429
+ is_fused_checkpoint_transposed: bool,
430
) -> torch.Tensor:
428
- """Normalise a fused expert tensor from either checkpoint orientation
429
- to (intermediate, hidden) for w1/w3 and (hidden, intermediate) for w2.
430
-
431
- Only transposes when the hidden dim is definitively on the wrong axis,
432
- leaving tensors that have no hidden dim alone (e.g. a fused per-channel
433
- scale, (2 * intermediate, 1)).
434
- """
435
- if shard_id == "w2":
436
- hidden_axis, intermediate_axis = -2, -1
437
- else:
438
- hidden_axis, intermediate_axis = -1, -2
439
- if (
440
- fused_weight.shape[hidden_axis] != unpadded_hidden
441
- and fused_weight.shape[intermediate_axis] == unpadded_hidden
442
- ):
431
+ """Normalise a fused expert tensor to the vLLM weight layout."""
432
+ if is_fused_checkpoint_transposed:
433
return fused_weight.transpose(-1, -2)
434
return fused_weight
435
@@ -882,9 +872,6 @@ class RoutedExperts(PluggableLayer):
872
self, weights: Iterable[tuple[str, torch.Tensor]]
873
) -> Iterable[str]:
874
expert_mapping = self.get_expert_mapping(include_fused=True)
885
- unpadded_hidden = (
886
- self.moe_config.hidden_dim_unpadded or self.moe_config.hidden_dim
887
- )
875
for expert_name, loaded_weight in weights:
876
qual_name = f"{self.layer_name}.{expert_name}"
877
# Fused expert weights can be identified by their 3D tensors
@@ -915,11 +902,19 @@ class RoutedExperts(PluggableLayer):
902
f"for checkpoint weight {qual_name!r}"
903
)
904
if is_fused:
905
+ quant_method = getattr(param, "quant_method", None)
906
+ # Block scales share the weight's two-dimensional layout.
907
+ # Other quantization metadata can use independent layouts.
908
+ uses_weight_layout = (
909
+ "scale" not in weight_name
910
+ or quant_method == FusedMoeWeightScaleSupported.BLOCK.value
911
+ )
912
# w1 and w3 share one fused tensor; use a local copy so the
913
# transpose below doesn't mutate loaded_weight across
914
# iterations (else w3 is transposed twice and wrongly chunked)
915
fused_weight = self._orient_fused_weight(
922
- loaded_weight, shard_id, unpadded_hidden
916
+ loaded_weight,
917
+ self.is_fused_checkpoint_transposed and uses_weight_layout,
918
)
919
if shard_id in {"w1", "w3"}:
920
# Repurpose expert_id for deconcatenating w1 and w3
vllm/model_executor/models/qwen3_moe.py
+11
-2
@@ -132,6 +132,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
132
self,
133
vllm_config: VllmConfig,
134
prefix: str = "",
135
+ is_fused_checkpoint_transposed: bool = False,
136
):
137
super().__init__()
138
@@ -208,6 +209,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
209
enable_eplb=self.enable_eplb,
210
num_redundant_experts=self.n_redundant_experts,
211
is_sequence_parallel=self.is_sequence_parallel,
212
+ is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
213
)
214
215
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -339,7 +341,12 @@ class Qwen3MoeAttention(nn.Module):
341
342
343
class Qwen3MoeDecoderLayer(nn.Module):
342
- def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
344
+ def __init__(
345
+ self,
346
+ vllm_config: VllmConfig,
347
+ prefix: str = "",
348
+ is_fused_checkpoint_transposed: bool = False,
349
+ ) -> None:
350
super().__init__()
351
352
config = vllm_config.model_config.hf_text_config
@@ -375,7 +382,9 @@ class Qwen3MoeDecoderLayer(nn.Module):
382
config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
383
):
384
self.mlp = Qwen3MoeSparseMoeBlock(
378
- vllm_config=vllm_config, prefix=f"{prefix}.mlp"
385
+ vllm_config=vllm_config,
386
+ prefix=f"{prefix}.mlp",
387
+ is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
388
)
389
else:
390
self.mlp = Qwen3MoeMLP(
vllm/model_executor/models/qwen3_vl_moe.py
+23
@@ -43,6 +43,7 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config
43
44
from .interfaces import MixtureOfExperts
45
from .qwen3_moe import (
46
+ Qwen3MoeDecoderLayer,
47
Qwen3MoeForCausalLM,
48
Qwen3MoeModel,
49
Qwen3MoeSparseMoeBlock,
@@ -64,6 +65,15 @@ class Qwen3VLMoeProcessingInfo(Qwen3VLProcessingInfo):
65
return self.ctx.get_hf_config(Qwen3VLMoeConfig)
66
67
68
+class Qwen3VLMoeDecoderLayer(Qwen3MoeDecoderLayer):
69
+ def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
70
+ super().__init__(
71
+ vllm_config=vllm_config,
72
+ prefix=prefix,
73
+ is_fused_checkpoint_transposed=True,
74
+ )
75
+
76
+
77
@support_torch_compile(
78
dynamic_arg_dims={
79
"input_ids": 0,
@@ -77,6 +87,19 @@ class Qwen3VLMoeProcessingInfo(Qwen3VLProcessingInfo):
87
}
88
)
89
class Qwen3MoeLLMModel(Qwen3MoeModel):
90
+ def __init__(
91
+ self,
92
+ *,
93
+ vllm_config: VllmConfig,
94
+ prefix: str = "",
95
+ decoder_layer_type: type[torch.nn.Module] = Qwen3VLMoeDecoderLayer,
96
+ ):
97
+ super().__init__(
98
+ vllm_config=vllm_config,
99
+ prefix=prefix,
100
+ decoder_layer_type=decoder_layer_type,
101
+ )
102
+
103
def forward(
104
self,
105
input_ids: torch.Tensor | None,