[Bugfix] Remove race in fused groupwise RMSNorm quantization (#54111)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Misha Goin committed
Aug 27, 2026 at 22:01 UTC
5bbc58c0ad0fec0e836f2ace6c61659f60328339
2 files changed
+35
-38
csrc/libtorch_stable/quantization/fused_kernels/layernorm_utils.cuh
+21
-37
@@ -47,27 +47,13 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
47
*rms = s_rms;
48
}
49
50
-__device__ float warpReduceMaxSpecialized(volatile float* val, int64_t tid,
51
- int64_t thread_in_warp,
52
- int64_t reduced_elems) {
50
+__device__ float warpReduceMax(float val) {
51
static_assert(WARP_SIZE == 32 || WARP_SIZE == 64);
54
- if constexpr (WARP_SIZE == 64) {
55
- if (thread_in_warp + 64 < reduced_elems)
56
- val[tid] = fmaxf(val[tid], val[tid + 64]);
52
+#pragma unroll
53
+ for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
54
+ val = fmaxf(val, VLLM_SHFL_DOWN_SYNC(val, offset));
55
}
58
- if (thread_in_warp + 32 < reduced_elems)
59
- val[tid] = fmaxf(val[tid], val[tid + 32]);
60
- if (thread_in_warp + 16 < reduced_elems)
61
- val[tid] = fmaxf(val[tid], val[tid + 16]);
62
- if (thread_in_warp + 8 < reduced_elems)
63
- val[tid] = fmaxf(val[tid], val[tid + 8]);
64
- if (thread_in_warp + 4 < reduced_elems)
65
- val[tid] = fmaxf(val[tid], val[tid + 4]);
66
- if (thread_in_warp + 2 < reduced_elems)
67
- val[tid] = fmaxf(val[tid], val[tid + 2]);
68
- if (thread_in_warp + 1 < reduced_elems)
69
- val[tid] = fmaxf(val[tid], val[tid + 1]);
70
- return val[tid];
56
+ return val;
57
}
58
59
template <typename scalar_t, typename scalar_out_t, bool has_residual = false,
@@ -115,16 +101,15 @@ __device__ void compute_dynamic_per_token_scales(
101
for (auto i = 0; i < groups_per_warp; ++i) {
102
int64_t const group_id = i * num_warps + warp_id;
103
if (group_id < num_groups) {
118
- int64_t warp_start = group_id * threads_per_group;
119
- int64_t const start = warp_start + thread_in_warp;
120
- int64_t const warp_end = min(warp_start + threads_per_group,
121
- static_cast<int64_t>(hidden_size));
122
- for (auto j = start; j + warp_size < warp_end; j += warp_size) {
123
- s_max_vals[start] =
124
- fmaxf(s_max_vals[start], s_max_vals[j + warp_size]);
104
+ int64_t const warp_start = group_id * threads_per_group;
105
+ float warp_max = 0.0f;
106
+ for (auto j = thread_in_warp; j < threads_per_group; j += warp_size) {
107
+ warp_max = fmaxf(warp_max, s_max_vals[warp_start + j]);
108
+ }
109
+ warp_max = warpReduceMax(warp_max);
110
+ if (thread_in_warp == 0) {
111
+ s_max_vals[warp_start] = warp_max;
112
}
126
- warpReduceMaxSpecialized(s_max_vals, start, thread_in_warp,
127
- min(warp_end - warp_start, warp_size));
113
}
114
}
115
__syncthreads();
@@ -375,16 +360,15 @@ __device__ void compute_dynamic_per_token_scales(
360
for (auto i = 0; i < groups_per_warp; ++i) {
361
int64_t const group_id = i * num_warps + warp_id;
362
if (group_id < num_groups) {
378
- int64_t warp_start = group_id * threads_per_group;
379
- int64_t const start = warp_start + thread_in_warp;
380
- int64_t const warp_end = min(warp_start + threads_per_group,
381
- static_cast<int64_t>(hidden_size));
382
- for (auto j = start; j + warp_size < warp_end; j += warp_size) {
383
- s_max_vals[start] =
384
- fmaxf(s_max_vals[start], s_max_vals[j + warp_size]);
363
+ int64_t const warp_start = group_id * threads_per_group;
364
+ float warp_max = 0.0f;
365
+ for (auto j = thread_in_warp; j < threads_per_group; j += warp_size) {
366
+ warp_max = fmaxf(warp_max, s_max_vals[warp_start + j]);
367
+ }
368
+ warp_max = warpReduceMax(warp_max);
369
+ if (thread_in_warp == 0) {
370
+ s_max_vals[warp_start] = warp_max;
371
}
386
- warpReduceMaxSpecialized(s_max_vals, start, thread_in_warp,
387
- min(warp_end - warp_start, warp_size));
372
}
373
}
374
__syncthreads();
tests/kernels/core/test_fused_quant_layernorm.py
+14
-1
@@ -330,6 +330,12 @@ def test_rms_norm(
330
and dtype == torch.bfloat16
331
and quant_dtype == current_platform.fp8_dtype()
332
)
333
+ allow_cuda_fp8_rounding_outliers = (
334
+ current_platform.is_cuda()
335
+ and group_size is None
336
+ and dtype == torch.bfloat16
337
+ and quant_dtype == current_platform.fp8_dtype()
338
+ )
339
340
def scales_close(rtol: float, atol: float) -> bool:
341
if torch.allclose(ref_scales, ops_scales, rtol=rtol, atol=atol):
@@ -358,8 +364,15 @@ def test_rms_norm(
364
# Valid gfx950 reduction trees can straddle an E4M3 boundary.
365
ok = fp8_allclose(ops_out, ref_out, rtol=0.125, atol=2e-3)
366
ok = ok and int(fp8_ulp_distance(ops_out, ref_out).max()) <= 1
367
+ elif allow_cuda_fp8_rounding_outliers:
368
+ # A valid BF16 reduction can cross an E4M3 boundary for
369
+ # isolated values.
370
+ ulp = fp8_ulp_distance(ref_out, ops_out)
371
+ max_outliers = ulp.numel() // 100_000 + 8
372
+ ok = int(ulp.max()) <= 1
373
+ ok = ok and int((ulp > 0).sum().item()) <= max_outliers
374
else:
362
- # CUDA (& non-bf16): compare dequantized values with relaxed tolerance.
375
+ # Compare dequantized values with relaxed tolerance.
376
if group_size is None:
377
a_deq = a * ref_scales.view(-1, 1)
378
b_deq = b * ops_scales.view(-1, 1)