#version 450
// Subgroup-reduced fused grouped Q8_0 matvec for MoE decode: ONE dispatch computes every routed slot's
// output y[m, nout] where y[slot][col] = dequant(bank[ids[slot]])[col, :] . x_flat[slot, :]. This is the
// subgroup variant of mul_mat_vec_id_q8.comp: instead of one INVOCATION per (slot,col) reading that
// expert-row's blocks serially (adjacent invocations stride distinct rows -> uncoalesced reads), one
// SUBGROUP owns each output and its lanes stride that single row's blocks (consecutive lanes touch
// consecutive blocks -> coalesced), then fuse the partials with a single subgroupAdd. The Q8_0 block
// decode below is COPIED VERBATIM from mul_mat_vec_id_q8.comp / mul_mat_vec_q8_sg.comp and MUST stay
// bit-identical (a wrong reduction silently garbles all MoE output). Dispatched ONLY when the device
// advertises subgroup ARITHMETIC support (gated host-side in vulkan_backend.rs); the scalar
// mul_mat_vec_id_q8 kernel is the fallback. Bank is [E,n,k] Q8_0 blocks (9 u32 / 32-block); expert e
// starts at woff = ids[slot] * per_expert_words u32 words. k is a multiple of 32.
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_arithmetic : require
// 64 invocations/workgroup: on RDNA (subgroup size 32 or 64) that is 1-2 subgroups, each computing one
// output element. gl_NumSubgroups outputs are produced per workgroup.
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer W { uint w[]; }; // Q8_0 bank, 9 u32 / 32-block
layout(set = 0, binding = 1) readonly buffer X { float x[]; }; // x_flat [m, k] row-major
layout(set = 0, binding = 2) writeonly buffer Y { float y[]; }; // out [m, nout] row-major
layout(set = 0, binding = 3) readonly buffer IDS { uint ids[]; }; // per-slot expert id, length m
layout(push_constant) uniform Pc { uint nout; uint k; uint per_expert_words; uint m; };
void main() {
// One subgroup per output element. Global subgroup index g -> (slot, col), row-major over [m, nout].
uint g = gl_WorkGroupID.x * gl_NumSubgroups + gl_SubgroupID;
uint slot = g / nout;
uint col = g % nout;
if (slot >= m) {
return;
}
uint nblocks = k / 32u;
uint woff = ids[slot] * per_expert_words; // expert id[slot]'s block base in the bank
uint base = woff + col * nblocks * 9u; // u32 offset of row col within that expert
uint xbase = slot * k; // this slot's activation row in x_flat
uint lane = gl_SubgroupInvocationID;
uint lanes = gl_SubgroupSize;
// Each lane sums a strided subset of the 32-blocks; partials reduced across the subgroup.
float acc = 0.0;
for (uint b = lane; b < nblocks; b += lanes) {
uint off = base + b * 9u;
float scale = unpackHalf2x16(w[off]).x;
float bsum = 0.0;
uint xb = xbase + b * 32u;
for (uint j = 0u; j < 8u; j++) {
uint word = w[off + 1u + j];
uint xo = xb + j * 4u;
// bitfieldExtract on a signed int sign-extends the 8-bit lane.
bsum += float(bitfieldExtract(int(word), 0, 8)) * x[xo + 0u];
bsum += float(bitfieldExtract(int(word), 8, 8)) * x[xo + 1u];
bsum += float(bitfieldExtract(int(word), 16, 8)) * x[xo + 2u];
bsum += float(bitfieldExtract(int(word), 24, 8)) * x[xo + 3u];
}
acc += scale * bsum;
}
// Single subgroup reduction over all lanes' partials; lane 0 writes the output element.
float total = subgroupAdd(acc);
if (subgroupElect()) {
y[slot * nout + col] = total;
}
}