iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix][MoE] Fix fused block-scale orientation (#50727)

Browse Source 
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
Andreas Karatzas
2026-08-11 16:05:38 -04:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent fd04edef3e
commit 36f4630d89
5 changed files with 80 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/test_moe_weight_loading_padded.py
+27 -7
View File
@@ -73,43 +73,63 @@ class TestOrientFusedWeight:
def test_w13_standard_orientation_is_untouched(self):
weight = torch.randn(8, 2048, self.HIDDEN)
result = RoutedExperts._orient_fused_weight(weight, "w1", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(weight, False)
assert result.shape == (8, 2048, self.HIDDEN)
def test_w13_transposed_checkpoint_is_normalised(self):
# e.g. Qwen3 VL MoE stores [experts, hidden, 2 * intermediate]
weight = torch.randn(8, self.HIDDEN, 2048)
result = RoutedExperts._orient_fused_weight(weight, "w3", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(weight, True)
assert result.shape == (8, 2048, self.HIDDEN)
def test_w2_standard_orientation_is_untouched(self):
weight = torch.randn(8, self.HIDDEN, 1024)
result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(weight, False)
assert result.shape == (8, self.HIDDEN, 1024)
def test_w2_transposed_checkpoint_is_normalised(self):
weight = torch.randn(8, 1024, self.HIDDEN)
result = RoutedExperts._orient_fused_weight(weight, "w2", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(weight, True)
assert result.shape == (8, self.HIDDEN, 1024)
def test_w13_per_channel_scale_is_untouched(self):
# A fused per-channel scale has no hidden dim, so transposing it would
# leave chunk()/TP sharding operating on the wrong axis.
scale = torch.randn(8, 2048, 1)
result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(scale, False)
assert result.shape == (8, 2048, 1)
assert result.chunk(2, dim=1)[0].shape == (8, 1024, 1)
def test_w2_per_channel_scale_is_untouched(self):
scale = torch.randn(8, self.HIDDEN, 1)
result = RoutedExperts._orient_fused_weight(scale, "w2", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(scale, False)
assert result.shape == (8, self.HIDDEN, 1)
def test_block_scale_is_untouched(self):
# Block scales are [experts, 2 * intermediate / block, hidden / block]
scale = torch.randn(8, 16, 24)
result = RoutedExperts._orient_fused_weight(scale, "w1", self.HIDDEN)
result = RoutedExperts._orient_fused_weight(scale, False)
assert result.shape == (8, 16, 24)
assert result.data_ptr() == scale.data_ptr()
@pytest.mark.parametrize(
"checkpoint_shape",
[
# Qwen3-VL stores block scales in checkpoint weight orientation.
(8, 16, 12),
(8, 6, 16),
(8, 16, 16),
],
)
def test_qwen3_vl_transposed_block_scale_uses_explicit_layout(
self,
checkpoint_shape: tuple[int, ...],
):
scale = torch.arange(torch.tensor(checkpoint_shape).prod()).reshape(
checkpoint_shape
)
result = RoutedExperts._orient_fused_weight(scale, True)
torch.testing.assert_close(result, scale.transpose(-1, -2))
class TestNarrowExpertDataForPadding:
vllm/model_executor/layers/fused_moe/layer.py
+4
View File
@@ -130,6 +130,7 @@ def FusedMoEFactory(
is_sequence_parallel: bool = False,
reduce_results: bool = True,
ckpt_names: tuple[str, str, str] = ("gate_proj", "down_proj", "up_proj"),
is_fused_checkpoint_transposed: bool = False,
n_shared_experts: int | None = None,
router_logits_dtype: torch.dtype | None = None,
gate: torch.nn.Module | None = None,
@@ -193,6 +194,8 @@ def FusedMoEFactory(
the late-AR path.
ckpt_names: Checkpoint parameter name tuple (gate_proj, down_proj,
up_proj) used for weight loading
is_fused_checkpoint_transposed: Whether fused checkpoint weights and
block scales use transposed storage.
n_shared_experts: Number of shared experts to fuse into the routed
grouped GEMM (ROCm; requires aiter FSE or the router-append path)
router_logits_dtype: Data type for router logits buffers
@@ -378,6 +381,7 @@ def FusedMoEFactory(
ckpt_gate_proj_name=ckpt_names[0],
ckpt_down_proj_name=ckpt_names[1],
ckpt_up_proj_name=ckpt_names[2],
is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
# Extra params that are needed by quant_methods, pass along for now
# Prefer getting these from other sources, e.g. moe_config or
# router object
vllm/model_executor/layers/fused_moe/routed_experts.py
+15 -20
View File
@@ -61,6 +61,7 @@ class RoutedExperts(PluggableLayer):
ckpt_gate_proj_name: str = "gate_proj",
ckpt_down_proj_name: str = "down_proj",
ckpt_up_proj_name: str = "up_proj",
is_fused_checkpoint_transposed: bool = False,
#
# Extra params that are needed by quant_methods, pass along for now
# Prefer getting these from other sources, e.g. moe_config or
@@ -86,6 +87,7 @@ class RoutedExperts(PluggableLayer):
self.ckpt_gate_proj_name = ckpt_gate_proj_name
self.ckpt_down_proj_name = ckpt_down_proj_name
self.ckpt_up_proj_name = ckpt_up_proj_name
self.is_fused_checkpoint_transposed = is_fused_checkpoint_transposed
self.expert_map_manager = expert_map_manager
self.hidden_size = moe_config.hidden_dim
self.global_num_experts = moe_config.num_experts
@@ -423,23 +425,11 @@ class RoutedExperts(PluggableLayer):
@staticmethod
def _orient_fused_weight(
fused_weight: torch.Tensor, shard_id: str, unpadded_hidden: int
fused_weight: torch.Tensor,
is_fused_checkpoint_transposed: bool,
) -> torch.Tensor:
"""Normalise a fused expert tensor from either checkpoint orientation
to (intermediate, hidden) for w1/w3 and (hidden, intermediate) for w2.
Only transposes when the hidden dim is definitively on the wrong axis,
leaving tensors that have no hidden dim alone (e.g. a fused per-channel
scale, (2 * intermediate, 1)).
"""
if shard_id == "w2":
hidden_axis, intermediate_axis = -2, -1
else:
hidden_axis, intermediate_axis = -1, -2
if (
fused_weight.shape[hidden_axis] != unpadded_hidden
and fused_weight.shape[intermediate_axis] == unpadded_hidden
):
"""Normalise a fused expert tensor to the vLLM weight layout."""
if is_fused_checkpoint_transposed:
return fused_weight.transpose(-1, -2)
return fused_weight
@@ -882,9 +872,6 @@ class RoutedExperts(PluggableLayer):
self, weights: Iterable[tuple[str, torch.Tensor]]
) -> Iterable[str]:
expert_mapping = self.get_expert_mapping(include_fused=True)
unpadded_hidden = (
self.moe_config.hidden_dim_unpadded or self.moe_config.hidden_dim
)
for expert_name, loaded_weight in weights:
qual_name = f"{self.layer_name}.{expert_name}"
# Fused expert weights can be identified by their 3D tensors
@@ -915,11 +902,19 @@ class RoutedExperts(PluggableLayer):
f"for checkpoint weight {qual_name!r}"
)
if is_fused:
quant_method = getattr(param, "quant_method", None)
# Block scales share the weight's two-dimensional layout.
# Other quantization metadata can use independent layouts.
uses_weight_layout = (
"scale" not in weight_name
or quant_method == FusedMoeWeightScaleSupported.BLOCK.value
)
# w1 and w3 share one fused tensor; use a local copy so the
# transpose below doesn't mutate loaded_weight across
# iterations (else w3 is transposed twice and wrongly chunked)
fused_weight = self._orient_fused_weight(
loaded_weight, shard_id, unpadded_hidden
loaded_weight,
self.is_fused_checkpoint_transposed and uses_weight_layout,
)
if shard_id in {"w1", "w3"}:
# Repurpose expert_id for deconcatenating w1 and w3
vllm/model_executor/models/qwen3_moe.py
+11 -2
View File
@@ -132,6 +132,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
self,
vllm_config: VllmConfig,
prefix: str = "",
is_fused_checkpoint_transposed: bool = False,
):
super().__init__()
@@ -208,6 +209,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
enable_eplb=self.enable_eplb,
num_redundant_experts=self.n_redundant_experts,
is_sequence_parallel=self.is_sequence_parallel,
is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
@@ -339,7 +341,12 @@ class Qwen3MoeAttention(nn.Module):
class Qwen3MoeDecoderLayer(nn.Module):
def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
def __init__(
self,
vllm_config: VllmConfig,
prefix: str = "",
is_fused_checkpoint_transposed: bool = False,
) -> None:
super().__init__()
config = vllm_config.model_config.hf_text_config
@@ -375,7 +382,9 @@ class Qwen3MoeDecoderLayer(nn.Module):
config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
):
self.mlp = Qwen3MoeSparseMoeBlock(
vllm_config=vllm_config, prefix=f"{prefix}.mlp"
vllm_config=vllm_config,
prefix=f"{prefix}.mlp",
is_fused_checkpoint_transposed=is_fused_checkpoint_transposed,
)
else:
self.mlp = Qwen3MoeMLP(
vllm/model_executor/models/qwen3_vl_moe.py
+23
View File
@@ -43,6 +43,7 @@ from vllm.tokenizers.registry import cached_tokenizer_from_config
from .interfaces import MixtureOfExperts
from .qwen3_moe import (
Qwen3MoeDecoderLayer,
Qwen3MoeForCausalLM,
Qwen3MoeModel,
Qwen3MoeSparseMoeBlock,
@@ -64,6 +65,15 @@ class Qwen3VLMoeProcessingInfo(Qwen3VLProcessingInfo):
return self.ctx.get_hf_config(Qwen3VLMoeConfig)
class Qwen3VLMoeDecoderLayer(Qwen3MoeDecoderLayer):
def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
super().__init__(
vllm_config=vllm_config,
prefix=prefix,
is_fused_checkpoint_transposed=True,
)
@support_torch_compile(
dynamic_arg_dims={
"input_ids": 0,
@@ -77,6 +87,19 @@ class Qwen3VLMoeProcessingInfo(Qwen3VLProcessingInfo):
}
)
class Qwen3MoeLLMModel(Qwen3MoeModel):
def __init__(
self,
*,
vllm_config: VllmConfig,
prefix: str = "",
decoder_layer_type: type[torch.nn.Module] = Qwen3VLMoeDecoderLayer,
):
super().__init__(
vllm_config=vllm_config,
prefix=prefix,
decoder_layer_type=decoder_layer_type,
)
def forward(
self,
input_ids: torch.Tensor | None,