diff --git a/csrc/quantization/fused_kernels/layernorm_utils.cuh b/csrc/quantization/fused_kernels/layernorm_utils.cuh index d633cb9a6e..8cab1824f2 100644 --- a/csrc/quantization/fused_kernels/layernorm_utils.cuh +++ b/csrc/quantization/fused_kernels/layernorm_utils.cuh @@ -516,6 +516,11 @@ __device__ void norm_and_quant( q8x4_t out; + // Group quant multiplies by the reciprocal of the scale (one div per + // group), matching per_token_group_quant_8bit. Per-token fp8 must keep + // exact division for FBGemm parity, so it is excluded here. + constexpr bool mul_by_recip = (group_size > 0) || is_scale_inverted; + float scale_val; if constexpr (group_size > 0) { @@ -528,14 +533,13 @@ __device__ void norm_and_quant( } else { scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size; } - scale_val = - is_scale_inverted ? 1.0f / scale[scale_idx] : scale[scale_idx]; + scale_val = 1.0f / scale[scale_idx]; } else { scale_val = *scale; } #pragma unroll for (int j = 0; j < VEC_SIZE; ++j) { - out.val[j] = ScaledQuant::quant_fn( + out.val[j] = ScaledQuant::quant_fn( static_cast(x.val[j] * rms) * w.val[j], scale_val); } vec_output[i] = out; diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index d2a8886438..0571cc9077 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -78,9 +78,13 @@ __device__ __forceinline__ void QuantizeGroup( const float y_s, const float min_8bit, const float max_8bit) { constexpr int vec_size = 16 / sizeof(T); + // One div.rn per group instead of per element; the register-resident + // kernel does the same, keeping both schedules bit-identical. + const float inv_y = 1.0f / y_s; + // quantize shared -> global 8-bit auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) { - float q = fminf(fmaxf(static_cast(src) / y_s, min_8bit), max_8bit); + float q = fminf(fmaxf(static_cast(src) * inv_y, min_8bit), max_8bit); dst = DST_DTYPE(q); }; @@ -272,16 +276,14 @@ __global__ void per_token_group_quant_8bit_register_kernel( } } - // A UE8M0 scale is an exact power of two, so multiplying by the reciprocal - // matches dividing. Otherwise the reciprocal costs an ulp and can move - // values that land exactly on an 8-bit boundary, so divide instead to stay - // bit-exact with the smem kernel. + // One div.rn per group; the smem path (QuantizeGroup) does the same, + // keeping both schedules bit-identical. const float inv_y = 1.0f / y_s; uint32_t packed[VEC_SIZE / 4] = {}; #pragma unroll for (int i = 0; i < VEC_SIZE; ++i) { const float v = static_cast(regs[i]); - const float scaled = SCALE_UE8M0 ? (v * inv_y) : (v / y_s); + const float scaled = v * inv_y; const float q = fminf(fmaxf(scaled, min_8bit), max_8bit); DST_DTYPE qb = DST_DTYPE(q); const uint8_t byte = *reinterpret_cast(&qb);