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 2, 08:30 UTC
[Bugfix] Fix RMSNorm kernels to multiply in weight's native dtype (#42379)
Signed-off-by: Lanze Liu <lanzetech@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Lanze Liu committed
May 29, 2026 at 23:16 UTC
124fac10cb0ea83aee2ffeabac0b413d6b759b26
2 files changed
+10
-23
csrc/libtorch_stable/layernorm_kernels.cu
+4
-6
@@ -78,8 +78,7 @@ __global__ void rms_norm_kernel(
78
#pragma unroll
79
for (int j = 0; j < VEC_SIZE; j++) {
80
float x = static_cast<float>(src1.val[j]);
81
- float w = static_cast<float>(src2.val[j]);
82
- dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
81
+ dst.val[j] = static_cast<scalar_t>(x * s_variance) * src2.val[j];
82
}
83
v_out[i] = dst;
84
}
@@ -143,8 +142,7 @@ fused_add_rms_norm_kernel(
142
#pragma unroll
143
for (int j = 0; j < width; ++j) {
144
float x = Converter::convert(res.data[j]);
146
- float wf = Converter::convert(w.data[j]);
147
- out.data[j] = Converter::convert(x * s_variance * wf);
145
+ out.data[j] = Converter::convert(x * s_variance) * w.data[j];
146
}
147
input_v[strided_id] = out;
148
}
@@ -183,8 +181,8 @@ fused_add_rms_norm_kernel(
181
182
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
183
float x = (float)residual[blockIdx.x * hidden_size + idx];
186
- float w = (float)weight[idx];
187
- input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
184
+ input[blockIdx.x * input_stride + idx] =
185
+ (scalar_t)(x * s_variance) * weight[idx];
186
}
187
}
188
csrc/libtorch_stable/layernorm_quant_kernels.cu
+6
-17
@@ -66,13 +66,8 @@ __global__ void rms_norm_static_fp8_quant_kernel(
66
#pragma unroll
67
for (int j = 0; j < VEC_SIZE; j++) {
68
float x = static_cast<float>(src1.val[j]);
69
- float w = static_cast<float>(src2.val[j]);
70
- // Round normalized result through scalar_t to match the precision of the
71
- // unfused composite (rms_norm writes scalar_t, then
72
- // static_scaled_fp8_quant re-loads it as float before FP8 conversion).
73
- // Without this round, the fused path is strictly more accurate and
74
- // disagrees with the composite at exact E4M3 quantization tie boundaries.
75
- scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
69
+ // Multiply in weight's native dtype to match rms_norm_kernel.
70
+ scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
71
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
72
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
73
scale_inv);
@@ -142,12 +137,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
137
#pragma unroll
138
for (int i = 0; i < width; ++i) {
139
float x = Converter::convert(res.data[i]);
145
- float wf = Converter::convert(w.data[i]);
146
- // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
147
- // to match the unfused composite path at FP8 boundaries. We use the
148
- // backend's hip_type for the intermediate since c10::Half/BFloat16 has
149
- // ambiguous conversions on CUDA and no implicit conversion on ROCm.
150
- HipT out_norm_h = Converter::convert(x * s_variance * wf);
140
+ // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
141
+ HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i];
142
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
143
Converter::convert(out_norm_h), scale_inv);
144
}
@@ -192,10 +183,8 @@ fused_add_rms_norm_static_fp8_quant_kernel(
183
184
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
185
float x = (float)residual[blockIdx.x * hidden_size + idx];
195
- float w = (float)weight[idx];
196
- // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t
197
- // to match the unfused composite path at FP8 boundaries.
198
- scalar_t out_norm = static_cast<scalar_t>(x * s_variance * w);
186
+ // Multiply in weight's native dtype to match fused_add_rms_norm_kernel.
187
+ scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * weight[idx];
188
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
189
static_cast<float>(out_norm), scale_inv);
190
}