[Bugfix] Make shared NVFP4 MoE scales writable (#49489)
Signed-off-by: S1ro1 <matej.sirovatka@gmail.com>
Matej Sirovatka committed
Jul 22, 2026 at 19:17 UTC
b07ec92faa2d534ed97fe97f44ca71c816404849
2 files changed
+46
-4
tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py
+42
@@ -1,13 +1,55 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
import torch
7
8
+from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend
9
+from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe
10
+from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
11
+ prepare_nvfp4_moe_layer_for_fi_or_cutlass,
12
+)
13
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
14
align_trtllm_fp4_moe_hidden_dim_for_fi,
15
)
16
17
18
+def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch):
19
+ monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x)
20
+
21
+ num_experts = 3
22
+ layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False))
23
+ w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
24
+ w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8)
25
+ w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
26
+ w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn)
27
+ weight_scale = torch.ones(num_experts)
28
+
29
+ outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass(
30
+ backend=NvFp4MoeBackend.FLASHINFER_CUTLASS,
31
+ layer=layer,
32
+ w13=w13,
33
+ w13_scale=w13_scale,
34
+ w13_scale_2=weight_scale,
35
+ a13_scale=torch.tensor([1.0, 2.0, 3.0]),
36
+ w2=w2,
37
+ w2_scale=w2_scale,
38
+ w2_scale_2=weight_scale,
39
+ a2_scale=torch.tensor([4.0, 5.0, 6.0]),
40
+ is_act_and_mul=False,
41
+ )
42
+ a13_scale, a2_scale = outputs[3], outputs[7]
43
+
44
+ torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0))
45
+ torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0))
46
+ distinct_values = torch.arange(num_experts, dtype=torch.float32)
47
+ a13_scale.copy_(distinct_values)
48
+ a2_scale.copy_(distinct_values)
49
+ torch.testing.assert_close(a13_scale, distinct_values)
50
+ torch.testing.assert_close(a2_scale, distinct_values)
51
+
52
+
53
def test_align_trtllm_fp4_moe_hidden_dim_noop():
54
w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256)
55
w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+4
-4
@@ -109,8 +109,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl(
109
110
# Global scaling factors (same as other FlashInfer backends).
111
num_experts = w13.shape[0]
112
- a13_scale = a13_scale.max().to(torch.float32).expand(num_experts)
113
- a2_scale = a2_scale.max().to(torch.float32).expand(num_experts)
112
+ a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
113
+ a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
114
115
half = w13.shape[1] // 2
116
w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1)
@@ -338,8 +338,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass(
338
# For some FI kernels, the input scales are shared by all experts.
339
if is_global_sf_supported_for_nvfp4_backend(backend):
340
num_experts = w13.shape[0]
341
- a13_scale = a13_scale.max().to(torch.float32).expand(num_experts)
342
- a2_scale = a2_scale.max().to(torch.float32).expand(num_experts)
341
+ a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts)
342
+ a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts)
343
else:
344
a13_scale = a13_scale.max(dim=1).values.to(torch.float32)
345