Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions csrc/quantization/fused_kernels/layernorm_utils.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -516,6 +516,11 @@ __device__ void norm_and_quant(

q8x4_t<scalar_out_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) {
Expand All @@ -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<scalar_out_t, is_scale_inverted>::quant_fn(
out.val[j] = ScaledQuant<scalar_out_t, mul_by_recip>::quant_fn(
static_cast<scalar_t>(x.val[j] * rms) * w.val[j], scale_val);
}
vec_output[i] = out;
Expand Down
14 changes: 8 additions & 6 deletions csrc/quantization/w8a8/fp8/per_token_group_quant.cu
Original file line number Diff line number Diff line change
Expand Up @@ -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<float>(src) / y_s, min_8bit), max_8bit);
float q = fminf(fmaxf(static_cast<float>(src) * inv_y, min_8bit), max_8bit);
dst = DST_DTYPE(q);
};

Expand Down Expand Up @@ -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<float>(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<uint8_t*>(&qb);
Expand Down
Loading