#version 450
// Q6_K matrix-vector product (decode/memory-bound path): y[n] = sum_k W[n,k]*x[k], W stored as
// GGUF Q6_K super-blocks. One Q6_K super-block packs 256 weights into 210 bytes:
// ql[0..128] = low 4 bits, two weights/byte
// qh[128..192] = high 2 bits, four weights/byte
// scales[192..208] = 16 signed i8 per-16 scales
// d[208..210] = f16 super-block scale
// 210 is not u32-aligned, so the host repacks each block into a PADDED 212-byte (53 u32) stride
// (`BLK_U32` below); the trailing 2 bytes are unused. Reading ~6.5 bits/weight (incl. pad) instead
// of 32 cuts decode memory traffic ~5x on this ~256 GB/s APU. One invocation computes one output
// element. Decode MUST match the CPU k_quants BlockQ6K::to_float.
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer W { uint w[]; }; // Q6_K blocks, 53 u32 (padded) / 256-block
layout(set = 0, binding = 1) readonly buffer X { float x[]; }; // activation vector, length k
layout(set = 0, binding = 2) writeonly buffer Y { float y[]; }; // output, length nout
// woff: u32 offset into w[] where this weight matrix starts (0 for a plain 2D weight; e*n*(k/256)*53
// to select expert e of a resident MoE bank, 53 u32 = padded Q6_K block). k is a multiple of 256.
layout(push_constant) uniform Pc { uint nout; uint k; uint woff; };
const uint QK_K = 256u;
const uint BLK_U32 = 53u; // 212 bytes (210 padded to a u32 multiple) / 4
const uint QL_BYTE = 0u; // byte offset of ql[128]
const uint QH_BYTE = 128u; // byte offset of qh[64]
const uint SC_BYTE = 192u; // byte offset of scales[16] (signed i8)
const uint D_BYTE = 208u; // byte offset of d (f16)
// Unsigned byte `b` (0-based) within the block at u32 index `base`.
uint byte_u(uint base, uint b) {
uint word = w[base + (b >> 2u)];
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
// Signed byte `b` (0-based) within the block (sign-extended).
int byte_s(uint base, uint b) {
uint word = w[base + (b >> 2u)];
return bitfieldExtract(int(word), int((b & 3u) * 8u), 8);
}
void main() {
uint n = gl_GlobalInvocationID.x;
if (n >= nout) {
return;
}
uint nblocks = k / QK_K;
uint rowbase = woff + n * nblocks * BLK_U32; // u32 offset of row n within the selected matrix
float acc = 0.0;
for (uint blk = 0u; blk < nblocks; blk++) {
uint base = rowbase + blk * BLK_U32;
// d is the f16 at byte 208; it sits in the low half of u32[52].
float d = float(unpackHalf2x16(w[base + (D_BYTE >> 2u)]).x);
uint xblk = blk * QK_K;
// 2 super-chunks of 128 weights (QK_K/128). Chunk idx uses scales[8*idx..], ql[64*idx..],
// qh[32*idx..]; the 4 lanes within a byte map to outputs l, l+32, l+64, l+96.
for (uint idx = 0u; idx < 2u; idx++) {
uint scoff = SC_BYTE + 8u * idx;
uint qloff = QL_BYTE + 64u * idx;
uint qhoff = QH_BYTE + 32u * idx;
uint xbase = xblk + idx * 128u;
for (uint l = 0u; l < 32u; l++) {
uint is = l >> 4u; // 0 for l<16, 1 otherwise
uint qll = byte_u(base, qloff + l);
uint qlh = byte_u(base, qloff + l + 32u);
uint qhv = byte_u(base, qhoff + l);
int q1 = int((qll & 0x0Fu) | ((qhv & 3u) << 4u)) - 32;
int q2 = int((qlh & 0x0Fu) | (((qhv >> 2u) & 3u) << 4u)) - 32;
int q3 = int((qll >> 4u) | (((qhv >> 4u) & 3u) << 4u)) - 32;
int q4 = int((qlh >> 4u) | (((qhv >> 6u) & 3u) << 4u)) - 32;
float s1 = d * float(byte_s(base, scoff + is));
float s2 = d * float(byte_s(base, scoff + is + 2u));
float s3 = d * float(byte_s(base, scoff + is + 4u));
float s4 = d * float(byte_s(base, scoff + is + 6u));
acc += s1 * float(q1) * x[xbase + l];
acc += s2 * float(q2) * x[xbase + l + 32u];
acc += s3 * float(q3) * x[xbase + l + 64u];
acc += s4 * float(q4) * x[xbase + l + 96u];
}
}
}
y[n] = acc;
}