// dense_gemm.metal — Dense F16 matrix multiply: C = A * B^T
//
// Portions of this file are modeled after llama.cpp's kernel_mul_mv_f16_f32
// (MIT licensed). Copyright the llama.cpp Authors. See LICENSE-MIT-llamacpp.
//
// Two kernels:
//
// 1. `dense_matvec_f16` — Specialized M=1 mat-vec (decode path).
// Modeled after llama.cpp's kernel_mul_mv_f16_f32 pattern.
// Each SIMD group (32 threads) processes N_DST rows using vectorized
// half4 loads and simd_sum reduction.
//
// 2. `dense_gemm_f16` — Fallback tiled GEMM for M>1.
// Uses shared memory tiling with simdgroup_matrix MMA (8x8 hardware
// accumulator) for high throughput.
//
// Buffer layout (both kernels):
// buffer(0): A — half [M, K]
// buffer(1): B — half [N, K]
// buffer(2): C — half [M, N] (output)
// buffer(3): params — DenseGemmParams {M, N, K}
#include <metal_stdlib>
#include <metal_simdgroup>
using namespace metal;
struct DenseGemmParams {
uint M;
uint N;
uint K;
};
// ============================================================
// Kernel 1: Specialized M=1 mat-vec — the hot path for lm_head decode
// ============================================================
//
// Architecture (matches llama.cpp pattern):
// - N_DST=4 rows per SIMD group, N_SIMDGROUP=2 per threadgroup -> 8 rows/tg
// - Each thread in a simdgroup handles K/(32*4) vectorized loads (half4)
// - simd_sum reduces across 32 lanes
// - Dispatch: threadgroups=(ceil(N/8), 1, 1), threads_per_tg=(32, 2, 1)
// (32 lanes x 2 simdgroups)
//
// For lm_head [1, 2816] x [262144, 2816]^T:
// 32768 threadgroups, 64 threads each, each tg outputs 8 values.
constant uint MV_N_DST = 4; // rows per SIMD group
constant uint MV_N_SIMDWIDTH = 32; // Apple GPU SIMD width
kernel void dense_matvec_f16(
device const half* A [[buffer(0)]],
device const half* B [[buffer(1)]],
device half* C [[buffer(2)]],
constant DenseGemmParams& params [[buffer(3)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint sgitg [[simdgroup_index_in_threadgroup]],
uint tiisg [[thread_index_in_simdgroup]]
) {
const uint K = params.K;
const uint N = params.N;
// Which row block this threadgroup handles
// N_SIMDGROUP = 2 (from dispatch: threads_per_tg.y = 2)
const uint row_base = tgpig.x * (MV_N_DST * 2);
// Which rows this simdgroup handles
const uint row0 = row_base + sgitg * MV_N_DST;
// Early exit if all rows are out of bounds
if (row0 >= N) return;
// The single input row (M=1)
device const half* a_ptr = A;
// Accumulators for N_DST rows (float for precision)
float acc0 = 0.0f;
float acc1 = 0.0f;
float acc2 = 0.0f;
float acc3 = 0.0f;
// Pointers to weight rows
device const half* b0 = B + (row0 + 0) * K;
device const half* b1 = B + min(row0 + 1, N - 1) * K;
device const half* b2 = B + min(row0 + 2, N - 1) * K;
device const half* b3 = B + min(row0 + 3, N - 1) * K;
// Main K-loop: vectorized half4 loads
// Each of 32 threads handles every 32nd group of 4 elements
const uint k_stride = MV_N_SIMDWIDTH * 4; // 128
uint k = tiisg * 4;
// Aligned loop (K divisible by k_stride handled automatically)
for (; k + 3 < K; k += k_stride) {
half4 av = *((device const half4*)(a_ptr + k));
float4 afv = float4(av);
half4 bv0 = *((device const half4*)(b0 + k));
acc0 += dot(afv, float4(bv0));
if (row0 + 1 < N) {
half4 bv1 = *((device const half4*)(b1 + k));
acc1 += dot(afv, float4(bv1));
}
if (row0 + 2 < N) {
half4 bv2 = *((device const half4*)(b2 + k));
acc2 += dot(afv, float4(bv2));
}
if (row0 + 3 < N) {
half4 bv3 = *((device const half4*)(b3 + k));
acc3 += dot(afv, float4(bv3));
}
}
// Scalar remainder — at most 3 elements per thread can remain,
// but since we stepped by k_stride (128) the remainder is at most
// K % 128 elements spread across 32 threads. Each thread handles
// at most ceil(remainder/32) elements.
for (; k < K; k += MV_N_SIMDWIDTH) {
if (k < K) {
float av = float(a_ptr[k]);
acc0 += av * float(b0[k]);
if (row0 + 1 < N) acc1 += av * float(b1[k]);
if (row0 + 2 < N) acc2 += av * float(b2[k]);
if (row0 + 3 < N) acc3 += av * float(b3[k]);
}
}
// SIMD reduction across 32 lanes
acc0 = simd_sum(acc0);
acc1 = simd_sum(acc1);
acc2 = simd_sum(acc2);
acc3 = simd_sum(acc3);
// Lane 0 writes results
if (tiisg == 0) {
C[row0 + 0] = half(acc0);
if (row0 + 1 < N) C[row0 + 1] = half(acc1);
if (row0 + 2 < N) C[row0 + 2] = half(acc2);
if (row0 + 3 < N) C[row0 + 3] = half(acc3);
}
}
// ============================================================
// Kernel 1b: Mixed-precision mat-vec — F16 weights × F32 input → F32 output
// ============================================================
//
// Identical to dense_matvec_f16 but:
// - A (input) is float* instead of half*
// - C (output) is float* instead of half*
// - B (weights) remains half*
// Eliminates the F32→F16 cast on input and F16→F32 cast on output.
kernel void dense_matvec_f16w_f32io(
device const float* A [[buffer(0)]],
device const half* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant DenseGemmParams& params [[buffer(3)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint sgitg [[simdgroup_index_in_threadgroup]],
uint tiisg [[thread_index_in_simdgroup]]
) {
const uint K = params.K;
const uint N = params.N;
const uint row_base = tgpig.x * (MV_N_DST * 2);
const uint row0 = row_base + sgitg * MV_N_DST;
if (row0 >= N) return;
device const float* a_ptr = A;
float acc0 = 0.0f;
float acc1 = 0.0f;
float acc2 = 0.0f;
float acc3 = 0.0f;
device const half* b0 = B + (row0 + 0) * K;
device const half* b1 = B + min(row0 + 1, N - 1) * K;
device const half* b2 = B + min(row0 + 2, N - 1) * K;
device const half* b3 = B + min(row0 + 3, N - 1) * K;
// Main K-loop: vectorized float4 loads for A, half4 for B weights
const uint k_stride = MV_N_SIMDWIDTH * 4; // 128
uint k = tiisg * 4;
for (; k + 3 < K; k += k_stride) {
float4 afv = *((device const float4*)(a_ptr + k));
half4 bv0 = *((device const half4*)(b0 + k));
acc0 += dot(afv, float4(bv0));
if (row0 + 1 < N) {
half4 bv1 = *((device const half4*)(b1 + k));
acc1 += dot(afv, float4(bv1));
}
if (row0 + 2 < N) {
half4 bv2 = *((device const half4*)(b2 + k));
acc2 += dot(afv, float4(bv2));
}
if (row0 + 3 < N) {
half4 bv3 = *((device const half4*)(b3 + k));
acc3 += dot(afv, float4(bv3));
}
}
// Scalar remainder
for (; k < K; k += MV_N_SIMDWIDTH) {
if (k < K) {
float av = a_ptr[k];
acc0 += av * float(b0[k]);
if (row0 + 1 < N) acc1 += av * float(b1[k]);
if (row0 + 2 < N) acc2 += av * float(b2[k]);
if (row0 + 3 < N) acc3 += av * float(b3[k]);
}
}
acc0 = simd_sum(acc0);
acc1 = simd_sum(acc1);
acc2 = simd_sum(acc2);
acc3 = simd_sum(acc3);
if (tiisg == 0) {
C[row0 + 0] = acc0;
if (row0 + 1 < N) C[row0 + 1] = acc1;
if (row0 + 2 < N) C[row0 + 2] = acc2;
if (row0 + 3 < N) C[row0 + 3] = acc3;
}
}
// ============================================================
// Kernel 2: Fallback tiled GEMM for M>1
// ============================================================
//
// Uses simdgroup_matrix for hardware-accelerated 8x8 MMA.
// 32x32 output tile, BK=16, 4 simdgroups (WM=2, WN=2).
//
// This is a simplified port of the MLX Steel GEMM for C = A * B^T.
// Attribution: Algorithm based on MLX (Apache-2.0), see
// https://github.com/ml-explore/mlx/blob/main/mlx/backend/metal/kernels/steel/gemm
constant uint BM = 32;
constant uint BN = 32;
constant uint BK = 16;
constant uint WM = 2; // simdgroups along M
constant uint WN = 2; // simdgroups along N
constant uint TGP_SIZE = WM * WN * 32; // 128 threads
// Simdgroup matrix tile strides
constant short GM_TM_stride = 8 * WM; // = 16
constant short GM_TN_stride = 8 * WN; // = 16
// Threadgroup memory sizes with padding to avoid bank conflicts
constant uint TGP_PAD = 8; // 16 / sizeof(half)
constant uint LDA_TGP = BK + TGP_PAD; // A: [BM, BK+pad] (A is not transposed)
constant uint LDB_TGP = BK + TGP_PAD; // B^T: [BN, BK+pad] (B is transposed)
constant uint TGP_MEM_A = BM * LDA_TGP;
constant uint TGP_MEM_B = BN * LDB_TGP;
// Number of simdgroup matrices per simdgroup sub-tile
constant uint GM_TM = 2; // 16/8 along M
constant uint GM_TN = 2; // 16/8 along N
kernel void dense_gemm_f16(
device const half* A [[buffer(0)]],
device const half* B [[buffer(1)]],
device half* C [[buffer(2)]],
constant DenseGemmParams& params [[buffer(3)]],
uint2 group_pos [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]],
uint tid_in_tg [[thread_index_in_threadgroup]]
) {
const uint M = params.M;
const uint N = params.N;
const uint K = params.K;
// Block position in output
const uint c_row = group_pos.y * BM;
const uint c_col = group_pos.x * BN;
// Early exit for out-of-bounds threadgroups
if (c_row >= M && c_col >= N) return;
// Threadgroup shared memory
threadgroup half As[TGP_MEM_A];
threadgroup half Bs[TGP_MEM_B];
// Advance A and B pointers to this block's starting position
device const half* A_block = A + c_row * K;
device const half* B_block = B + c_col * K;
// Simdgroup position within the WM x WN grid
const uint wm_id = sgid / WN; // 0 or 1
const uint wn_id = sgid % WN; // 0 or 1
// Determine thread position within simdgroup matrix layout
short qid = slid / 4;
short sm = (qid & 4) + (slid / 2) % 4;
short sn = (qid & 2) * 2 + (slid % 2) * 2;
// Simdgroup tile offsets
short tm = 8 * wm_id;
short tn = 8 * wn_id;
// Initialize result accumulators
simdgroup_matrix<float, 8, 8> results[GM_TM * GM_TN] = {
simdgroup_matrix<float, 8, 8>(0),
simdgroup_matrix<float, 8, 8>(0),
simdgroup_matrix<float, 8, 8>(0),
simdgroup_matrix<float, 8, 8>(0)
};
// Cooperative tile loading constants
const uint a_loads_per_thread = (BM * BK) / TGP_SIZE; // 4
const uint b_loads_per_thread = (BN * BK) / TGP_SIZE; // 4
const uint k_iterations = (K + BK - 1) / BK;
for (uint k_iter = 0; k_iter < k_iterations; k_iter++) {
const uint k_off = k_iter * BK;
const uint k_remain = min(BK, K - k_off);
threadgroup_barrier(mem_flags::mem_threadgroup);
// Load A tile [BM, BK] into shared memory
for (uint i = 0; i < a_loads_per_thread; i++) {
uint flat_idx = tid_in_tg + i * TGP_SIZE;
uint local_row = flat_idx / BK;
uint local_col = flat_idx % BK;
half val = 0;
if (c_row + local_row < M && local_col < k_remain) {
val = A_block[local_row * K + k_off + local_col];
}
As[local_row * LDA_TGP + local_col] = val;
}
// Load B tile [BN, BK] into shared memory
for (uint i = 0; i < b_loads_per_thread; i++) {
uint flat_idx = tid_in_tg + i * TGP_SIZE;
uint local_row = flat_idx / BK;
uint local_col = flat_idx % BK;
half val = 0;
if (c_col + local_row < N && local_col < k_remain) {
val = B_block[local_row * K + k_off + local_col];
}
Bs[local_row * LDB_TGP + local_col] = val;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// MMA: iterate over BK in steps of 8
simdgroup_matrix<float, 8, 8> Asimd[GM_TM];
simdgroup_matrix<float, 8, 8> Bsimd[GM_TN];
// A offsets: not transposed
short As_off = sn + (tm + sm) * (short)LDA_TGP;
// B offsets: transposed
short Bs_off = (tn + sn) * (short)LDB_TGP + sm;
threadgroup const half* As_ptr = As + As_off;
threadgroup const half* Bs_ptr = Bs + Bs_off;
for (uint kk = 0; kk < BK; kk += 8) {
simdgroup_barrier(mem_flags::mem_none);
// Load A simdgroup matrices
for (uint i = 0; i < GM_TM; i++) {
Asimd[i].thread_elements()[0] =
float(As_ptr[i * GM_TM_stride * (short)LDA_TGP + 0]);
Asimd[i].thread_elements()[1] =
float(As_ptr[i * GM_TM_stride * (short)LDA_TGP + 1]);
}
simdgroup_barrier(mem_flags::mem_none);
// Load B simdgroup matrices
for (uint j = 0; j < GM_TN; j++) {
Bsimd[j].thread_elements()[0] =
float(Bs_ptr[j * GM_TN_stride * (short)LDB_TGP + 0]);
Bsimd[j].thread_elements()[1] =
float(Bs_ptr[j * GM_TN_stride * (short)LDB_TGP + (short)LDB_TGP]);
}
simdgroup_barrier(mem_flags::mem_none);
// MMA with serpentine access pattern
for (uint i = 0; i < GM_TM; i++) {
for (uint j = 0; j < GM_TN; j++) {
short j_serp = (i % 2) ? (GM_TN - 1 - j) : j;
simdgroup_multiply_accumulate(
results[i * GM_TN + j_serp],
Asimd[i],
Bsimd[j_serp],
results[i * GM_TN + j_serp]);
}
}
As_ptr += 8;
Bs_ptr += 8;
}
}
// Store results
short tgp_bm = min((short)BM, (short)(M - c_row));
short tgp_bn = min((short)BN, (short)(N - c_col));
device half* D = C + (uint)(sm + tm + c_row) * N + (uint)(tn + sn + c_col);
for (uint i = 0; i < GM_TM; i++) {
for (uint j = 0; j < GM_TN; j++) {
thread const auto& accum = results[i * GM_TN + j].thread_elements();
int offset = (int)(i * GM_TM_stride) * (int)N + (int)(j * GM_TN_stride);
short out_row = sm + tm + (short)(i * GM_TM_stride);
short out_col = sn + tn + (short)(j * GM_TN_stride);
if (out_row < tgp_bm && out_col < tgp_bn) {
D[offset] = half(accum[0]);
}
if (out_row < tgp_bm && out_col + 1 < tgp_bn) {
D[offset + 1] = half(accum[1]);
}
}
}
}