#version 450
// Q5_K matrix-matrix product (the prefill/compute path): y[m,n] = sum_k x[m,k]*W[n,k] with W stored
// as GGUF Q5_K super-blocks kept verbatim in VRAM, SAME layout and decode as mul_mat_vec_q5k (256
// weights / 176 bytes = 44 u32 per super-block: u32[0]={d,dmin} f16 pair; u32[1..4]=12 packed
// scale/min bytes; u32[4..12]=32 qh high-bit bytes; u32[12..44]=128 qs low-nibble bytes). Prefill
// has M>1 activation rows sharing one weight matrix, so the old path dequantized the whole weight to
// f32 every forward (massive VRAM + bandwidth). This kernel keeps W quantized and decodes each weight
// value ONCE per output column, reusing it across all M rows: weight traffic equals a single matvec
// while M rows are produced. One invocation computes one output column n for up to MAX_M rows (the
// host tiles M by MAX_M). Decode MUST match the CPU k_quants BlockQ5K::to_float (identical to
// mul_mat_vec_q5k).
#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[]; }; // Q5_K blocks, 44 u32 / 256-block
layout(set = 0, binding = 1) readonly buffer X { float x[]; }; // activations, row-major [M, k]
layout(set = 0, binding = 2) writeonly buffer Y { float y[]; }; // output, row-major [M, nout]
// m0: first activation row this dispatch handles; mcount: rows in this tile (1..MAX_M); nout: total
// columns; k: contraction dim (multiple of 256). x row stride is k, y row stride is nout. woff: u32
// offset into w[] where this weight matrix starts (0 for a plain 2D weight; e*nout*(k/256)*44 to
// select expert e of a resident MoE bank).
layout(push_constant) uniform Pc { uint m0; uint mcount; uint nout; uint k; uint woff; };
const uint QK_K = 256u;
const uint BLK_U32 = 44u; // 176 bytes / 4
const uint QH_U32 = 4u; // u32 index where the 32 qh bytes start
const uint QS_U32 = 12u; // u32 index where the 128 qs bytes start
const uint MAX_M = 8u; // register accumulators per invocation; host tiles M by this
// Byte `b` (0-based) out of the block's scale region (u32[1..4], i.e. the 12 scale bytes).
uint scale_byte(uint base, uint b) {
uint word = w[base + 1u + (b >> 2u)];
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
// get_scale_min_k4 from quantized/utils.rs: returns packed (sc<<8 | m), each 6-bit.
uint get_scale_min_k4(uint base, uint j) {
uint sc, m;
if (j < 4u) {
sc = scale_byte(base, j) & 63u;
m = scale_byte(base, j + 4u) & 63u;
} else {
uint s_j4 = scale_byte(base, j + 4u);
uint s_jm4 = scale_byte(base, j - 4u);
uint s_j = scale_byte(base, j);
sc = (s_j4 & 0x0Fu) | ((s_jm4 >> 6u) << 4u);
m = (s_j4 >> 4u) | ((s_j >> 6u) << 4u);
}
return (sc << 8u) | m;
}
// One byte (0-based) out of the 32 qh high-bit bytes (u32[4..12]).
uint qh_byte(uint base, uint b) {
uint word = w[base + QH_U32 + (b >> 2u)];
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
// One byte (0-based) out of the 128 qs low-nibble bytes (u32[12..44]).
uint qs_byte(uint base, uint b) {
uint word = w[base + QS_U32 + (b >> 2u)];
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
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 weight row n within the selected matrix
float acc[MAX_M];
for (uint r = 0u; r < MAX_M; r++) {
acc[r] = 0.0;
}
for (uint blk = 0u; blk < nblocks; blk++) {
uint base = rowbase + blk * BLK_U32;
float d = float(unpackHalf2x16(w[base]).x);
float dmin = float(unpackHalf2x16(w[base]).y);
uint xblk = blk * QK_K; // activation offset for this super-block
// 4 chunks of 64 weights (QK_K/64). Chunk ci uses scale-pair (2ci, 2ci+1), qs bytes
// [ci*32, ci*32+32), the shared 32 qh bytes, and high-bit masks u1=1<<(2ci), u2=2<<(2ci).
for (uint ci = 0u; ci < 4u; ci++) {
uint is = ci * 2u;
uint sm1 = get_scale_min_k4(base, is);
uint sm2 = get_scale_min_k4(base, is + 1u);
float d1 = d * float(sm1 >> 8u);
float m1 = dmin * float(sm1 & 0xFFu);
float d2 = d * float(sm2 >> 8u);
float m2 = dmin * float(sm2 & 0xFFu);
uint qoff = ci * 32u; // byte offset into qs for this chunk
uint u1 = 1u << (ci * 2u); // high-bit mask for the low nibble
uint u2 = 2u << (ci * 2u); // high-bit mask for the high nibble
uint xlo = xblk + ci * 64u; // lower-nibble outputs land here
uint xhi = xlo + 32u; // upper-nibble outputs
for (uint l = 0u; l < 32u; l++) {
uint qval = qs_byte(base, qoff + l);
uint hbit = qh_byte(base, l);
float addlo = ((hbit & u1) != 0u) ? 16.0 : 0.0;
float addhi = ((hbit & u2) != 0u) ? 16.0 : 0.0;
// Decode the two weights of this byte ONCE, then reuse across all rows.
float wlo = d1 * (float(qval & 0x0Fu) + addlo) - m1;
float whi = d2 * (float(qval >> 4u) + addhi) - m2;
for (uint r = 0u; r < mcount; r++) {
uint xrow = (m0 + r) * k;
acc[r] += wlo * x[xrow + xlo + l];
acc[r] += whi * x[xrow + xhi + l];
}
}
}
}
for (uint r = 0u; r < mcount; r++) {
y[(m0 + r) * nout + n] = acc[r];
}
}