vulkan: FA MMQ should use fp32 for Q quantization calculations (#27413)

Codex found that qd could be a denorm and 1/qd would overflow.
This commit is contained in:
Jeff Bolz
2026-08-20 09:18:11 -05:00
committed by GitHub
parent 63b64a50a3
commit 78ec4c3780
@@ -121,13 +121,13 @@ void main() {
const uint buf_ib = r * qf_stride + d / 8; const uint buf_ib = r * qf_stride + d / 8;
const uint buf_iqs = d % 8; const uint buf_iqs = d % 8;
FLOAT_TYPEV4 vals = is_in_bounds ? FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale) : FLOAT_TYPEV4(0.0f); vec4 vals = is_in_bounds ? data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale : vec4(0.0f);
const FLOAT_TYPEV4 abs_vals = abs(vals); const vec4 abs_vals = abs(vals);
const FLOAT_TYPE thread_max = max(max(abs_vals.x, abs_vals.y), max(abs_vals.z, abs_vals.w)); const float thread_max = max(max(abs_vals.x, abs_vals.y), max(abs_vals.z, abs_vals.w));
const FLOAT_TYPE amax = subgroupClusteredMax(thread_max, 8); const float amax = subgroupClusteredMax(thread_max, 8);
const FLOAT_TYPE qd = amax / FLOAT_TYPE(127.0); const float qd = amax / 127.0f;
const FLOAT_TYPE qd_inv = qd != FLOAT_TYPE(0.0) ? FLOAT_TYPE(1.0) / qd : FLOAT_TYPE(0.0); const float qd_inv = qd != 0.0f ? 1.0f / qd : 0.0f;
vals = round(vals * qd_inv); vals = round(vals * qd_inv);
Qf[buf_ib].qs[buf_iqs] = pack32(i8vec4(vals)); Qf[buf_ib].qs[buf_iqs] = pack32(i8vec4(vals));
@@ -136,11 +136,11 @@ void main() {
// the row-sum scaled by qd, used in k_dot_correction. // the row-sum scaled by qd, used in k_dot_correction.
if (FaTypeK == FA_TYPE_Q8_0) { if (FaTypeK == FA_TYPE_Q8_0) {
if (buf_iqs == 0) { if (buf_iqs == 0) {
Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0); Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f);
} }
} else { } else {
const FLOAT_TYPE thread_sum = vals.x + vals.y + vals.z + vals.w; const float thread_sum = vals.x + vals.y + vals.z + vals.w;
const FLOAT_TYPE sum = subgroupClusteredAdd(thread_sum, 8); const float sum = subgroupClusteredAdd(thread_sum, 8);
if (buf_iqs == 0) { if (buf_iqs == 0) {
Qf[buf_ib].ds = FLOAT_TYPEV2(qd, sum * qd); Qf[buf_ib].ds = FLOAT_TYPEV2(qd, sum * qd);