Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
@@ -81,11 +81,11 @@ __global__ void rms_norm_kernel(
|
||||
#pragma unroll
|
||||
for (int j = 0; j < VEC_SIZE; j++) {
|
||||
float x = static_cast<float>(src1.val[j]);
|
||||
scalar_t normalized = static_cast<scalar_t>(x * s_variance);
|
||||
if constexpr (HasWeight) {
|
||||
dst.val[j] = normalized * src2.val[j];
|
||||
float w = static_cast<float>(src2.val[j]);
|
||||
dst.val[j] = static_cast<scalar_t>(x * s_variance * w);
|
||||
} else {
|
||||
dst.val[j] = normalized;
|
||||
dst.val[j] = static_cast<scalar_t>(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);
|
||||
}
|
||||
|
||||
@@ -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<float>(src1.val[j]);
|
||||
// Multiply in weight's native dtype to match rms_norm_kernel.
|
||||
scalar_t out_norm = static_cast<scalar_t>(x * s_variance) * src2.val[j];
|
||||
float w = static_cast<float>(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<scalar_t>(x * s_variance * w);
|
||||
out[blockIdx.x * hidden_size + idx * VEC_SIZE + j] =
|
||||
scaled_fp8_conversion<true, fp8_type>(static_cast<float>(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<true, fp8_type>(
|
||||
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<scalar_t>(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<scalar_t>(x * s_variance * w);
|
||||
out[blockIdx.x * hidden_size + idx] = scaled_fp8_conversion<true, fp8_type>(
|
||||
static_cast<float>(out_norm), scale_inv);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user