#version 450
// Subgroup-reduced Q6_K matrix-vector product: 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[128] low 4 bits; qh[64]
// high 2 bits; scales[16] i8; d f16); 210 is not u32-aligned, so the host repacks each block into a
// PADDED 212-byte (53 u32) stride (trailing 2 bytes unused). Decode is memory-bandwidth bound on this
// ~256 GB/s APU. This variant fuses the per-row reduction into one subgroup arithmetic reduce
// (subgroupAdd) instead of one thread per row, so a whole subgroup streams one row's super-blocks in
// parallel — more memory-level parallelism per row, no shared-memory barrier. Dispatched ONLY when
// the device advertises subgroup ARITHMETIC support (gated host-side); the scalar mul_mat_vec_q6k
// kernel is the fallback. Decode is bit-identical to mul_mat_vec_q6k (CPU k_quants BlockQ6K::to_float).
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#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, 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() {
// One subgroup per output row. Row = global subgroup index.
uint n = gl_WorkGroupID.x * gl_NumSubgroups + gl_SubgroupID;
if (n >= nout) {
return;
}
uint nblocks = k / QK_K;
uint rowbase = woff + n * nblocks * BLK_U32; // u32 offset of this row within the selected matrix
uint lane = gl_SubgroupInvocationID;
uint lanes = gl_SubgroupSize;
float acc = 0.0;
for (uint blk = lane; blk < nblocks; blk += lanes) {
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];
}
}
}
float total = subgroupAdd(acc);
if (subgroupElect()) {
y[n] = total;
}
}