Skip to content
Merged
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
45 changes: 36 additions & 9 deletions runtime/metal/kernels/decode/linear_gguf_sgmatrix.metal
Original file line number Diff line number Diff line change
Expand Up @@ -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<half2>((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<half2>((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
Expand Down Expand Up @@ -186,9 +194,19 @@ template <> struct NibbleChunk<FmtAffine<4, 1>> { static constexpr constant bool
template <> struct NibbleChunk<FmtQ5K> { static constexpr constant bool value = true; };
template <> struct NibbleChunk<FmtQ51> { static constexpr constant bool value = true; };
template <> struct NibbleChunk<ZeroAsMin<FmtQ50>> { 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 <class F> struct CenteredChunk { static constexpr constant bool value = false; };
template <> struct CenteredChunk<FmtQ6K> { static constexpr constant bool value = true; };
template <> struct CenteredChunk<FmtQ6K> {
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<FmtQ3K> {
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 <class F>
Expand Down Expand Up @@ -332,7 +350,16 @@ inline void decode(device const TableElem<F32> *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<F>(src[i], u, JS * (e / NC));
if (e < NC * S::J / JS) {
C k = coefficient<F>(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<F>::value) k = float2(64 * k.x, fma(64.0f, k.x, k.y));
// The centered chains' Scale s.
if constexpr (F32 && CenteredChunk<F>::value) k = CenteredChunk<F>::Scale * k;
cu[JS * (e / NC) * NC + e % NC] = k;
}
}
simdgroup_barrier(mem_flags::mem_threadgroup);
#pragma unroll
Expand Down Expand Up @@ -404,7 +431,7 @@ inline void decode(device const TableElem<F32> *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<F>::operand(cur[nf][q], f);
else if constexpr (F32) a[nf][f] = operand_f32<F>(cur[nf][q], f, lut);
else a[nf][f] = operand<F>(cur[nf][q], f, lut);
}
Expand Down Expand Up @@ -432,15 +459,15 @@ inline void decode(device const TableElem<F32> *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]);
}
Expand Down
Loading