[Bugfix] FusedMoE: coerce shape-(1,) per-tensor scales to 0-D scalar … (#43362)
Signed-off-by: Varshith <kvarshithgowda@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Varshith committed
Jun 22, 2026 at 15:26 UTC
c0b2d8f471699a0619ab9cc2e99719937eba503f
2 files changed
+33
-4
tests/kernels/moe/test_moe_weight_loading_padded.py
+21
@@ -325,3 +325,24 @@ class TestWeightLoadingWithPaddedHiddenSize:
325
shard_id="w2",
326
expert_id=0,
327
)
328
+
329
+
330
+class TestPerTensorScaleCoercion:
331
+ """Regression test for shape-(1,) per-tensor scales (issue #43297).
332
+
333
+ llm-compressor NVFP4 emits per-tensor weight and input scales as
334
+ shape-(1,) tensors. `_to_scalar` collapses them to a 0-D scalar so the
335
+ scalar-slot assignments in the weight loader neither broadcast nor raise.
336
+ """
337
+
338
+ def test_collapses_to_scalar(self):
339
+ # shape-(1,) and 0-D both reduce to a 0-D scalar.
340
+ for loaded_weight in (torch.tensor([0.5]), torch.tensor(0.5)):
341
+ scalar = RoutedExperts._to_scalar(loaded_weight)
342
+ assert scalar.shape == ()
343
+ assert scalar.item() == pytest.approx(0.5)
344
+
345
+ def test_rejects_non_scalar(self):
346
+ # numel > 1 must fail loudly instead of silently picking an element.
347
+ with pytest.raises(RuntimeError):
348
+ RoutedExperts._to_scalar(torch.tensor([0.1, 0.2]))
vllm/model_executor/layers/fused_moe/routed_experts.py
+12
-4
@@ -275,6 +275,12 @@ class RoutedExperts(PluggableLayer):
275
# Weight Loading Methods
276
#
277
278
+ @staticmethod
279
+ def _to_scalar(loaded_weight: torch.Tensor) -> torch.Tensor:
280
+ # Per-tensor scales arrive 0-D or as shape-(1,) (llm-compressor NVFP4);
281
+ # reduce to a 0-D scalar. numel > 1 raises instead of broadcasting.
282
+ return loaded_weight.reshape(())
283
+
284
def _load_per_tensor_weight_scale(
285
self,
286
shard_id: str,
@@ -288,10 +294,10 @@ class RoutedExperts(PluggableLayer):
294
# We have to keep the weight scales of w1 and w3 because
295
# we need to re-quantize w1/w3 weights after weight loading.
296
idx = 0 if shard_id == "w1" else 1
291
- param_data[expert_id][idx] = loaded_weight
297
+ param_data[expert_id][idx] = self._to_scalar(loaded_weight)
298
# If we are in the row parallel case (down_proj)
299
elif shard_id == "w2":
294
- param_data[expert_id] = loaded_weight
300
+ param_data[expert_id] = self._to_scalar(loaded_weight)
301
302
def _load_combined_w13_weight_scale(
303
self,
@@ -525,7 +531,7 @@ class RoutedExperts(PluggableLayer):
531
param_data = param.data
532
533
# Input scales can be loaded directly and should be equal.
528
- param_data[expert_id] = loaded_weight
534
+ param_data[expert_id] = self._to_scalar(loaded_weight)
535
536
def _load_g_idx(
537
self,
@@ -692,7 +698,9 @@ class RoutedExperts(PluggableLayer):
698
):
699
scale_expert_id = global_expert_id if use_global_sf else expert_id
700
scale_shard_id = 0 if shard_id == "w1" else 1
695
- param.data[scale_expert_id][scale_shard_id] = loaded_weight.reshape(())
701
+ param.data[scale_expert_id][scale_shard_id] = self._to_scalar(
702
+ loaded_weight
703
+ )
704
return True if return_success else None
705
706
if (