/* SPDX-License-Identifier: AGPL-3.0-or-later */ /* Copyright (C) 2026 David (Corey) Herrera */ /** * @file fm_kernel_matvec_q4_k.m * @brief Host dispatch and MSL source for the GGUF Q4_K kernels: fused * dequant+matvec, the tiled simdgroup batched GEMM, the dequantize-to-F16 * helper, and the F16 MPS prefill GEMM built on top of it. * * @status None * @issues None * @todo None */ #include "metal/kernels/fm_kernel_matvec_q4_k.h" #import #include "metal/fm_metal_internal.h" #include "metal/kernels/fm_kernel_matvec_lane_split.h" #include "metal/kernels/fm_kernel_registry.h" /* Lanes one output row gets, for whichever of the two matvec kernels is being * encoded. Both call this: their bit-exactness contract is that they share a * decomposition, so they must share the arithmetic that defines it. */ static int fm_q4_k_lanes_per_row(int cols) { return fm_matvec_lanes_per_row(cols / FM_Q4_K_BLOCK_WEIGHTS); } void fm_matvec_q4_k_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; /* A row of `cols` weights is only cols/256 super-blocks, so at the MoE * expert shapes it cannot fill a SIMD group on its own. Give the row just * the lanes it can keep busy and pack several rows into the group instead; * the decode path is untouched, so no lane re-reads another lane's header. */ int lanes = fm_q4_k_lanes_per_row(cols); fm_fc_values fc; fm_fc_reset(&fc); if (fm_fc_set_int(&fc, FM_Q4_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 }; /* The weight offset is applied by the binding, so the kernel's own row * indexing is unchanged whether the tensor sits at the start of its own * buffer or deep inside a shared mmap region. */ fm_dispatch_threadgroups_refs(@"matvec_q4_k", &fc, refs, 3, consts, 3, fm_matvec_threadgroup_count(rows, lanes), FM_MATVEC_SIMD_LANES); } void fm_matvec_q4_k(fm_buffer *w, fm_buffer *x, fm_buffer *out, int rows, int cols, int row_stride_bytes) { if (!w || !x || !out) return; fm_matvec_q4_k_at(fm_bufref_at(w, 0), fm_bufref_at(x, 0), fm_bufref_at(out, 0), rows, cols, row_stride_bytes); } /* Encodes matvec_q4_k_nx for `tile_width` columns starting at `col_base`, * with the tile width bound as a function constant so the kernel's * accumulator loop has a compile-time bound. */ static void fm_matvec_q4_k_nx_tiles(fm_bufref w, fm_bufref x, fm_bufref out, int rows, int cols, int row_stride_bytes, int col_base, int tile_width, int tile_count) { int lanes = fm_q4_k_lanes_per_row(cols); fm_fc_values fc; fm_fc_reset(&fc); fm_fc_set_int(&fc, FM_Q4_K_NX_TILE_CONSTANT_INDEX, tile_width); if (fm_fc_set_int(&fc, FM_Q4_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, (uint32_t)col_base, (uint32_t)tile_count }; /* One threadgroup per (row group, column tile). Splitting the batch * across threadgroups rather than looping inside one keeps the grid * large enough to fill the GPU when rows is small. The row group is the * same lane split the single-column kernel uses — that shared * decomposition is what makes a batched column bit-identical to it. */ fm_dispatch_threadgroups_refs(@"matvec_q4_k_nx", &fc, refs, 3, consts, 5, fm_matvec_threadgroup_count(rows, lanes) * tile_count, FM_MATVEC_SIMD_LANES); } void fm_matvec_q4_k_nx(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 (rows <= 0 || cols <= 0 || n <= 0) return; /* At most two dispatches: the full-width tiles, then the remainder as * one narrower tile of its own width. Specializing the remainder * rather than padding it means no lane ever accumulates a column that * is thrown away. */ int full_tiles = n / FM_Q4_K_NX_MAX_TILE; int remainder = n % FM_Q4_K_NX_MAX_TILE; if (full_tiles > 0) fm_matvec_q4_k_nx_tiles(w, x, out, rows, cols, row_stride_bytes, 0, FM_Q4_K_NX_MAX_TILE, full_tiles); if (remainder > 0) fm_matvec_q4_k_nx_tiles(w, x, out, rows, cols, row_stride_bytes, full_tiles * FM_Q4_K_NX_MAX_TILE, remainder, 1); } void fm_sgemm_q4_k_at(fm_bufref w, fm_bufref x, fm_bufref out, int rows, int cols, int row_stride_bytes, int B) { if (!fm_metal_is_initialized() || !w.buffer || !x.buffer || !out.buffer) return; if (B <= 0) return; fm_bufref refs[] = { w, x, out }; uint32_t consts[] = { (uint32_t)rows, (uint32_t)cols, (uint32_t)row_stride_bytes, (uint32_t)B }; /* Tiled simdgroup_matrix GEMM: one threadgroup per 64-row tile, 256 threads * (8 simdgroups). Threadgroup memory is statically declared in the kernel. */ int ntiles = (rows + 63) / 64; fm_dispatch_threadgroups_refs(@"gemm_q4_k", NULL, refs, 3, consts, 4, ntiles, 256); } void fm_sgemm_q4_k(fm_buffer *w, fm_buffer *x, fm_buffer *out, int rows, int cols, int row_stride_bytes, int B) { if (!w || !x || !out) return; fm_sgemm_q4_k_at(fm_bufref_at(w, 0), fm_bufref_at(x, 0), fm_bufref_at(out, 0), rows, cols, row_stride_bytes, B); } /* Persistent F16 staging for the prefill GEMM (dequantized weight, converted * activations, F16 result), grown on demand and reused across calls. */ static fm_buffer *g_pf_w = NULL; static size_t g_pf_w_cap = 0; static fm_buffer *g_pf_x = NULL; static size_t g_pf_x_cap = 0; static fm_buffer *g_pf_o = NULL; static size_t g_pf_o_cap = 0; static fm_buffer *pf_stage(fm_buffer **b, size_t *cap, size_t bytes) { if (*b && *cap >= bytes) return *b; if (*b) fm_buffer_free(*b); *b = fm_buffer_alloc(bytes); *cap = *b ? bytes : 0; return *b; } void fm_gemm_q4_k_f16_at(fm_bufref w, fm_bufref x_f32, fm_bufref out_f32, int M, int N, int K, int row_stride_bytes) { if (!fm_metal_is_initialized() || !w.buffer || !x_f32.buffer || !out_f32.buffer) return; if (M <= 0 || N <= 0 || K <= 0) return; fm_buffer *wf = pf_stage(&g_pf_w, &g_pf_w_cap, (size_t)N * K * sizeof(uint16_t)); fm_buffer *xf = pf_stage(&g_pf_x, &g_pf_x_cap, (size_t)M * K * sizeof(uint16_t)); fm_buffer *of = pf_stage(&g_pf_o, &g_pf_o_cap, (size_t)M * N * sizeof(uint16_t)); if (!wf || !xf || !of) return; /* 1. Dequantize the Q4_K weight to F16 [N, K]. */ { fm_bufref refs[] = { w, fm_bufref_at(wf, 0) }; uint32_t c[] = { (uint32_t)N, (uint32_t)K, (uint32_t)row_stride_bytes }; fm_dispatch_1d_refs(@"dequant_q4_k_f16", NULL, refs, 2, c, 3, N * (K / 256)); } /* 2. Convert activations F32 [M, K] -> F16. Encoded here rather than * through fm_convert_f32_to_f16 only because that entry point binds at * offset 0, and x may be a slice of a larger staging buffer. */ { fm_bufref refs[] = { x_f32, fm_bufref_at(xf, 0) }; uint32_t c[] = { (uint32_t)(M * K) }; fm_dispatch_1d_refs(@"cvt_f32_f16", NULL, refs, 2, c, 1, M * K); } /* 3. F16 MPS matmul: of[M, N] = xf[M, K] @ wf[N, K]^T. */ MPSMatrixDescriptor *dA = [MPSMatrixDescriptor matrixDescriptorWithRows:M columns:K rowBytes:K * sizeof(uint16_t) dataType:MPSDataTypeFloat16]; MPSMatrixDescriptor *dB = [MPSMatrixDescriptor matrixDescriptorWithRows:N columns:K rowBytes:K * sizeof(uint16_t) dataType:MPSDataTypeFloat16]; MPSMatrixDescriptor *dC = [MPSMatrixDescriptor matrixDescriptorWithRows:M columns:N rowBytes:N * sizeof(uint16_t) dataType:MPSDataTypeFloat16]; MPSMatrix *mA = [[MPSMatrix alloc] initWithBuffer:xf->mtl descriptor:dA]; MPSMatrix *mB = [[MPSMatrix alloc] initWithBuffer:wf->mtl descriptor:dB]; MPSMatrix *mC = [[MPSMatrix alloc] initWithBuffer:of->mtl descriptor:dC]; MPSMatrixMultiplication *mm = [[MPSMatrixMultiplication alloc] initWithDevice:fm_metal_device() transposeLeft:NO transposeRight:YES resultRows:M resultColumns:N interiorColumns:K alpha:1.0 beta:0.0]; id cmdbuf = fm_get_command_buffer(); [mm encodeToCommandBuffer:cmdbuf leftMatrix:mA rightMatrix:mB resultMatrix:mC]; fm_submit_if_not_batching(); /* 4. Convert the F16 result back to F32 [M, N]. */ { fm_bufref refs[] = { fm_bufref_at(of, 0), out_f32 }; uint32_t c[] = { (uint32_t)(M * N) }; fm_dispatch_1d_refs(@"cvt_f16_f32", NULL, refs, 2, c, 1, M * N); } } void fm_gemm_q4_k_f16(fm_buffer *w, fm_buffer *x_f32, fm_buffer *out_f32, int M, int N, int K, int row_stride_bytes) { if (!w || !x_f32 || !out_f32) return; fm_gemm_q4_k_f16_at(fm_bufref_at(w, 0), fm_bufref_at(x_f32, 0), fm_bufref_at(out_f32, 0), M, N, K, row_stride_bytes); } NSString *fm_kernel_msl_matvec_q4_k(void) { return @ "/* ============================================================\n" " * Q4_K fused dequant + matvec (no scratch roundtrip)\n" " *\n" " * One threadgroup is one SIMD group (32 lanes) and covers\n" " * 32/FMQ4K_LANES output rows: a row gets as many lanes as it has\n" " * super-blocks, capped at 32. Weights stay GPU-resident as raw Q4_K\n" " * super-blocks (144 bytes / 256 weights), so only ~0.5625 bytes/weight are\n" " * read. Each lane dequantizes whole 256-element super-blocks (strided by\n" " * its row's lane count) and FMA-accumulates against x; a shuffle butterfly\n" " * reduces the row's partials. Header decode stays once per super-block —\n" " * splitting a row across MORE lanes than it has blocks is what makes eight\n" " * lanes re-decode one header, and that is the cost this avoids.\n" " * Mirrors pn_dequantize_block_q4_k exactly (value = d*sc*q - dmin*mn).\n" " * ============================================================ */\n" "\n" "/* Lanes one output row gets; 32 (one row per SIMD group) unspecialized.\n" " * Shared by matvec_q4_k and matvec_q4_k_nx — the two are bit-identical\n" " * per column only while they split a row the same way. */\n" "constant int fmq4k_lanes_in [[function_constant(81)]];\n" "constant uint FMQ4K_LANES =\n" " is_function_constant_defined(fmq4k_lanes_in) ? uint(fmq4k_lanes_in) : 32u;\n" "\n" "kernel void matvec_q4_k(\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 group [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]]\n" ") {\n" " const uint lanes = FMQ4K_LANES;\n" " uint lane_in_row = tid & (lanes - 1u);\n" " uint row = group * (32u / lanes) + tid / lanes;\n" " uint n_blocks = cols / 256u;\n" "\n" " // Lanes owning a row past the end accumulate nothing but still enter the\n" " // reduction, so all 32 lanes are active at the shuffles.\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 < n_blocks; b += lanes) {\n" " device const uchar* blk = base + (ulong)b * 144u;\n" " device const half* dh = (device const half*)blk;\n" " float d = float(dh[0]);\n" " float dmin = float(dh[1]);\n" " device const uchar* sr = blk + 4; // scales[12]\n" " device const uchar* qs = blk + 16; // qs[128]\n" " uchar sc[8]; uchar mn[8];\n" " for (uint j = 0; j < 4; j++) { sc[j] = sr[j] & 63; mn[j] = sr[j+4] & 63; }\n" " for (uint j = 4; j < 8; j++) {\n" " sc[j] = (sr[j+4] & 0xF) | ((sr[j-4] >> 6) << 4);\n" " mn[j] = (sr[j+4] >> 4) | ((sr[j] >> 6) << 4);\n" " }\n" " device const float* xb = x + (ulong)b * 256u;\n" " for (uint p = 0; p < 4; p++) { // 4 pairs share 32 qs bytes\n" " float d1 = d * float(sc[2*p]);\n" " float m1 = dmin * float(mn[2*p]);\n" " float d2 = d * float(sc[2*p+1]);\n" " float m2 = dmin * float(mn[2*p+1]);\n" " device const uchar* q = qs + p * 32u;\n" " uint o = p * 64u;\n" " for (uint l = 0; l < 32; l++) {\n" " uchar qb = q[l];\n" " acc += (d1 * float(qb & 0xF) - m1) * xb[o + l];\n" " acc += (d2 * float(qb >> 4) - m2) * xb[o + 32u + l];\n" " }\n" " }\n" " }\n" " }\n" " float row_sum = acc;\n" " if (lanes == 32u) {\n" " row_sum = simd_sum(acc);\n" " } else {\n" " for (uint off = lanes >> 1; 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" "/* ============================================================\n" " * Q4_K matvec batched over n activation columns.\n" " *\n" " * out[t][rows] = dequant(W) . X[t], with X = [n, cols] and\n" " * out = [n, rows]. The batch axis is the ACTIVATION COLUMNS; rows are\n" " * already the parallel axis. One threadgroup per (row group, column\n" " * tile); a lane decodes its super-blocks ONCE and FMAs each dequantized\n" " * weight into TILE accumulators, so the dequant amortizes across the\n" " * tile.\n" " *\n" " * TILE is a function constant, not an argument. That is the whole\n" " * point: with a runtime bound the accumulator loop cannot unroll, so\n" " * acc[] would live in thread-local memory instead of registers and\n" " * the batching would cost more than it saves. The host emits full\n" " * FM_Q4_K_NX_MAX_TILE tiles and, if n does not divide evenly, one\n" " * remainder tile specialized to its own exact width — so every column\n" " * a lane accumulates is a column that gets stored.\n" " *\n" " * The lane split, the block walk, the inner (p, l) order and the final\n" " * reduction are character-for-character the single-column kernel's, which\n" " * is what makes each column bit-identical to fm_matvec_q4_k_at on that\n" " * column. FMQ4K_LANES is shared with it for exactly that reason.\n" " * ============================================================ */\n" "\n" "constant int fmq4k_nx_tile_in [[function_constant(80)]];\n" "constant int FMQ4K_NX_TILE =\n" " is_function_constant_defined(fmq4k_nx_tile_in) ? fmq4k_nx_tile_in : 1;\n" "\n" "kernel void matvec_q4_k_nx(\n" " device const uchar* W [[buffer(0)]],\n" " device const float* X [[buffer(1)]],\n" // [n, cols] " device float* out [[buffer(2)]],\n" // [n, rows] " constant uint& rows [[buffer(3)]],\n" " constant uint& cols [[buffer(4)]],\n" " constant uint& row_stride [[buffer(5)]],\n" " constant uint& col_base [[buffer(6)]],\n" " constant uint& tile_count [[buffer(7)]],\n" " uint tgid [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]]\n" ") {\n" " // acc[] is declared at the widest tile so its size is a literal;\n" " // every loop over it stops at the specialized FMQ4K_NX_TILE, which\n" " // is a compile-time bound, so the array stays in registers.\n" " const uint TILE = (uint)FMQ4K_NX_TILE;\n" " const uint lanes = FMQ4K_LANES;\n" " uint lane_in_row = tid & (lanes - 1u);\n" " uint row = (tgid / tile_count) * (32u / lanes) + tid / lanes;\n" " uint t0 = col_base + (tgid % tile_count) * TILE;\n" "\n" " uint n_blocks = cols / 256u;\n" " device const float* xcol[8];\n" " for (uint t = 0; t < TILE; t++) xcol[t] = X + (ulong)(t0 + t) * cols;\n" "\n" " float acc[8];\n" " for (uint t = 0; t < TILE; t++) acc[t] = 0.0f;\n" "\n" " // Lanes owning a row past the end accumulate nothing but still enter the\n" " // reduction, so all 32 lanes are active at the shuffles.\n" " if (row < rows) {\n" " device const uchar* base = W + (ulong)row * row_stride;\n" " for (uint b = lane_in_row; b < n_blocks; b += lanes) {\n" " device const uchar* blk = base + (ulong)b * 144u;\n" " device const half* dh = (device const half*)blk;\n" " float d = float(dh[0]);\n" " float dmin = float(dh[1]);\n" " device const uchar* sr = blk + 4; // scales[12]\n" " device const uchar* qs = blk + 16; // qs[128]\n" " uchar sc[8]; uchar mn[8];\n" " for (uint j = 0; j < 4; j++) { sc[j] = sr[j] & 63; mn[j] = sr[j+4] & 63; }\n" " for (uint j = 4; j < 8; j++) {\n" " sc[j] = (sr[j+4] & 0xF) | ((sr[j-4] >> 6) << 4);\n" " mn[j] = (sr[j+4] >> 4) | ((sr[j] >> 6) << 4);\n" " }\n" " uint bcol = b * 256u;\n" " for (uint p = 0; p < 4; p++) { // 4 pairs share 32 qs bytes\n" " float d1 = d * float(sc[2*p]);\n" " float m1 = dmin * float(mn[2*p]);\n" " float d2 = d * float(sc[2*p+1]);\n" " float m2 = dmin * float(mn[2*p+1]);\n" " device const uchar* q = qs + p * 32u;\n" " uint o = bcol + p * 64u;\n" " for (uint l = 0; l < 32; l++) {\n" " uchar qb = q[l];\n" " for (uint t = 0; t < TILE; t++) {\n" " acc[t] += (d1 * float(qb & 0xF) - m1) * xcol[t][o + l];\n" " acc[t] += (d2 * float(qb >> 4) - m2) * xcol[t][o + 32u + l];\n" " }\n" " }\n" " }\n" " }\n" " }\n" "\n" " for (uint t = 0; t < TILE; t++) {\n" " float col_sum = acc[t];\n" " if (lanes == 32u) {\n" " col_sum = simd_sum(acc[t]);\n" " } else {\n" " for (uint off = lanes >> 1; off > 0u; off >>= 1)\n" " col_sum += simd_shuffle_xor(col_sum, ushort(off));\n" " }\n" " if (lane_in_row == 0u && row < rows)\n" " out[(ulong)(t0 + t) * rows + row] = col_sum;\n" " }\n" "}\n" "\n" "/* ============================================================\n" " * Q4_K x B GEMM (batched decode). out[B, rows] = X[B, cols] @ dequant(W)^T.\n" " * One SIMD group (32 lanes) per output row; each lane dequantizes its\n" " * assigned 256-blocks ONCE and FMAs every element into B accumulators, so\n" " * the Q4_K dequant is amortized across the whole batch (the GPU analog of\n" " * the CPU AMX panel). B is capped at 32 (accumulator registers).\n" " * ============================================================ */\n" "/* Tiled simdgroup_matrix GEMM.\n" " * - Each threadgroup owns ROWS_TILE=64 output rows and the full batch B.\n" " * - 256 threads = 8 simdgroups; simdgroup sg owns rows [sg*8, sg*8+8).\n" " * - Per KCHUNK=32 slice of cols: stage X into Xt[B][32] (f32, coalesced,\n" " * reused by all 8 simdgroups) and dequantize W into Wt[32][64] (f32,\n" " * transposed to [k][r], reused by all B tiles), then accumulate with\n" " * simdgroup_float8x8 (F32 accumulation, exact enough for err/RMS<1e-3).\n" " * - out[b,r] = sum_k X[b,k] * Wdeq[r,k]; C[b][r] = A[b][k] * Bm[k][r]. */\n" "kernel void gemm_q4_k(\n" " device const uchar* W [[buffer(0)]],\n" " device const float* X [[buffer(1)]],\n" // [B, cols] " device float* out [[buffer(2)]],\n" // [B, rows] " constant uint& rows [[buffer(3)]],\n" " constant uint& cols [[buffer(4)]],\n" " constant uint& row_stride [[buffer(5)]],\n" " constant uint& B [[buffer(6)]],\n" " uint tgid [[threadgroup_position_in_grid]],\n" " uint tid [[thread_position_in_threadgroup]],\n" " uint sgid [[simdgroup_index_in_threadgroup]],\n" " uint lane [[thread_index_in_simdgroup]]\n" ") {\n" " const uint ROWS_TILE = 64u;\n" " const uint KCHUNK = 64u;\n" " threadgroup float Xt[32u * 64u]; // [b][k], ld = KCHUNK\n" " threadgroup float Wt[64u * 64u]; // [k][r], ld = ROWS_TILE\n" " threadgroup float scratch[8u * 64u]; // per-simdgroup 8x8 spill\n" "\n" " uint rtile = tgid * ROWS_TILE;\n" " uint NB = (B + 7u) / 8u; // batch tiles (<=4)\n" " uint nchunks = cols / KCHUNK;\n" "\n" " simdgroup_float8x8 C[4];\n" " for (uint i = 0; i < NB; i++) C[i] = simdgroup_float8x8(0.0f);\n" "\n" " for (uint c = 0; c < nchunks; c++) {\n" " uint chunkbase = c * KCHUNK;\n" " // --- stage activations Xt[b][k] (coalesced device reads) ---\n" " for (uint i = tid; i < 32u * KCHUNK; i += 256u) {\n" " uint b = i / KCHUNK; uint k = i % KCHUNK;\n" " Xt[i] = (b < B) ? X[(ulong)b * cols + chunkbase + k] : 0.0f;\n" " }\n" " // --- dequantize W into Wt[k][r] (4 threads per row cover 64 k) ---\n" " {\n" " uint row = tid & 63u; // 0..63 -> output row in tile\n" " uint kg = tid >> 6; // 0..3 -> k group of 16\n" " uint grow = rtile + row;\n" " uint sb = chunkbase / 256u; // super-block index\n" " uint p = (chunkbase % 256u) / 64u; // sub-block pair 0..3\n" " if (grow < rows) {\n" " device const uchar* blk = W + (ulong)grow * row_stride + (ulong)sb * 144u;\n" " device const half* dh = (device const half*)blk;\n" " float d = float(dh[0]);\n" " float dmin = float(dh[1]);\n" " device const uchar* sr = blk + 4;\n" " device const uchar* qs = blk + 16;\n" " uint i0 = 2u * p, i1 = 2u * p + 1u;\n" " uchar s0, n0, s1, n1;\n" " if (i0 < 4u) { s0 = sr[i0] & 63u; n0 = sr[i0+4u] & 63u; }\n" " else { s0 = (sr[i0+4u] & 0xFu) | ((sr[i0-4u] >> 6) << 4);\n" " n0 = (sr[i0+4u] >> 4) | ((sr[i0] >> 6) << 4); }\n" " if (i1 < 4u) { s1 = sr[i1] & 63u; n1 = sr[i1+4u] & 63u; }\n" " else { s1 = (sr[i1+4u] & 0xFu) | ((sr[i1-4u] >> 6) << 4);\n" " n1 = (sr[i1+4u] >> 4) | ((sr[i1] >> 6) << 4); }\n" " float d1 = d * float(s0), m1 = dmin * float(n0);\n" " float d2 = d * float(s1), m2 = dmin * float(n1);\n" " device const uchar* q = qs + p * 32u;\n" " for (uint kk = 0; kk < 16u; kk++) {\n" " uint k = kg * 16u + kk;\n" " float val;\n" " if (k < 32u) { uchar qb = q[k]; val = d1 * float(qb & 0xF) - m1; }\n" " else { uchar qb = q[k - 32u]; val = d2 * float(qb >> 4) - m2; }\n" " Wt[k * 64u + row] = val;\n" " }\n" " } else {\n" " for (uint kk = 0; kk < 16u; kk++)\n" " Wt[(kg * 16u + kk) * 64u + row] = 0.0f;\n" " }\n" " }\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " // --- simdgroup matmul: accumulate this chunk ---\n" " for (uint kk = 0; kk < KCHUNK; kk += 8u) {\n" " simdgroup_float8x8 bm;\n" " simdgroup_load(bm, &Wt[kk * 64u + sgid * 8u], 64);\n" " for (uint nb = 0; nb < NB; nb++) {\n" " simdgroup_float8x8 a;\n" " simdgroup_load(a, &Xt[nb * 8u * KCHUNK + kk], KCHUNK);\n" " simdgroup_multiply_accumulate(C[nb], a, bm, C[nb]);\n" " }\n" " }\n" " threadgroup_barrier(mem_flags::mem_threadgroup);\n" " }\n" " // --- store C[b][r] -> out[b*rows + r] ---\n" " uint rbase = rtile + sgid * 8u;\n" " for (uint nb = 0; nb < NB; nb++) {\n" " uint bbase = nb * 8u;\n" " if (bbase + 8u <= B && rbase + 8u <= rows) {\n" " simdgroup_store(C[nb], out + (ulong)bbase * rows + rbase, rows);\n" " } else {\n" " simdgroup_store(C[nb], &scratch[sgid * 64u], 8);\n" " simdgroup_barrier(mem_flags::mem_threadgroup);\n" " for (uint idx = lane; idx < 64u; idx += 32u) {\n" " uint bi = idx / 8u, rj = idx % 8u;\n" " if (bbase + bi < B && rbase + rj < rows)\n" " out[(ulong)(bbase + bi) * rows + rbase + rj] = scratch[sgid * 64u + idx];\n" " }\n" " }\n" " }\n" "}\n" "\n" "/* Dequantize a Q4_K weight [rows, cols] to F16 (one thread per super-block). */\n" "kernel void dequant_q4_k_f16(\n" " device const uchar* W [[buffer(0)]],\n" " device half* out [[buffer(1)]],\n" " constant uint& rows [[buffer(2)]],\n" " constant uint& cols [[buffer(3)]],\n" " constant uint& row_stride [[buffer(4)]],\n" " uint gid [[thread_position_in_grid]]\n" ") {\n" " uint n_blocks = cols / 256u;\n" " if (gid >= rows * n_blocks) return;\n" " uint row = gid / n_blocks, b = gid % n_blocks;\n" " device const uchar* blk = W + (ulong)row * row_stride + (ulong)b * 144u;\n" " device const half* dh = (device const half*)blk;\n" " float d = float(dh[0]), dmin = float(dh[1]);\n" " device const uchar* sr = blk + 4; device const uchar* qs = blk + 16;\n" " uchar sc[8], mn[8];\n" " for (uint j=0;j<4;j++){ sc[j]=sr[j]&63; mn[j]=sr[j+4]&63; }\n" " for (uint j=4;j<8;j++){ sc[j]=(sr[j+4]&0xF)|((sr[j-4]>>6)<<4); mn[j]=(sr[j+4]>>4)|((sr[j]>>6)<<4); }\n" " device half* o = out + (ulong)row * cols + (ulong)b * 256u;\n" " for (uint p=0;p<4;p++){\n" " float d1=d*float(sc[2*p]), m1=dmin*float(mn[2*p]);\n" " float d2=d*float(sc[2*p+1]), m2=dmin*float(mn[2*p+1]);\n" " device const uchar* q = qs + p*32u; uint o0 = p*64u;\n" " for (uint l=0;l<32;l++){ uchar qb=q[l];\n" " o[o0+l] = half(d1*float(qb & 0xF) - m1);\n" " o[o0+32u+l] = half(d2*float(qb >> 4) - m2);\n" " }\n" " }\n" "}\n" "\n" ; }