From c0b2d8f471699a0619ab9cc2e99719937eba503f Mon Sep 17 00:00:00 2001 From: Varshith Date: Mon, 22 Jun 2026 15:26:53 -0500 Subject: [PATCH] =?UTF-8?q?[Bugfix]=20FusedMoE:=20coerce=20shape-(1,)=20pe?= =?UTF-8?q?r-tensor=20scales=20to=200-D=20scalar=20=E2=80=A6=20(#43362)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Varshith Co-authored-by: Claude Co-authored-by: Michael Goin --- .../moe/test_moe_weight_loading_padded.py | 21 +++++++++++++++++++ .../layers/fused_moe/routed_experts.py | 16 ++++++++++---- 2 files changed, 33 insertions(+), 4 deletions(-) diff --git a/tests/kernels/moe/test_moe_weight_loading_padded.py b/tests/kernels/moe/test_moe_weight_loading_padded.py index d4939c79e5..2fd4e0fed5 100644 --- a/tests/kernels/moe/test_moe_weight_loading_padded.py +++ b/tests/kernels/moe/test_moe_weight_loading_padded.py @@ -325,3 +325,24 @@ class TestWeightLoadingWithPaddedHiddenSize: shard_id="w2", expert_id=0, ) + + +class TestPerTensorScaleCoercion: + """Regression test for shape-(1,) per-tensor scales (issue #43297). + + llm-compressor NVFP4 emits per-tensor weight and input scales as + shape-(1,) tensors. `_to_scalar` collapses them to a 0-D scalar so the + scalar-slot assignments in the weight loader neither broadcast nor raise. + """ + + def test_collapses_to_scalar(self): + # shape-(1,) and 0-D both reduce to a 0-D scalar. + for loaded_weight in (torch.tensor([0.5]), torch.tensor(0.5)): + scalar = RoutedExperts._to_scalar(loaded_weight) + assert scalar.shape == () + assert scalar.item() == pytest.approx(0.5) + + def test_rejects_non_scalar(self): + # numel > 1 must fail loudly instead of silently picking an element. + with pytest.raises(RuntimeError): + RoutedExperts._to_scalar(torch.tensor([0.1, 0.2])) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 669d1d3769..421cb396a6 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -275,6 +275,12 @@ class RoutedExperts(PluggableLayer): # Weight Loading Methods # + @staticmethod + def _to_scalar(loaded_weight: torch.Tensor) -> torch.Tensor: + # Per-tensor scales arrive 0-D or as shape-(1,) (llm-compressor NVFP4); + # reduce to a 0-D scalar. numel > 1 raises instead of broadcasting. + return loaded_weight.reshape(()) + def _load_per_tensor_weight_scale( self, shard_id: str, @@ -288,10 +294,10 @@ class RoutedExperts(PluggableLayer): # We have to keep the weight scales of w1 and w3 because # we need to re-quantize w1/w3 weights after weight loading. idx = 0 if shard_id == "w1" else 1 - param_data[expert_id][idx] = loaded_weight + param_data[expert_id][idx] = self._to_scalar(loaded_weight) # If we are in the row parallel case (down_proj) elif shard_id == "w2": - param_data[expert_id] = loaded_weight + param_data[expert_id] = self._to_scalar(loaded_weight) def _load_combined_w13_weight_scale( self, @@ -525,7 +531,7 @@ class RoutedExperts(PluggableLayer): param_data = param.data # Input scales can be loaded directly and should be equal. - param_data[expert_id] = loaded_weight + param_data[expert_id] = self._to_scalar(loaded_weight) def _load_g_idx( self, @@ -692,7 +698,9 @@ class RoutedExperts(PluggableLayer): ): scale_expert_id = global_expert_id if use_global_sf else expert_id scale_shard_id = 0 if shard_id == "w1" else 1 - param.data[scale_expert_id][scale_shard_id] = loaded_weight.reshape(()) + param.data[scale_expert_id][scale_shard_id] = self._to_scalar( + loaded_weight + ) return True if return_success else None if (