#version 450
// Coalesced multi-thread Q6_K matrix-vector product: y[n] = sum_k W[n,k]*x[k], W stored as GGUF Q6_K
// super-blocks (256 weights / 210 bytes; host-repacked to a padded 212-byte, 53-u32 stride).
//
// This is the coalesced-load, cooperative-thread decode matvec (the memory-BW lever for decode on the
// gfx1151 UMA APU). Structure follows llama.cpp's mul_mat_vec_q6_k: a fixed group of THREADS_PER_BLOCK
// threads cooperate on ONE super-block, each thread owning a fixed 16-weight sub-position, and each
// workgroup computes ROWS_PER_WG output rows so adjacent threads read adjacent u32 weight words (a
// coalesced 32-byte burst) instead of one thread streaming a whole 212-byte block strided from its
// neighbours (uncoalesced, and idle whenever a row has fewer super-blocks than the subgroup is wide).
// A super-block's 256 weights split into 16 sub-positions x 16 weights; THREADS_PER_BLOCK=16 tiles it.
// ROWS_PER_WG rows amortise the activation loads. The per-row partials reduce with a subgroup add plus
// a small shared-memory sum across subgroups. Decode is bit-identical to mul_mat_vec_q6k / the scalar
// mul_mat_vec_q6k_sg (CPU k_quants BlockQ6K::to_float). Dispatched when the device advertises subgroup
// ARITHMETIC (gated host-side); VK_Q6K_CM_OFF falls back to mul_mat_vec_q6k_sg.
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_control_flow_attributes : enable
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_arithmetic : 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). k is a multiple of 256.
layout(push_constant) uniform Pc { uint nout; uint k; uint woff; };
const uint BLK_U32 = 53u; // padded Q6_K block: 212 bytes / 4
const uint TPB = 16u; // threads cooperating on one super-block
const uint ITS = gl_WorkGroupSize.x / TPB; // super-blocks processed in parallel per workgroup
const uint NR = 2u; // output rows per workgroup (llama rm_kq on RDNA)
const uint MAX_SG = 8u; // >= workgroup_size / min_subgroup_size
shared float sh[NR][MAX_SG];
// Signed int8 scale `j` (0..15) of the block at u32 base `b` (scales occupy bytes 192..207 = u32 48..51).
int scale_s(uint b, uint j) {
uint word = w[b + 48u + (j >> 2u)];
return bitfieldExtract(int(word), int((j & 3u) * 8u), 8);
}
void main() {
uint first_row = NR * gl_WorkGroupID.x;
uint nblocks = k / 256u;
uint tid = gl_LocalInvocationID.x;
uint itid = tid & 15u; // 0..15: sub-position within a super-block
uint ix = tid >> 4u; // 0..ITS-1: which super-block slot this thread group owns
uint v_im = itid >> 3u; // 0 or 1: which 128-weight half
uint v_in = itid & 7u; // 0..7
uint l0 = v_in << 2u; // 0,4,...,28
uint is = v_in >> 2u; // 0 or 1: scale sub-index
uint s_off = 8u * v_im + is; // first of the 4 scales this thread applies
uint qlw = 16u * v_im + v_in; // ql u32 index (low plane)
uint qhw = 32u + 8u * v_im + v_in;
uint yb0 = 128u * v_im + l0; // activation base offset within the block
float acc[NR];
[[unroll]] for (uint n = 0u; n < NR; n++) acc[n] = 0.0;
[[unroll]] for (uint n = 0u; n < NR; n++) {
uint row = first_row + n;
if (row >= nout) break;
uint rowbase = woff + row * nblocks * BLK_U32;
for (uint blk = ix; blk < nblocks; blk += ITS) {
uint b = rowbase + blk * BLK_U32;
uint ql0 = w[b + qlw];
uint ql32 = w[b + qlw + 8u];
uint qh = w[b + qhw];
float d = float(unpackHalf2x16(w[b + 52u]).x);
// ggml Q6_K decode: q = (ql_nibble | (qh_bitpair << 4)) - 32, four packed values per u32.
uint ql0lo = ql0 & 0x0F0F0F0Fu, ql0hi = (ql0 >> 4u) & 0x0F0F0F0Fu;
uint q32lo = ql32 & 0x0F0F0F0Fu, q32hi = (ql32 >> 4u) & 0x0F0F0F0Fu;
uint h0 = (qh & 0x03030303u) << 4u;
uint h2 = (qh & 0x0C0C0C0Cu) << 2u;
uint h4 = (qh & 0x30303030u);
uint h6 = (qh & 0xC0C0C0C0u) >> 2u;
uint Q0 = ql0lo | h0, Q1 = q32lo | h2, Q2 = ql0hi | h4, Q3 = q32hi | h6;
uint yb = blk * 256u + yb0;
float s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0;
[[unroll]] for (uint l = 0u; l < 4u; l++) {
uint sh8 = l * 8u;
s0 = fma(x[yb + l], float(int((Q0 >> sh8) & 0xFFu) - 32), s0);
s1 = fma(x[yb + 32u + l], float(int((Q1 >> sh8) & 0xFFu) - 32), s1);
s2 = fma(x[yb + 64u + l], float(int((Q2 >> sh8) & 0xFFu) - 32), s2);
s3 = fma(x[yb + 96u + l], float(int((Q3 >> sh8) & 0xFFu) - 32), s3);
}
float g = s0 * float(scale_s(b, s_off))
+ s1 * float(scale_s(b, s_off + 2u))
+ s2 * float(scale_s(b, s_off + 4u))
+ s3 * float(scale_s(b, s_off + 6u));
acc[n] = fma(d, g, acc[n]);
}
}
uint nsg = gl_NumSubgroups;
uint sgid = gl_SubgroupID;
[[unroll]] for (uint n = 0u; n < NR; n++) {
float v = subgroupAdd(acc[n]);
if (subgroupElect()) sh[n][sgid] = v;
}
barrier();
if (tid == 0u) {
[[unroll]] for (uint n = 0u; n < NR; n++) {
uint row = first_row + n;
if (row >= nout) continue;
float t = 0.0;
for (uint s = 0u; s < nsg; s++) t += sh[n][s];
y[row] = t;
}
}
}