#version 450
// Fused MoE grouped quant matvec (IQ3_S): y[s, r] = sum_k W[ids[s], r, k] * x[s, k], reading the
// per-expert slice from a GGML IQ3_S bank [E, n, k] in VRAM. Router gather + per-expert GEMV in one
// dispatch; the decode is the dense `mul_mat_vec_iq3s` one. Super-block = 110 bytes held at a padded
// 112-byte stride: d (f16), qs[64], qh[8], signs[32], scales[4]. Per sub-block e,
// db = d*(1+2*nibble); a 9-bit grid index from qs and qh picks 4 magnitudes, signs flip them
// (BlockIQ3s::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, IQ3_S at a 112-byte stride
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 a multiple of 256
uint rdbyte(uint bo) { return bitfieldExtract(w[bo >> 2u], int((bo & 3u) * 8u), 8); }
float rdscale(uint bo) { return unpackHalf2x16(rdbyte(bo) | (rdbyte(bo + 1u) << 8u)).x; }
const uint IQ3S_GRID[512] = uint[512](
0x01010101u, 0x01010103u, 0x01010105u, 0x0101010bu, 0x0101010fu, 0x01010301u, 0x01010303u, 0x01010305u, 0x01010309u,
0x0101030du, 0x01010501u, 0x01010503u, 0x0101050bu, 0x01010707u, 0x01010901u, 0x01010905u, 0x0101090bu, 0x0101090fu,
0x01010b03u, 0x01010b07u, 0x01010d01u, 0x01010d05u, 0x01010f03u, 0x01010f09u, 0x01010f0fu, 0x01030101u, 0x01030103u,
0x01030105u, 0x01030109u, 0x01030301u, 0x01030303u, 0x0103030bu, 0x01030501u, 0x01030507u, 0x0103050fu, 0x01030703u,
0x0103070bu, 0x01030909u, 0x01030d03u, 0x01030d0bu, 0x01030f05u, 0x01050101u, 0x01050103u, 0x0105010bu, 0x0105010fu,
0x01050301u, 0x01050307u, 0x0105030du, 0x01050503u, 0x0105050bu, 0x01050701u, 0x01050709u, 0x01050905u, 0x0105090bu,
0x0105090fu, 0x01050b03u, 0x01050b07u, 0x01050f01u, 0x01050f07u, 0x01070107u, 0x01070303u, 0x0107030bu, 0x01070501u,
0x01070505u, 0x01070703u, 0x01070707u, 0x0107070du, 0x01070909u, 0x01070b01u, 0x01070b05u, 0x01070d0fu, 0x01070f03u,
0x01070f0bu, 0x01090101u, 0x01090307u, 0x0109030fu, 0x01090503u, 0x01090509u, 0x01090705u, 0x01090901u, 0x01090907u,
0x01090b03u, 0x01090f01u, 0x010b0105u, 0x010b0109u, 0x010b0501u, 0x010b0505u, 0x010b050du, 0x010b0707u, 0x010b0903u,
0x010b090bu, 0x010b090fu, 0x010b0d0du, 0x010b0f07u, 0x010d010du, 0x010d0303u, 0x010d0307u, 0x010d0703u, 0x010d0b05u,
0x010d0f03u, 0x010f0101u, 0x010f0105u, 0x010f0109u, 0x010f0501u, 0x010f0505u, 0x010f050du, 0x010f0707u, 0x010f0b01u,
0x010f0b09u, 0x03010101u, 0x03010103u, 0x03010105u, 0x03010109u, 0x03010301u, 0x03010303u, 0x03010307u, 0x0301030bu,
0x0301030fu, 0x03010501u, 0x03010505u, 0x03010703u, 0x03010709u, 0x0301070du, 0x03010b09u, 0x03010b0du, 0x03010d03u,
0x03010f05u, 0x03030101u, 0x03030103u, 0x03030107u, 0x0303010du, 0x03030301u, 0x03030309u, 0x03030503u, 0x03030701u,
0x03030707u, 0x03030903u, 0x03030b01u, 0x03030b05u, 0x03030f01u, 0x03030f0du, 0x03050101u, 0x03050305u, 0x0305030bu,
0x0305030fu, 0x03050501u, 0x03050509u, 0x03050705u, 0x03050901u, 0x03050907u, 0x03050b0bu, 0x03050d01u, 0x03050f05u,
0x03070103u, 0x03070109u, 0x0307010fu, 0x03070301u, 0x03070307u, 0x03070503u, 0x0307050fu, 0x03070701u, 0x03070709u,
0x03070903u, 0x03070d05u, 0x03070f01u, 0x03090107u, 0x0309010bu, 0x03090305u, 0x03090309u, 0x03090703u, 0x03090707u,
0x03090905u, 0x0309090du, 0x03090b01u, 0x03090b09u, 0x030b0103u, 0x030b0301u, 0x030b0307u, 0x030b0503u, 0x030b0701u,
0x030b0705u, 0x030b0b03u, 0x030d0501u, 0x030d0509u, 0x030d050fu, 0x030d0909u, 0x030d090du, 0x030f0103u, 0x030f0107u,
0x030f0301u, 0x030f0305u, 0x030f0503u, 0x030f070bu, 0x030f0903u, 0x030f0d05u, 0x030f0f01u, 0x05010101u, 0x05010103u,
0x05010107u, 0x0501010bu, 0x0501010fu, 0x05010301u, 0x05010305u, 0x05010309u, 0x0501030du, 0x05010503u, 0x05010507u,
0x0501050fu, 0x05010701u, 0x05010705u, 0x05010903u, 0x05010907u, 0x0501090bu, 0x05010b01u, 0x05010b05u, 0x05010d0fu,
0x05010f01u, 0x05010f07u, 0x05010f0bu, 0x05030101u, 0x05030105u, 0x05030301u, 0x05030307u, 0x0503030fu, 0x05030505u,
0x0503050bu, 0x05030703u, 0x05030709u, 0x05030905u, 0x05030b03u, 0x05050103u, 0x05050109u, 0x0505010fu, 0x05050503u,
0x05050507u, 0x05050701u, 0x0505070fu, 0x05050903u, 0x05050b07u, 0x05050b0fu, 0x05050f03u, 0x05050f09u, 0x05070101u,
0x05070105u, 0x0507010bu, 0x05070303u, 0x05070505u, 0x05070509u, 0x05070703u, 0x05070707u, 0x05070905u, 0x05070b01u,
0x05070d0du, 0x05090103u, 0x0509010fu, 0x05090501u, 0x05090507u, 0x05090705u, 0x0509070bu, 0x05090903u, 0x05090f05u,
0x05090f0bu, 0x050b0109u, 0x050b0303u, 0x050b0505u, 0x050b070fu, 0x050b0901u, 0x050b0b07u, 0x050b0f01u, 0x050d0101u,
0x050d0105u, 0x050d010fu, 0x050d0503u, 0x050d0b0bu, 0x050d0d03u, 0x050f010bu, 0x050f0303u, 0x050f050du, 0x050f0701u,
0x050f0907u, 0x050f0b01u, 0x07010105u, 0x07010303u, 0x07010307u, 0x0701030bu, 0x0701030fu, 0x07010505u, 0x07010703u,
0x07010707u, 0x0701070bu, 0x07010905u, 0x07010909u, 0x0701090fu, 0x07010b03u, 0x07010d07u, 0x07010f03u, 0x07030103u,
0x07030107u, 0x0703010bu, 0x07030309u, 0x07030503u, 0x07030507u, 0x07030901u, 0x07030d01u, 0x07030f05u, 0x07030f0du,
0x07050101u, 0x07050305u, 0x07050501u, 0x07050705u, 0x07050709u, 0x07050b01u, 0x07070103u, 0x07070301u, 0x07070309u,
0x07070503u, 0x07070507u, 0x0707050fu, 0x07070701u, 0x07070903u, 0x07070907u, 0x0707090fu, 0x07070b0bu, 0x07070f07u,
0x07090107u, 0x07090303u, 0x0709030du, 0x07090505u, 0x07090703u, 0x07090b05u, 0x07090d01u, 0x07090d09u, 0x070b0103u,
0x070b0301u, 0x070b0305u, 0x070b050bu, 0x070b0705u, 0x070b0909u, 0x070b0b0du, 0x070b0f07u, 0x070d030du, 0x070d0903u,
0x070f0103u, 0x070f0107u, 0x070f0501u, 0x070f0505u, 0x070f070bu, 0x09010101u, 0x09010109u, 0x09010305u, 0x09010501u,
0x09010509u, 0x0901050fu, 0x09010705u, 0x09010903u, 0x09010b01u, 0x09010f01u, 0x09030105u, 0x0903010fu, 0x09030303u,
0x09030307u, 0x09030505u, 0x09030701u, 0x0903070bu, 0x09030907u, 0x09030b03u, 0x09030b0bu, 0x09050103u, 0x09050107u,
0x09050301u, 0x0905030bu, 0x09050503u, 0x09050707u, 0x09050901u, 0x09050b0fu, 0x09050d05u, 0x09050f01u, 0x09070109u,
0x09070303u, 0x09070307u, 0x09070501u, 0x09070505u, 0x09070703u, 0x0907070bu, 0x09090101u, 0x09090105u, 0x09090509u,
0x0909070fu, 0x09090901u, 0x09090f03u, 0x090b010bu, 0x090b010fu, 0x090b0503u, 0x090b0d05u, 0x090d0307u, 0x090d0709u,
0x090d0d01u, 0x090f0301u, 0x090f030bu, 0x090f0701u, 0x090f0907u, 0x090f0b03u, 0x0b010105u, 0x0b010301u, 0x0b010309u,
0x0b010505u, 0x0b010901u, 0x0b010909u, 0x0b01090fu, 0x0b010b05u, 0x0b010d0du, 0x0b010f09u, 0x0b030103u, 0x0b030107u,
0x0b03010bu, 0x0b030305u, 0x0b030503u, 0x0b030705u, 0x0b030f05u, 0x0b050101u, 0x0b050303u, 0x0b050507u, 0x0b050701u,
0x0b05070du, 0x0b050b07u, 0x0b070105u, 0x0b07010fu, 0x0b070301u, 0x0b07050fu, 0x0b070909u, 0x0b070b03u, 0x0b070d0bu,
0x0b070f07u, 0x0b090103u, 0x0b090109u, 0x0b090501u, 0x0b090705u, 0x0b09090du, 0x0b0b0305u, 0x0b0b050du, 0x0b0b0b03u,
0x0b0b0b07u, 0x0b0d0905u, 0x0b0f0105u, 0x0b0f0109u, 0x0b0f0505u, 0x0d010303u, 0x0d010307u, 0x0d01030bu, 0x0d010703u,
0x0d010707u, 0x0d010d01u, 0x0d030101u, 0x0d030501u, 0x0d03050fu, 0x0d030d09u, 0x0d050305u, 0x0d050709u, 0x0d050905u,
0x0d050b0bu, 0x0d050d05u, 0x0d050f01u, 0x0d070101u, 0x0d070309u, 0x0d070503u, 0x0d070901u, 0x0d09050bu, 0x0d090907u,
0x0d090d05u, 0x0d0b0101u, 0x0d0b0107u, 0x0d0b0709u, 0x0d0b0d01u, 0x0d0d010bu, 0x0d0d0901u, 0x0d0f0303u, 0x0d0f0307u,
0x0f010101u, 0x0f010109u, 0x0f01010fu, 0x0f010501u, 0x0f010505u, 0x0f01070du, 0x0f010901u, 0x0f010b09u, 0x0f010d05u,
0x0f030105u, 0x0f030303u, 0x0f030509u, 0x0f030907u, 0x0f03090bu, 0x0f050103u, 0x0f050109u, 0x0f050301u, 0x0f05030du,
0x0f050503u, 0x0f050701u, 0x0f050b03u, 0x0f070105u, 0x0f070705u, 0x0f07070bu, 0x0f070b07u, 0x0f090103u, 0x0f09010bu,
0x0f090307u, 0x0f090501u, 0x0f090b01u, 0x0f0b0505u, 0x0f0b0905u, 0x0f0d0105u, 0x0f0d0703u, 0x0f0f0101u
);
void main() {
uint gid = gl_GlobalInvocationID.x;
if (gid >= nrows * n) {
return;
}
uint s = gid / n;
uint r = gid - s * n;
uint nb = k / 256u;
uint rowbase = (ids[s] * n + r) * nb * 112u; // byte offset of row r of the routed expert
uint xbase = s * k;
float acc = 0.0;
for (uint b = 0u; b < nb; b++) {
uint bb = rowbase + b * 112u;
float d = rdscale(bb);
uint xb = xbase + b * 256u;
for (uint e = 0u; e < 8u; e++) {
uint sc = rdbyte(bb + 106u + (e >> 1u));
uint nib = ((e & 1u) != 0u) ? (sc >> 4u) : (sc & 0xFu);
uint qh = rdbyte(bb + 66u + e);
float blk = 0.0;
for (uint l = 0u; l < 4u; l++) {
uint signs = rdbyte(bb + 74u + 4u * e + l);
uint xo = xb + e * 32u + l * 8u;
for (uint m = 0u; m < 8u; m++) {
uint hf = m >> 2u;
uint shift = (m < 4u) ? (8u - 2u * l) : (7u - 2u * l);
uint idx = rdbyte(bb + 2u + 8u * e + 2u * l + hf) | ((qh << shift) & 256u);
float g = float((IQ3S_GRID[idx] >> (8u * (m & 3u))) & 0xFFu);
blk += (((signs >> m) & 1u) != 0u ? -g : g) * x[xo + m];
}
}
acc += d * (1.0 + 2.0 * float(nib)) * blk;
}
}
y[gid] = acc;
}