[Bugfix] Make shared NVFP4 MoE scales writable (#49489)

Signed-off-by: S1ro1 <[email protected]>

tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py

@@ -1,13 +1,55 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import torch
 
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
+from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
+    prepare_nvfp4_moe_layer_for_fi_or_cutlass,
+)
 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
     align_trtllm_fp4_moe_hidden_dim_for_fi,
 )
 
 
+def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
+    monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
+
+    num_experts = 3
+    layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
+    w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
+    w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
+    w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
+    w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
+    weight_scale = torch.ones(num_experts)
+
+    outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(
+        backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,
+        layer=layer,
+        w13=w13,
+        w13_scale=w13_scale,
+        w13_scale_2=weight_scale,
+        a13_scale=torch.tensor([1.0, 2.0, 3.0]),
+        w2=w2,
+        w2_scale=w2_scale,
+        w2_scale_2=weight_scale,
+        a2_scale=torch.tensor([4.0, 5.0, 6.0]),
+        is_act_and_mul=False,
+    )
+    a13_scale, a2_scale = outputs[3], outputs[7]
+
+    torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0))
+    torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0))
+    distinct_values = torch.arange(num_experts, dtype=torch.float32)
+    a13_scale.copy_(distinct_values)
+    a2_scale.copy_(distinct_values)
+    torch.testing.assert_close(a13_scale, distinct_values)
+    torch.testing.assert_close(a2_scale, distinct_values)
+
+
 def test_align_trtllm_fp4_moe_hidden_dim_noop():
     w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
     w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

@@ -109,8 +109,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
 
     # Global scaling factors (same as other FlashInfer backends).
     num_experts = w13.shape[0]
-    a13_scale = a13_scale.max().to(torch.float32).expand(num_experts)
-    a2_scale = a2_scale.max().to(torch.float32).expand(num_experts)
+    a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
+    a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
 
     half = w13.shape[1] // 2
     w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1)
@@ -338,8 +338,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
     # For some FI kernels, the input scales are shared by all experts.
     if is_global_sf_supported_for_nvfp4_backend(backend):
         num_experts = w13.shape[0]
-        a13_scale = a13_scale.max().to(torch.float32).expand(num_experts)
-        a2_scale = a2_scale.max().to(torch.float32).expand(num_experts)
+        a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
+        a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
     else:
         a13_scale = a13_scale.max(dim=1).values.to(torch.float32)