diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index aa283419bf..d3716a196b 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -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: diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index cfbadba7ac..28b994057f 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -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 diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 718ab9126c..1eb96d5594 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -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 diff --git a/vllm/model_executor/models/qwen3_moe.py b/vllm/model_executor/models/qwen3_moe.py index 41698bfd4d..c33166469a 100644 --- a/vllm/model_executor/models/qwen3_moe.py +++ b/vllm/model_executor/models/qwen3_moe.py @@ -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( diff --git a/vllm/model_executor/models/qwen3_vl_moe.py b/vllm/model_executor/models/qwen3_vl_moe.py index f4a36f2897..1360edfdd7 100644 --- a/vllm/model_executor/models/qwen3_vl_moe.py +++ b/vllm/model_executor/models/qwen3_vl_moe.py @@ -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,