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 13, 14:30 UTC
[Bugfix] Fix test fused quant layernorm tests (#27865)
Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
ElizaWszola committed
Nov 8, 2025 at 23:31 UTC
171133f929f2e896af767ca6e6402990a5c2814e
2 files changed
+28
-10
csrc/quantization/w8a8/int8/scaled_quant.cu
+3
@@ -1,5 +1,6 @@
1
#include <ATen/cuda/CUDAContext.h>
2
#include <torch/all.h>
3
+#include <c10/cuda/CUDAGuard.h>
4
5
#include <cmath>
6
@@ -275,6 +276,7 @@ void static_scaled_int8_quant(torch::Tensor& out, // [..., hidden_size]
276
int const num_tokens = input.numel() / hidden_size;
277
dim3 const grid(num_tokens);
278
dim3 const block(std::min(hidden_size, 256));
279
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
280
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
281
VLLM_DISPATCH_FLOATING_TYPES(
282
input.scalar_type(), "static_scaled_int8_quant_kernel", [&] {
@@ -306,6 +308,7 @@ void dynamic_scaled_int8_quant(
308
int const num_tokens = input.numel() / hidden_size;
309
dim3 const grid(num_tokens);
310
dim3 const block(std::min(hidden_size, 256));
311
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
312
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
313
VLLM_DISPATCH_FLOATING_TYPES(
314
input.scalar_type(), "dynamic_scaled_int8_quant_kernel", [&] {
tests/kernels/core/test_fused_quant_layernorm.py
+25
-10
@@ -11,7 +11,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm
11
12
DTYPES = [torch.bfloat16, torch.float]
13
QUANT_DTYPES = [torch.int8, torch.float8_e4m3fn]
14
-VEC_HIDDEN_SIZES = range(1024, 1030)
14
+VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
15
# Avoid combinatorial explosion with full Cartesian product
16
NUM_TOKENS_HIDDEN_SIZES = [
17
*[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
@@ -65,7 +65,7 @@ def ref_dynamic_per_token_quant(
65
)
66
else:
67
assert quant_dtype == torch.int8
68
- torch_out, scales = ops.scaled_int8_quant(torch_out)
68
+ torch_out, scales, _ = ops.scaled_int8_quant(torch_out)
69
70
return torch_out, scales, residual
71
@@ -109,7 +109,7 @@ def ops_impl(
109
110
@pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
111
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
112
-@pytest.mark.parametrize("scale_ub", SCALE_UBS)
112
+@pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
113
@pytest.mark.parametrize("dtype", DTYPES)
114
@pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
115
@pytest.mark.parametrize("seed", SEEDS)
@@ -119,7 +119,7 @@ def test_rms_norm(
119
num_tokens: int,
120
hidden_size: int,
121
add_residual: bool,
122
- scale_ub: bool,
122
+ has_scale_ub: bool,
123
dtype: torch.dtype,
124
quant_dtype: torch.dtype,
125
seed: int,
@@ -130,7 +130,7 @@ def test_rms_norm(
130
torch.cuda.manual_seed(seed)
131
torch.set_default_device(device)
132
133
- if scale_ub is not None and quant_dtype != torch.float8_e4m3fn:
133
+ if has_scale_ub and quant_dtype != torch.float8_e4m3fn:
134
# skip
135
return
136
@@ -143,9 +143,11 @@ def test_rms_norm(
143
scale = 1 / (hidden_size)
144
x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
145
residual = torch.randn_like(x) * scale if add_residual else None
146
- if scale_ub is not None:
146
+ if has_scale_ub:
147
rms_x, _ = ref_rms_norm(layer, x, residual)
148
scale_ub = torch.mean(rms_x).to(dtype=torch.float32, device="cuda")
149
+ else:
150
+ scale_ub = None
151
152
ref_out, ref_scales, ref_residual = ref_impl(
153
layer, x, quant_dtype, residual, scale_ub
@@ -156,14 +158,27 @@ def test_rms_norm(
158
159
assert ref_out.dtype == quant_dtype
160
assert ops_out.dtype == quant_dtype
159
- assert torch.allclose(ref_scales, ops_scales)
161
if quant_dtype == torch.int8:
162
+ assert torch.allclose(ref_scales, ops_scales, atol=1e-6)
163
# big atol to account for round-off errors.
164
assert torch.allclose(ref_out, ops_out, atol=1)
165
else:
164
- assert torch.allclose(
165
- ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
166
- )
166
+ assert torch.allclose(ref_scales, ops_scales)
167
+ a = ref_out.to(dtype=torch.float32)
168
+ b = ops_out.to(dtype=torch.float32)
169
+ ok = torch.allclose(a, b)
170
+ if not ok:
171
+ # fallback: compare dequantized values with relaxed tolerance
172
+ a_deq = a * ref_scales.view(-1, 1)
173
+ b_deq = b * ops_scales.view(-1, 1)
174
+ # NOTE: It is possible that some future test cases trigger this
175
+ # max diff due to precision issues. If such an error is
176
+ # encountered, it's recommended to inspect the differences between
177
+ # all corresponding elements from each tensor (e.g. by looping over
178
+ # them) and checking how many the max diff error shows up on (just
179
+ # a few bad elements should still be considered acceptable).
180
+ ok = torch.allclose(a_deq, b_deq, rtol=5e-2, atol=5e-2)
181
+ assert ok
182
if add_residual:
183
assert torch.allclose(ref_residual, ops_residual)
184