#include <metal_stdlib>
using namespace metal;
// GEMM for Q5_K weights using simdgroup matrix multiply-accumulate.
//
// Same tiled `kernel_mul_mm` framework as gemm_q4_k (64×32 output tile, 8×8
// simdgroup ops, 128 threads / 4 simdgroups), specialized to the Q5_K
// super-block: nl = QK_K/16 = 16 sixteen-element tiles per 256-element block.
// Q5_K is Q4_K plus the `qh` high-bit plane and reuses its 6-bit scale/min
// packing verbatim, so the only differences from gemm_q4_k are the block struct,
// the 176-byte row stride, and the extra `qh` term in the tile decode.
// `dequantize_q5_K` decodes tile `il` (0..15) with the same bit layout as cera's
// `dequantize_q5_k_block` (quant.rs). Like gemm_q4_0/gemm_q8_0, the dequantized
// weights are rounded to `half` before the simdgroup matmul, so this matches the
// f32 per-token `gemv_q5_k` / CPU path within f16 tolerance, not bit-for-bit.
// The sibling comments add "argmax preserved"; that is established for them by
// model-level logit parity tests, and no Q5_K model test exists yet, so the
// claim is deliberately not repeated here.
//
// Dispatch: (ceil(n/32), ceil(m/64)) TGs × 128 threads. Threadgroup memory: 8 KB.
#define BLOCK_SIZE_M 64
#define BLOCK_SIZE_N 32
#define BLOCK_SIZE_K 32
#define THREAD_MAT_M 4
#define THREAD_MAT_N 2
#define THREAD_PER_ROW 2
#define THREAD_PER_COL 4
#define SG_MAT_SIZE 64
#define Q5K_NL 16
struct GemmParams {
uint m;
uint k;
uint n;
uint x_stride;
uint y_stride;
uint _pad;
};
struct block_q5_K {
half d;
half dmin;
uchar scales[12];
uchar qh[32];
uchar qs[128];
};
// (scale, min) for sub-block `j + k` — port of `decode_q4km_scales` in quant.rs,
// the same `get_scale_min_k4` used by llama.cpp's dequantize.
static inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {
return j < 4 ? uchar2{uchar(q[j + k] & 63), uchar(q[j + 4 + k] & 63)}
: uchar2{uchar((q[j + 4 + k] & 0xF) | ((q[j - 4 + k] >> 6) << 4)),
uchar((q[j + 4 + k] >> 4) | ((q[j + k] >> 6) << 4))};
}
// Decode the 16-element tile `il` (0..15) of a Q5_K super-block into `reg`.
// Tile il covers output elements [16*il, 16*il+16): sub-block
// `2*(il/4) + (il%4)/2`, qs base `32*(il/4) + 16*(il%2)`, low nibble for il%4<2
// else high. Matches `out[64j+l] = d*sc[2j]*((qs[32j+l]&0xF) + 16*bit) - dmin*mn[2j]`
// and the high-nibble sibling in `dequantize_q5_k_block`.
//
// The `qh` plane is indexed by `l` alone — the same 32 bytes serve all four `j`
// iterations, consuming bit pairs (2j, 2j+1) — so its offset is `16*(il%2)`,
// *without* the `32*(il/4)` term the `qs` pointer carries. The bit selector
// `1 << (il/2)` is taken before `il` is reduced mod 4, and reproduces the
// `u1`/`u2 <<= 2` walk of the CPU reference: il/2 == 2j + (nibble half).
//
// On the high-nibble tiles the mask is 0xF0 with `d/16`, so the high bit must be
// pre-scaled to match: 256/16 == 16, hence `qh_val` of 256 there against 16 on
// the low-nibble tiles.
void dequantize_q5_K(device const block_q5_K * xb, short il, thread half4x4 & reg) {
device const uchar * q = xb->qs;
device const uchar * qh = xb->qh;
short is = (il / 4) * 2;
q = q + (il / 4) * 32 + 16 * (il & 1);
qh = qh + 16 * (il & 1);
const uchar ul = 1 << (il / 2);
il = il & 3;
const uchar2 sc = get_scale_min_k4_just2(is, il / 2, xb->scales);
// High nibbles (il>=2) use mask 0xF0 with d/16 so dl*(q&0xF0) == d*sc*(q>>4).
const float d = il < 2 ? float(xb->d) : float(xb->d) / 16.0f;
const float mn = float(xb->dmin);
const float dl = d * float(sc[0]);
const float ml = mn * float(sc[1]);
const ushort mask = il < 2 ? 0x0F : 0xF0;
const float qh_val = il < 2 ? 16.0f : 256.0f;
float4x4 reg_f;
for (int i = 0; i < 16; i++) {
const float hb = (qh[i] & ul) ? qh_val : 0.0f;
reg_f[i / 4][i % 4] = dl * (float(q[i] & mask) + hb) - ml;
}
reg = (half4x4) reg_f;
}
kernel void gemm_q5_k(
const device uchar * src0 [[buffer(0)]],
const device float * src1 [[buffer(1)]],
device float * dst [[buffer(2)]],
constant GemmParams & params [[buffer(3)]],
threadgroup char * shmem [[threadgroup(0)]],
uint3 tgpig [[threadgroup_position_in_grid]],
ushort tiitg [[thread_index_in_threadgroup]],
ushort sgitg [[simdgroup_index_in_threadgroup]]
) {
const uint m = params.m;
const uint k = params.k;
const uint n = params.n;
const uint x_stride = params.x_stride;
const uint y_stride = params.y_stride;
const uint nb = k / 256;
const uint row_bytes = nb * 176;
const short nl = Q5K_NL;
threadgroup half * sa = (threadgroup half *)(shmem);
threadgroup float * sb = (threadgroup float *)(shmem + 4096);
const int r0 = tgpig.y;
const int r1 = tgpig.x;
const short n_rows = min((int)m - r0 * BLOCK_SIZE_M, BLOCK_SIZE_M);
const short n_cols = min((int)n - r1 * BLOCK_SIZE_N, BLOCK_SIZE_N);
const short thread_row = min((short)(tiitg / THREAD_PER_ROW), (short)(n_rows - 1));
const short thread_col = min((short)(tiitg / THREAD_PER_COL), (short)(n_cols - 1));
simdgroup_half8x8 ma[4];
simdgroup_float8x8 mb[2];
simdgroup_float8x8 mc[8];
for (short i = 0; i < 8; i++) {
mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);
}
short il = (tiitg % THREAD_PER_ROW);
device const block_q5_K * x = (device const block_q5_K *)(src0
+ row_bytes * (r0 * BLOCK_SIZE_M + thread_row)) + il / nl;
device const float * y = src1
+ x_stride * (r1 * BLOCK_SIZE_N + thread_col)
+ (BLOCK_SIZE_K / THREAD_PER_COL * (tiitg % THREAD_PER_COL));
for (uint loop_k = 0; loop_k < k; loop_k += BLOCK_SIZE_K) {
half4x4 temp_a;
dequantize_q5_K(x, il, temp_a);
threadgroup_barrier(mem_flags::mem_threadgroup);
#pragma unroll(16)
for (short i = 0; i < 16; i++) {
*(sa + SG_MAT_SIZE * ((tiitg / THREAD_PER_ROW / 8)
+ (tiitg % THREAD_PER_ROW) * 16 + (i / 8) * 8)
+ (tiitg / THREAD_PER_ROW) % 8 + (i & 7) * 8) = temp_a[i/4][i%4];
}
*(threadgroup float2x4 *)(sb + 32 * 8 * (tiitg % THREAD_PER_COL) + 8 * (tiitg / THREAD_PER_COL)) = *((device float2x4 *) y);
il = (il + 2 < nl) ? il + 2 : il % 2;
x = (il < 2) ? x + (2 + nl - 1) / nl : x;
y += BLOCK_SIZE_K;
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup const half * lsma = (sa + THREAD_MAT_M * SG_MAT_SIZE * (sgitg % 2));
threadgroup const float * lsmb = (sb + THREAD_MAT_N * SG_MAT_SIZE * (sgitg / 2));
#pragma unroll(4)
for (short ik = 0; ik < BLOCK_SIZE_K / 8; ik++) {
simdgroup_barrier(mem_flags::mem_none);
#pragma unroll(4)
for (short i = 0; i < 4; i++) {
simdgroup_load(ma[i], lsma + SG_MAT_SIZE * i);
}
#pragma unroll(2)
for (short i = 0; i < 2; i++) {
simdgroup_load(mb[i], lsmb + SG_MAT_SIZE * i);
}
simdgroup_barrier(mem_flags::mem_none);
#pragma unroll(8)
for (short i = 0; i < 8; i++) {
simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);
}
lsma += (BLOCK_SIZE_M / 8) * SG_MAT_SIZE;
lsmb += (BLOCK_SIZE_N / 8) * SG_MAT_SIZE;
}
}
if ((r0 + 1) * BLOCK_SIZE_M <= (int)m && (r1 + 1) * BLOCK_SIZE_N <= (int)n) {
// Fast path: full tile, no accumulate — direct simdgroup store.
device float * C = dst
+ (BLOCK_SIZE_M * r0 + 32 * (sgitg & 1))
+ (BLOCK_SIZE_N * r1 + 16 * (sgitg >> 1)) * y_stride;
for (short i = 0; i < 8; i++) {
simdgroup_store(mc[i], C + 8 * (i % 4) + 8 * y_stride * (i / 4), y_stride);
}
} else {
threadgroup_barrier(mem_flags::mem_threadgroup);
threadgroup float * temp_str = ((threadgroup float *) shmem)
+ 32 * (sgitg & 1) + (16 * (sgitg >> 1)) * BLOCK_SIZE_M;
for (short i = 0; i < 8; i++) {
simdgroup_store(mc[i], temp_str + 8 * (i % 4) + 8 * BLOCK_SIZE_M * (i / 4), BLOCK_SIZE_M);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (sgitg == 0) {
for (int j = tiitg; j < n_cols; j += BLOCK_SIZE_N) {
device float * D = dst + (r0 * BLOCK_SIZE_M) + (r1 * BLOCK_SIZE_N + j) * y_stride;
threadgroup float * S = ((threadgroup float *) shmem) + (j * BLOCK_SIZE_M);
device float4 * D4 = (device float4 *) D;
threadgroup float4 * S4 = (threadgroup float4 *) S;
int i = 0;
for (; i < n_rows / 4; i++) {
*(D4 + i) = *(S4 + i);
}
i *= 4;
for (; i < n_rows; i++) {
*(D + i) = *(S + i);
}
}
}
}
}