[Bugfix] Fix NVFP4 fused SiLU+mul scale allocation and global scale direction (#55643)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: khluu <khluu000@gmail.com> Co-authored-by: jiahao <jxia77@terpmail.umd.edu> Co-authored-by: Kimi Code <noreply@moonshot.cn>
Jiangyun Zhu committed
Sep 8, 2026 at 23:26 UTC
8ebc5b0a182b3351a1c62fe73e97be2c489ed2a6
5 files changed
+67
-18
.buildkite/test_areas/compile.yaml
+3
@@ -46,6 +46,7 @@ steps:
46
source_file_dependencies:
47
- csrc/quantization/fp4/
48
- vllm/model_executor/layers/quantization/
49
+ - vllm/model_executor/layers/fusion/
50
- vllm/model_executor/layers/layernorm.py
51
- vllm/model_executor/layers/activation.py
52
- vllm/model_executor/layers/attention/attention.py
@@ -54,6 +55,7 @@ steps:
55
- tests/compile/passes/test_fusion_attn.py
56
- tests/compile/passes/test_mla_attn_quant_fusion.py
57
- tests/compile/passes/test_silu_mul_quant_fusion.py
58
+ - tests/compile/passes/test_silu_mul_quant_manual_fusion.py
59
- tests/compile/passes/distributed/test_fusion_all_reduce.py
60
- tests/compile/fullgraph/test_full_graph.py
61
commands:
@@ -62,6 +64,7 @@ steps:
64
- pytest -v -s tests/compile/passes/test_fusion_attn.py -k FLASHINFER
65
- pytest -v -s tests/compile/passes/test_mla_attn_quant_fusion.py
66
- pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py
67
+ - pytest -v -s tests/compile/passes/test_silu_mul_quant_manual_fusion.py
68
# this runner has 2 GPUs available even though num_devices=2 is not set
69
- pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
70
.buildkite/test_areas/kernels.yaml
+4
@@ -312,8 +312,10 @@ steps:
312
# optional: true
313
source_file_dependencies:
314
- csrc/quantization/fp4/
315
+ - csrc/libtorch_stable/quantization/fp4/
316
- csrc/attention/mla/
317
- csrc/quantization/cutlass_w8a8/moe/
318
+ - vllm/model_executor/layers/fusion/
319
- vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
320
- vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
321
- vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
@@ -349,6 +351,8 @@ steps:
351
- tests/kernels/attention/test_minimax_m3.py
352
- tests/kernels/test_ll_bf16_gemm.py
353
- tests/kernels/test_top_k_per_row.py
354
+ - tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
355
+ - tests/models/quantization/test_nvfp4.py
356
commands:
357
- nvidia-smi
358
- |
tests/compile/passes/test_silu_mul_quant_manual_fusion.py
+39
-8
@@ -47,13 +47,21 @@ from vllm.platforms import current_platform
47
class MockLinearForFusion(torch.nn.Module):
48
"""Mock linear layer that exposes input_quant_key for fusion testing."""
49
50
- def __init__(self, quant_key, input_scale=None, input_global_scale=None):
50
+ def __init__(
51
+ self,
52
+ quant_key,
53
+ input_scale=None,
54
+ input_global_scale=None,
55
+ input_global_scale_inv=None,
56
+ ):
57
super().__init__()
58
self.input_quant_key = quant_key
59
if input_scale is not None:
60
self.input_scale = input_scale
61
if input_global_scale is not None:
62
self.input_global_scale = input_global_scale
63
+ if input_global_scale_inv is not None:
64
+ self.input_global_scale_inv = input_global_scale_inv
65
66
67
ROCM_KERNELS = [ROCmFP8ScaledMMLinearKernel, PerTensorTorchFP8ScaledMMLinearKernel]
@@ -206,6 +214,10 @@ def test_manual_fusion_fp8_dynamic_128(dtype: torch.dtype):
214
215
216
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
217
+# Cover M < 128 (padded swizzled SF tiles) and hidden_size that makes the
218
+# kernel launch grid.y > 1 (num_packed_cols = hidden_size / 16 > 512).
219
+@pytest.mark.parametrize("num_tokens", [1, 127, 128])
220
+@pytest.mark.parametrize("hidden_size", [256, 14336])
221
@pytest.mark.skipif(not current_platform.is_cuda(), reason="NVFP4 CUDA only")
222
@pytest.mark.skipif(
223
not current_platform.has_device_capability(100), reason="NVFP4 requires SM100+"
@@ -213,7 +225,9 @@ def test_manual_fusion_fp8_dynamic_128(dtype: torch.dtype):
225
@pytest.mark.skipif(
226
envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"], reason="Only test on CUDA and ROCm"
227
)
216
-def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype):
228
+def test_manual_fusion_nvfp4_dynamic(
229
+ dtype: torch.dtype, num_tokens: int, hidden_size: int
230
+):
231
"""Test kNvfp4Dynamic fusion path.
232
233
Compares fused (silu_and_mul_nvfp4_quant) vs unfused (silu_and_mul)
@@ -228,10 +242,12 @@ def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype):
242
torch.set_default_dtype(dtype)
243
244
# NVFP4 requires hidden_size divisible by 16 (block size) and by 2 (packing).
231
- # M (num_tokens) must be >= 128 for the 128x4 swizzled scale layout.
232
- num_tokens, hidden_size = 128, 256
245
x = torch.rand(num_tokens, hidden_size * 2)
234
- input_global_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda")
246
+ # Non-1.0 global scale so the test is sensitive to the scale direction:
247
+ # the fused producer must quantize with input_global_scale_inv (the GEMM's
248
+ # alpha divides by input_global_scale), matching the unfused path.
249
+ input_global_scale = torch.tensor([0.5], dtype=torch.float32, device="cuda")
250
+ input_global_scale_inv = 1.0 / input_global_scale
251
252
config = VllmConfig(
253
compilation_config=CompilationConfig(custom_ops=["none"]),
@@ -247,7 +263,9 @@ def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype):
263
264
# Fused path: apply silu_and_mul + NVFP4 quantization
265
mock_linear_with_key = MockLinearForFusion(
250
- kNvfp4Dynamic, input_global_scale=input_global_scale
266
+ kNvfp4Dynamic,
267
+ input_global_scale=input_global_scale,
268
+ input_global_scale_inv=input_global_scale_inv,
269
)
270
result_fused = maybe_fused_act_quant(silu_and_mul, x, mock_linear_with_key)
271
@@ -260,11 +278,24 @@ def test_manual_fusion_nvfp4_dynamic(dtype: torch.dtype):
278
assert result_fused.orig_dtype == dtype
279
assert result_fused.orig_shape == (num_tokens, hidden_size)
280
263
- # Dequantize fused result and compare with unfused
281
+ # The scale tensor must use the padded 128x4-tile swizzled layout,
282
+ # same as scaled_fp4_quant's allocator; the kernel writes a full
283
+ # 128-row tile even when M is smaller.
284
+ from vllm.utils.math_utils import round_up
285
+
286
+ assert result_fused.scale.dtype == current_platform.fp8_dtype()
287
+ assert result_fused.scale.shape == (
288
+ round_up(num_tokens, 128),
289
+ round_up(hidden_size // 16, 4),
290
+ )
291
+
292
+ # Dequantize fused result and compare with unfused. The block scales
293
+ # must have been folded with input_global_scale_inv, so dequantize
294
+ # with the same inverse scale to recover the activation.
295
dequant_result = dequantize_nvfp4_to_dtype(
296
tensor_fp4=result_fused.data,
297
tensor_sf=result_fused.scale,
267
- global_scale=input_global_scale,
298
+ global_scale=input_global_scale_inv,
299
dtype=dtype,
300
device="cuda",
301
block_size=16,
tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
+1
-1
@@ -23,7 +23,7 @@ FP4_DTYPE = torch.uint8
23
FP8_DTYPE = current_platform.fp8_dtype()
24
25
DTYPES = [torch.float16, torch.bfloat16]
26
-SHAPES = [(128, 256), (128, 128), (256, 256), (256, 128)]
26
+SHAPES = [(128, 256), (128, 128), (256, 256), (256, 128), (1, 256), (64, 256)]
27
BLOCK_SIZE = 16
28
29
vllm/model_executor/layers/fusion/fused_act_quant.py
+20
-9
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
28
kNvfp4Dynamic,
29
)
30
from vllm.platforms import current_platform
31
+from vllm.utils.math_utils import round_up
32
33
FP8_DTYPE = current_platform.fp8_dtype()
34
FP4_DTYPE = torch.uint8
@@ -105,20 +106,30 @@ def _silu_and_mul_nvfp4_dynamic(
106
# NVFP4 packs 2 values into 1 byte
107
result = torch.empty((num_tokens, d // 2), dtype=FP4_DTYPE, device=x.device)
108
108
- # Block scale output shape: swizzled layout for tensor cores
109
- # Each group of 16 elements shares one FP8 scale
110
- num_k_tiles = (d + 63) // 64
109
+ # Block scale output: swizzled tensor-core layout
110
+ # [num_m_tiles, num_k_tiles, 32, 4, 4] of int32-packed FP8 scales, so the
111
+ # row/col extents must be padded to the 128x4 tile like
112
+ # scaled_fp4_quant's allocator (create_fp4_scale_tensor). Each group of 16
113
+ # elements shares one FP8 scale.
114
+ rounded_m = round_up(num_tokens, 128)
115
+ rounded_n = round_up(d // 16, 4)
116
block_scale = torch.empty(
112
- (num_tokens, num_k_tiles * 4), dtype=FP8_DTYPE, device=x.device
117
+ (rounded_m, rounded_n // 4), dtype=torch.int32, device=x.device
118
+ ).view(FP8_DTYPE)
119
+
120
+ # The kernel folds the global scale into the block scales, and the
121
+ # consumer GEMM's alpha (= input_global_scale * weight_global_scale)
122
+ # divides it back out, so quantize with the reciprocal like the unfused
123
+ # scaled_fp4_quant path does.
124
+ input_global_scale_inv = getattr(linear, "input_global_scale_inv", None)
125
+ assert input_global_scale_inv is not None, (
126
+ "input_global_scale_inv is required for NVFP4 quantization"
127
)
128
115
- input_global_scale = getattr(linear, "input_global_scale", None)
116
- assert input_global_scale is not None, (
117
- "input_global_scale is required for NVFP4 quantization"
129
+ torch.ops._C.silu_and_mul_nvfp4_quant(
130
+ result, block_scale, x, input_global_scale_inv
131
)
132
120
- torch.ops._C.silu_and_mul_nvfp4_quant(result, block_scale, x, input_global_scale)
121
-
133
return QuantizedActivation(
134
data=result.view(out_shape[:-1] + (d // 2,)),
135
scale=block_scale,