A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Fix unstable silu_mul+nvfp4 quant fusion test (#24370)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
elvischenv committed
Sep 7, 2025 at 04:39 UTC
e68dc2f014369771de4892a008d994c33317834f
2 files changed
+38
-16
tests/compile/test_silu_mul_quant_fusion.py
+30
-16
@@ -1,9 +1,12 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import cast
4
+
5
import pytest
6
import torch
7
8
import vllm.envs as envs
9
+from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor
10
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
11
# yapf conflicts with isort for this block
12
# yapf: disable
@@ -64,24 +67,27 @@ class TestSiluMulFp8QuantModel(torch.nn.Module):
67
68
class TestSiluMulNvfp4QuantModel(torch.nn.Module):
69
67
- def __init__(self, hidden_size: int, **kwargs):
70
+ def __init__(self, hidden_size: int, x: torch.Tensor, **kwargs):
71
super().__init__()
72
self.silu_and_mul = SiluAndMul()
70
- self.w = torch.randint(256, (hidden_size, hidden_size // 2),
71
- dtype=FP4_DTYPE)
72
- self.wscale = torch.randn(hidden_size,
73
- hidden_size // 16).to(dtype=FP8_DTYPE)
74
- self.wscale2 = torch.rand(1, dtype=torch.float32)
75
- self.scale = torch.rand(1, dtype=torch.float32)
73
+
74
+ # create nvfp4 weight
75
+ w = torch.rand((hidden_size, hidden_size))
76
+ self.w, self.w_block_scale, self.w_global_scale = quant_nvfp4_tensor(w)
77
+
78
+ # get global scale offline
79
+ _, _, self.y_global_scale = quant_nvfp4_tensor(self.silu_and_mul(x))
80
+
81
+ self.alpha = 1.0 / (self.w_global_scale * self.y_global_scale)
82
83
def forward(self, x):
84
y = self.silu_and_mul(x)
79
- y_quant, y_block_scale = scaled_fp4_quant(y, 1 / self.scale)
85
+ y_quant, y_block_scale = scaled_fp4_quant(y, self.y_global_scale)
86
out = cutlass_scaled_fp4_mm(a=y_quant,
87
b=self.w,
88
block_scale_a=y_block_scale,
83
- block_scale_b=self.wscale,
84
- alpha=self.scale * self.wscale2,
89
+ block_scale_b=self.w_block_scale,
90
+ alpha=self.alpha,
91
out_dtype=y.dtype)
92
return out
93
@@ -95,8 +101,9 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
101
@pytest.mark.parametrize("num_tokens", [64])
102
@pytest.mark.parametrize("hidden_size", [128])
103
@pytest.mark.parametrize(
98
- "model_class", [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
99
- if is_nvfp4_supported() else [TestSiluMulFp8QuantModel])
104
+ "model_class",
105
+ cast(list[type], [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
106
+ if is_nvfp4_supported() else [TestSiluMulFp8QuantModel]))
107
# cuda_force_torch used to test torch code path on platforms that
108
# cutlass_fp8_supported() == True.
109
@pytest.mark.parametrize("cuda_force_torch",
@@ -111,6 +118,8 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, model_class,
118
torch.set_default_device("cuda")
119
torch.set_default_dtype(torch.float16)
120
121
+ x = torch.rand(num_tokens, hidden_size * 2)
122
+
123
# Reshape pass is needed for the fusion pass to work
124
config = VllmConfig()
125
config.compilation_config = CompilationConfig(
@@ -119,10 +128,10 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, model_class,
128
129
backend = TestBackend(NoOpEliminationPass(config), fusion_pass)
130
model = model_class(hidden_size=hidden_size,
122
- cuda_force_torch=cuda_force_torch)
131
+ cuda_force_torch=cuda_force_torch,
132
+ x=x)
133
134
# First dimension dynamic
125
- x = torch.rand(num_tokens, hidden_size * 2)
135
torch._dynamo.mark_dynamic(x, 0)
136
137
result = model(x)
@@ -131,10 +140,15 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, model_class,
140
result2 = model2(x)
141
142
# Check that it gives the same answer
143
+ if model_class == TestSiluMulFp8QuantModel:
144
+ atol, rtol = 1e-3, 1e-3
145
+ elif model_class == TestSiluMulNvfp4QuantModel:
146
+ atol, rtol = 1e-1, 1e-1
147
+
148
torch.testing.assert_close(result[0].to(dtype=torch.float16),
149
result2[0].to(dtype=torch.float16),
136
- atol=1e-3,
137
- rtol=1e-3)
150
+ atol=atol,
151
+ rtol=rtol)
152
153
# In pre-nodes, quant op should be present and fused kernels should not
154
backend.check_before_ops(model.ops_in_model_before())
tests/kernels/quantization/nvfp4_utils.py
+8
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import torch
4
5
+from vllm._custom_ops import scaled_fp4_quant
6
from vllm.scalar_type import scalar_types
7
8
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
@@ -65,3 +66,10 @@ def break_fp4_bytes(a, dtype):
66
67
# Reshape to final form
68
return values.reshape(m, n * 2).to(dtype=dtype)
69
+
70
+
71
+def quant_nvfp4_tensor(a: torch.Tensor):
72
+ a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
73
+ torch.abs(a).max().to(torch.float32))
74
+ a_quant, a_block_scale = scaled_fp4_quant(a, a_global_scale)
75
+ return a_quant, a_block_scale, a_global_scale