#version 450
// Q4_K TILED prefill matmul (L1): y[m,n] = sum_k x[m,k]*W[n,k], BIT-IDENTICAL to mul_mat_q4k (same f32
// decode + same accumulation order: for blk -> for chunk c -> for byte l -> (low nibble then high)), but
// the column's weight is staged into SHARED MEMORY ONCE per workgroup and reused across ALL M rows. So a
// column's quantized weight streams from VRAM exactly once for the whole prefill, vs mul_mat_q4k's
// ceil(M/MAX_M) re-streams (the 44x prefill gap at M=512 on gfx1151). One workgroup == one output column
// n; its WG_SIZE threads cooperatively load the column's nblocks super-blocks (raw qs bytes + decoded
// per-sub-block d1/m1/d2/m2) into LDS, barrier, then stride over the M rows (thread t owns rows
// t, t+WG_SIZE, ...) doing the identical per-byte f32 MAC out of LDS. Bit-exact gate: mul_mm_q4k_shared
// vs mul_mat_q4k vs the CPU to_float oracle.
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
const uint WG_SIZE = 256u;
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer W { uint w[]; }; // Q4_K blocks, 36 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]
layout(push_constant) uniform Pc { uint m0; uint mcount; uint nout; uint k; uint woff; };
const uint QK_K = 256u;
const uint BLK_U32 = 36u; // 144 bytes / 4
const uint QS_U32 = 32u; // 128 qs bytes / 4, per super-block
// Per super-block LDS: 32 qs u32 + 16 f32 (d1,m1,d2,m2 per chunk c=0..3) = 192 B. MAX_BLOCKS=64 ->
// 12 KiB, far inside gfx1151's 64 KiB LDS; k<=64*256=16384 covers every Qwen3 projection (max k=14336).
// The host gates nblocks<=MAX_BLOCKS (else falls back to mul_mat_q4k).
const uint MAX_BLOCKS = 64u;
shared uint s_qs[MAX_BLOCKS * QS_U32];
shared float s_d1[MAX_BLOCKS * 4u];
shared float s_m1[MAX_BLOCKS * 4u];
shared float s_d2[MAX_BLOCKS * 4u];
shared float s_m2[MAX_BLOCKS * 4u];
uint scale_byte(uint base, uint b) {
uint word = w[base + 1u + (b >> 2u)];
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
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;
}
void main() {
uint n = gl_WorkGroupID.x; // one workgroup per output column
if (n >= nout) {
return;
}
uint tid = gl_LocalInvocationID.x;
uint nblocks = k / QK_K;
uint rowbase = woff + n * nblocks * BLK_U32;
// Phase 1: cooperatively stage this column into LDS (thread takes strided super-blocks). Raw qs
// bytes verbatim + the per-sub-block d1/m1/d2/m2 decoded exactly as mul_mat_q4k.
for (uint blk = tid; blk < nblocks; blk += WG_SIZE) {
uint base = rowbase + blk * BLK_U32;
float d = float(unpackHalf2x16(w[base]).x);
float dmin = float(unpackHalf2x16(w[base]).y);
uint qbase = base + 4u;
for (uint i = 0u; i < QS_U32; i++) {
s_qs[blk * QS_U32 + i] = w[qbase + i];
}
for (uint c = 0u; c < 4u; c++) {
uint sm1 = get_scale_min_k4(base, c * 2u);
uint sm2 = get_scale_min_k4(base, c * 2u + 1u);
s_d1[blk * 4u + c] = d * float(sm1 >> 8u);
s_m1[blk * 4u + c] = dmin * float(sm1 & 0xFFu);
s_d2[blk * 4u + c] = d * float(sm2 >> 8u);
s_m2[blk * 4u + c] = dmin * float(sm2 & 0xFFu);
}
}
barrier();
// Phase 2: each thread strides over the M rows; identical per-byte f32 MAC to mul_mat_q4k, reading
// the staged weight from LDS (zero VRAM weight traffic here).
for (uint r = tid; r < mcount; r += WG_SIZE) {
uint xrow = (m0 + r) * k;
float acc = 0.0;
for (uint blk = 0u; blk < nblocks; blk++) {
uint xblk = blk * QK_K;
for (uint c = 0u; c < 4u; c++) {
float d1 = s_d1[blk * 4u + c];
float m1 = s_m1[blk * 4u + c];
float d2 = s_d2[blk * 4u + c];
float m2 = s_m2[blk * 4u + c];
uint qoff = c * 32u;
uint xlo = xblk + c * 64u;
uint xhi = xlo + 32u;
for (uint l = 0u; l < 32u; l++) {
uint qb = qoff + l;
uint qword = s_qs[blk * QS_U32 + (qb >> 2u)];
uint q = (qword >> ((qb & 3u) * 8u)) & 0xFFu;
float wlo = d1 * float(q & 0x0Fu) - m1;
float whi = d2 * float(q >> 4u) - m2;
acc += wlo * x[xrow + xlo + l];
acc += whi * x[xrow + xhi + l];
}
}
}
y[(m0 + r) * nout + n] = acc;
}
}