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 accuracy issue for silu_mul + nvfp4 quant fusion kernel (#24833)
Signed-off-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
elvischenv committed
Sep 18, 2025 at 07:37 UTC
e6585ddb451ba6056e044184f7fc88dcc13f8cfe
5 files changed
+110
-226
.buildkite/test-pipeline.yaml
+1
-1
@@ -796,7 +796,7 @@ steps:
796
# Quantization
797
- pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'
798
- pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py
799
- - pytest -v -s tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
799
+ - pytest -v -s tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
800
- pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py
801
- pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py
802
- pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
+27
-93
@@ -30,109 +30,41 @@
30
31
namespace vllm {
32
33
-template <class Type>
34
-__inline__ __device__ PackedVec<Type> compute_silu(PackedVec<Type>& vec,
35
- PackedVec<Type>& vec2) {
36
- PackedVec<Type> result;
37
-#pragma unroll
38
- for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
39
- if constexpr (std::is_same_v<Type, half>) {
40
- half2 val(0.5f, 0.5f);
41
- half2 t0 = __hmul2(vec.elts[i], val);
42
- half2 t1 = __hfma2(h2tanh(t0), val, val);
43
- half2 t2 = __hmul2(vec.elts[i], t1);
44
- result.elts[i] = __hmul2(t2, vec2.elts[i]);
45
- } else {
46
- __nv_bfloat162 val(0.5f, 0.5f);
47
- __nv_bfloat162 t0 = __hmul2(vec.elts[i], val);
48
- __nv_bfloat162 t1 = __hfma2(h2tanh(t0), val, val);
49
- __nv_bfloat162 t2 = __hmul2(vec.elts[i], t1);
50
- result.elts[i] = __hmul2(t2, vec2.elts[i]);
51
- }
52
- }
53
- return result;
33
+// silu in float32
34
+__device__ __forceinline__ float silu(float x) {
35
+ return __fdividef(x, (1.f + __expf(-x)));
36
}
37
56
-// Quantizes the provided PackedVec into the uint32_t output
57
-template <class Type, bool UE8M0_SF = false>
58
-__device__ uint32_t silu_and_cvt_warp_fp16_to_fp4(PackedVec<Type>& vec,
59
- PackedVec<Type>& vec2,
60
- float SFScaleVal,
61
- uint8_t* SFout) {
62
- PackedVec<Type> out_silu = compute_silu(vec, vec2);
63
- // Get absolute maximum values among the local 8 values.
64
- auto localMax = __habs2(out_silu.elts[0]);
65
-
66
-// Local maximum value.
67
-#pragma unroll
68
- for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
69
- localMax = __hmax2(localMax, __habs2(out_silu.elts[i]));
70
- }
71
-
72
- // Get the absolute maximum among all 16 values (two threads).
73
- localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
74
- // Get the final absolute maximum values.
75
- float vecMax = float(__hmax(localMax.x, localMax.y));
76
-
77
- // Get the SF (max value of the vector / max value of e2m1).
78
- // maximum value of e2m1 = 6.0.
79
- // TODO: use half as compute data type.
80
- float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
81
- // 8 bits representation of the SF.
82
- uint8_t fp8SFVal;
83
- // Write the SF to global memory (STG.8).
84
- if constexpr (UE8M0_SF) {
85
- // Extract the 8 exponent bits from float32.
86
- // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
87
- uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
88
- fp8SFVal = tmp & 0xff;
89
- // Convert back to fp32.
90
- reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
91
- } else {
92
- // Here SFValue is always positive, so E4M3 is the same as UE4M3.
93
- __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
94
- reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
95
- // Convert back to fp32.
96
- SFValue = float(tmp);
97
- }
98
- // Get the output scale.
99
- // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
100
- // reciprocal(SFScaleVal))
101
- float outputScale =
102
- SFValue != 0 ? reciprocal_approximate_ftz(
103
- SFValue * reciprocal_approximate_ftz(SFScaleVal))
104
- : 0.0f;
105
-
106
- if (SFout) {
107
- // Write the SF to global memory (STG.8).
108
- *SFout = fp8SFVal;
109
- }
38
+__device__ __forceinline__ float2 silu2(float2 x) {
39
+ return make_float2(silu(x.x), silu(x.y));
40
+}
41
111
- // Convert the input to float.
112
- float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
42
+template <class Type>
43
+__inline__ __device__ PackedVec<Type> compute_silu_mul(PackedVec<Type>& vec,
44
+ PackedVec<Type>& vec2) {
45
+ PackedVec<Type> result;
46
+ using packed_type = typename TypeConverter<Type>::Type;
47
48
#pragma unroll
115
- for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
49
+ for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
50
+ // silu_mul in float32
51
if constexpr (std::is_same_v<Type, half>) {
117
- fp2Vals[i] = __half22float2(out_silu.elts[i]);
52
+ float2 silu_vec = silu2(__half22float2(vec.elts[i]));
53
+ result.elts[i] =
54
+ __float22half2_rn(__fmul2_rn(silu_vec, __half22float2(vec2.elts[i])));
55
} else {
119
- fp2Vals[i] = __bfloat1622float2(out_silu.elts[i]);
56
+ float2 silu_vec = silu2(__bfloat1622float2(vec.elts[i]));
57
+ result.elts[i] = __float22bfloat162_rn(
58
+ __fmul2_rn(silu_vec, __bfloat1622float2(vec2.elts[i])));
59
}
121
- fp2Vals[i].x *= outputScale;
122
- fp2Vals[i].y *= outputScale;
60
}
124
-
125
- // Convert to e2m1 values.
126
- uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
127
-
128
- // Write the e2m1 values to global memory.
129
- return e2m1Vec;
61
+ return result;
62
}
63
64
// Use UE4M3 by default.
65
template <class Type, bool UE8M0_SF = false>
66
__global__ void __launch_bounds__(1024, 4)
135
- silu_and_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
67
+ silu_mul_cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
68
float const* SFScale, uint32_t* out,
69
uint32_t* SFout) {
70
using PackedVec = PackedVec<Type>;
@@ -160,16 +92,18 @@ __global__ void __launch_bounds__(1024, 4)
92
// Get the output tensor offset.
93
// Same as inOffset because 8 elements are packed into one uint32_t.
94
int64_t outOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
163
- ;
95
auto& out_pos = out[outOffset];
96
97
+ // Compute silu and mul
98
+ PackedVec out_silu_mul = compute_silu_mul(in_vec, in_vec2);
99
+
100
auto sf_out =
101
cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
102
CVT_FP4_NUM_THREADS_PER_SF>(
103
rowIdx, colIdx, numCols, SFout);
104
171
- out_pos = silu_and_cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(
172
- in_vec, in_vec2, SFScaleVal, sf_out);
105
+ out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(out_silu_mul, SFScaleVal,
106
+ sf_out);
107
}
108
}
109
}
@@ -204,7 +138,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d]
138
input.scalar_type(), "silu_and_mul_nvfp4_quant_kernel", [&] {
139
using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
140
auto input_ptr = static_cast<cuda_type const*>(input.data_ptr());
207
- vllm::silu_and_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
141
+ vllm::silu_mul_cvt_fp16_to_fp4<cuda_type><<<grid, block, 0, stream>>>(
142
m, n, input_ptr, input_sf_ptr,
143
reinterpret_cast<uint32_t*>(output_ptr),
144
reinterpret_cast<uint32_t*>(sf_out));
tests/compile/test_silu_mul_quant_fusion.py
+7
-6
@@ -98,8 +98,9 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
98
return [FUSED_OPS[kNvfp4Quant]]
99
100
101
-@pytest.mark.parametrize("num_tokens", [64])
102
-@pytest.mark.parametrize("hidden_size", [128])
101
+@pytest.mark.parametrize("num_tokens", [32, 64])
102
+@pytest.mark.parametrize("hidden_size", [128, 256])
103
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
104
@pytest.mark.parametrize(
105
"model_class",
106
cast(list[type], [TestSiluMulFp8QuantModel, TestSiluMulNvfp4QuantModel]
@@ -110,13 +111,13 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
111
[True, False] if cutlass_fp8_supported() else [True])
112
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
113
reason="Only test on CUDA and ROCm")
113
-def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, model_class,
114
+def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, dtype, model_class,
115
cuda_force_torch):
116
if model_class == TestSiluMulNvfp4QuantModel and cuda_force_torch:
117
pytest.skip("Duplicate tests for NVFP4")
118
119
torch.set_default_device("cuda")
119
- torch.set_default_dtype(torch.float16)
120
+ torch.set_default_dtype(dtype)
121
122
x = torch.rand(num_tokens, hidden_size * 2)
123
@@ -145,8 +146,8 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size, model_class,
146
elif model_class == TestSiluMulNvfp4QuantModel:
147
atol, rtol = 1e-1, 1e-1
148
148
- torch.testing.assert_close(result[0].to(dtype=torch.float16),
149
- result2[0].to(dtype=torch.float16),
149
+ torch.testing.assert_close(result[0].to(dtype=dtype),
150
+ result2[0].to(dtype=dtype),
151
atol=atol,
152
rtol=rtol)
153
tests/kernels/quantization/test_silu_mul_nvfp4_quant.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
7
+ FLOAT8_E4M3_MAX,
8
+ dequantize_nvfp4_to_dtype)
9
+from vllm._custom_ops import scaled_fp4_quant
10
+from vllm.model_executor.layers.activation import SiluAndMul
11
+from vllm.platforms import current_platform
12
+
13
+if not current_platform.has_device_capability(100):
14
+ pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
15
+ allow_module_level=True)
16
+
17
+FP4_DTYPE = torch.uint8
18
+FP8_DTYPE = current_platform.fp8_dtype()
19
+
20
+DTYPES = [torch.float16, torch.bfloat16]
21
+SHAPES = [(128, 256), (128, 128), (256, 256), (256, 128)]
22
+BLOCK_SIZE = 16
23
+
24
+
25
+@pytest.mark.parametrize("dtype", DTYPES)
26
+@pytest.mark.parametrize("shape", SHAPES)
27
+@torch.inference_mode()
28
+def test_silu_mul_nvfp4_quant(
29
+ dtype: torch.dtype,
30
+ shape: tuple[int, int],
31
+) -> None:
32
+ current_platform.seed_everything(42)
33
+ device = 'cuda:0'
34
+ torch.set_default_device(device)
35
+
36
+ x = torch.randn(shape, dtype=dtype)
37
+
38
+ # ref op
39
+ ref_output = SiluAndMul().forward_native(x)
40
+ ref_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
41
+ torch.abs(ref_output).max().to(torch.float32))
42
+ ref_output_quant, ref_block_scale = scaled_fp4_quant(
43
+ ref_output, ref_global_scale)
44
+
45
+ # fused op
46
+ fused_output_quant = torch.empty_like(ref_output_quant)
47
+ fused_block_scale = torch.empty_like(ref_block_scale)
48
+ torch.ops._C.silu_and_mul_nvfp4_quant(fused_output_quant,
49
+ fused_block_scale, x,
50
+ ref_global_scale)
51
+
52
+ # check dtype
53
+ assert ref_output_quant.dtype == FP4_DTYPE
54
+ assert fused_output_quant.dtype == FP4_DTYPE
55
+ assert ref_output_quant.shape == fused_output_quant.shape
56
+
57
+ assert ref_block_scale.dtype == FP8_DTYPE
58
+ assert fused_block_scale.dtype == FP8_DTYPE
59
+ assert ref_block_scale.shape == fused_block_scale.shape
60
+
61
+ # check dequantized output
62
+ ref_output_dequant = dequantize_nvfp4_to_dtype(ref_output_quant,
63
+ ref_block_scale,
64
+ ref_global_scale, dtype,
65
+ device)
66
+ fused_output_dequant = dequantize_nvfp4_to_dtype(fused_output_quant,
67
+ fused_block_scale,
68
+ ref_global_scale, dtype,
69
+ device)
70
+
71
+ atol, rtol = 3e-1, 3e-1
72
+ torch.testing.assert_close(ref_output_dequant,
73
+ fused_output_dequant,
74
+ atol=atol,
75
+ rtol=rtol)
tests/kernels/quantization/test_silu_nvfp4_quant_fusion.py
deleted
-126
@@ -1,126 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import pytest
4
-import torch
5
-
6
-from tests.kernels.utils import opcheck
7
-from vllm.model_executor.layers.activation import SiluAndMul
8
-from vllm.platforms import current_platform
9
-from vllm.scalar_type import scalar_types
10
-
11
-if not current_platform.has_device_capability(100):
12
- pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
13
- allow_module_level=True)
14
-
15
-DTYPES = [torch.float16, torch.bfloat16]
16
-SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)]
17
-SEEDS = [42]
18
-CUDA_DEVICES = ['cuda:0']
19
-
20
-FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
21
-FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
22
-
23
-BLOCK_SIZE = 16
24
-
25
-
26
-def ref_impl(silu_and_mul: SiluAndMul, x: torch.Tensor,
27
- global_scale: torch.Tensor,
28
- ref_output_scale: torch.Tensor) -> torch.Tensor:
29
- silu_and_mul_out = silu_and_mul.forward_native(x)
30
- assert not current_platform.is_rocm()
31
- assert silu_and_mul_out.ndim >= 1, (
32
- f'input.ndim needs to be >= 1, but got {silu_and_mul_out.ndim}.')
33
- other_dims = 1 if silu_and_mul_out.ndim == 1 else -1
34
- silu_and_mul_out = silu_and_mul_out.reshape(other_dims,
35
- silu_and_mul_out.shape[-1])
36
- m, n = silu_and_mul_out.shape
37
- device = silu_and_mul_out.device
38
-
39
- # Two fp4 values will be packed into an uint8.
40
- out = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
41
-
42
- output_scale = ref_output_scale
43
-
44
- torch.ops._C.scaled_fp4_quant(out, silu_and_mul_out, output_scale,
45
- global_scale)
46
-
47
- return out, output_scale
48
-
49
-
50
-def ops_impl(x: torch.Tensor, global_scale: torch.Tensor,
51
- ref_output_scale: torch.Tensor) -> torch.Tensor:
52
- out_shape = (x.shape[0], x.shape[1] // 4)
53
- output_scale = ref_output_scale
54
- out = torch.empty(out_shape, dtype=torch.uint8, device=x.device)
55
- torch.ops._C.silu_and_mul_nvfp4_quant(out, output_scale, x, global_scale)
56
- return out, output_scale
57
-
58
-
59
-@pytest.mark.parametrize("dtype", DTYPES)
60
-@pytest.mark.parametrize("shape", SHAPES)
61
-@pytest.mark.parametrize("seed", SEEDS)
62
-@pytest.mark.parametrize("device", CUDA_DEVICES)
63
-@torch.inference_mode()
64
-def test_quantize_to_fp4(
65
- dtype: torch.dtype,
66
- shape: tuple[int, int],
67
- seed: int,
68
- device: str,
69
-) -> None:
70
- current_platform.seed_everything(seed)
71
- torch.set_default_device(device)
72
-
73
- m, n = shape
74
-
75
- x = torch.randn((m, n), dtype=dtype)
76
- tensor_amax = torch.abs(x).max().to(torch.float32)
77
- global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
78
-
79
- block_size = 16
80
-
81
- assert n % block_size == 0, (
82
- f'last dim has to be multiple of 16, but got {n}.')
83
- assert x.dtype in (torch.float16, torch.bfloat16), (
84
- f'input.dtype needs to be fp16 or bf16 but got {x.dtype}.')
85
-
86
- round_up = lambda x, y: (x + y - 1) // y * y
87
- rounded_m = round_up(x.shape[0], 128)
88
- scale_n = x.shape[1] // (2 * block_size)
89
- rounded_n = round_up(scale_n, 4)
90
- output_scale = torch.empty((rounded_m, rounded_n // 4),
91
- device=x.device,
92
- dtype=torch.int32)
93
-
94
- layer = SiluAndMul()
95
-
96
- ref_out, ref_out_scale = ref_impl(layer, x, global_scale, output_scale)
97
-
98
- fusion_out, fusion_out_scale = ops_impl(x, global_scale, output_scale)
99
-
100
- assert ref_out.dtype == torch.uint8
101
- assert fusion_out.dtype == torch.uint8
102
- assert ref_out.shape == fusion_out.shape
103
-
104
- assert ref_out_scale.dtype == torch.int32
105
- assert fusion_out_scale.dtype == torch.int32
106
- assert ref_out_scale.shape == fusion_out_scale.shape
107
-
108
- # Allow up to 2% of mismatched values since BF16 has accuracy issues.
109
- mis_threshold = 0.02
110
- atol = 0.4
111
- rtol = 0.4
112
- ref_logits = ref_out[-1]
113
- fusion_logits = fusion_out[-1]
114
-
115
- mis_count = torch.sum(
116
- torch.abs(fusion_logits - ref_logits) > (atol +
117
- rtol * torch.abs(ref_logits)))
118
- mis_ratio = mis_count / fusion_logits.numel()
119
-
120
- assert mis_ratio < mis_threshold, \
121
- f"Mismatch ratio {mis_ratio} exceeds threshold {mis_threshold}"
122
-
123
- torch.testing.assert_close(ref_out_scale, fusion_out_scale)
124
-
125
- opcheck(torch.ops._C.silu_and_mul_nvfp4_quant,
126
- (fusion_out, fusion_out_scale, x, global_scale))