#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
// ALL WG lanes participate, two lanes per column (BN=128, WG=256), each decoding 4 of the 8 q-words.
// The Q6_K coopmat op is dequant-ALU-bound (ql-nibble + qh-2bit combine per value), so decode ALU --
// not LDS feed -- gates it; vectorise the extraction (mask each 32-bit ql/qh word once for four lanes)
// and split the lanes across the whole workgroup. QL0/QH0 are 4-aligned so a K-step's ql and qh bytes
// are 8 contiguous u32 each; word wi covers kk in [4*wi,4*wi+4), never crossing the 16-wide sub-block
// boundary, so the scale is constant (dsc0 wi<4, dsc1 wi>=4). Stored f16 values are bit-identical to
// the per-byte form (same nibble+2bit assembly, same affine, same layout).
for (uint idx = tid; idx < BN * 2u; idx += WG) {
uint col = idx >> 1u;
uint w0 = (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 + 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
uint qlwb = base + (QL0 >> 2u);
uint qhwb = base + (QH0 >> 2u);
[[unroll]] for (uint u = 0u; u < 4u; u++) {
uint wi = w0 + u;
uint qlw = w[qlwb + wi];
uint qhw = w[qhwb + wi];
float dsc = (wi < 4u) ? dsc0 : dsc1;
uint nib4 = use_high ? ((qlw >> 4u) & 0x0F0F0F0Fu) : (qlw & 0x0F0F0F0Fu);
uint hi4 = ((qhw >> qh_shift) & 0x03030303u) << 4u;
uint q6x4 = nib4 | hi4; // four 6-bit codes, one per byte
float q0 = float(int( q6x4 & 0xFFu) - 32);
float q1 = float(int((q6x4 >> 8u) & 0xFFu) - 32);
float q2 = float(int((q6x4 >> 16u) & 0xFFu) - 32);
float q3 = float(int((q6x4 >> 24u) & 0xFFu) - 32);
s_b[col * SB_STRIDE + 2u * wi] = f16vec2(float16_t(dsc * q0), float16_t(dsc * q1));
s_b[col * SB_STRIDE + 2u * wi + 1u] = f16vec2(float16_t(dsc * q2), float16_t(dsc * q3));
}
} else {
[[unroll]] for (uint u = 0u; u < 4u; u++) {
uint wi = w0 + u;
s_b[col * SB_STRIDE + 2u * wi] = f16vec2(0.0);
s_b[col * SB_STRIDE + 2u * wi + 1u] = f16vec2(0.0);
}
}
}
}
// Square 2x2 warp grid (see mul_mm_q4k_coopmat.comp): each of the four warps owns a WMT*T x WNT*T (64x64)
// region so they split BOTH s_a rows and s_b columns two ways -- halving the redundant cold-weight s_b LDS
// reads vs a wide 4x1 stacking. 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) {
[[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);
}
}
}
}