#version 450
// Q4_K coopmat prefill matmul, faithful to llama's mul_mm.comp COOPMAT path (RADV gfx1151, wave64):
// y[m,n] = sum_k x[m,k] * W[n,k] on the f16 matrix cores (VK_KHR_cooperative_matrix). Structure that
// feeds the cores:
// * decode the Q4_K scale/min ONCE per 32-element sub-block, then unpack that sub-block's 32 nibbles;
// * BK=32 aligned to the Q4_K sub-block; BM=128 so each weight sub-block is decoded once per 128 rows;
// * WG=256 = 4 wave64 subgroups in a SQUARE 2x2 warp grid, each owning a WMT*T x WNT*T (64x64) region
// of the 128x128 block. The four warps split BOTH the s_a rows and the s_b columns two ways, so each
// warp reads only its 64 s_b columns -- vs a wide 4x1 stacking where every warp re-reads all 128 s_b
// columns, quadrupling the cold-weight LDS traffic that feeds the matrix cores (the measured bound);
// * SINGLE-BUFFERED K loop (llama's `shared FLOAT_TYPEV2 buf_a[...]`, no double buffer): the weight
// streams from GTT at a small fraction of DRAM peak, so a double buffer's global-read/compute
// overlap hides latency that is already hidden -- while its 2x LDS caps occupancy. Halving the LDS
// lets more subgroups reside per SIMD, and those extra waves (not intra-wave overlap) hide the
// decode + LDS-load latency. Two barriers per K-step (stage; compute; overwrite) hidden behind the
// added occupancy.
// * PACKED f16vec2 LDS: two adjacent-k f16 share one 32-bit LDS word (llama's FLOAT_TYPEV2). The
// coopMatLoad that feeds the cores then issues one 32-bit LDS read per pair (ds_load_b32) instead of
// two 16-bit reads (ds_load_u16), halving the LDS-read traffic that gates the matrix cores; the
// decode likewise writes one 32-bit store per pair. s_a is [M, K/2] loaded RowMajor; s_b is the
// transposed weight [N, K/2] loaded ColumnMajor -- so a K-subtile is contiguous in each row and the
// pair-packing lands on the load's fast axis.
// Per-element decode value is byte-identical to mul_mat_q4k / mul_mm_q4k_tiled_dp4a, so the correctness
// gate is the f16-rounding f64-oracle tolerance. Shapes: m,n multiples of 16, k a multiple of 256.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_control_flow_attributes : enable
// Tunable schedule (the autotune genome). Each knob is `#ifndef`-defaulted to the shipped value so a
// no-`-D` compile is byte-identical to the committed SPIR-V (the tuner's byte-identical gate); the
// tuner overrides them with `glslc -DHK_*` to compile a variant. RM/RN/NWARP set the accumulator tile
// geometry (BM = NWARP*RM*T rows, BN = RN*T cols, WG = NWARP*64 lanes); HK_PAD is the LDS
// anti-bank-aliasing stride pad. HK_WG is derived (NWARP*64) but named separately because a layout
// qualifier needs a bare literal. BK stays 32 (the Q4_K sub-block width the decode below is written
// against) and single-buffering stays fixed -- both are refuted knobs, carried only as autotune priors
// (deny-list), never realized here.
#ifndef HK_RM
#define HK_RM 2u
#endif
#ifndef HK_RN
#define HK_RN 8u
#endif
#ifndef HK_NWARP
#define HK_NWARP 4u
#endif
#ifndef HK_WG
#define HK_WG 256
#endif
#ifndef HK_PAD
#define HK_PAD 4u
#endif
const uint T = 16u; // cooperative-matrix tile edge
const uint RM = HK_RM; // row tiles per warp
const uint RN = HK_RN; // column tiles per warp (BN = 128)
const uint NWARP = HK_NWARP; // subgroups per workgroup; wave64 => WG/64 == NWARP
const uint BM = NWARP * RM * T; // 128
const uint BN = RN * T; // 128
const uint BK = 32u; // one Q4_K sub-block per K-step
const uint WG = NWARP * 64u; // 256
layout(local_size_x = HK_WG, local_size_y = 1, local_size_z = 1) in;
layout(set = 0, binding = 0) readonly buffer W { uint w[]; }; // Q4_K blocks, 36 u32 / 256-block
layout(set = 0, binding = 1) readonly buffer X { float x[]; }; // activations, row-major [M, k]
layout(set = 0, binding = 2) writeonly buffer Y { float y[]; }; // output, row-major [M, nout]
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;
// f16vec2-packed LDS: each element is a (k, k+1) pair. Row stride padded to BK/2 + 4 pairs so the
// coopMatLoad row stride is not a power of two (a stride of BK/2 = 16 pairs aliases every row onto the
// same bank set, throttling the loads that feed the matrix cores; +4 breaks the aliasing -- llama's
// mul_mm.comp uses the same SHMEM_STRIDE = BK/2 + 4). s_a holds activations [M, K/2]; s_b holds the
// transposed weight [N, K/2].
const uint SA_STRIDE = BK / 2u + HK_PAD; // 20 pairs
const uint SB_STRIDE = BK / 2u + HK_PAD; // 20 pairs
shared f16vec2 s_a[BM * SA_STRIDE]; // [M=128, K/2=16(pad20)] row-major, single-buffered
shared f16vec2 s_b[BN * SB_STRIDE]; // [N=128, K/2=16(pad20)] row-major (transposed W), single
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;
}
// Stage activation [BM,BK] and decode weight sub-block [BK,BN] for K-offset `koff` into the LDS buffer.
void stage(uint koff, uint row0, uint col0, uint nblocks) {
uint tid = gl_LocalInvocationID.x;
// Activation: pack 2 adjacent k into one f16vec2, so BM*(BK/2) pairs.
for (uint idx = tid; idx < BM * (BK / 2u); idx += WG) {
uint r = idx / (BK / 2u);
uint kp = idx - r * (BK / 2u);
uint mrow = row0 + r;
if (mrow < mcount) {
float lo = x[mrow * k + koff + 2u * kp];
float hi = x[mrow * k + koff + 2u * kp + 1u];
s_a[r * SA_STRIDE + kp] = f16vec2(float16_t(lo), float16_t(hi));
} else {
s_a[r * SA_STRIDE + kp] = f16vec2(0.0);
}
}
uint blk = koff / QK_K;
uint sub = (koff / BK) & 7u; // 0..7 within the 256-superblock (BK == 32)
uint c = sub >> 1u;
uint shift = ((sub & 1u) == 0u) ? 0u : 4u; // low or high nibble of every byte
// Weight decode: ALL WG lanes participate, two lanes per column (BN=128, WG=256), each decoding 4 of
// the 8 q-words (16 of the sub-block's 32 nibbles). The prior one-lane-per-column form left half the
// workgroup idle through the (dominant) decode phase; a prefill Q4_K GEMM is compute/feed-bound (the
// weight streams from GTT well below DRAM peak), so decode ALU competes with the matrix cores and
// spreading it across all lanes shortens the K-step. Nibble extraction is vectorised: one mask over
// the whole 32-bit q-word yields four nibbles, split into a vec4 -- no per-nibble branch. The stored
// f16 values are bit-identical to the one-lane form (same mask, same affine, same layout).
for (uint idx = tid; idx < BN * 2u; idx += WG) {
uint col = idx >> 1u;
uint u0 = (idx & 1u) * 4u; // this lane's q-word half: words 0..3 or 4..7
uint n = col0 + col;
if (n < nout) {
uint base = woff + n * nblocks * BLK_U32 + blk * BLK_U32;
float d = float(unpackHalf2x16(w[base]).x);
float dmin = float(unpackHalf2x16(w[base]).y);
uint sm = get_scale_min_k4(base, sub);
float ds = d * float(sm >> 8u);
float ms = dmin * float(sm & 0xFFu);
[[unroll]] for (uint u = 0u; u < 4u; u++) {
uint qword = w[base + 4u + c * 8u + u0 + u];
uint q4 = (qword >> shift) & 0x0F0F0F0Fu; // four nibbles, one mask
float q0 = float( q4 & 0xFFu);
float q1 = float((q4 >> 8u) & 0xFFu);
float q2 = float((q4 >> 16u) & 0xFFu);
float q3 = float((q4 >> 24u) & 0xFFu);
uint p = (u0 + u) * 2u;
s_b[col * SB_STRIDE + p] = f16vec2(float16_t(ds * q0 - ms), float16_t(ds * q1 - ms));
s_b[col * SB_STRIDE + p + 1u] = f16vec2(float16_t(ds * q2 - ms), float16_t(ds * q3 - ms));
}
} else {
[[unroll]] for (uint u = 0u; u < 4u; u++) {
uint p = (u0 + u) * 2u;
s_b[col * SB_STRIDE + p] = f16vec2(0.0);
s_b[col * SB_STRIDE + p + 1u] = f16vec2(0.0);
}
}
}
}
// Square warp grid: each warp owns a WMT*T x WNT*T region of the BM x BN block. With BM=BN=128 and 4
// warps this is a 2x2 grid of 64x64 warps, so the four warps split BOTH the s_a rows and the s_b columns
// two ways -- unlike the wide 4x1 stacking where every warp re-reads all BN s_b columns. Halves the
// redundant s_b (cold weight) LDS traffic that feeds the matrix cores. WMT*WNT == RM*RN accumulators.
const uint WMT = 4u;
const uint WNT = 4u;
const uint WGR = BM / (WMT * T); // warp-grid rows (2)
void main() {
uint nblocks = k / QK_K;
uint row0 = gl_WorkGroupID.y * BM;
uint col0 = gl_WorkGroupID.x * BN;
uint warp = gl_SubgroupID; // 0..NWARP-1 on wave64
uint warp_r = warp % WGR; // 0..1
uint warp_c = warp / WGR; // 0..1
uint nsteps = k / BK;
coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator> acc[WMT][WNT];
[[unroll]] for (uint i = 0u; i < WMT; i++)
[[unroll]] for (uint j = 0u; j < WNT; j++)
acc[i][j] = coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>(0.0);
for (uint s = 0u; s < nsteps; s++) {
stage(s * BK, row0, col0, nblocks);
barrier();
if (warp < NWARP) {
// Per K-subtile: load this warp's WMT A row-tiles once, then stream the WNT B col-tiles one at
// a time (one live B), reusing each across the WMT rows. s_a is [M, K/2] RowMajor; s_b is the
// transposed weight [N, K/2] presented ColumnMajor; K-subtile kc lands at pair kc*(T/2).
[[unroll]] for (uint kc = 0u; kc < 2u; kc++) {
coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA> ma[WMT];
[[unroll]] for (uint i = 0u; i < WMT; i++)
coopMatLoad(ma[i], s_a, (warp_r * WMT + i) * T * SA_STRIDE + kc * (T / 2u), SA_STRIDE,
gl_CooperativeMatrixLayoutRowMajor);
[[unroll]] for (uint j = 0u; j < WNT; j++) {
coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB> mb;
coopMatLoad(mb, s_b, (warp_c * WNT + j) * T * SB_STRIDE + kc * (T / 2u), SB_STRIDE,
gl_CooperativeMatrixLayoutColumnMajor);
[[unroll]] for (uint i = 0u; i < WMT; i++)
acc[i][j] = coopMatMulAdd(ma[i], mb, acc[i][j]);
}
}
}
barrier();
}
if (warp < NWARP) {
[[unroll]] for (uint i = 0u; i < WMT; i++) {
uint orow = row0 + (warp_r * WMT + i) * T;
[[unroll]] for (uint j = 0u; j < WNT; j++) {
uint ocol = col0 + (warp_c * WNT + j) * T;
if (orow < mcount && ocol < nout)
coopMatStore(acc[i][j], y, orow * nout + ocol, nout,
gl_CooperativeMatrixLayoutRowMajor);
}
}
}
}