268 lines
12 KiB
Mathematica
268 lines
12 KiB
Mathematica
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file fm_kernel_matvec_q3_k.m
|
|
* @brief Host dispatch and MSL source for the Q3_K matvec kernels.
|
|
*
|
|
* Written from a behaviour specification (docs/cleanroom/f_metal_spec.md); no ggml or
|
|
* llama.cpp source was consulted.
|
|
*
|
|
* @status None
|
|
* @issues _nx leaves most lanes idle at narrow cols (one row per group, lanes stride
|
|
* blocks by 32).
|
|
* @todo None
|
|
*/
|
|
|
|
#include "metal/kernels/fm_kernel_matvec_q3_k.h"
|
|
|
|
#include "metal/fm_metal_internal.h"
|
|
#include "metal/kernels/fm_kernel_matvec_lane_split.h"
|
|
#include "metal/kernels/fm_kernel_registry.h"
|
|
|
|
/* Shape checks shared by both entry points. The stride must hold a whole row of
|
|
* blocks; anything wider is padding the kernel skips. */
|
|
static bool fm_q3k_shape_ok(int rows, int cols, int row_stride_bytes) {
|
|
if (rows <= 0 || cols <= 0 || (cols % FM_Q3_K_BLOCK_WEIGHTS) != 0)
|
|
return false;
|
|
return row_stride_bytes >= (cols / FM_Q3_K_BLOCK_WEIGHTS) * FM_Q3_K_BLOCK_BYTES;
|
|
}
|
|
|
|
void fm_matvec_q3k_at(fm_bufref w, fm_bufref x, fm_bufref out,
|
|
int rows, int cols, int row_stride_bytes) {
|
|
if (!fm_metal_is_initialized() || !w.buffer || !x.buffer || !out.buffer)
|
|
return;
|
|
if (!fm_q3k_shape_ok(rows, cols, row_stride_bytes))
|
|
return;
|
|
|
|
/* Lanes per row is a function constant so the lane mask, rows per group and
|
|
* the butterfly fold at pipeline-build time. Set even when it is 32. */
|
|
int lanes = fm_matvec_lanes_per_row(cols / FM_Q3_K_BLOCK_WEIGHTS);
|
|
fm_fc_values fc;
|
|
fm_fc_reset(&fc);
|
|
if (fm_fc_set_int(&fc, FM_Q3_K_FC_LANES_PER_ROW, lanes) != 0)
|
|
return;
|
|
|
|
fm_bufref refs[] = { w, x, out };
|
|
uint32_t consts[] = { (uint32_t)rows, (uint32_t)cols, (uint32_t)row_stride_bytes };
|
|
fm_dispatch_threadgroups_refs(@"matvec_q3k", &fc, refs, 3, consts, 3,
|
|
fm_matvec_threadgroup_count(rows, lanes),
|
|
FM_MATVEC_SIMD_LANES);
|
|
}
|
|
|
|
void fm_matvec_q3k(fm_buffer *w, fm_buffer *x, fm_buffer *out,
|
|
int rows, int cols, int row_stride_bytes) {
|
|
if (!w || !x || !out)
|
|
return;
|
|
fm_matvec_q3k_at(fm_bufref_at(w, 0), fm_bufref_at(x, 0), fm_bufref_at(out, 0),
|
|
rows, cols, row_stride_bytes);
|
|
}
|
|
|
|
/* One _nx dispatch: `tiles` tiles of `tile_cols` columns starting at column col0,
|
|
* one threadgroup per (row, tile). */
|
|
static void fm_q3k_nx_dispatch(const fm_bufref *refs, int rows, int cols,
|
|
int row_stride_bytes, int n, int tile_cols,
|
|
int col0, int tiles) {
|
|
fm_fc_values fc;
|
|
fm_fc_reset(&fc);
|
|
if (fm_fc_set_int(&fc, FM_Q3_K_FC_NX_COLS, tile_cols) != 0)
|
|
return;
|
|
|
|
uint32_t consts[] = { (uint32_t)rows, (uint32_t)cols, (uint32_t)row_stride_bytes,
|
|
(uint32_t)n, (uint32_t)col0, (uint32_t)tiles };
|
|
fm_dispatch_threadgroups_refs(@"matvec_q3k_nx", &fc, refs, 3, consts, 6,
|
|
rows * tiles, FM_MATVEC_SIMD_LANES);
|
|
}
|
|
|
|
void fm_matvec_q3k_nx_at(fm_bufref w, fm_bufref x, fm_bufref out,
|
|
int rows, int cols, int row_stride_bytes, int n) {
|
|
if (!fm_metal_is_initialized() || !w.buffer || !x.buffer || !out.buffer)
|
|
return;
|
|
if (!fm_q3k_shape_ok(rows, cols, row_stride_bytes) || n <= 0)
|
|
return;
|
|
|
|
fm_bufref refs[] = { w, x, out };
|
|
int full_tiles = n / FM_Q3_K_NX_MAX_COLS;
|
|
int remainder = n % FM_Q3_K_NX_MAX_COLS;
|
|
|
|
if (full_tiles > 0)
|
|
fm_q3k_nx_dispatch(refs, rows, cols, row_stride_bytes, n,
|
|
FM_Q3_K_NX_MAX_COLS, 0, full_tiles);
|
|
if (remainder > 0)
|
|
fm_q3k_nx_dispatch(refs, rows, cols, row_stride_bytes, n,
|
|
remainder, full_tiles * FM_Q3_K_NX_MAX_COLS, 1);
|
|
}
|
|
|
|
void fm_matvec_q3k_nx(fm_buffer *w, fm_buffer *x, fm_buffer *out,
|
|
int rows, int cols, int row_stride_bytes, int n) {
|
|
if (!w || !x || !out)
|
|
return;
|
|
fm_matvec_q3k_nx_at(fm_bufref_at(w, 0), fm_bufref_at(x, 0), fm_bufref_at(out, 0),
|
|
rows, cols, row_stride_bytes, n);
|
|
}
|
|
|
|
NSString *fm_kernel_msl_matvec_q3_k(void)
|
|
{
|
|
return @
|
|
"/* ============================================================\n"
|
|
" * Q3_K fused dequant + matvec\n"
|
|
" *\n"
|
|
" * Block (110 bytes, 256 weights): hmask[32] | qs[64] | scales[12] | half d.\n"
|
|
" * Sub-block s (16 weights) = half u, level v, group g with s = 8u + 2v + g.\n"
|
|
" * Weight l of s: low two bits from qs[32u + 16g + l] at shift 2v, minus 4\n"
|
|
" * unless hmask[16g + l] bit (4u + v) is set; multiplier d * (S_s - 32).\n"
|
|
" *\n"
|
|
" * Every group is one SIMD group; no threadgroup memory. Fast math makes\n"
|
|
" * the bits depend on lane ownership, visit order and the reduction\n"
|
|
" * primitive, so those are fixed: simd_sum at 32 lanes per row, an xor\n"
|
|
" * butterfly below it.\n"
|
|
" * ============================================================ */\n"
|
|
"\n"
|
|
"constant int fm_q3k_nx_cols_fc [[function_constant(30)]];\n"
|
|
"constant int fm_q3k_nx_cols = is_function_constant_defined(fm_q3k_nx_cols_fc)\n"
|
|
" ? fm_q3k_nx_cols_fc : 8;\n"
|
|
"constant int fm_q3k_lanes_fc [[function_constant(31)]];\n"
|
|
"constant int fm_q3k_lanes = is_function_constant_defined(fm_q3k_lanes_fc)\n"
|
|
" ? fm_q3k_lanes_fc : 32;\n"
|
|
"\n"
|
|
"/* Six-bit scale of sub-block s, with i = s & 3 and quarter q = s >> 2: low\n"
|
|
" * four bits from scales[i + 4(q & 1)] (low nibble for q < 2, high nibble\n"
|
|
" * otherwise), top two bits from scales[8 + i] at shift 2q. */\n"
|
|
"static inline int fm_q3k_six_bit_scale(device const uchar* sc, uint s) {\n"
|
|
" uint i = s & 3u;\n"
|
|
" uint q = s >> 2;\n"
|
|
" int lo_byte = int(sc[i + 4u * (q & 1u)]);\n"
|
|
" int lo = (q < 2u) ? (lo_byte & 15) : (lo_byte >> 4);\n"
|
|
" int hi = (int(sc[8u + i]) >> (2u * q)) & 3;\n"
|
|
" return lo | (hi << 4);\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* The 16 sub-block multipliers d * (S_s - 32) of one block. */\n"
|
|
"static inline void fm_q3k_block_multipliers(device const uchar* blk, thread float* dl) {\n"
|
|
" float d = float(*(device const half*)(blk + 108));\n"
|
|
" for (uint s = 0; s < 16u; s++)\n"
|
|
" dl[s] = d * float(fm_q3k_six_bit_scale(blk + 96, s) - 32);\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* The 16 dequantized weights of sub-block s. */\n"
|
|
"static inline void fm_q3k_sub_block_weights(device const uchar* blk, uint s, float dl,\n"
|
|
" thread float* w) {\n"
|
|
" uint u = s >> 3;\n"
|
|
" uint v = (s & 7u) >> 1;\n"
|
|
" uint g = s & 1u;\n"
|
|
" device const uchar* qs = blk + 32u + 32u * u + 16u * g;\n"
|
|
" device const uchar* hm = blk + 16u * g;\n"
|
|
" uint shift = 2u * v;\n"
|
|
" uint mask = 1u << (4u * u + v);\n"
|
|
" for (uint l = 0; l < 16u; l++) {\n"
|
|
" int q2 = int((uint(qs[l]) >> shift) & 3u);\n"
|
|
" int quant = (uint(hm[l]) & mask) ? q2 : q2 - 4;\n"
|
|
" w[l] = dl * float(quant);\n"
|
|
" }\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* One block against 256 activations into a running chain: per sub-block a\n"
|
|
" * fresh 16-term partial, then added to acc. */\n"
|
|
"static inline float fm_q3k_block_dot(device const uchar* blk, device const float* xb,\n"
|
|
" float acc) {\n"
|
|
" float dl[16];\n"
|
|
" fm_q3k_block_multipliers(blk, dl);\n"
|
|
" for (uint s = 0; s < 16u; s++) {\n"
|
|
" float w[16];\n"
|
|
" fm_q3k_sub_block_weights(blk, s, dl[s], w);\n"
|
|
" float part = 0.0f;\n"
|
|
" for (uint l = 0; l < 16u; l++)\n"
|
|
" part += w[l] * xb[16u * s + l];\n"
|
|
" acc += part;\n"
|
|
" }\n"
|
|
" return acc;\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* L lanes per row, 32 / L rows per SIMD group; lane k of a row takes blocks\n"
|
|
" * k, k + L, ... Rows past the end accumulate 0 but still join the reduction. */\n"
|
|
"kernel void matvec_q3k(\n"
|
|
" device const uchar* W [[buffer(0)]],\n"
|
|
" device const float* x [[buffer(1)]],\n"
|
|
" device float* out [[buffer(2)]],\n"
|
|
" constant uint& rows [[buffer(3)]],\n"
|
|
" constant uint& cols [[buffer(4)]],\n"
|
|
" constant uint& row_stride [[buffer(5)]],\n"
|
|
" uint tg [[threadgroup_position_in_grid]],\n"
|
|
" uint tid [[thread_position_in_threadgroup]]\n"
|
|
") {\n"
|
|
" const uint L = uint(fm_q3k_lanes);\n"
|
|
" uint lane_in_row = tid & (L - 1u);\n"
|
|
" uint row = tg * (32u / L) + tid / L;\n"
|
|
" uint nb = cols / 256u;\n"
|
|
"\n"
|
|
" float acc = 0.0f;\n"
|
|
" if (row < rows) {\n"
|
|
" device const uchar* base = W + (ulong)row * row_stride;\n"
|
|
" for (uint b = lane_in_row; b < nb; b += L)\n"
|
|
" acc = fm_q3k_block_dot(base + (ulong)b * 110u, x + (ulong)b * 256u, acc);\n"
|
|
" }\n"
|
|
"\n"
|
|
" float row_sum;\n"
|
|
" if (L == 32u) {\n"
|
|
" row_sum = simd_sum(acc);\n"
|
|
" } else {\n"
|
|
" row_sum = acc;\n"
|
|
" for (uint off = L / 2u; off > 0u; off >>= 1)\n"
|
|
" row_sum += simd_shuffle_xor(row_sum, (ushort)off);\n"
|
|
" }\n"
|
|
" if (lane_in_row == 0u && row < rows) out[row] = row_sum;\n"
|
|
"}\n"
|
|
"\n"
|
|
"/* One SIMD group per (row, tile of NC token columns). Each block is decoded\n"
|
|
" * once per sub-block and applied to every live column of the tile. */\n"
|
|
"kernel void matvec_q3k_nx(\n"
|
|
" device const uchar* W [[buffer(0)]],\n"
|
|
" device const float* X [[buffer(1)]],\n"
|
|
" device float* out [[buffer(2)]],\n"
|
|
" constant uint& rows [[buffer(3)]],\n"
|
|
" constant uint& cols [[buffer(4)]],\n"
|
|
" constant uint& row_stride [[buffer(5)]],\n"
|
|
" constant uint& n [[buffer(6)]],\n"
|
|
" constant uint& col0 [[buffer(7)]],\n"
|
|
" constant uint& tiles [[buffer(8)]],\n"
|
|
" uint tg [[threadgroup_position_in_grid]],\n"
|
|
" uint tid [[thread_position_in_threadgroup]],\n"
|
|
" uint T [[threads_per_threadgroup]]\n"
|
|
") {\n"
|
|
" const uint NC = uint(fm_q3k_nx_cols);\n"
|
|
" uint row = tg / tiles;\n"
|
|
" if (row >= rows) return;\n"
|
|
" uint c0 = col0 + (tg - row * tiles) * NC;\n"
|
|
" uint nb = cols / 256u;\n"
|
|
"\n"
|
|
" float acc[8];\n"
|
|
" for (uint c = 0; c < NC; c++) acc[c] = 0.0f;\n"
|
|
"\n"
|
|
" device const uchar* base = W + (ulong)row * row_stride;\n"
|
|
" for (uint b = tid; b < nb; b += T) {\n"
|
|
" device const uchar* blk = base + (ulong)b * 110u;\n"
|
|
" float dl[16];\n"
|
|
" fm_q3k_block_multipliers(blk, dl);\n"
|
|
" for (uint s = 0; s < 16u; s++) {\n"
|
|
" float w[16];\n"
|
|
" fm_q3k_sub_block_weights(blk, s, dl[s], w);\n"
|
|
" for (uint c = 0; c < NC; c++) {\n"
|
|
" if (c0 + c >= n) continue;\n"
|
|
" device const float* xs = X + (ulong)(c0 + c) * cols + (ulong)b * 256u\n"
|
|
" + 16u * s;\n"
|
|
" float part = 0.0f;\n"
|
|
" for (uint l = 0; l < 16u; l++)\n"
|
|
" part += w[l] * xs[l];\n"
|
|
" acc[c] += part;\n"
|
|
" }\n"
|
|
" }\n"
|
|
" }\n"
|
|
"\n"
|
|
" for (uint c = 0; c < NC; c++) {\n"
|
|
" float col_sum = simd_sum(acc[c]);\n"
|
|
" if (tid == 0u && c0 + c < n) out[(ulong)(c0 + c) * rows + row] = col_sum;\n"
|
|
" }\n"
|
|
"}\n"
|
|
"\n"
|
|
;
|
|
}
|