diff --git a/runtime/metal/kernels/decode/linear_gguf_sgmatrix.metal b/runtime/metal/kernels/decode/linear_gguf_sgmatrix.metal index b79679cff..4853051ed 100644 --- a/runtime/metal/kernels/decode/linear_gguf_sgmatrix.metal +++ b/runtime/metal/kernels/decode/linear_gguf_sgmatrix.metal @@ -149,6 +149,14 @@ inline half2 nibble6_operand(uint2 chunk, uint f) { const uint high = f == 3 ? chunk.y >> 4 : chunk.y << (8 - 4 * f); return as_type((at & 0x00F000F0u) | (high & 0x03000300u) | 0x3C003C00u) - half2(1.5h); } +// The exact half operand (code - 4) / 8 of pair f of a Q3_K chunk (uint2: the low bits as quant_spread2 places them, +// pair p's at bits 4p and 16 + 4p, then the hmask bits as quant_spread1 does, at 2p and 16 + 2p): the low bits at bits +// 7-8 and the hmask bit at 9 of each halfword, 1 + code / 8, less 1.5, so a chain of Q3_K's zero point 4 needs no +// seed. On an M2 Max this takes 2% off the 27B's Q3_K projections at one and two lanes and 6-10% at three and four. +inline half2 q3_operand(uint2 chunk, uint f) { + const uint low = f == 0 ? chunk.x << 7 : f == 1 ? chunk.x << 3 : chunk.x >> (4 * f - 7); + return as_type((low & 0x01800180u) | ((chunk.y << (9 - 2 * f)) & 0x02000200u) | 0x3C003C00u) - half2(1.5h); +} // Apple7/8's (F32) view of Q4_0 and Q5_0: the zero point taken as a min, m = -zero d, so their chains close once per // 32 inputs against the table's sums, as Q4_1's and Q5_1's do, instead of once per 16 against the seeds. On an M2 @@ -186,9 +194,19 @@ template <> struct NibbleChunk> { static constexpr constant bool template <> struct NibbleChunk { static constexpr constant bool value = true; }; template <> struct NibbleChunk { static constexpr constant bool value = true; }; template <> struct NibbleChunk> { static constexpr constant bool value = true; }; -// Q6_K, whose F32 chains per 16 inputs take nibble6_operand, unseeded. +// Q6_K and Q3_K, whose F32 chains per 16 inputs take nibble6_operand and q3_operand, (code - zero) / Scale, +// unseeded. template struct CenteredChunk { static constexpr constant bool value = false; }; -template <> struct CenteredChunk { static constexpr constant bool value = true; }; +template <> struct CenteredChunk { + static constexpr constant bool value = true; + static constexpr constant float Scale = 64; + static half2 operand(uint2 chunk, uint f) { return nibble6_operand(chunk, f); } +}; +template <> struct CenteredChunk { + static constexpr constant bool value = true; + static constexpr constant float Scale = 8; + static half2 operand(uint2 chunk, uint f) { return q3_operand(chunk, f); } +}; // The exact bf16 operand of pair f of a chunk. template @@ -332,7 +350,16 @@ inline void decode(device const TableElem *table, device const float *sums, #pragma unroll for (uint i = 0; i < I; ++i) { const uint e = lane + 32 * i; - if (e < NC * S::J / JS) cu[JS * (e / NC) * NC + e % NC] = coefficient(src[i], u, JS * (e / NC)); + if (e < NC * S::J / JS) { + C k = coefficient(src[i], u, JS * (e / NC)); + // The F32 per-32 nibble chains' epilogue factors, 64 s and m - 64 s (code = 64 (operand - 1)), once per + // coefficient rather than per chain: on an M2 Max, 1-2% off Q4_K's, Q4_0's and Q5_K's 27B projections at one to + // three lanes (Q4_K and Q4_0 0.3% more at four). + if constexpr (F32 && NibbleChunk::value) k = float2(64 * k.x, fma(64.0f, k.x, k.y)); + // The centered chains' Scale s. + if constexpr (F32 && CenteredChunk::value) k = CenteredChunk::Scale * k; + cu[JS * (e / NC) * NC + e % NC] = k; + } } simdgroup_barrier(mem_flags::mem_threadgroup); #pragma unroll @@ -404,7 +431,7 @@ inline void decode(device const TableElem *table, device const float *sums, #pragma unroll for (uint f = 0; f < 4; ++f) { if constexpr (nibbles) a[nf][f] = nibble_operand(cur[nf][q], f); - else if constexpr (centered) a[nf][f] = nibble6_operand(cur[nf][q], f); + else if constexpr (centered) a[nf][f] = CenteredChunk::operand(cur[nf][q], f); else if constexpr (F32) a[nf][f] = operand_f32(cur[nf][q], f, lut); else a[nf][f] = operand(cur[nf][q], f, lut); } @@ -432,15 +459,15 @@ inline void decode(device const TableElem *table, device const float *sums, for (uint f = h * FG; f < (h + 1) * FG; ++f) sgmatrix::mma_mixed(dot, a[nf][f], b[f]); if constexpr (nibbles) { - // code = 64 (operand - 1): s code + m = 64 s operand + m - 64 s, cs.y being m - 128 s. - acc[r][nf] = fma(dot, 64 * cs[h][nf].x, acc[r][nf]); - acc[r][nf] = fma(sum, fma(64.0f, cs[h][nf].x, cs[h][nf].y), acc[r][nf]); + // code = 64 (operand - 1): s code + m = 64 s operand + m - 64 s. + acc[r][nf] = fma(dot, cs[h][nf].x, acc[r][nf]); + acc[r][nf] = fma(sum, cs[h][nf].y, acc[r][nf]); } else if constexpr (S::HasMin) { acc[r][nf] = fma(dot, cs[h][nf].x, acc[r][nf]); acc[r][nf] = fma(S::Seeded ? seed[h] : sum, cs[h][nf].y, acc[r][nf]); } else if constexpr (centered) { - // code - 32 = 64 operand. - acc[r][nf] = fma(dot, 64 * float(cs[h][nf]), acc[r][nf]); + // code - zero = Scale operand, the coefficient being Scale s. + acc[r][nf] = fma(dot, float(cs[h][nf]), acc[r][nf]); } else { acc[r][nf] = fma(dot, float(cs[h][nf]), acc[r][nf]); }