[CI Bug] Revert #42379 to fix CI `Multi-Modal Models (Extended Generation 1)` (#46070)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jun 18, 2026 at 15:37 UTC
225936a1dd10586798c0181696d628e7b609ea90
2 files changed
+24
-12
csrc/libtorch_stable/layernorm_kernels.cu
+7
-6
@@ -81,11 +81,11 @@ __global__ void rms_norm_kernel(
81
#pragma unroll
82
for (int j = 0; j < VEC_SIZE; j++) {
83
float x = static_cast<float>(src1.val[j]);
84
- scalar_t normalized = static_cast<scalar_t>(x * s_variance);
84
if constexpr (HasWeight) {
86
- dst.val[j] = normalized * src2.val[j];
85
+ float w = static_cast<float>(src2.val[j]);
86
+ dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
87
} else {
88
- dst.val[j] = normalized;
88
+ dst.val[j] = static_cast<scalar_t>(x * s_variance);
89
}
90
}
91
v_out[i] = dst;
@@ -151,7 +151,8 @@ fused_add_rms_norm_kernel(
151
#pragma unroll
152
for (int j = 0; j < width; ++j) {
153
float x = Converter::convert(res.data[j]);
154
- out.data[j] = Converter::convert(x * s_variance) * w.data[j];
154
+ float wf = Converter::convert(w.data[j]);
155
+ out.data[j] = Converter::convert(x * s_variance * wf);
156
}
157
} else {
158
#pragma unroll
@@ -198,8 +199,8 @@ fused_add_rms_norm_kernel(
199
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
200
float x = (float)residual[blockIdx.x * hidden_size + idx];
201
if constexpr (HasWeight) {
201
- input[blockIdx.x * input_stride + idx] =
202
- (scalar_t)(x * s_variance) * weight[idx];
202
+ float w = (float)weight[idx];
203
+ input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w);
204
} else {
205
input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance);
206
}
csrc/libtorch_stable/layernorm_quant_kernels.cu
+17
-6
@@ -66,8 +66,13 @@ __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
- // 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];
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);
76
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
77
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(out_norm),
78
scale_inv);
@@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel(
142
#pragma unroll
143
for (int i = 0; i < width; ++i) {
144
float x = Converter::convert(res.data[i]);
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];
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);
151
out[id * width + i] = scaled_fp8_conversion<true, fp8_type>(
152
Converter::convert(out_norm_h), scale_inv);
153
}
@@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
192
193
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
194
float x = (float)residual[blockIdx.x * hidden_size + idx];
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];
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);
199
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
200
static_cast<float>(out_norm), scale_inv);
201
}