#version 450
// Q4_K int8-dp4a decode matvec, TWO output rows per invocation (ILP + activation reuse). Same math as
// mul_mat_q4k_dp4a (the column dp4a, 1.8x) but each thread computes rows n0 and n0+1: the q8_1 int8
// activation (xq) + its scales are loaded ONCE per super-block and reused across both rows, and the two
// rows' dp4a chains are independent -> more memory latency hidden per thread. Decode == mul_mat_q4k.
// One invocation per output-row PAIR; host dispatches ceil(nout/2 / WG1D) workgroups. mcount fixed at 1
// (decode); woff selects an MoE bank slice. Tests whether the column dp4a is latency-bound (ILP helps)
// or pure-bandwidth-bound (flat).
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_spirv_intrinsics : require
spirv_instruction(extensions = ["SPV_KHR_integer_dot_product"], capabilities = [6018, 6019], id = 4453)
int sdot_accsat(int a, int b, int acc, spirv_literal int fmt);
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer W { uint w[]; };
layout(set = 0, binding = 1) readonly buffer XQ { uint xq[]; };
layout(set = 0, binding = 2) readonly buffer XS { float xs[]; };
layout(set = 0, binding = 3) readonly buffer XSUM { float xsum[]; };
layout(set = 0, binding = 4) writeonly buffer Y { float y[]; };
layout(push_constant) uniform Pc { uint m0; uint mcount; uint nout; uint k; uint woff; };
const uint QK_K = 256u;
const uint BLK_U32 = 36u;
const int FMT4x8 = 0;
uint scale_byte(uint base, uint b) {
uint word = w[base + 1u + (b >> 2u)];
return (word >> ((b & 3u) * 8u)) & 0xFFu;
}
uint get_scale_min_k4(uint base, uint j) {
uint sc, m;
if (j < 4u) {
sc = scale_byte(base, j) & 63u;
m = scale_byte(base, j + 4u) & 63u;
} else {
uint s_j4 = scale_byte(base, j + 4u);
uint s_jm4 = scale_byte(base, j - 4u);
uint s_j = scale_byte(base, j);
sc = (s_j4 & 0x0Fu) | ((s_jm4 >> 6u) << 4u);
m = (s_j4 >> 4u) | ((s_j >> 6u) << 4u);
}
return (sc << 8u) | m;
}
void main() {
uint n0 = gl_GlobalInvocationID.x * 2u; // first of the two output rows
if (n0 >= nout) { return; }
uint nr = min(2u, nout - n0); // 1 if nout is odd and this is the last pair
uint nblocks = k / QK_K;
float acc0 = 0.0;
float acc1 = 0.0;
uint rb0 = woff + n0 * nblocks * BLK_U32;
uint rb1 = rb0 + nblocks * BLK_U32;
for (uint blk = 0u; blk < nblocks; blk++) {
uint base0 = rb0 + blk * BLK_U32;
uint base1 = rb1 + blk * BLK_U32;
float d0 = float(unpackHalf2x16(w[base0]).x);
float dm0 = float(unpackHalf2x16(w[base0]).y);
float d1 = float(unpackHalf2x16(w[base1]).x);
float dm1 = float(unpackHalf2x16(w[base1]).y);
uint qb0 = base0 + 4u;
uint qb1 = base1 + 4u;
uint ablk = blk * 8u;
for (uint c = 0u; c < 4u; c++) {
uint is = c * 2u;
// per-row sub-block scales (lo = 2c, hi = 2c+1)
uint s0a = get_scale_min_k4(base0, is);
uint s0b = get_scale_min_k4(base0, is + 1u);
uint s1a = get_scale_min_k4(base1, is);
uint s1b = get_scale_min_k4(base1, is + 1u);
uint qc0 = qb0 + c * 8u;
uint qc1 = qb1 + c * 8u;
uint blo = ablk + is;
uint xqlo = blo * 8u;
uint xqhi = (blo + 1u) * 8u;
int d0lo = 0; int d0hi = 0;
int d1lo = 0; int d1hi = 0;
for (uint j = 0u; j < 8u; j++) {
uint a_lo = xq[xqlo + j]; // activation int8 codes, shared across both rows
uint a_hi = xq[xqhi + j];
uint w0 = w[qc0 + j];
uint w1 = w[qc1 + j];
d0lo = sdot_accsat(int(w0 & 0x0F0F0F0Fu), int(a_lo), d0lo, FMT4x8);
d0hi = sdot_accsat(int((w0 >> 4u) & 0x0F0F0F0Fu), int(a_hi), d0hi, FMT4x8);
d1lo = sdot_accsat(int(w1 & 0x0F0F0F0Fu), int(a_lo), d1lo, FMT4x8);
d1hi = sdot_accsat(int((w1 >> 4u) & 0x0F0F0F0Fu), int(a_hi), d1hi, FMT4x8);
}
float xslo = xs[blo]; float xsumlo = xsum[blo];
float xshi = xs[blo + 1u]; float xsumhi = xsum[blo + 1u];
acc0 += d0 * float(s0a >> 8u) * xslo * float(d0lo) - dm0 * float(s0a & 0xFFu) * xsumlo;
acc0 += d0 * float(s0b >> 8u) * xshi * float(d0hi) - dm0 * float(s0b & 0xFFu) * xsumhi;
acc1 += d1 * float(s1a >> 8u) * xslo * float(d1lo) - dm1 * float(s1a & 0xFFu) * xsumlo;
acc1 += d1 * float(s1b >> 8u) * xshi * float(d1hi) - dm1 * float(s1b & 0xFFu) * xsumhi;
}
}
y[n0] = acc0;
if (nr == 2u) { y[n0 + 1u] = acc1; }
}