#version 450
// Q6_K coopmat prefill matmul, the Q6_K twin of mul_mm_q4k_coopmat: y[m,n] = sum_k x[m,k]*W[n,k] on the
// f16 matrix cores (VK_KHR_cooperative_matrix). Tile/buffering/padding/f16vec2-packing are IDENTICAL to
// the Q4_K coopmat kernel; only the weight decode differs. Q6_K is SYMMETRIC (w = d*sc*q6, q6 in
// [-32,31], no min term), so the decode-into-LDS is simpler than Q4_K's affine `d*sc*nibble - dmin*m`.
// Structure that feeds the cores:
// * 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;
// * BK=32 K-step = TWO Q6_K 16-wide sub-blocks (Q6_K scales are per-16 vs Q4_K per-32), so a K-step
// fetches two signed scales: the first 16 k use scales[s0], the second 16 use scales[s0+1];
// * SINGLE-BUFFERED K loop (like the Q4_K kernel): the weight streams from GTT well below DRAM peak,
// so the double buffer's global-read/compute overlap is worthless while its 2x LDS caps occupancy.
// Q6_K's heavier decode makes it register-limited sooner, so halving LDS lifts the resident-wave
// count that hides its decode/LDS latency across waves rather than within one;
// * PACKED f16vec2 LDS: two adjacent-k f16 share one 32-bit LDS word (llama's FLOAT_TYPEV2), so the
// coopMatLoad feeding the cores issues one ds_load_b32 per pair instead of two ds_load_u16.
// Q6_K block = 210 B padded to 212 (53 u32): ql[128] low nibbles (bytes 0..127), qh[64] high 2 bits
// (128..191), scales[16] signed i8 (192..207), d f16 (byte 208 = low half of u32[52]). Decode position
// p (contiguous k) is byte-identical to mul_mat_q6k / mul_mm_q6k_tiled_dp4a and to BlockQ6K::to_float, so
// the correctness gate is the f16-rounding 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; // 2 Q6_K sub-blocks 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[]; }; // Q6_K blocks, 53 u32 (padded) / 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 = 53u; // 212 bytes (210 padded) / 4
const uint SC_BYTE = 192u; // scales[16] (signed i8)
const uint D_U32 = 52u; // d f16 at byte 208 -> low half of u32[52]
// 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 (breaks bank aliasing -- llama's SHMEM_STRIDE = BK/2 + 4).
const uint SA_STRIDE = BK / 2u + 4u; // 20 pairs
const uint SB_STRIDE = BK / 2u + 4u; // 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
// Signed byte `b` (0-based) within the block (sign-extended) -- the i8 scales.
int byte_s(uint base, uint b) {
uint word = w[base + (b >> 2u)];
return bitfieldExtract(int(word), int((b & 3u) * 8u), 8);
}
// 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. IDENTICAL to Q4_K coopmat.
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);
}
}
// Q6_K decode geometry for the 32 contiguous k of this K-step (BK == 32). koff is a 32-multiple, so
// the 32 k span two 16-wide sub-blocks s0 (even), s0+1 within one super-block; both share the 128-chunk
// `hh` and the quadrant, differing only in the is=0/1 half. Unified byte addressing over kk in [0,32):
// ql byte = 64*hh + (quad&1)*32 + kk (region [0,128))
// qh byte = 128 + 32*hh + kk (region [128,192))
// scale = scales[s0] for kk<16, scales[s0+1] for kk>=16
uint blk = koff / QK_K;
uint local = koff & (QK_K - 1u); // koff % 256, a 32-multiple
uint s0 = local >> 4u; // first sub-block index (even), 0..14
uint hh = s0 >> 3u; // 128-chunk: 0 or 1
uint quad = (s0 & 7u) >> 1u; // quadrant 0..3
bool use_high = quad >= 2u;
uint qh_shift = 2u * quad;
uint QL0 = 64u * hh + (quad & 1u) * 32u; // ql byte base at kk=0
uint QH0 = 128u + 32u * hh; // qh byte base at kk=0
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 + D_U32]).x);
float dsc0 = d * float(byte_s(base, SC_BYTE + s0)); // scale for kk 0..15
float dsc1 = d * float(byte_s(base, SC_BYTE + s0 + 1u)); // scale for kk 16..31
// WORD-load the ql/qh planes (llama's qword decode): QL0/QH0 are 4-aligned, so the 32 ql
// bytes and 32 qh bytes of this K-step are 8 contiguous u32 each. Load one ql word + one qh
// word per 4 k-values (vs a per-byte load), then extract the 4 lanes from registers -- fewer
// global loads and less register pressure than the per-element path. Each word covers kk in
// [4*wi, 4*wi+4), which never straddles the 16-wide sub-block boundary, so the scale is
// constant across the word: dsc0 for wi<4, dsc1 for wi>=4. Pair 2*wi holds (kk,kk+1).
uint qlwb = base + (QL0 >> 2u);
uint qhwb = base + (QH0 >> 2u);
[[unroll]] for (uint wi = 0u; wi < 8u; wi++) {
uint qlw = w[qlwb + wi];
uint qhw = w[qhwb + wi];
float dsc = (wi < 4u) ? dsc0 : dsc1;
float16_t v[4];
[[unroll]] for (uint b = 0u; b < 4u; b++) {
uint qlb = (qlw >> (b * 8u)) & 0xFFu;
uint qhb = (qhw >> (b * 8u)) & 0xFFu;
uint nib = use_high ? (qlb >> 4u) : (qlb & 0x0Fu);
uint hi = ((qhb >> qh_shift) & 3u) << 4u;
int q6 = int(nib | hi) - 32;
v[b] = float16_t(dsc * float(q6));
}
s_b[col * SB_STRIDE + 2u * wi] = f16vec2(v[0], v[1]);
s_b[col * SB_STRIDE + 2u * wi + 1u] = f16vec2(v[2], v[3]);
}
} else {
[[unroll]] for (uint kp = 0u; kp < BK / 2u; kp++)
s_b[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);
for (uint s = 0u; s < nsteps; s++) {
stage(s * BK, row0, col0, nblocks);
barrier();
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, (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 without a scratch spill. 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, (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);
}
}
}
}