From ea5dbffa7f968e80802e861e121f78236de3b039 Mon Sep 17 00:00:00 2001 From: Niko1221 Date: Fri, 30 Oct 2026 20:09:02 +0100 Subject: [PATCH] Metal: Q6_K multi-column dense matvec (ds4_q6k_mvn) --- metal/src/engine/ds4_finish.cpp | 1 + metal/src/engine/msl/ds4_iq.metal | 55 +++++++++++++++++++++++++++++++ metal/tests/dsu_parity.cpp | 5 +-- 3 files changed, 59 insertions(+), 2 deletions(-) diff --git a/metal/src/engine/ds4_finish.cpp b/metal/src/engine/ds4_finish.cpp index a0fc63a4..d755e911 100644 --- a/metal/src/engine/ds4_finish.cpp +++ b/metal/src/engine/ds4_finish.cpp @@ -746,6 +746,7 @@ bool iq_mv(const void* jobs, int nj) { const char* tn = nullptr; switch (q.type) { case 13: tn = "ds4_q5k_mvn"; break; + case 14: tn = "ds4_q6k_mvn"; break; // (the Unsloth sets' Q6_K dense matrices) case 16: tn = "ds4_iq2xxs_mvn"; break; case 17: tn = "ds4_iq2xs_mvn"; break; case 18: tn = "ds4_iq3xxs_mvn"; break; diff --git a/metal/src/engine/msl/ds4_iq.metal b/metal/src/engine/msl/ds4_iq.metal index 4e5ad419..5bf41e04 100644 --- a/metal/src/engine/msl/ds4_iq.metal +++ b/metal/src/engine/msl/ds4_iq.metal @@ -953,6 +953,60 @@ kernel void ds4_q5k_mvn(device const uchar* jobs [[buffer(0)]], uint2 tgp [[thre ds4_mvn_store(b, row, lane, acc); } +// Q6_K (210-byte blocks; the Unsloth sets' attention, shared-expert, head and engram_kv matrices) for NC activations: four lanes a +// block, a lane one 128-value half and one of its two quarter pairs (quarters p and p + 2 share the 32 ql bytes at 32 p: the low +// nibbles are quarter p, the high ones p + 2; quarter q's 2 high bits are qh >> 2 q), eight blocks at once. The loads go through +// packed_uchar4: a block starts only 2-byte aligned. ggml's dequantize_row_q6_K arithmetic: d * scale * (q - 32). +template +kernel void ds4_q6k_mvn(device const uchar* jobs [[buffer(0)]], uint2 tgp [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], ushort lane [[thread_index_in_simdgroup]]) { + const device Ds4MvJob& b = *(const device Ds4MvJob*) (jobs + (size_t) tgp.y * sizeof(Ds4MvJob)); + const uint row = tgp.x * 8u + (tid >> 5); + if (int(row) >= b.n) return; + const uint nblk = uint(b.k) / 256u, nsb = uint(b.k) / 32u; + const device uchar* w = b.w + (size_t) row * b.ld; + const uint j = lane & 3u, h = j >> 1, pr = j & 1u; + const uint bsel = lane >> 2; + float acc[NC]; + for (uint c = 0; c < NC; ++c) acc[c] = 0.0f; + for (uint blk = bsel; blk < nblk; blk += 8u) { + const device uchar* p = w + (size_t) blk * 210u; + const float d = float(as_type(*(const device ushort*) (p + 208))); + const device uchar* ql = p + 64u * h + 32u * pr; + const device uchar* qh = p + 128u + 32u * h; + const device char* scl = (const device char*) (p + 192u + 8u * h); + uint qlw[8], qhw[8]; + for (uint t = 0; t < 8u; ++t) { + qlw[t] = as_type(uchar4(*(const device packed_uchar4*) (ql + 4u * t))); + qhw[t] = as_type(uchar4(*(const device packed_uchar4*) (qh + 4u * t))); + } + for (uint qi = 0; qi < 2u; ++qi) { + const uint q = pr + 2u * qi; + const uint sh = 4u * qi; + int cw[32]; + for (uint t = 0; t < 8u; ++t) { + const uint x = ((qlw[t] >> sh) & 0x0f0f0f0fu) | (((qhw[t] >> (2u * q)) & 0x03030303u) << 4); + for (uint i = 0; i < 4u; ++i) cw[t * 4u + i] = int((x >> (8u * i)) & 0xffu) - 32; + } + const float s0 = float(scl[2u * q]), s1 = float(scl[2u * q + 1u]); + const uint g = blk * 8u + 4u * h + q; + for (uint c = 0; c < NC; ++c) { + const device uint4* xq4 = (const device uint4*) (b.xq + (size_t) c * b.k + (size_t) blk * 256u + 128u * h + 32u * q); + const uint4 xa = xq4[0], xb = xq4[1]; + const uint xw[8] = {xa.x, xa.y, xa.z, xa.w, xb.x, xb.y, xb.z, xb.w}; + int sa = 0, sb = 0; + for (uint t = 0; t < 8u; ++t) { + int part = 0; + for (uint i = 0; i < 4u; ++i) part += cw[t * 4u + i] * int(char(xw[t] >> (8u * i))); + if (t < 4u) sa += part; else sb += part; + } + acc[c] = fma(b.xd[(size_t) c * nsb + g] * d, s0 * float(sa) + s1 * float(sb), acc[c]); + } + } + } + ds4_mvn_store(b, row, lane, acc); +} + // the grid-coded types: decode the sub-block's 32 weights (signs applied) into w[32] ... static inline void ds4_w32_iq2xxs(const device uchar* p, uint ib, thread int (&w)[32], thread float& db) { const float d = float(as_type(*(const device ushort*) p)); @@ -1117,6 +1171,7 @@ kernel void ds4_iqg_mvn(device const uchar* jobs [[buffer(0)]], uint2 tgp [[thre #define DS4_MVN_INST(NC) \ template [[host_name("ds4_q5k_mvn" #NC)]] kernel void ds4_q5k_mvn(device const uchar*, uint2, uint, ushort); \ + template [[host_name("ds4_q6k_mvn" #NC)]] kernel void ds4_q6k_mvn(device const uchar*, uint2, uint, ushort); \ template [[host_name("ds4_iq2xxs_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ template [[host_name("ds4_iq2xs_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ template [[host_name("ds4_iq2s_mvn" #NC)]] kernel void ds4_iqg_mvn(device const uchar*, uint2, uint, ushort); \ diff --git a/metal/tests/dsu_parity.cpp b/metal/tests/dsu_parity.cpp index 492ecb84..b78d55b5 100644 --- a/metal/tests/dsu_parity.cpp +++ b/metal/tests/dsu_parity.cpp @@ -47,9 +47,10 @@ int main(int argc, char** argv) { struct Case { const char* kern; int type; int nc; }; const std::vector dummy; const std::vector cases = {{"ds4_iq1s_mvn1", 19, 1}, {"ds4_iq1s_mvn2", 19, 2}, {"ds4_iq1s_mvn4", 19, 4}, {"ds4_iq1s_mvn8", 19, 8}, - {"ds4_iq1m_mvn4", 29, 4}, {"ds4_iq2xxs_mvn4", 16, 4}, {"ds4_iq_mv_k", 14, 8}, {"ds4_iq_mv_k", 14, 1}}; + {"ds4_iq1m_mvn4", 29, 4}, {"ds4_iq2xxs_mvn4", 16, 4}, {"ds4_iq_mv_k", 14, 8}, {"ds4_iq_mv_k", 14, 1}, + {"ds4_q6k_mvn1", 14, 1}, {"ds4_q6k_mvn2", 14, 2}, {"ds4_q6k_mvn4", 14, 4}, {"ds4_q6k_mvn8", 14, 8}}; if (bench) { // dense decode matvec speed on a real attention-sized matrix: Q5_K (the 2.00 set's) against Q6_K (the Unsloth sets') - for (const Case& cs : std::vector{{"ds4_q5k_mvn1", 13, 1}, {"ds4_iq_mv_k", 14, 1}, {"ds4_iq1s_mvn1", 19, 1}, {"ds4_iq2xxs_mvn1", 16, 1}, {"ds4_iq_mv_k", 14, 8}, {"ds4_q5k_mvn8", 13, 8}}) { + for (const Case& cs : std::vector{{"ds4_q5k_mvn1", 13, 1}, {"ds4_iq_mv_k", 14, 1}, {"ds4_iq1s_mvn1", 19, 1}, {"ds4_iq2xxs_mvn1", 16, 1}, {"ds4_iq_mv_k", 14, 8}, {"ds4_q5k_mvn8", 13, 8}, {"ds4_q6k_mvn1", 14, 1}, {"ds4_q6k_mvn8", 14, 8}}) { const int k = 4096, n = 16384, nc = cs.nc; const size_t rb = ref::row_bytes(cs.type, k); std::vector w(rb * n);