#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). The first port
// ran at ~1/17 of llama's f16-coopmat throughput because the matrix cores were starved by a per-element
// weight decode. 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, each owning RM row tiles x RN col tiles of accumulator fragments;
// * DOUBLE-BUFFERED K loop: stage sub-block s+1 into the alternate LDS buffer while the matrix cores
// consume sub-block s, so global weight reads overlap coopMatMulAdd and there is one barrier/step;
// * 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
const uint T = 16u; // cooperative-matrix tile edge
const uint RM = 2u; // row tiles per warp
const uint RN = 8u; // column tiles per warp (BN = 128)
const uint NWARP = 4u; // 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 = 256, 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 + 4u; // 20 pairs
const uint SB_STRIDE = BK / 2u + 4u; // 20 pairs
shared f16vec2 s_a[2u][BM * SA_STRIDE]; // [M=128, K/2=16(pad20)] row-major, double-buffered
shared f16vec2 s_b[2u][BN * SB_STRIDE]; // [N=128, K/2=16(pad20)] row-major (transposed W), double-buffered
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 LDS buffer `buf`.
void stage(uint buf, 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[buf][r * SA_STRIDE + kp] = f16vec2(float16_t(lo), float16_t(hi));
} else {
s_a[buf][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;
bool lo = (sub & 1u) == 0u;
for (uint col = tid; col < BN; col += WG) {
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);
// Decode the sub-block's 32 nibbles into 16 f16vec2 pairs of adjacent k, one 32-bit store each.
[[unroll]] for (uint u = 0u; u < 8u; u++) {
uint qword = w[base + 4u + c * 8u + u];
uint b0 = (qword >> 0u) & 0xFFu;
uint b1 = (qword >> 8u) & 0xFFu;
uint b2 = (qword >> 16u) & 0xFFu;
uint b3 = (qword >> 24u) & 0xFFu;
uint q0 = lo ? (b0 & 0x0Fu) : (b0 >> 4u);
uint q1 = lo ? (b1 & 0x0Fu) : (b1 >> 4u);
uint q2 = lo ? (b2 & 0x0Fu) : (b2 >> 4u);
uint q3 = lo ? (b3 & 0x0Fu) : (b3 >> 4u);
float16_t v0 = float16_t(ds * float(q0) - ms);
float16_t v1 = float16_t(ds * float(q1) - ms);
float16_t v2 = float16_t(ds * float(q2) - ms);
float16_t v3 = float16_t(ds * float(q3) - ms);
s_b[buf][col * SB_STRIDE + u * 2u] = f16vec2(v0, v1);
s_b[buf][col * SB_STRIDE + u * 2u + 1u] = f16vec2(v2, v3);
}
} else {
[[unroll]] for (uint kp = 0u; kp < BK / 2u; kp++)
s_b[buf][col * SB_STRIDE + kp] = f16vec2(0.0);
}
}
}
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 nsteps = k / BK;
coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator> acc[RM][RN];
[[unroll]] for (uint i = 0u; i < RM; i++)
[[unroll]] for (uint j = 0u; j < RN; j++)
acc[i][j] = coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>(0.0);
stage(0u, 0u, row0, col0, nblocks);
barrier();
for (uint s = 0u; s < nsteps; s++) {
uint cur = s & 1u;
if (s + 1u < nsteps)
stage((s + 1u) & 1u, (s + 1u) * BK, row0, col0, nblocks);
if (warp < NWARP) {
// A fragments: [M-tile][K-subtile]. s_a is [M, K/2] f16vec2 -> K-subtile kc lands at pair
// offset kc*(T/2) and the row stride is SA_STRIDE pairs; RowMajor.
coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA> ma[RM][2];
[[unroll]] for (uint i = 0u; i < RM; i++)
[[unroll]] for (uint kc = 0u; kc < 2u; kc++)
coopMatLoad(ma[i][kc], s_a[cur], (warp * RM + i) * T * SA_STRIDE + kc * (T / 2u), SA_STRIDE,
gl_CooperativeMatrixLayoutRowMajor);
// Load the RN B-fragments for one K-subtile at a time and reuse the registers across the two
// K-subtiles. Halves the peak live B-fragment count (RN vs 2*RN) so the 16 f32 accumulators
// fit in registers without a scratch spill; math is identical (the two K-subtiles accumulate
// into acc in sequence rather than after a joint load). s_b is the transposed weight [N, K/2]
// f16vec2; ColumnMajor presents it as the [K, N] B fragment, K-subtile kc at pair kc*(T/2).
[[unroll]] for (uint kc = 0u; kc < 2u; kc++) {
coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB> mb[RN];
[[unroll]] for (uint j = 0u; j < RN; j++)
coopMatLoad(mb[j], s_b[cur], (j * T) * SB_STRIDE + kc * (T / 2u), SB_STRIDE,
gl_CooperativeMatrixLayoutColumnMajor);
[[unroll]] for (uint i = 0u; i < RM; i++)
[[unroll]] for (uint j = 0u; j < RN; j++)
acc[i][j] = coopMatMulAdd(ma[i][kc], mb[j], acc[i][j]);
}
}
barrier();
}
if (warp < NWARP) {
[[unroll]] for (uint i = 0u; i < RM; i++) {
uint orow = row0 + (warp * RM + i) * T;
[[unroll]] for (uint j = 0u; j < RN; j++) {
uint ocol = col0 + j * T;
if (orow < mcount && ocol < nout)
coopMatStore(acc[i][j], y, orow * nout + ocol, nout,
gl_CooperativeMatrixLayoutRowMajor);
}
}
}
}