diff --git a/csrc/libtorch_stable/layernorm_kernels.cu b/csrc/libtorch_stable/layernorm_kernels.cu index eb121b0b88..f29734fc26 100644 --- a/csrc/libtorch_stable/layernorm_kernels.cu +++ b/csrc/libtorch_stable/layernorm_kernels.cu @@ -81,11 +81,11 @@ __global__ void rms_norm_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - scalar_t normalized = static_cast(x * s_variance); if constexpr (HasWeight) { - dst.val[j] = normalized * src2.val[j]; + float w = static_cast(src2.val[j]); + dst.val[j] = static_cast(x * s_variance * w); } else { - dst.val[j] = normalized; + dst.val[j] = static_cast(x * s_variance); } } v_out[i] = dst; @@ -151,7 +151,8 @@ fused_add_rms_norm_kernel( #pragma unroll for (int j = 0; j < width; ++j) { float x = Converter::convert(res.data[j]); - out.data[j] = Converter::convert(x * s_variance) * w.data[j]; + float wf = Converter::convert(w.data[j]); + out.data[j] = Converter::convert(x * s_variance * wf); } } else { #pragma unroll @@ -198,8 +199,8 @@ fused_add_rms_norm_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; if constexpr (HasWeight) { - input[blockIdx.x * input_stride + idx] = - (scalar_t)(x * s_variance) * weight[idx]; + float w = (float)weight[idx]; + input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance * w); } else { input[blockIdx.x * input_stride + idx] = (scalar_t)(x * s_variance); } diff --git a/csrc/libtorch_stable/layernorm_quant_kernels.cu b/csrc/libtorch_stable/layernorm_quant_kernels.cu index 32f3495f4e..26ffa76d6e 100644 --- a/csrc/libtorch_stable/layernorm_quant_kernels.cu +++ b/csrc/libtorch_stable/layernorm_quant_kernels.cu @@ -66,8 +66,13 @@ __global__ void rms_norm_static_fp8_quant_kernel( #pragma unroll for (int j = 0; j < VEC_SIZE; j++) { float x = static_cast(src1.val[j]); - // Multiply in weight's native dtype to match rms_norm_kernel. - scalar_t out_norm = static_cast(x * s_variance) * src2.val[j]; + float w = static_cast(src2.val[j]); + // Round normalized result through scalar_t to match the precision of the + // unfused composite (rms_norm writes scalar_t, then + // static_scaled_fp8_quant re-loads it as float before FP8 conversion). + // Without this round, the fused path is strictly more accurate and + // disagrees with the composite at exact E4M3 quantization tie boundaries. + scalar_t out_norm = static_cast(x * s_variance * w); out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] = scaled_fp8_conversion(static_cast(out_norm), scale_inv); @@ -137,8 +142,12 @@ fused_add_rms_norm_static_fp8_quant_kernel( #pragma unroll for (int i = 0; i < width; ++i) { float x = Converter::convert(res.data[i]); - // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. - HipT out_norm_h = Converter::convert(x * s_variance) * w.data[i]; + float wf = Converter::convert(w.data[i]); + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. We use the + // backend's hip_type for the intermediate since c10::Half/BFloat16 has + // ambiguous conversions on CUDA and no implicit conversion on ROCm. + HipT out_norm_h = Converter::convert(x * s_variance * wf); out[id * width + i] = scaled_fp8_conversion( Converter::convert(out_norm_h), scale_inv); } @@ -183,8 +192,10 @@ fused_add_rms_norm_static_fp8_quant_kernel( for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) { float x = (float)residual[blockIdx.x * hidden_size + idx]; - // Multiply in weight's native dtype to match fused_add_rms_norm_kernel. - scalar_t out_norm = static_cast(x * s_variance) * weight[idx]; + float w = (float)weight[idx]; + // See note in rms_norm_static_fp8_quant_kernel: round through scalar_t + // to match the unfused composite path at FP8 boundaries. + scalar_t out_norm = static_cast(x * s_variance * w); out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion( static_cast(out_norm), scale_inv); }