#version 450
// Subgroup-reduced Q8_0 matrix-vector product: y[n] = sum_k W[n,k]*x[k], W stored quantized
// (1 fp16 scale + 32 int8 per 32-block, 9 u32/block). Decode is memory-bandwidth bound on this
// ~256 GB/s APU; reading ~1.125 B/weight is the lever. This variant fuses the per-row reduction
// into a single subgroup arithmetic reduce (subgroupAdd) instead of one thread per row, so a whole
// subgroup streams one row's blocks in parallel — more memory-level parallelism per row and no
// shared-memory barrier. It is dispatched ONLY when the device advertises subgroup ARITHMETIC
// support (gated host-side in vulkan_backend.rs); the scalar mul_mat_vec_q8 kernel is the fallback.
#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 handling
// one output row. gl_NumSubgroups rows 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[]; }; // quantized weights, 9 u32 / 32-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/32)*9
// to select expert e of a resident MoE bank). k is a multiple of 32.
layout(push_constant) uniform Pc { uint nout; uint k; uint woff; };
void main() {
// One subgroup per output row. Row = global subgroup index.
uint row = gl_WorkGroupID.x * gl_NumSubgroups + gl_SubgroupID;
if (row >= nout) {
return;
}
uint nblocks = k / 32u;
uint base = woff + row * nblocks * 9u; // u32 offset of this row within the selected matrix
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 = 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 row.
float total = subgroupAdd(acc);
if (subgroupElect()) {
y[row] = total;
}
}