#version 450
// Fused MoE grouped quant matvec (Q6_K): y[s, r] = sum_k W[ids[s], r, k] * x[s, k], reading the
// per-expert slice from a resident Q6_K weight bank [E, n, k]. Router gather + per-expert matvec in
// one dispatch. Each 210-byte GGUF Q6_K super-block (ql[128] low 4 bits; qh[64] high 2 bits;
// scales[16] i8; d f16) is host-repacked to a PADDED 212-byte (53 u32) stride via quantize_q6k --
// 210 is not u32-aligned -- so this reads the SAME layout and decode as mul_mat_q6k / mul_mat_vec_q6k
// (byte-exact with BlockQ6K::to_float). One invocation = one y element.
#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[]; }; // expert bank, 53 u32 (padded) / 256-block
layout(set = 0, binding = 1) readonly buffer X { float x[]; }; // [S, k] activations
layout(set = 0, binding = 2) readonly buffer Ids { uint ids[]; }; // [S] expert id per slot
layout(set = 0, binding = 3) writeonly buffer Y { float y[]; }; // [S, n] outputs
layout(push_constant) uniform Pc { uint n; uint k; uint nrows; }; // k mult of 256
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 gid = gl_GlobalInvocationID.x;
uint total = nrows * n;
if (gid >= total) {
return;
}
uint s = gid / n;
uint r = gid - s * n;
uint expert = ids[s];
uint nblocks = k / QK_K;
uint rowbase = (expert * n + r) * nblocks * BLK_U32; // u32 offset of this expert's weight row r
uint xbase = s * k;
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 = xbase + blk * QK_K;
// 2 super-chunks of 128 weights. 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 xb = 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[xb + l];
acc += s2 * float(q2) * x[xb + l + 32u];
acc += s3 * float(q3) * x[xb + l + 64u];
acc += s4 * float(q4) * x[xb + l + 96u];
}
}
}
y[gid] = acc;
}