use crate::device::metal_device;
use metal::{Buffer, ComputePipelineState, Library, MTLResourceOptions};
use std::sync::OnceLock;
pub const RLX_KERNELS_MSL: &str = r#"
#include <metal_stdlib>
using namespace metal;
// Naive sgemm: one thread per output element, one dot product each.
// C[m,n] = A[m,k] @ B[k,n]. Good baseline; tiled version below for speed.
kernel void sgemm(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint row = gid.y;
uint col = gid.x;
if (row >= M || col >= N) return;
float sum = 0.0;
for (uint k = 0; k < K; ++k) {
sum += A[row * K + k] * B[k * N + col];
}
C[row * N + col] = sum;
}
// ── Half-precision (f16) variants ──────────────────────────────────────
// Apple Silicon supports simdgroup_half8x8 — same tensor unit pipeline
// but 2× peak FLOPs and ½ memory bandwidth vs simdgroup_float8x8.
// Tiled half-precision matmul: 32x32 output per TG, 16 simdgroups cooperate.
// Inputs A, B and output C all in f16; bias also f16 if provided.
kernel void hgemm_simd_4x4(
device const half* A [[buffer(0)]],
device const half* B [[buffer(1)]],
device half* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]]
) {
uint sg_row = sgid / 4;
uint sg_col = sgid % 4;
uint tg_row_base = tgid.y * 32;
uint tg_col_base = tgid.x * 32;
threadgroup half A_tg[32 * 32];
threadgroup half B_tg[32 * 32];
simdgroup_half8x8 a, b;
simdgroup_half8x8 c = simdgroup_half8x8(0.0h);
for (uint kk = 0; kk < K; kk += 32) {
uint linear = sgid * 32 + slid;
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint ar = idx / 32, ac = idx % 32;
A_tg[idx] = A[(tg_row_base + ar) * K + (kk + ac)];
}
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint br = idx / 32, bc = idx % 32;
B_tg[idx] = B[(kk + br) * N + (tg_col_base + bc)];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k_inner = 0; k_inner < 32; k_inner += 8) {
simdgroup_load(a, &A_tg[sg_row * 8 * 32 + k_inner], 32);
simdgroup_load(b, &B_tg[k_inner * 32 + sg_col * 8], 32);
simdgroup_multiply_accumulate(c, a, b, c);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
uint out_row = tg_row_base + sg_row * 8;
uint out_col = tg_col_base + sg_col * 8;
simdgroup_store(c, &C[out_row * N + out_col], N);
}
// Half-precision matmul + bias + activation fused.
kernel void hgemm_simd_4x4_bias(
device const half* A [[buffer(0)]],
device const half* B [[buffer(1)]],
device const half* bias [[buffer(2)]],
device half* C [[buffer(3)]],
constant uint& M [[buffer(4)]],
constant uint& K [[buffer(5)]],
constant uint& N [[buffer(6)]],
constant uint& act_kind [[buffer(7)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]]
) {
uint sg_row = sgid / 4;
uint sg_col = sgid % 4;
uint tg_row_base = tgid.y * 32;
uint tg_col_base = tgid.x * 32;
threadgroup half A_tg[32 * 32];
threadgroup half B_tg[32 * 32];
simdgroup_half8x8 a, b;
simdgroup_half8x8 c = simdgroup_half8x8(0.0h);
for (uint kk = 0; kk < K; kk += 32) {
uint linear = sgid * 32 + slid;
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint ar = idx / 32, ac = idx % 32;
A_tg[idx] = A[(tg_row_base + ar) * K + (kk + ac)];
}
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint br = idx / 32, bc = idx % 32;
B_tg[idx] = B[(kk + br) * N + (tg_col_base + bc)];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k_inner = 0; k_inner < 32; k_inner += 8) {
simdgroup_load(a, &A_tg[sg_row * 8 * 32 + k_inner], 32);
simdgroup_load(b, &B_tg[k_inner * 32 + sg_col * 8], 32);
simdgroup_multiply_accumulate(c, a, b, c);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
threadgroup half tile[16 * 64];
simdgroup_store(c, &tile[sgid * 64], 8);
threadgroup_barrier(mem_flags::mem_threadgroup);
uint out_row_base = tg_row_base + sg_row * 8;
uint out_col_base = tg_col_base + sg_col * 8;
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint r = idx / 8;
uint cc = idx % 8;
// Promote to fp32 for activation math (more accurate)
float v = float(tile[sgid * 64 + idx]) + float(bias[out_col_base + cc]);
if (act_kind == 1) {
float arg = v * 0.7071067811865475f;
float sign = arg >= 0.0f ? 1.0f : -1.0f;
float xa = abs(arg);
float t = 1.0f / (1.0f + 0.3275911f * xa);
float y = t * (0.254829592f + t * (-0.284496736f + t * (1.421413741f
+ t * (-1.453152027f + t * 1.061405429f))));
float erf_val = sign * (1.0f - y * exp(-xa * xa));
v = v * 0.5f * (1.0f + erf_val);
} else if (act_kind == 2) {
v = v / (1.0f + exp(-v));
}
C[(out_row_base + r) * N + (out_col_base + cc)] = half(v);
}
}
// ── Half-precision element-wise + reduction kernels ─────────────────
kernel void bias_add_h(
device half* data [[buffer(0)]],
device const half* bias [[buffer(1)]],
constant uint& m [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]
) {
uint row = gid.y, col = gid.x;
if (row >= m || col >= n) return;
data[row * n + col] += bias[col];
}
kernel void gelu_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
// Promote to f32 for math (more accurate for f16 input)
float x = float(data[gid]);
float arg = x * 0.7071067811865475f;
float sign = arg >= 0.0f ? 1.0f : -1.0f;
float xa = abs(arg);
float t = 1.0f / (1.0f + 0.3275911f * xa);
float y = t * (0.254829592f + t * (-0.284496736f + t * (1.421413741f
+ t * (-1.453152027f + t * 1.061405429f))));
float erf_val = sign * (1.0f - y * exp(-xa * xa));
data[gid] = half(x * 0.5f * (1.0f + erf_val));
}
kernel void gelu_approx_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float x = float(data[gid]);
float inner = 0.7978845608f * (x + 0.044715f * x * x * x);
data[gid] = half(0.5f * x * (1.0f + tanh(inner)));
}
kernel void silu_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float x = float(data[gid]);
data[gid] = half(x / (1.0f + exp(-x)));
}
// f16 input, f32 reduction, f16 output (mixed precision LayerNorm)
kernel void layer_norm_h(
device const half* input [[buffer(0)]],
device const half* gamma [[buffer(1)]],
device const half* beta [[buffer(2)]],
device half* output [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant float& eps [[buffer(5)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
float local_sum = 0.0f, local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = float(input[row * h + i]);
local_sum += v;
local_sumsq += v * v;
}
partial_sum[tid] = local_sum;
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sum[tid] += partial_sum[tid + stride];
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float mean = partial_sum[0] / float(h);
float var = partial_sumsq[0] / float(h) - mean * mean;
float inv_std = rsqrt(var + eps);
for (uint i = tid; i < h; i += tsize) {
float v = float(input[row * h + i]);
output[row * h + i] = half((v - mean) * inv_std * float(gamma[i]) + float(beta[i]));
}
}
kernel void fused_residual_ln_h(
device const half* x [[buffer(0)]],
device const half* res [[buffer(1)]],
device const half* gamma [[buffer(2)]],
device const half* beta [[buffer(3)]],
device half* out [[buffer(4)]],
constant uint& h [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
float local_sum = 0.0f, local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = float(x[row * h + i]) + float(res[row * h + i]);
local_sum += v;
local_sumsq += v * v;
}
partial_sum[tid] = local_sum;
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sum[tid] += partial_sum[tid + stride];
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float mean = partial_sum[0] / float(h);
float var = partial_sumsq[0] / float(h) - mean * mean;
float inv_std = rsqrt(var + eps);
for (uint i = tid; i < h; i += tsize) {
float v = float(x[row * h + i]) + float(res[row * h + i]);
out[row * h + i] = half((v - mean) * inv_std * float(gamma[i]) + float(beta[i]));
}
}
kernel void fused_residual_rms_norm_h(
device const half* x [[buffer(0)]],
device const half* res [[buffer(1)]],
device const half* gamma [[buffer(2)]],
device const half* beta [[buffer(3)]],
device half* out [[buffer(4)]],
constant uint& h [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial_sumsq[256];
float local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = float(x[row * h + i]) + float(res[row * h + i]);
local_sumsq += v * v;
}
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_rms = rsqrt(partial_sumsq[0] / float(h) + eps);
for (uint i = tid; i < h; i += tsize) {
float v = float(x[row * h + i]) + float(res[row * h + i]);
out[row * h + i] = half(v * inv_rms * float(gamma[i]) + float(beta[i]));
}
}
kernel void elem_add_h(
device const half* a [[buffer(0)]],
device const half* b [[buffer(1)]],
device half* c [[buffer(2)]],
constant uint& len [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
c[gid] = a[gid] + b[gid];
}
kernel void elem_mul_h(
device const half* a [[buffer(0)]],
device const half* b [[buffer(1)]],
device half* c [[buffer(2)]],
constant uint& len [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
c[gid] = a[gid] * b[gid];
}
kernel void gather_axis0_h(
device const half* table [[buffer(0)]],
device const half* idx [[buffer(1)]],
device half* out [[buffer(2)]],
constant uint& num_idx [[buffer(3)]],
constant uint& trailing [[buffer(4)]],
uint2 gid [[thread_position_in_grid]]
) {
uint i = gid.y, j = gid.x;
if (i >= num_idx || j >= trailing) return;
uint row = uint(float(idx[i]));
out[i * trailing + j] = table[row * trailing + j];
}
kernel void narrow_lastax_h(
device const char* arena_src [[buffer(0)]],
device char* arena_dst [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis [[buffer(3)]],
constant uint& start [[buffer(4)]],
constant uint& len [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
constant ulong& dst_byte_off [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
device const half* src = (device const half*)(arena_src + src_byte_off);
device half* dst = (device half*)(arena_dst + dst_byte_off);
uint i = gid.y, j = gid.x;
if (i >= outer || j >= len) return;
dst[i * len + j] = src[i * src_axis + start + j];
}
kernel void sdpa_h(
device const half* Q [[buffer(0)]],
device const half* K [[buffer(1)]],
device const half* V [[buffer(2)]],
device const half* M [[buffer(3)]],
device half* OUT [[buffer(4)]],
constant uint& batch [[buffer(5)]],
constant uint& seq [[buffer(6)]],
constant uint& heads [[buffer(7)]],
constant uint& head_dim [[buffer(8)]],
constant uint& seq_stride [[buffer(9)]],
constant uint& mask_kind [[buffer(10)]],
constant uint& seq_k [[buffer(11)]], // unused; mirrors sdpa signature
constant uint& k_stride [[buffer(12)]], // unused; mirrors sdpa signature
constant uint& bhsd [[buffer(13)]], // unused; mirrors sdpa signature
constant uint& window [[buffer(14)]],
uint tgid_x [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
(void)seq_k; (void)k_stride; (void)bhsd; // accepted to share encode_sdpa layout
threadgroup float scores[64 * 64];
threadgroup float row_max;
threadgroup float row_sum;
uint bi = tgid_x / heads;
uint hi = tgid_x % heads;
if (bi >= batch) return;
uint hs = heads * head_dim;
float scale = 1.0f / precise::sqrt(float(head_dim));
uint per_batch_stride = seq_stride * hs;
uint total = seq * seq;
for (uint idx = tid; idx < total; idx += tsize) {
uint qi = idx / seq;
uint ki = idx % seq;
float dot = 0.0f;
uint q_base = bi * per_batch_stride + qi * hs + hi * head_dim;
uint k_base = bi * per_batch_stride + ki * hs + hi * head_dim;
for (uint d = 0; d < head_dim; ++d) {
dot += float(Q[q_base + d]) * float(K[k_base + d]);
}
float s = dot * scale;
if (mask_kind == 1u) {
if (ki > qi) s = -1e9f;
} else if (mask_kind == 2u) {
if (float(M[bi * seq_stride + ki]) < 0.5f) s = -1e9f;
} else if (mask_kind == 4u) {
// Lq == Lk here (sdpa_h is prefill-only), so abs_q == qi.
uint lo = qi > window ? qi - window : 0u;
if (ki < lo || ki > qi) s = -1e9f;
}
scores[qi * seq + ki] = s;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint qi = 0; qi < seq; ++qi) {
if (tid == 0) {
float mx = -1e30f;
for (uint ki = 0; ki < seq; ++ki) {
mx = max(mx, scores[qi * seq + ki]);
}
row_max = mx;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid == 0) {
float sum = 0.0f;
for (uint ki = 0; ki < seq; ++ki) {
float e = precise::exp(scores[qi * seq + ki] - row_max);
scores[qi * seq + ki] = e;
sum += e;
}
row_sum = sum;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint ki = tid; ki < seq; ki += tsize) {
scores[qi * seq + ki] /= row_sum;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
uint out_total = seq * head_dim;
for (uint idx = tid; idx < out_total; idx += tsize) {
uint qi = idx / head_dim;
uint d = idx % head_dim;
float acc = 0.0f;
for (uint ki = 0; ki < seq; ++ki) {
uint v_base = bi * per_batch_stride + ki * hs + hi * head_dim;
acc += scores[qi * seq + ki] * float(V[v_base + d]);
}
uint o_base = bi * per_batch_stride + qi * hs + hi * head_dim;
OUT[o_base + d] = half(acc);
}
}
kernel void rope_h(
device const half* x [[buffer(0)]],
device const half* cos [[buffer(1)]],
device const half* sin [[buffer(2)]],
device half* out [[buffer(3)]],
constant uint& batch [[buffer(4)]],
constant uint& seq [[buffer(5)]],
constant uint& hidden [[buffer(6)]],
constant uint& head_dim [[buffer(7)]],
constant uint& src_row_stride [[buffer(8)]],
constant uint& seq_stride [[buffer(9)]],
constant uint& n_rot [[buffer(10)]],
uint3 gid [[thread_position_in_grid]]
) {
uint half_dh = head_dim / 2;
uint rot_half = n_rot / 2;
if (gid.x >= head_dim) return;
uint bs = gid.z;
uint bi = bs / seq;
uint si = bs % seq;
if (bi >= batch || si >= seq) return;
uint nh = hidden / head_dim;
uint hi = gid.y;
if (hi >= nh) return;
// PLAN L1 — `seq_stride` is the compile-time full extent for buffer
// offsets; `seq` is the (possibly scaled) iteration bound.
uint src_base = bi * seq_stride * src_row_stride + si * src_row_stride + hi * head_dim;
uint dst_base = bi * seq_stride * hidden + si * hidden + hi * head_dim;
uint d = gid.x;
if (d < rot_half) {
float x1 = float(x[src_base + d]);
float x2 = float(x[src_base + rot_half + d]);
float c = float(cos[si * half_dh + d]);
float s = float(sin[si * half_dh + d]);
out[dst_base + d] = half(x1 * c - x2 * s);
out[dst_base + rot_half + d] = half(x2 * c + x1 * s);
} else if (d >= n_rot) {
out[dst_base + d] = x[src_base + d];
}
}
// Cast f32 → f16 (used at I/O boundary)
kernel void cast_f32_to_f16(
device const float* src [[buffer(0)]],
device half* dst [[buffer(1)]],
constant uint& len [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
dst[gid] = half(src[gid]);
}
// Cast f16 → f32 (used at I/O boundary)
kernel void cast_f16_to_f32(
device const half* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& len [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
dst[gid] = float(src[gid]);
}
// Plain f32 buffer copy — used for Reshape/Expand thunks when we want
// to stay on the shared compute encoder instead of switching to a blit
// encoder (encoder-switch overhead dominates for small ops).
kernel void copy_f32(
device const char* arena [[buffer(0)]],
constant ulong& src_byte_off [[buffer(1)]],
constant ulong& dst_byte_off [[buffer(2)]],
constant uint& len [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device const float* src = (device const float*)(arena + src_byte_off);
device float* dst = (device float*)(arena + dst_byte_off);
dst[gid] = src[gid];
}
kernel void copy4(
device const char* arena [[buffer(0)]],
constant ulong& src_byte_off [[buffer(1)]],
constant ulong& dst_byte_off [[buffer(2)]],
constant uint& len4 [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* src = (device const packed_float4*)(arena + src_byte_off);
device packed_float4* dst = (device packed_float4*)(arena + dst_byte_off);
dst[gid] = src[gid];
}
// SIMD-group matrix sgemm: uses Apple Silicon's dedicated tensor units.
// One simdgroup computes an 8x8 output tile via simdgroup_multiply_accumulate.
// Threadgroup has 32 threads = 1 simdgroup, computing one 8x8 tile of C.
// For larger output, dispatch more threadgroups.
//
// All dimensions must be multiples of 8 for this kernel. Caller is responsible
// for routing non-multiple-of-8 cases to the scalar tiled fallback.
kernel void sgemm_simd(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]]
) {
uint row_base = tgid.y * 8;
uint col_base = tgid.x * 8;
if (row_base >= M || col_base >= N) return;
simdgroup_float8x8 a;
simdgroup_float8x8 b;
simdgroup_float8x8 c;
c = simdgroup_float8x8(0.0f);
for (uint k = 0; k < K; k += 8) {
simdgroup_load(a, A + row_base * K + k, K);
simdgroup_load(b, B + k * N + col_base, N);
simdgroup_multiply_accumulate(c, a, b, c);
}
simdgroup_store(c, C + row_base * N + col_base, N);
}
// High-throughput simdgroup matmul: 32x32 output per threadgroup,
// 4x4 = 16 simdgroups cooperate through threadgroup memory.
// Each B element is reused 4× across rows of simdgroups; each A element 4× across cols.
// K loaded in 32-wide stripes into threadgroup memory.
//
// Requires M%32==K%32==N%32==0. Falls back to sgemm_simd for smaller dims.
kernel void sgemm_simd_4x4(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]]
) {
// 4x4 simdgroup grid within threadgroup
uint sg_row = sgid / 4; // 0..3
uint sg_col = sgid % 4; // 0..3
uint tg_row_base = tgid.y * 32;
uint tg_col_base = tgid.x * 32;
threadgroup float A_tg[32 * 32]; // 4 KB
threadgroup float B_tg[32 * 32]; // 4 KB
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint kk = 0; kk < K; kk += 32) {
// Cooperative load: 16 simdgroups × 32 threads = 512 threads
// load 32×32 A tile and 32×32 B tile (1024 floats each = 4 elements per thread)
uint linear = sgid * 32 + slid; // 0..511
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint ar = idx / 32;
uint ac = idx % 32;
A_tg[idx] = A[(tg_row_base + ar) * K + (kk + ac)];
}
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint br = idx / 32;
uint bc = idx % 32;
B_tg[idx] = B[(kk + br) * N + (tg_col_base + bc)];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// 4 inner-K steps of 8 each, accumulating into c
for (uint k_inner = 0; k_inner < 32; k_inner += 8) {
simdgroup_load(a, &A_tg[sg_row * 8 * 32 + k_inner], 32);
simdgroup_load(b, &B_tg[k_inner * 32 + sg_col * 8], 32);
simdgroup_multiply_accumulate(c, a, b, c);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
uint out_row = tg_row_base + sg_row * 8;
uint out_col = tg_col_base + sg_col * 8;
simdgroup_store(c, &C[out_row * N + out_col], N);
}
// 32x32-tiled with bias + optional activation fused.
kernel void sgemm_simd_4x4_bias(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device const float* bias [[buffer(2)]],
device float* C [[buffer(3)]],
constant uint& M [[buffer(4)]],
constant uint& K [[buffer(5)]],
constant uint& N [[buffer(6)]],
constant uint& act_kind [[buffer(7)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]]
) {
uint sg_row = sgid / 4;
uint sg_col = sgid % 4;
uint tg_row_base = tgid.y * 32;
uint tg_col_base = tgid.x * 32;
threadgroup float A_tg[32 * 32];
threadgroup float B_tg[32 * 32];
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint kk = 0; kk < K; kk += 32) {
uint linear = sgid * 32 + slid;
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint ar = idx / 32;
uint ac = idx % 32;
A_tg[idx] = A[(tg_row_base + ar) * K + (kk + ac)];
}
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint br = idx / 32;
uint bc = idx % 32;
B_tg[idx] = B[(kk + br) * N + (tg_col_base + bc)];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k_inner = 0; k_inner < 32; k_inner += 8) {
simdgroup_load(a, &A_tg[sg_row * 8 * 32 + k_inner], 32);
simdgroup_load(b, &B_tg[k_inner * 32 + sg_col * 8], 32);
simdgroup_multiply_accumulate(c, a, b, c);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
// Stage 8x8 output, apply bias + activation per element
threadgroup float tile[16 * 64]; // 16 simdgroups × 64 elements each
simdgroup_store(c, &tile[sgid * 64], 8);
threadgroup_barrier(mem_flags::mem_threadgroup);
uint out_row_base = tg_row_base + sg_row * 8;
uint out_col_base = tg_col_base + sg_col * 8;
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint r = idx / 8;
uint cc = idx % 8;
float v = tile[sgid * 64 + idx] + bias[out_col_base + cc];
if (act_kind == 1) {
float arg = v * 0.7071067811865475;
float sign = arg >= 0.0 ? 1.0 : -1.0;
float xa = abs(arg);
float t = 1.0 / (1.0 + 0.3275911 * xa);
float y = t * (0.254829592 + t * (-0.284496736 + t * (1.421413741
+ t * (-1.453152027 + t * 1.061405429))));
float erf_val = sign * (1.0 - y * exp(-xa * xa));
v = v * 0.5 * (1.0 + erf_val);
} else if (act_kind == 2) {
v = v / (1.0 + exp(-v));
}
C[(out_row_base + r) * N + (out_col_base + cc)] = v;
}
}
// sgemm + bias (broadcast per column) fused into one kernel.
// Dispatched same as sgemm_simd: 1 threadgroup per 8x8 output tile.
kernel void sgemm_simd_bias(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device const float* bias [[buffer(2)]],
device float* C [[buffer(3)]],
constant uint& M [[buffer(4)]],
constant uint& K [[buffer(5)]],
constant uint& N [[buffer(6)]],
constant uint& act_kind [[buffer(7)]], // 0=none, 1=gelu, 2=silu
uint2 tgid [[threadgroup_position_in_grid]],
uint slid [[thread_index_in_simdgroup]]
) {
uint row_base = tgid.y * 8;
uint col_base = tgid.x * 8;
if (row_base >= M || col_base >= N) return;
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint k = 0; k < K; k += 8) {
simdgroup_load(a, A + row_base * K + k, K);
simdgroup_load(b, B + k * N + col_base, N);
simdgroup_multiply_accumulate(c, a, b, c);
}
// Stage tile in threadgroup memory, then apply bias + activation per element
threadgroup float tile[64];
simdgroup_store(c, tile, 8);
threadgroup_barrier(mem_flags::mem_threadgroup);
// 32 threads × 2 elements each cover the 8x8 tile
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint r = idx / 8;
uint cc = idx % 8;
float v = tile[idx] + bias[col_base + cc];
if (act_kind == 1) {
// GELU (Abramowitz & Stegun erf approx)
float arg = v * 0.7071067811865475;
float sign = arg >= 0.0 ? 1.0 : -1.0;
float xa = abs(arg);
float t = 1.0 / (1.0 + 0.3275911 * xa);
float y = t * (0.254829592 + t * (-0.284496736 + t * (1.421413741
+ t * (-1.453152027 + t * 1.061405429))));
float erf_val = sign * (1.0 - y * exp(-xa * xa));
v = v * 0.5 * (1.0 + erf_val);
} else if (act_kind == 2) {
v = v / (1.0 + exp(-v));
}
C[(row_base + r) * N + (col_base + cc)] = v;
}
}
// Padded variant: arbitrary M with bounds-checked stores + bias + optional act.
kernel void sgemm_simd_padded_bias(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device const float* bias [[buffer(2)]],
device float* C [[buffer(3)]],
constant uint& M [[buffer(4)]],
constant uint& K [[buffer(5)]],
constant uint& N [[buffer(6)]],
constant uint& act_kind [[buffer(7)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint slid [[thread_index_in_simdgroup]]
) {
uint row_base = tgid.y * 8;
uint col_base = tgid.x * 8;
threadgroup float A_pad[64];
threadgroup float B_pad[64];
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint k = 0; k < K; k += 8) {
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint ar = idx / 8, ac = idx % 8;
uint sr = row_base + ar, sc = k + ac;
A_pad[idx] = (sr < M && sc < K) ? A[sr * K + sc] : 0.0f;
}
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint br = idx / 8, bc = idx % 8;
uint sr = k + br, sc = col_base + bc;
B_pad[idx] = (sr < K && sc < N) ? B[sr * N + sc] : 0.0f;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
simdgroup_load(a, A_pad, 8);
simdgroup_load(b, B_pad, 8);
simdgroup_multiply_accumulate(c, a, b, c);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
threadgroup float C_pad[64];
simdgroup_store(c, C_pad, 8);
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint r = idx / 8;
uint cc = idx % 8;
uint dst_row = row_base + r;
uint dst_col = col_base + cc;
if (dst_row < M && dst_col < N) {
float v = C_pad[idx] + bias[dst_col];
if (act_kind == 1) {
float arg = v * 0.7071067811865475;
float sign = arg >= 0.0 ? 1.0 : -1.0;
float xa = abs(arg);
float t = 1.0 / (1.0 + 0.3275911 * xa);
float y = t * (0.254829592 + t * (-0.284496736 + t * (1.421413741
+ t * (-1.453152027 + t * 1.061405429))));
float erf_val = sign * (1.0 - y * exp(-xa * xa));
v = v * 0.5 * (1.0 + erf_val);
} else if (act_kind == 2) {
v = v / (1.0 + exp(-v));
}
C[dst_row * N + dst_col] = v;
}
}
}
// Padded simdgroup sgemm: handles arbitrary M/K/N by zero-padding.
// Reads A row-by-row with bounds checks, computes 8x8 simdgroup tiles,
// writes back row-by-row with bounds checks. Slower than sgemm_simd for
// aligned dims but works for the common batch=1 case (m=6).
//
// Strategy: pre-stage A's relevant rows into threadgroup memory (zero-pad
// missing rows), then use simdgroup ops on the padded tile. Same for B's
// columns.
kernel void sgemm_simd_padded(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]]
) {
uint row_base = tgid.y * 8;
uint col_base = tgid.x * 8;
// Per-tile staging in threadgroup memory: 8x8 A tile, 8x8 B tile.
// 32 threads collaborate to stage; reuse the simdgroup_load API for
// the multiply once data is in threadgroup or device memory.
threadgroup float A_pad[64];
threadgroup float B_pad[64];
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint k = 0; k < K; k += 8) {
// Stage 8x8 A tile with bounds-checked loads (32 threads cover 64 elements: 2 each)
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint ar = idx / 8;
uint ac = idx % 8;
uint src_row = row_base + ar;
uint src_col = k + ac;
float v = (src_row < M && src_col < K) ? A[src_row * K + src_col] : 0.0f;
A_pad[idx] = v;
}
// Stage 8x8 B tile
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint br = idx / 8;
uint bc = idx % 8;
uint src_row = k + br;
uint src_col = col_base + bc;
float v = (src_row < K && src_col < N) ? B[src_row * N + src_col] : 0.0f;
B_pad[idx] = v;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
simdgroup_load(a, A_pad, 8);
simdgroup_load(b, B_pad, 8);
simdgroup_multiply_accumulate(c, a, b, c);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
// Bounds-checked store of the 8x8 C tile (32 threads × 2 elements each)
threadgroup float C_pad[64];
simdgroup_store(c, C_pad, 8);
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
uint cr = idx / 8;
uint cc = idx % 8;
uint dst_row = row_base + cr;
uint dst_col = col_base + cc;
if (dst_row < M && dst_col < N) {
C[dst_row * N + dst_col] = C_pad[idx];
}
}
}
// Tiled sgemm: TILExTILE output blocks, K loaded in TILE-wide stripes
// into threadgroup memory. Used for non-multiple-of-8 dimensions.
constant uint TILE = 16;
kernel void sgemm_tiled(
device const float* A [[buffer(0)]],
device const float* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint2 gid [[thread_position_in_grid]],
uint2 tid [[thread_position_in_threadgroup]],
uint2 tgid [[threadgroup_position_in_grid]]
) {
threadgroup float Asub[16][16];
threadgroup float Bsub[16][16];
uint row = tgid.y * TILE + tid.y;
uint col = tgid.x * TILE + tid.x;
float sum = 0.0;
uint num_tiles = (K + TILE - 1) / TILE;
for (uint t = 0; t < num_tiles; ++t) {
uint a_col = t * TILE + tid.x;
uint b_row = t * TILE + tid.y;
Asub[tid.y][tid.x] = (row < M && a_col < K) ? A[row * K + a_col] : 0.0;
Bsub[tid.y][tid.x] = (b_row < K && col < N) ? B[b_row * N + col] : 0.0;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k = 0; k < TILE; ++k) {
sum += Asub[tid.y][k] * Bsub[k][tid.x];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (row < M && col < N) {
C[row * N + col] = sum;
}
}
// out = bias_add(data, bias, m, n)
kernel void bias_add(
device float* data [[buffer(0)]],
device const float* bias [[buffer(1)]],
constant uint& m [[buffer(2)]],
constant uint& n [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]
) {
uint row = gid.y;
uint col = gid.x;
if (row >= m || col >= n) return;
data[row * n + col] += bias[col];
}
// in-place GELU using Abramowitz & Stegun erf approximation
// (matches CPU NEON kernel for parity)
kernel void gelu_inplace(
device char* arena [[buffer(0)]],
constant ulong& data_byte_off [[buffer(1)]],
constant uint& len [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device float* data = (device float*)(arena + data_byte_off);
float x = data[gid];
float arg = x * 0.7071067811865475; // x / sqrt(2)
float sign = arg >= 0.0 ? 1.0 : -1.0;
float xa = abs(arg);
float t = 1.0 / (1.0 + 0.3275911 * xa);
float y = t * (0.254829592 + t * (-0.284496736 + t * (1.421413741
+ t * (-1.453152027 + t * 1.061405429))));
float erf_val = sign * (1.0 - y * exp(-xa * xa));
data[gid] = x * 0.5 * (1.0 + erf_val);
}
// Tanh-approximation GELU — matches CPU `scalar_gelu_approx` / PyTorch default:
// y = 0.5 · x · (1 + tanh(√(2/π) · (x + 0.044715 · x³)))
// Routed from `Activation::GeluApprox` (Gemma 4 MLP, ViT, DINOv2).
kernel void gelu_approx_inplace(
device char* arena [[buffer(0)]],
constant ulong& data_byte_off [[buffer(1)]],
constant uint& len [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device float* data = (device float*)(arena + data_byte_off);
float x = data[gid];
float inner = 0.7978845608f * (x + 0.044715f * x * x * x);
data[gid] = 0.5f * x * (1.0f + tanh(inner));
}
kernel void gelu_approx_inplace4(
device char* arena [[buffer(0)]],
constant ulong& data_byte_off [[buffer(1)]],
constant uint& len4 [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device packed_float4* data = (device packed_float4*)(arena + data_byte_off);
packed_float4 px = data[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float x = px[c];
float inner = 0.7978845608f * (x + 0.044715f * x * x * x);
out[c] = 0.5f * x * (1.0f + tanh(inner));
}
data[gid] = out;
}
kernel void gelu_approx_out4(
device const char* arena [[buffer(0)]],
constant ulong& src_off [[buffer(1)]],
constant ulong& dst_off [[buffer(2)]],
constant uint& len4 [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* src = (device const packed_float4*)(arena + src_off);
device packed_float4* dst = (device packed_float4*)(arena + dst_off);
packed_float4 px = src[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float x = px[c];
float inner = 0.7978845608f * (x + 0.044715f * x * x * x);
out[c] = 0.5f * x * (1.0f + tanh(inner));
}
dst[gid] = out;
}
kernel void gelu_inplace4(
device char* arena [[buffer(0)]],
constant ulong& data_byte_off [[buffer(1)]],
constant uint& len4 [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device packed_float4* data = (device packed_float4*)(arena + data_byte_off);
packed_float4 px = data[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float xv = px[c];
float arg = xv * 0.7071067811865475;
float sign = arg >= 0.0 ? 1.0 : -1.0;
float xa = abs(arg);
float t = 1.0 / (1.0 + 0.3275911 * xa);
float y = t * (0.254829592 + t * (-0.284496736 + t * (1.421413741
+ t * (-1.453152027 + t * 1.061405429))));
float erf_val = sign * (1.0 - y * exp(-xa * xa));
out[c] = xv * 0.5 * (1.0 + erf_val);
}
data[gid] = out;
}
kernel void silu_inplace4(
device char* arena [[buffer(0)]],
constant ulong& data_byte_off [[buffer(1)]],
constant uint& len4 [[buffer(2)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device packed_float4* data = (device packed_float4*)(arena + data_byte_off);
packed_float4 px = data[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float xv = px[c];
out[c] = xv / (1.0 + exp(-xv));
}
data[gid] = out;
}
// rhs [cols] broadcast across rows: out[m, n] = lhs[m*cols+n] op rhs[n]
kernel void binary_broadcast_rhs_col_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
constant uint& op [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint m = gid.y;
uint n = gid.x;
if (m >= rows || n >= cols) return;
float lv = lhs[m * cols + n];
float rv = rhs[n];
float out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst[m * cols + n] = out;
}
kernel void binary_broadcast_rhs_col4(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols4 [[buffer(4)]],
constant uint& op [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint m = gid.y;
uint n4 = gid.x;
uint cols = cols4 * 4u;
if (m >= rows || n4 >= cols4) return;
device const packed_float4* lhs4 =
(device const packed_float4*)(lhs + m * cols);
device const packed_float4* rhs4 = (device const packed_float4*)(rhs);
device packed_float4* dst4 = (device packed_float4*)(dst + m * cols);
packed_float4 lv = lhs4[n4];
packed_float4 rv = rhs4[n4];
packed_float4 out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst4[n4] = out;
}
// Dense lhs + rhs row vector broadcast on the last axis (e.g. `[rows, cols] op [rows, 1]`).
kernel void binary_broadcast_rhs_row_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
constant uint& op [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint m = gid.y;
uint n = gid.x;
if (m >= rows || n >= cols) return;
float lv = lhs[m * cols + n];
float rv = rhs[m];
float out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst[m * cols + n] = out;
}
kernel void binary_broadcast_rhs_row4(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols4 [[buffer(4)]],
constant uint& op [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint m = gid.y;
uint n4 = gid.x;
uint cols = cols4 * 4u;
if (m >= rows || n4 >= cols4) return;
device const packed_float4* lhs4 =
(device const packed_float4*)(lhs + m * cols);
float rv = rhs[m];
packed_float4 rv4 = packed_float4(rv);
device packed_float4* dst4 = (device packed_float4*)(dst + m * cols);
packed_float4 lv = lhs4[n4];
packed_float4 out;
switch (op) {
case 0: out = lv + rv4; break;
case 1: out = lv - rv4; break;
case 2: out = lv * rv4; break;
case 3: out = lv / rv4; break;
case 4: out = max(lv, rv4); break;
case 5: out = min(lv, rv4); break;
default: out = pow(lv, rv4); break;
}
dst4[n4] = out;
}
// Dense lhs + scalar rhs (all broadcast strides zero).
kernel void binary_broadcast_rhs_scalar_f32(
device const char* arena [[buffer(0)]],
constant ulong& lhs_off [[buffer(1)]],
constant ulong& rhs_off [[buffer(2)]],
constant ulong& dst_off [[buffer(3)]],
constant uint& len [[buffer(4)]],
constant uint& op [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device const float* lhs = (device const float*)(arena + lhs_off);
device const float* rhs = (device const float*)(arena + rhs_off);
device float* dst = (device float*)(arena + dst_off);
float lv = lhs[gid];
float rv = rhs[0];
float out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst[gid] = out;
}
kernel void binary_broadcast_rhs_scalar4(
device const char* arena [[buffer(0)]],
constant ulong& lhs_off [[buffer(1)]],
constant ulong& rhs_off [[buffer(2)]],
constant ulong& dst_off [[buffer(3)]],
constant uint& len4 [[buffer(4)]],
constant uint& op [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* lhs = (device const packed_float4*)(arena + lhs_off);
device const float* rhs = (device const float*)(arena + rhs_off);
device packed_float4* dst = (device packed_float4*)(arena + dst_off);
float rv = rhs[0];
packed_float4 rv4 = packed_float4(rv);
packed_float4 lv = lhs[gid];
packed_float4 out;
switch (op) {
case 0: out = lv + rv4; break;
case 1: out = lv - rv4; break;
case 2: out = lv * rv4; break;
case 3: out = lv / rv4; break;
case 4: out = max(lv, rv4); break;
case 5: out = min(lv, rv4); break;
default: out = pow(lv, rv4); break;
}
dst[gid] = out;
}
// Dense lhs + rhs broadcast on exactly one axis (e.g. `[B, T, H] op [B, 1, H]`).
kernel void binary_broadcast_1ax_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
constant uint& mid [[buffer(5)]],
constant uint& op [[buffer(6)]],
uint2 gid [[thread_position_in_grid]]
) {
uint col = gid.x;
uint row = gid.y;
if (col >= cols || row >= rows) return;
uint pre_i = row / mid;
uint li = row * cols + col;
uint ri = pre_i * cols + col;
float lv = lhs[li];
float rv = rhs[ri];
float out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst[li] = out;
}
kernel void binary_broadcast_1ax4(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols4 [[buffer(4)]],
constant uint& mid [[buffer(5)]],
constant uint& op [[buffer(6)]],
uint2 gid [[thread_position_in_grid]]
) {
uint col4 = gid.x;
uint row = gid.y;
uint cols = cols4 * 4u;
if (col4 >= cols4 || row >= rows) return;
uint pre_i = row / mid;
device const packed_float4* lhs4 =
(device const packed_float4*)(lhs + row * cols);
device const float* rhs_row = rhs + pre_i * cols;
device const packed_float4* rhs4 =
(device const packed_float4*)(rhs_row);
device packed_float4* dst4 =
(device packed_float4*)(dst + row * cols);
packed_float4 lv = lhs4[col4];
packed_float4 rv = rhs4[col4];
packed_float4 out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst4[col4] = out;
}
inline float fused_bin(float lv, float rv, uint op) {
switch (op) {
case 0: return lv + rv;
case 1: return lv - rv;
case 2: return lv * rv;
case 3: return lv / rv;
case 4: return max(lv, rv);
case 5: return min(lv, rv);
default: return pow(lv, rv);
}
}
inline float fused_act(float x, uint act) {
switch (act) {
case 0: {
float arg = x * 0.7071067811865475;
float sign = arg >= 0.0 ? 1.0 : -1.0;
float xa = abs(arg);
float t = 1.0 / (1.0 + 0.3275911 * xa);
float y = t * (0.254829592 + t * (-0.284496736 + t * (1.421413741
+ t * (-1.453152027 + t * 1.061405429))));
float erf_val = sign * (1.0 - y * exp(-xa * xa));
return x * 0.5 * (1.0 + erf_val);
}
case 1: return x / (1.0 + exp(-x));
case 2: return max(x, 0.0);
case 3: return 1.0 / (1.0 + exp(-x));
case 4: return tanh(x);
default: return x;
}
}
kernel void fused_binary_activation_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& len [[buffer(3)]],
constant uint& bin_op [[buffer(4)]],
constant uint& act_op [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float v = fused_bin(lhs[gid], rhs[gid], bin_op);
dst[gid] = fused_act(v, act_op);
}
kernel void fused_binary_activation4(
device const packed_float4* lhs [[buffer(0)]],
device const packed_float4* rhs [[buffer(1)]],
device packed_float4* dst [[buffer(2)]],
constant uint& len4 [[buffer(3)]],
constant uint& bin_op [[buffer(4)]],
constant uint& act_op [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
packed_float4 lv = lhs[gid];
packed_float4 rv = rhs[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float v = fused_bin(lv[c], rv[c], bin_op);
out[c] = fused_act(v, act_op);
}
dst[gid] = out;
}
kernel void fused_ternary_activation_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs0 [[buffer(1)]],
device const float* rhs1 [[buffer(2)]],
device float* dst [[buffer(3)]],
constant uint& len [[buffer(4)]],
constant uint& bin_op0 [[buffer(5)]],
constant uint& bin_op1 [[buffer(6)]],
constant uint& act_op [[buffer(7)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float v = fused_bin(lhs[gid], rhs0[gid], bin_op0);
v = fused_bin(v, rhs1[gid], bin_op1);
dst[gid] = fused_act(v, act_op);
}
kernel void fused_ternary_activation4(
device const packed_float4* lhs [[buffer(0)]],
device const packed_float4* rhs0 [[buffer(1)]],
device const packed_float4* rhs1 [[buffer(2)]],
device packed_float4* dst [[buffer(3)]],
constant uint& len4 [[buffer(4)]],
constant uint& bin_op0 [[buffer(5)]],
constant uint& bin_op1 [[buffer(6)]],
constant uint& act_op [[buffer(7)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
packed_float4 lv = lhs[gid];
packed_float4 r0 = rhs0[gid];
packed_float4 r1 = rhs1[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float v = fused_bin(lv[c], r0[c], bin_op0);
v = fused_bin(v, r1[c], bin_op1);
out[c] = fused_act(v, act_op);
}
dst[gid] = out;
}
// Element-wise add: c = a + b (same length)
kernel void elem_add(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device const float* a = (device const float*)(arena + a_off);
device const float* b = (device const float*)(arena + b_off);
device float* c = (device float*)(arena + c_off);
c[gid] = a[gid] + b[gid];
}
kernel void elem_add4(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len4 [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* a = (device const packed_float4*)(arena + a_off);
device const packed_float4* b = (device const packed_float4*)(arena + b_off);
device packed_float4* c = (device packed_float4*)(arena + c_off);
c[gid] = a[gid] + b[gid];
}
kernel void elem_sub4(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len4 [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* a = (device const packed_float4*)(arena + a_off);
device const packed_float4* b = (device const packed_float4*)(arena + b_off);
device packed_float4* c = (device packed_float4*)(arena + c_off);
c[gid] = a[gid] - b[gid];
}
// Rank-2 broadcast without a per-element rank loop (fallback after vec fast paths).
kernel void binary_broadcast_rank2_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& len [[buffer(3)]],
constant uint& dim0 [[buffer(4)]],
constant uint& dim1 [[buffer(5)]],
constant uint& lhs_stride0 [[buffer(6)]],
constant uint& lhs_stride1 [[buffer(7)]],
constant uint& rhs_stride0 [[buffer(8)]],
constant uint& rhs_stride1 [[buffer(9)]],
constant uint& op [[buffer(10)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
uint j = gid % dim1;
uint i = gid / dim1;
uint li = i * lhs_stride0 + j * lhs_stride1;
uint ri = i * rhs_stride0 + j * rhs_stride1;
float lv = lhs[li];
float rv = rhs[ri];
float out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst[gid] = out;
}
kernel void binary_broadcast_rank24(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& len4 [[buffer(3)]],
constant uint& dim0 [[buffer(4)]],
constant uint& dim1 [[buffer(5)]],
constant uint& lhs_stride0 [[buffer(6)]],
constant uint& lhs_stride1 [[buffer(7)]],
constant uint& rhs_stride0 [[buffer(8)]],
constant uint& rhs_stride1 [[buffer(9)]],
constant uint& op [[buffer(10)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
uint cols4 = dim1 / 4u;
uint j4 = gid % cols4;
uint i = gid / cols4;
uint j = j4 * 4u;
device const packed_float4* lhs4 =
(device const packed_float4*)(lhs + i * lhs_stride0 + j * lhs_stride1);
device const packed_float4* rhs4 =
(device const packed_float4*)(rhs + i * rhs_stride0 + j * rhs_stride1);
device packed_float4* dst4 = (device packed_float4*)(dst + i * dim1 + j);
packed_float4 lv = *lhs4;
packed_float4 rv = *rhs4;
packed_float4 out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
*dst4 = out;
}
// Shape-aware broadcast binary op. Each thread computes one output
// element by decomposing gid into coords against `out_dims` (row-major)
// and walking `lhs_strides`/`rhs_strides` (stride 0 ⇒ broadcast).
// Op encoding matches `rlx_ir::op::BinaryOp` discriminant order:
// 0=Add, 1=Sub, 2=Mul, 3=Div, 4=Max, 5=Min, 6=Pow. Rank capped at 8.
kernel void binary_broadcast_f32(
device const float* lhs [[buffer(0)]],
device const float* rhs [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& len [[buffer(3)]],
constant uint& rank [[buffer(4)]],
constant uint* out_dims [[buffer(5)]],
constant uint* lhs_strides [[buffer(6)]],
constant uint* rhs_strides [[buffer(7)]],
constant uint& op [[buffer(8)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
uint rem = gid;
uint li = 0;
uint ri = 0;
// Walk from innermost dim to outermost (matches row-major decomposition).
for (uint ax_rev = 0; ax_rev < rank; ++ax_rev) {
uint ax = rank - 1 - ax_rev;
uint sz = out_dims[ax];
uint coord = rem % sz;
rem /= sz;
li += coord * lhs_strides[ax];
ri += coord * rhs_strides[ax];
}
float lv = lhs[li];
float rv = rhs[ri];
float out;
switch (op) {
case 0: out = lv + rv; break;
case 1: out = lv - rv; break;
case 2: out = lv * rv; break;
case 3: out = lv / rv; break;
case 4: out = max(lv, rv); break;
case 5: out = min(lv, rv); break;
default: out = pow(lv, rv); break;
}
dst[gid] = out;
}
// Element-wise multiply: c = a * b
kernel void elem_mul(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device const float* a = (device const float*)(arena + a_off);
device const float* b = (device const float*)(arena + b_off);
device float* c = (device float*)(arena + c_off);
c[gid] = a[gid] * b[gid];
}
kernel void elem_mul4(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len4 [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* a = (device const packed_float4*)(arena + a_off);
device const packed_float4* b = (device const packed_float4*)(arena + b_off);
device packed_float4* c = (device packed_float4*)(arena + c_off);
c[gid] = a[gid] * b[gid];
}
kernel void elem_div4(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len4 [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len4) return;
device const packed_float4* a = (device const packed_float4*)(arena + a_off);
device const packed_float4* b = (device const packed_float4*)(arena + b_off);
device packed_float4* c = (device packed_float4*)(arena + c_off);
c[gid] = a[gid] / b[gid];
}
// Element-wise subtract: c = a - b
kernel void elem_sub(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device const float* a = (device const float*)(arena + a_off);
device const float* b = (device const float*)(arena + b_off);
device float* c = (device float*)(arena + c_off);
c[gid] = a[gid] - b[gid];
}
// Element-wise divide: c = a / b
kernel void elem_div(
device const char* arena [[buffer(0)]],
constant ulong& a_off [[buffer(1)]],
constant ulong& b_off [[buffer(2)]],
constant ulong& c_off [[buffer(3)]],
constant uint& len [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
device const float* a = (device const float*)(arena + a_off);
device const float* b = (device const float*)(arena + b_off);
device float* c = (device float*)(arena + c_off);
c[gid] = a[gid] / b[gid];
}
kernel void elem_max(
device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& len [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; c[gid] = max(a[gid], b[gid]); }
kernel void elem_min(
device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& len [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; c[gid] = min(a[gid], b[gid]); }
kernel void elem_pow(
device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& len [[buffer(3)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; c[gid] = pow(a[gid], b[gid]); }
// Element-wise compare: writes 1.0 / 0.0 per element. `op_kind` selects:
// 0=Eq 1=Ne 2=Lt 3=Le 4=Gt 5=Ge
// One kernel for all six variants keeps the binary-shaped dispatch path
// uniform — the encoder picks op_kind at submit time.
kernel void elem_compare(
device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device float* c [[buffer(2)]],
constant uint& len [[buffer(3)]],
constant uint& op_kind [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float x = a[gid], y = b[gid];
bool r = false;
if (op_kind == 0) r = (x == y);
else if (op_kind == 1) r = (x != y);
else if (op_kind == 2) r = (x < y);
else if (op_kind == 3) r = (x <= y);
else if (op_kind == 4) r = (x > y);
else r = (x >= y);
c[gid] = r ? 1.0f : 0.0f;
}
// 2D convolution (naive direct, NCHW input). One thread per output
// element. Supports groups, dilation. Bias is a separate Op (matches the
// IR's two-input Conv shape). Two u32-arrays of dims pack into one
// constant buffer; an `aux` buffer carries the param triplets.
kernel void conv2d(
device const float* src [[buffer(0)]],
device const float* wt [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint4& nch [[buffer(3)]], // [N, C_in, H, W]
constant uint4& out_dims [[buffer(4)]], // [C_out, H_out, W_out, groups]
constant uint4& kshape [[buffer(5)]], // [kh, kw, sh, sw]
constant uint4& padd [[buffer(6)]], // [ph, pw, dh, dw]
uint3 gid [[thread_position_in_grid]]
) {
uint nco = gid.z; // n * c_out + co
uint ho = gid.y;
uint wo = gid.x;
uint c_out = out_dims.x;
uint h_out = out_dims.y;
uint w_out = out_dims.z;
uint groups = out_dims.w;
if (ho >= h_out || wo >= w_out || nco >= nch.x * c_out) return;
uint n = nco / c_out;
uint co = nco % c_out;
uint c_in = nch.y;
uint h = nch.z;
uint w = nch.w;
uint c_in_per_g = c_in / groups;
uint c_out_per_g = c_out / groups;
uint g = co / c_out_per_g;
uint ci_start = g * c_in_per_g;
uint kh = kshape.x; uint kw = kshape.y;
uint sh = kshape.z; uint sw = kshape.w;
uint ph = padd.x; uint pw = padd.y;
uint dh = padd.z; uint dw = padd.w;
float acc = 0.0f;
for (uint ci_off = 0; ci_off < c_in_per_g; ++ci_off) {
uint ci = ci_start + ci_off;
uint in_chan = ((n * c_in) + ci) * h * w;
uint wt_chan = ((co * c_in_per_g) + ci_off) * kh * kw;
for (uint ki = 0; ki < kh; ++ki) {
for (uint kj = 0; kj < kw; ++kj) {
int hi = (int)(ho * sh + ki * dh) - (int)ph;
int wi = (int)(wo * sw + kj * dw) - (int)pw;
if (hi < 0 || wi < 0 || hi >= (int)h || wi >= (int)w) continue;
acc += src[in_chan + (uint)hi * w + (uint)wi]
* wt[wt_chan + ki * kw + kj];
}
}
}
dst[((n * c_out) + co) * h_out * w_out + ho * w_out + wo] = acc;
}
// 1-D conv (W_in = W_out = 1) — Voxtral codec layout `[N,C,T,1]`.
kernel void conv2d_w1(
device const float* src [[buffer(0)]],
device const float* wt [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint4& nch [[buffer(3)]], // [N, C_in, H, 1]
constant uint4& out_dims [[buffer(4)]], // [C_out, H_out, 1, groups]
constant uint4& kshape [[buffer(5)]],
constant uint4& padd [[buffer(6)]],
uint3 gid [[thread_position_in_grid]]
) {
uint nco = gid.z;
uint ho = gid.y;
uint h_out = out_dims.y;
uint c_out = out_dims.x;
uint groups = out_dims.w;
if (ho >= h_out || nco >= nch.x * c_out) return;
uint n = nco / c_out;
uint co = nco % c_out;
uint c_in = nch.y;
uint h = nch.z;
uint c_in_per_g = c_in / groups;
uint c_out_per_g = c_out / groups;
uint g = co / c_out_per_g;
uint ci_start = g * c_in_per_g;
uint kh = kshape.x;
uint kw = kshape.y;
uint sh = kshape.z;
uint ph = padd.x;
uint pw = padd.y;
uint dh = padd.z;
uint dw = padd.w;
float acc = 0.0f;
for (uint ci_off = 0; ci_off < c_in_per_g; ++ci_off) {
uint ci = ci_start + ci_off;
uint in_chan = ((n * c_in) + ci) * h;
uint wt_chan = ((co * c_in_per_g) + ci_off) * kh * kw;
for (uint ki = 0; ki < kh; ++ki) {
for (uint kj = 0; kj < kw; ++kj) {
int hi = (int)(ho * sh + ki * dh) - (int)ph;
int wi = (int)(kj * dw) - (int)pw;
if (hi < 0 || wi < 0 || hi >= (int)h || wi >= 1) continue;
acc += src[in_chan + (uint)hi] * wt[wt_chan + ki * kw + kj];
}
}
}
dst[((n * c_out) + co) * h_out + ho] = acc;
}
// LayerNorm2d (candle / SAM semantics): normalize across channels at each
// spatial position. One thread per (batch, ho, wo). gamma/beta are [C].
kernel void layer_norm2d(
device const float* src [[buffer(0)]],
device const float* gamma [[buffer(1)]],
device const float* beta [[buffer(2)]],
device float* dst [[buffer(3)]],
constant uint4& nchw [[buffer(4)]], // [N, C, H, W]
constant float& eps [[buffer(5)]],
uint3 gid [[thread_position_in_grid]]
) {
uint n = gid.z;
uint ho = gid.y;
uint wo = gid.x;
uint batch = nchw.x;
uint c = nchw.y;
uint h = nchw.z;
uint w = nchw.w;
if (n >= batch || ho >= h || wo >= w) return;
float mean = 0.0f;
for (uint ch = 0; ch < c; ++ch) {
mean += src[((n * c + ch) * h + ho) * w + wo];
}
mean /= (float)c;
float var = 0.0f;
for (uint ch = 0; ch < c; ++ch) {
float d = src[((n * c + ch) * h + ho) * w + wo] - mean;
var += d * d;
}
var /= (float)c;
float inv = rsqrt(var + eps);
for (uint ch = 0; ch < c; ++ch) {
uint idx = ((n * c + ch) * h + ho) * w + wo;
float v = (src[idx] - mean) * inv;
dst[idx] = v * gamma[ch] + beta[ch];
}
}
// Transposed 2D convolution (NCHW, PyTorch ConvTranspose2d, no bias).
// Weight layout [C_in, C_out/groups, kH, kW]. One thread per output
// element; accumulates in-register (no output zero pass).
kernel void conv_transpose2d(
device const float* src [[buffer(0)]],
device const float* wt [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint4& nch [[buffer(3)]], // [N, C_in, H, W]
constant uint4& out_dims [[buffer(4)]], // [C_out, H_out, W_out, groups]
constant uint4& kshape [[buffer(5)]], // [kh, kw, sh, sw]
constant uint4& padd [[buffer(6)]], // [ph, pw, dh, dw]
uint3 gid [[thread_position_in_grid]]
) {
uint nco = gid.z;
uint ho = gid.y;
uint wo = gid.x;
uint c_out = out_dims.x;
uint h_out = out_dims.y;
uint w_out = out_dims.z;
uint groups = out_dims.w;
if (ho >= h_out || wo >= w_out || nco >= nch.x * c_out) return;
uint n = nco / c_out;
uint co = nco % c_out;
uint c_in = nch.y;
uint h = nch.z;
uint w = nch.w;
uint c_in_per_g = c_in / groups;
uint c_out_per_g = c_out / groups;
uint g = co / c_out_per_g;
uint oc_off = co % c_out_per_g;
uint kh = kshape.x; uint kw = kshape.y;
uint sh = kshape.z; uint sw = kshape.w;
uint ph = padd.x; uint pw = padd.y;
uint dh = padd.z; uint dw = padd.w;
float acc = 0.0f;
for (uint ci_off = 0; ci_off < c_in_per_g; ++ci_off) {
uint ci = g * c_in_per_g + ci_off;
for (uint ky = 0; ky < kh; ++ky) {
int t_h = (int)ho + (int)ph - (int)ky * (int)dh;
if (t_h < 0 || t_h % (int)sh != 0) continue;
int iy = t_h / (int)sh;
if (iy < 0 || iy >= (int)h) continue;
for (uint kx = 0; kx < kw; ++kx) {
int t_w = (int)wo + (int)pw - (int)kx * (int)dw;
if (t_w < 0 || t_w % (int)sw != 0) continue;
int ix = t_w / (int)sw;
if (ix < 0 || ix >= (int)w) continue;
uint w_idx = ((ci * c_out_per_g + oc_off) * kh + ky) * kw + kx;
float v = src[((n * c_in + ci) * h + (uint)iy) * w + (uint)ix];
acc += v * wt[w_idx];
}
}
}
dst[((n * c_out) + co) * h_out * w_out + ho * w_out + wo] = acc;
}
// NCHW group norm: normalize each (C/G)×H×W block. One threadgroup per
// (batch, group); 256-wide reduction then normalize.
kernel void group_norm(
device const float* src [[buffer(0)]],
device const float* gamma [[buffer(1)]],
device const float* beta [[buffer(2)]],
device float* dst [[buffer(3)]],
constant uint4& nchw [[buffer(4)]], // [N, C, H, W]
constant uint& num_groups [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint ng [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
uint batch = nchw.x;
uint c = nchw.y;
uint h = nchw.z;
uint w = nchw.w;
if (ng >= batch * num_groups) return;
uint n = ng / num_groups;
uint g = ng % num_groups;
uint cpg = c / num_groups;
uint c0 = g * cpg;
uint plane = h * w;
uint count = cpg * plane;
float local_sum = 0.0f;
float local_sumsq = 0.0f;
for (uint i = tid; i < count; i += tsize) {
uint c_off = i / plane;
uint s = i % plane;
uint ch = c0 + c_off;
float v = src[((n * c + ch) * plane) + s];
local_sum += v;
local_sumsq += v * v;
}
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
partial_sum[tid] = local_sum;
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sum[tid] += partial_sum[tid + stride];
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float mean = partial_sum[0] / float(count);
float var = partial_sumsq[0] / float(count) - mean * mean;
float inv = rsqrt(var + eps);
for (uint i = tid; i < count; i += tsize) {
uint c_off = i / plane;
uint s = i % plane;
uint ch = c0 + c_off;
uint idx = ((n * c + ch) * plane) + s;
float v = (src[idx] - mean) * inv;
dst[idx] = v * gamma[ch] + beta[ch];
}
}
// Nearest-neighbor 2× upsample on planar NCHW. One thread per output pixel.
kernel void resize_nearest_2x(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint4& nchw [[buffer(2)]], // [N, C, H, W] input
uint3 gid [[thread_position_in_grid]]
) {
uint wo = gid.x;
uint ho = gid.y;
uint nc = gid.z;
uint n = nchw.x;
uint c = nchw.y;
uint h = nchw.z;
uint w = nchw.w;
uint h2 = h * 2u;
uint w2 = w * 2u;
if (nc >= n * c || ho >= h2 || wo >= w2) return;
uint ni = nc / c;
uint ci = nc % c;
uint hi = ho / 2u;
uint wi = wo / 2u;
float v = src[((ni * c + ci) * h + hi) * w + wi];
dst[((ni * c + ci) * h2 + ho) * w2 + wo] = v;
}
// 2D pooling. One thread per output element (n, c, ho, wo). Padding is
// implicit-zero; Mean divides by the full kernel area to match torch's
// `count_include_pad=True`. `kind`: 0=Mean (catch-all), 2=Max.
kernel void pool2d(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint4& nchw [[buffer(2)]], // [N, C, H, W]
constant uint2& hw_out [[buffer(3)]], // [H_out, W_out]
constant uint4& khsw [[buffer(4)]], // [kh, kw, sh, sw]
constant uint2& pad [[buffer(5)]], // [ph, pw]
constant uint& kind [[buffer(6)]],
uint3 gid [[thread_position_in_grid]]
) {
uint nc = gid.z;
uint ho = gid.y;
uint wo = gid.x;
uint n_total = nchw.x;
uint c_total = nchw.y;
if (nc >= n_total * c_total || ho >= hw_out.x || wo >= hw_out.y) return;
uint n = nc / c_total;
uint c = nc % c_total;
uint h = nchw.z;
uint w = nchw.w;
uint h_out = hw_out.x;
uint w_out = hw_out.y;
uint kh = khsw.x; uint kw = khsw.y;
uint sh = khsw.z; uint sw = khsw.w;
uint ph = pad.x; uint pw = pad.y;
float acc = (kind == 2) ? -INFINITY : 0.0f;
uint in_chan = ((n * c_total) + c) * h * w;
for (uint ki = 0; ki < kh; ++ki) {
for (uint kj = 0; kj < kw; ++kj) {
int hi = (int)(ho * sh + ki) - (int)ph;
int wi = (int)(wo * sw + kj) - (int)pw;
if (hi < 0 || wi < 0 || hi >= (int)h || wi >= (int)w) continue;
float v = src[in_chan + (uint)hi * w + (uint)wi];
if (kind == 2) acc = max(acc, v); else acc += v;
}
}
if (kind == 0 || kind == 1) acc /= (float)(kh * kw); // Mean
dst[((n * c_total) + c) * h_out * w_out + ho * w_out + wo] = acc;
}
// Gather along an arbitrary axis. One thread per output element. Output
// is laid out as [outer, num_idx, trailing]; source as [outer, axis_dim, trailing].
kernel void gather_axis(
device const float* table [[buffer(0)]],
device const float* idx [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& outer [[buffer(3)]],
constant uint& axis_dim [[buffer(4)]],
constant uint& num_idx [[buffer(5)]],
constant uint& trailing [[buffer(6)]],
uint3 gid [[thread_position_in_grid]]
) {
uint o = gid.z;
uint k = gid.y;
uint t = gid.x;
if (o >= outer || k >= num_idx || t >= trailing) return;
uint row = (uint)(idx[k]);
dst[(o * num_idx + k) * trailing + t] =
table[(o * axis_dim + row) * trailing + t];
}
// General N-D transpose. One thread per output element. The encoder packs
// out_dims and in_strides into a single u32 buffer of length 2*rank:
// buffer = [out_dim_0, ..., out_dim_{r-1}, in_stride_0, ..., in_stride_{r-1}]
// Rank is bounded at 8 (sufficient for current models).
kernel void transpose_nd(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& rank [[buffer(2)]],
constant uint& total [[buffer(3)]],
constant uint* meta [[buffer(4)]], // [out_dims..., in_strides...]
uint gid [[thread_position_in_grid]]
) {
if (gid >= total) return;
uint src_idx = 0;
uint remaining = gid;
// Decompose flat output index into multi-dim coords (outer-to-inner)
// using stride math, then accumulate the source index from in_strides.
// Compute denominators on the fly to avoid a separate divisor table.
uint stride_rem = total;
for (uint d = 0; d < rank; ++d) {
uint dim = meta[d];
stride_rem /= dim;
uint coord = remaining / stride_rem;
remaining = remaining - coord * stride_rem;
src_idx += coord * meta[rank + d];
}
dst[gid] = src[src_idx];
}
// Rank-2 swap (row-major [rows, cols] → [cols, rows]). Cheaper than transpose_nd
// for attention/layout reshapes — 2D thread grid, coalesced reads.
kernel void transpose_2d_f32(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& rows [[buffer(2)]],
constant uint& cols [[buffer(3)]],
uint2 gid [[thread_position_in_grid]]
) {
uint r = gid.x;
uint c = gid.y;
if (r >= rows || c >= cols) return;
dst[c * rows + r] = src[r * cols + c];
}
// Tiled transpose for large matrices (32x32 tile staged through threadgroup memory).
// Threadgroup size: (32, 8, 1). Each thread loads 4 rows of the tile.
kernel void transpose_2d_tiled_f32(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& rows [[buffer(2)]],
constant uint& cols [[buffer(3)]],
ushort2 tid [[thread_position_in_threadgroup]],
ushort2 tgp [[threadgroup_position_in_grid]]
) {
threadgroup float tile[32][33];
uint r0 = (uint)tgp.x * 32u + (uint)tid.x;
uint c0 = (uint)tgp.y * 32u + (uint)tid.y;
// Load 32x32 tile from src.
for (uint i = 0; i < 32; i += 8) {
uint c = c0 + i;
if (r0 < rows && c < cols) {
tile[tid.x][tid.y + i] = src[r0 * cols + c];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// Store transposed tile to dst.
uint rr0 = (uint)tgp.y * 32u + (uint)tid.x;
uint cc0 = (uint)tgp.x * 32u + (uint)tid.y;
for (uint i = 0; i < 32; i += 8) {
uint rr = rr0;
uint cc = cc0 + i;
if (rr < cols && cc < rows) {
dst[rr * rows + cc] = tile[tid.y + i][tid.x];
}
}
}
// Batched swap of the last two dims: src [batch, rows, cols] -> dst [batch, cols, rows]
kernel void transpose_last2_batched_f32(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& batch [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
uint3 gid [[thread_position_in_grid]]
) {
uint b = gid.z;
uint r = gid.x;
uint c = gid.y;
if (b >= batch || r >= rows || c >= cols) return;
uint src_base = b * rows * cols;
uint dst_base = b * rows * cols;
dst[dst_base + c * rows + r] = src[src_base + r * cols + c];
}
// Tiled batched last2 transpose. Dispatch threadgroups over (rows, cols, batch).
// Threadgroup size: (32, 8, 1).
kernel void transpose_last2_batched_tiled_f32(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& batch [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
uint3 tid [[thread_position_in_threadgroup]],
uint3 tgp [[threadgroup_position_in_grid]]
) {
threadgroup float tile[32][33];
uint b = tgp.z;
if (b >= batch) return;
uint r0 = tgp.x * 32u + tid.x;
uint c0 = tgp.y * 32u + tid.y;
uint base = b * rows * cols;
for (uint i = 0; i < 32; i += 8) {
uint c = c0 + i;
if (r0 < rows && c < cols) {
tile[tid.x][tid.y + i] = src[base + r0 * cols + c];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
uint rr0 = tgp.y * 32u + tid.x; // out row (col index)
uint cc0 = tgp.x * 32u + tid.y; // out col (row index)
for (uint i = 0; i < 32; i += 8) {
uint rr = rr0;
uint cc = cc0 + i;
if (rr < cols && cc < rows) {
dst[base + rr * rows + cc] = tile[tid.y + i][tid.x];
}
}
}
// `[B, A, C, D] → [B, C, A, D]` — swap axes 1 and 2 with trailing axis contiguous.
kernel void transpose_swap12_batched_trail_f32(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& batch [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
constant uint& trail [[buffer(5)]],
uint3 gid [[thread_position_in_grid]]
) {
uint bd = gid.z;
uint b = bd / trail;
uint d = bd % trail;
uint r = gid.x;
uint c = gid.y;
if (b >= batch || r >= rows || c >= cols) return;
uint block = rows * cols * trail;
uint src_idx = b * block + c * rows * trail + r * trail + d;
uint dst_idx = b * block + r * cols * trail + c * trail + d;
dst[dst_idx] = src[src_idx];
}
kernel void transpose_swap12_batched_trail_tiled_f32(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& batch [[buffer(2)]],
constant uint& rows [[buffer(3)]],
constant uint& cols [[buffer(4)]],
constant uint& trail [[buffer(5)]],
uint3 tid [[thread_position_in_threadgroup]],
uint3 tgp [[threadgroup_position_in_grid]]
) {
threadgroup float tile[32][33];
uint bd = tgp.z;
uint b = bd / trail;
uint d = bd % trail;
if (b >= batch) return;
uint block = rows * cols * trail;
uint plane = b * block + d;
uint r0 = tgp.x * 32u + tid.x;
uint c0 = tgp.y * 32u + tid.y;
for (uint i = 0; i < 32; i += 8) {
uint c = c0 + i;
if (r0 < rows && c < cols) {
tile[tid.x][tid.y + i] =
src[plane + c * rows * trail + r0 * trail];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
uint rr0 = tgp.y * 32u + tid.x;
uint cc0 = tgp.x * 32u + tid.y;
for (uint i = 0; i < 32; i += 8) {
uint rr = rr0; // output column index (in `cols` range)
uint cc = cc0 + i; // output row index (in `rows` range)
if (rr < cols && cc < rows) {
// Output layout [batch, rows, cols, trail]: element (row=cc, col=rr)
// lives at cc*cols*trail + rr*trail (+ plane carries batch & trail).
dst[plane + cc * cols * trail + rr * trail] = tile[tid.y + i][tid.x];
}
}
}
// Two-phase scatter-add: phase 0 zeros the output buffer, phase 1
// accumulates updates atomically. Atomic add is required because
// multiple updates may target the same destination row from different
// threads. `op_phase`: 0 = zero, 1 = accumulate.
//
// Each phase is a single dispatch: phase 0 runs over `out_total` threads,
// phase 1 over `num_updates * trailing` threads. The encoder fires both
// in sequence within one command buffer.
kernel void scatter_add_zero(
device atomic_uint* dst [[buffer(0)]], // bit-cast view of f32 buffer
constant uint& out_total [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= out_total) return;
atomic_store_explicit(&dst[gid], 0u, memory_order_relaxed);
}
kernel void scatter_add_accumulate(
device const float* updates [[buffer(0)]],
device const float* indices [[buffer(1)]],
device atomic_uint* dst [[buffer(2)]], // f32 reinterpreted as u32 atomic
constant uint& trailing [[buffer(3)]],
constant uint& num_updates [[buffer(4)]],
constant uint& out_dim [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint i = gid.y; // which update
uint j = gid.x; // which trailing element
if (i >= num_updates || j >= trailing) return;
uint row = (uint)indices[i];
if (row >= out_dim) return; // OOB safety
float v = updates[i * trailing + j];
// Compare-and-swap loop for atomic float-add. Metal lacks native
// atomic_add for float; reinterpret as uint, CAS the float bits.
uint dst_idx = row * trailing + j;
uint old_bits = atomic_load_explicit(&dst[dst_idx], memory_order_relaxed);
while (true) {
float old_f = as_type<float>(old_bits);
float new_f = old_f + v;
uint new_bits = as_type<uint>(new_f);
if (atomic_compare_exchange_weak_explicit(
&dst[dst_idx], &old_bits, new_bits,
memory_order_relaxed, memory_order_relaxed)) {
break;
}
// CAS failed → old_bits now holds the latest value; retry.
}
}
// Indexed batched matmul (MoE GEMM). One thread per output element
// (i, j). Token i looks up its expert via expert_idx, then dot-products
// the row of `input` against the column of `weight[expert_idx[i]]`.
kernel void grouped_matmul(
device const float* input [[buffer(0)]],
device const float* weight [[buffer(1)]],
device const float* expert_idx [[buffer(2)]],
device float* dst [[buffer(3)]],
constant uint& m [[buffer(4)]],
constant uint& k_dim [[buffer(5)]],
constant uint& n [[buffer(6)]],
constant uint& num_experts [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
uint i = gid.y;
uint j = gid.x;
if (i >= m || j >= n) return;
uint e = (uint)(expert_idx[i]);
if (e >= num_experts) return; // OOB safety
uint w_base = e * k_dim * n;
uint in_base = i * k_dim;
float acc = 0.0f;
for (uint kk = 0; kk < k_dim; ++kk) {
acc += input[in_base + kk] * weight[w_base + kk * n + j];
}
dst[i * n + j] = acc;
}
// Top-K indices along the last axis. One thread per output row. Repeated
// argmax with masking — O(k * axis_dim) per row; fine for small k (MoE
// typical k=2–8). Each thread maintains its own scratch space in private
// memory, no threadgroup coordination needed.
//
// Important: rlx writes float32-encoded indices; downstream Gather reads
// them via `(uint)idx[k]`. Cast on store mirrors that.
kernel void topk_lastax(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& axis_dim [[buffer(2)]],
constant uint& k [[buffer(3)]],
uint o [[thread_position_in_grid]]
) {
// Hard cap on axis_dim — guards the on-chip scratch. MoE expert
// counts top out around 256 in practice; raise this if a real
// workload needs more.
const uint MAX_AXIS = 1024;
if (axis_dim > MAX_AXIS) return;
float scratch[MAX_AXIS];
uint base = o * axis_dim;
for (uint i = 0; i < axis_dim; ++i) scratch[i] = src[base + i];
uint out_base = o * k;
for (uint ki = 0; ki < k; ++ki) {
float best_v = scratch[0];
uint best_i = 0;
for (uint i = 1; i < axis_dim; ++i) {
float v = scratch[i];
if (v > best_v) { best_v = v; best_i = i; }
}
dst[out_base + ki] = (float)best_i;
scratch[best_i] = -INFINITY;
}
}
// Reduce over a contiguous axis range. Input layout [outer, reduced, inner];
// output [outer, inner]. One thread per output element walks `reduced`
// values with stride `inner`. `op_kind`: 0=Sum 1=Mean 2=Max 3=Min 4=Prod.
//
// Trade-off: a serial reduction loop per thread is slower than threadgroup
// reduction when `reduced` is large, but it generalises trivially to any
// axis range and avoids the per-row threadgroup setup cost. For the shapes
// we care about (Reduce::Sum on 60×768 is 22 µs CPU vs 135 µs Metal — the
// wait latency dominates either way), kernel choice barely moves the
// needle. Revisit if a launch-bound reduction shows up.
kernel void reduce_axes(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& reduced [[buffer(2)]],
constant uint& inner [[buffer(3)]],
constant uint& op_kind [[buffer(4)]],
uint2 gid [[thread_position_in_grid]]
) {
uint i = gid.x; // inner axis index
uint o = gid.y; // outer axis index
if (i >= inner) return;
float acc;
if (op_kind == 2) acc = -INFINITY;
else if (op_kind == 3) acc = INFINITY;
else if (op_kind == 4) acc = 1.0f;
else acc = 0.0f; // Sum / Mean
uint base = o * reduced * inner + i;
for (uint r = 0; r < reduced; ++r) {
float v = src[base + r * inner];
if (op_kind == 0 || op_kind == 1) acc += v;
else if (op_kind == 2) acc = max(acc, v);
else if (op_kind == 3) acc = min(acc, v);
else acc *= v;
}
if (op_kind == 1) acc /= float(reduced);
dst[o * inner + i] = acc;
}
// Ternary select: cond != 0 ? a : b. cond is treated as bool via != 0.
kernel void elem_where(
device const float* cond [[buffer(0)]],
device const float* a [[buffer(1)]],
device const float* b [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& len [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
out[gid] = cond[gid] != 0.0f ? a[gid] : b[gid];
}
// In-place ReLU: data = max(0, data)
kernel void relu_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = max(0.0f, data[gid]);
}
// In-place sigmoid: 1 / (1 + exp(-x))
kernel void sigmoid_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = 1.0f / (1.0f + exp(-data[gid]));
}
// In-place tan
kernel void tan_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = tan(data[gid]);
}
// In-place atan
kernel void atan_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = atan(data[gid]);
}
// In-place sin
kernel void sin_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = sin(data[gid]);
}
// In-place cos
kernel void cos_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = cos(data[gid]);
}
// In-place tanh
kernel void tanh_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = tanh(data[gid]);
}
// In-place exp / log / sqrt / rsqrt / neg / abs — one kernel each so the
// dispatch path stays uniform with the existing `*_inplace` family.
kernel void exp_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = exp(data[gid]); }
kernel void log_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = log(data[gid]); }
kernel void sqrt_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = sqrt(data[gid]); }
kernel void rsqrt_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = rsqrt(data[gid]); }
kernel void neg_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = -data[gid]; }
kernel void abs_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = abs(data[gid]); }
kernel void round_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = round(data[gid]); }
// Standalone softmax along the last axis. One threadgroup per row,
// reduces max + exp-sum across the row, then normalizes. tg_size is
// the actual number of threads per group (passed via threads_per_threadgroup).
kernel void softmax_lastax(
device float* data [[buffer(0)]],
constant uint& cols [[buffer(1)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial[256];
uint base = row * cols;
// Pass 1: find row max for numerical stability.
float local_max = -INFINITY;
for (uint i = tid; i < cols; i += tsize) {
local_max = max(local_max, data[base + i]);
}
partial[tid] = (float)local_max;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial[tid] = max(partial[tid], partial[tid + stride]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float row_max = partial[0];
// Pass 2: exp(x - max) and sum.
float local_sum = 0.0f;
for (uint i = tid; i < cols; i += tsize) {
float e = exp(data[base + i] - row_max);
data[base + i] = e;
local_sum += e;
}
partial[tid] = (float)local_sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial[tid] += partial[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_sum = 1.0f / partial[0];
// Pass 3: normalize.
for (uint i = tid; i < cols; i += tsize) {
data[base + i] *= inv_sum;
}
}
// Embedding lookup: out[i, .] = table[idx[i], .]
// table: [vocab, trailing], idx: [num_idx], out: [num_idx, trailing]
kernel void gather_axis0(
device const float* table [[buffer(0)]],
device const float* idx [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& num_idx [[buffer(3)]],
constant uint& trailing [[buffer(4)]],
uint2 gid [[thread_position_in_grid]]
) {
uint i = gid.y;
uint j = gid.x;
if (i >= num_idx || j >= trailing) return;
uint row = uint(idx[i]);
out[i * trailing + j] = table[row * trailing + j];
}
// Narrow / slice along last axis. src is [outer, src_axis], dst is [outer, len].
// Each invocation copies one (outer, j) element.
kernel void narrow_lastax(
device const char* arena_src [[buffer(0)]],
device char* arena_dst [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis [[buffer(3)]],
constant uint& start [[buffer(4)]],
constant uint& len [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
constant ulong& dst_byte_off [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
// Task #50: > 4 GB activations need ulong byte offsets.
device const float* src = (device const float*)(arena_src + src_byte_off);
device float* dst = (device float*)(arena_dst + dst_byte_off);
uint i = gid.y;
uint j = gid.x;
if (i >= outer || j >= len) return;
dst[i * len + j] = src[i * src_axis + start + j];
}
// Vectorized narrow for aligned shapes: src/dst treated as packed_float4.
// Requirements (enforced in encoder): start, src_axis, len are divisible by 4.
kernel void narrow_lastax4(
device const char* arena_src [[buffer(0)]],
device char* arena_dst [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis4 [[buffer(3)]],
constant uint& start4 [[buffer(4)]],
constant uint& len4 [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
constant ulong& dst_byte_off [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
device const packed_float4* src = (device const packed_float4*)(arena_src + src_byte_off);
device packed_float4* dst = (device packed_float4*)(arena_dst + dst_byte_off);
uint i = gid.y;
uint j4 = gid.x;
if (i >= outer || j4 >= len4) return;
dst[i * len4 + j4] = src[i * src_axis4 + start4 + j4];
}
// `dst` widened to ulong (task #50: ≥4 GB models have activation byte
// offsets that exceed u32 — truncated wrap-around made K_rep narrows
// write to the wrong slot and SDPA saw all-zero K_rep).
struct NarrowSeg {
ulong dst;
uint start;
uint len;
};
// Concat VJP: multiple last-axis slices from one source in one dispatch.
kernel void split_lastax(
device const char* arena_src [[buffer(0)]],
device char* arena [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis [[buffer(3)]],
constant uint& num_seg [[buffer(4)]],
constant NarrowSeg* segs [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
uint3 gid [[thread_position_in_grid]]
) {
device const float* src = (device const float*)(arena_src + src_byte_off);
uint s = gid.z;
uint i = gid.y;
uint j = gid.x;
if (s >= num_seg) return;
NarrowSeg seg = segs[s];
if (i >= outer || j >= seg.len) return;
device float* dst = (device float*)(arena + seg.dst);
dst[i * seg.len + j] = src[i * src_axis + seg.start + j];
}
kernel void split_lastax4(
device const char* arena_src [[buffer(0)]],
device char* arena [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis4 [[buffer(3)]],
constant uint& num_seg [[buffer(4)]],
constant NarrowSeg* segs [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
uint3 gid [[thread_position_in_grid]]
) {
device const packed_float4* src = (device const packed_float4*)(arena_src + src_byte_off);
uint s = gid.z;
uint i = gid.y;
uint j4 = gid.x;
if (s >= num_seg) return;
NarrowSeg seg = segs[s];
uint len4 = seg.len / 4u;
if (i >= outer || j4 >= len4) return;
device packed_float4* dst = (device packed_float4*)(arena + seg.dst);
uint start4 = seg.start / 4u;
dst[i * len4 + j4] = src[i * src_axis4 + start4 + j4];
}
// Concat segment: copy one [outer, src_axis] tensor into [outer, dst_axis]
// at the column slice [dst_col .. dst_col + src_axis]. Multi-input concat
// = N dispatches of this kernel, one per source. Mirror of narrow_lastax.
kernel void concat_segment_lastax(
device const char* arena_src [[buffer(0)]],
device char* arena_dst [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis [[buffer(3)]],
constant uint& dst_axis [[buffer(4)]],
constant uint& dst_col [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
constant ulong& dst_byte_off [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
// Task #50: large set_buffer offsets silently lose kernel writes on
// M-series — bind arena base and apply byte offsets here.
device const float* src = (device const float*)(arena_src + src_byte_off);
device float* dst = (device float*)(arena_dst + dst_byte_off);
uint i = gid.y;
uint j = gid.x;
if (i >= outer || j >= src_axis) return;
dst[i * dst_axis + dst_col + j] = src[i * src_axis + j];
}
kernel void concat_segment_lastax4(
device const char* arena_src [[buffer(0)]],
device char* arena_dst [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis4 [[buffer(3)]],
constant uint& dst_axis4 [[buffer(4)]],
constant uint& dst_col4 [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
constant ulong& dst_byte_off [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
device const packed_float4* src = (device const packed_float4*)(arena_src + src_byte_off);
device packed_float4* dst = (device packed_float4*)(arena_dst + dst_byte_off);
uint i = gid.y;
uint j4 = gid.x;
if (i >= outer || j4 >= src_axis4) return;
dst[i * dst_axis4 + dst_col4 + j4] = src[i * src_axis4 + j4];
}
// `src` widened to ulong (task #50: ≥4 GB models have activation byte
// offsets > u32 — truncated wrap-around made `repeat_kv` write the wrong
// slot and SDPA saw all-zero K_rep).
struct ConcatSeg {
ulong src;
uint dst_col;
uint len;
};
kernel void concat_lastax_multi(
device char* arena [[buffer(0)]],
constant ulong& dst_byte [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& dst_axis [[buffer(3)]],
constant uint& num_seg [[buffer(4)]],
constant ConcatSeg* segs [[buffer(5)]],
uint3 gid [[thread_position_in_grid]]
) {
uint s = gid.z;
uint i = gid.y;
uint j = gid.x;
if (s >= num_seg) return;
ConcatSeg seg = segs[s];
if (i >= outer || j >= seg.len) return;
device const float* src = (device const float*)(arena + seg.src);
device float* dst = (device float*)(arena + dst_byte);
dst[i * dst_axis + seg.dst_col + j] = src[i * seg.len + j];
}
kernel void concat_lastax_multi4(
device char* arena [[buffer(0)]],
constant ulong& dst_byte [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& dst_axis4 [[buffer(3)]],
constant uint& num_seg [[buffer(4)]],
constant ConcatSeg* segs [[buffer(5)]],
uint3 gid [[thread_position_in_grid]]
) {
uint s = gid.z;
uint i = gid.y;
uint j4 = gid.x;
if (s >= num_seg) return;
ConcatSeg seg = segs[s];
uint len4 = seg.len / 4u;
if (i >= outer || j4 >= len4) return;
device const packed_float4* src =
(device const packed_float4*)(arena + seg.src);
device packed_float4* dst =
(device packed_float4*)(arena + dst_byte);
uint dst_col4 = seg.dst_col / 4u;
dst[i * dst_axis4 + dst_col4 + j4] = src[i * len4 + j4];
}
kernel void concat_segment_lastax_h(
device const char* arena_src [[buffer(0)]],
device char* arena_dst [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& src_axis [[buffer(3)]],
constant uint& dst_axis [[buffer(4)]],
constant uint& dst_col [[buffer(5)]],
constant ulong& src_byte_off [[buffer(6)]],
constant ulong& dst_byte_off [[buffer(7)]],
uint2 gid [[thread_position_in_grid]]
) {
device const half* src = (device const half*)(arena_src + src_byte_off);
device half* dst = (device half*)(arena_dst + dst_byte_off);
uint i = gid.y;
uint j = gid.x;
if (i >= outer || j >= src_axis) return;
dst[i * dst_axis + dst_col + j] = src[i * src_axis + j];
}
// Fused residual + LN: out = LN(x + residual + bias, gamma, beta)
// (bias is broadcast per row; pass empty/null offset for no-bias variant)
kernel void fused_residual_ln(
device const float* x [[buffer(0)]],
device const float* res [[buffer(1)]],
device const float* gamma [[buffer(2)]],
device const float* beta [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& h [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
// Pass 1: compute (x + res) on the fly, accumulate sum/sumsq
float local_sum = 0.0;
float local_sumsq = 0.0;
for (uint i = tid; i < h; i += tsize) {
float v = x[row * h + i] + res[row * h + i];
local_sum += v;
local_sumsq += v * v;
}
partial_sum[tid] = local_sum;
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sum[tid] += partial_sum[tid + stride];
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float mean = partial_sum[0] / float(h);
float var = partial_sumsq[0] / float(h) - mean * mean;
float inv_std = rsqrt(var + eps);
// Pass 2: write normalized output
for (uint i = tid; i < h; i += tsize) {
float v = x[row * h + i] + res[row * h + i];
out[row * h + i] = (v - mean) * inv_std * gamma[i] + beta[i];
}
}
// Fused residual + RMSNorm: out = RmsNorm(x + residual, gamma, beta)
kernel void fused_residual_rms_norm(
device const float* x [[buffer(0)]],
device const float* res [[buffer(1)]],
device const float* gamma [[buffer(2)]],
device const float* beta [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& h [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial_sumsq[256];
float local_sumsq = 0.0;
for (uint i = tid; i < h; i += tsize) {
float v = x[row * h + i] + res[row * h + i];
local_sumsq += v * v;
}
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_rms = rsqrt(partial_sumsq[0] / float(h) + eps);
for (uint i = tid; i < h; i += tsize) {
float v = x[row * h + i] + res[row * h + i];
out[row * h + i] = v * inv_rms * gamma[i] + beta[i];
}
}
// Packed SDPA byte offsets (task #50: lets the kernel reach activations
// that sit past 4 GB in the arena without needing `set_buffer(... ,large)`,
// which silently dropped kernel writes on M-series).
struct SdpaOffsets {
ulong q;
ulong k;
ulong v;
ulong m;
ulong o;
};
// Q/K/V offset helpers — BSNH [B, L, H*D] vs BHSD [B, H, L, D].
static inline uint qkv_q_offset(
uint bi, uint hi, uint qi,
uint heads, uint seq_q, uint head_dim, uint q_stride, uint bhsd
) {
if (bhsd != 0u) {
return bi * heads * seq_q * head_dim + hi * seq_q * head_dim + qi * head_dim;
}
uint hs = heads * head_dim;
return bi * q_stride * hs + qi * hs + hi * head_dim;
}
static inline uint qkv_kv_offset(
uint bi, uint hi, uint ki,
uint heads, uint seq_k, uint head_dim, uint k_stride, uint bhsd
) {
if (bhsd != 0u) {
return bi * heads * seq_k * head_dim + hi * seq_k * head_dim + ki * head_dim;
}
uint hs = heads * head_dim;
return bi * k_stride * hs + ki * hs + hi * head_dim;
}
// Multi-head SDPA: attention(Q, K, V, mask) → out
// Shapes: Q/out [batch, seq_q, heads*head_dim]; K/V [batch, seq_k, heads*head_dim]
// One threadgroup per (batch, head). Each TG computes [seq_q, seq_k] scores
// in threadgroup memory (seq_q * seq_k ≤ 64*64), applies softmax, then
// accumulates scores @ V.
kernel void sdpa(
device const float* arena_q [[buffer(0)]],
device const float* arena_k [[buffer(1)]],
device const float* arena_v [[buffer(2)]],
device const float* arena_m [[buffer(3)]],
device float* arena_o [[buffer(4)]],
constant uint& batch [[buffer(5)]],
constant uint& seq_q [[buffer(6)]],
constant uint& heads [[buffer(7)]],
constant uint& head_dim [[buffer(8)]],
constant uint& q_stride [[buffer(9)]],
constant uint& mask_kind [[buffer(10)]],
constant uint& seq_k [[buffer(11)]],
constant uint& k_stride [[buffer(12)]],
constant uint& bhsd [[buffer(13)]],
constant uint& window [[buffer(14)]],
constant float& score_scale [[buffer(15)]],
constant float& attn_softcap [[buffer(16)]],
// Byte offsets relative to the arena buffer, packed as one struct.
// For ≥4 GB models the activation byte offsets exceed u32 — and Metal
// silently drops kernel writes when `set_buffer` is called with
// `offset > 4 GB`. Binding all five buffers to `offset=0` and adding
// the offsets here is the workaround proven by the dequant kernel
// (works for offsets ≥ 14 GB). One inline-constant slot for all five
// (task #50).
constant SdpaOffsets& byte_offs [[buffer(17)]],
uint tgid_x [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
ulong q_byte_off = byte_offs.q;
ulong k_byte_off = byte_offs.k;
ulong v_byte_off = byte_offs.v;
ulong m_byte_off = byte_offs.m;
ulong o_byte_off = byte_offs.o;
device const float* Q = (device const float*)((device const char*)arena_q + q_byte_off);
device const float* K = (device const float*)((device const char*)arena_k + k_byte_off);
device const float* V = (device const float*)((device const char*)arena_v + v_byte_off);
device const float* M = (device const float*)((device const char*)arena_m + m_byte_off);
device float* OUT = (device float*)((device char*)arena_o + o_byte_off);
// mask_kind:
// 0 = None (no masking)
// 1 = Causal (mask ki > (seq_k - seq_q) + qi)
// 2 = Custom (column-wise binary mask buffer M; 0 = padded)
// 4 = SlidingWindow (visible range [abs_q - window, abs_q],
// absolute positions so decode w/ cached K/V works)
threadgroup float scores[64 * 64]; // up to seq_q * seq_k = 4096
threadgroup float row_max;
threadgroup float row_sum;
// Linearized: tgid_x = bi * heads + hi
uint bi = tgid_x / heads;
uint hi = tgid_x % heads;
if (bi >= batch) return;
// `score_scale` is the host-provided multiplier (Gemma 4 sets 1.0
// because Q is per-head RMS-normed before attention). Sentinel `0.0`
// means "use the relaxed-precision default `1/sqrt(head_dim)`".
float scale = (score_scale > 0.0f) ? score_scale : 1.0f / precise::sqrt(float(head_dim));
// Gemma 2 carries an attention-logit softcap (50.0); Gemma 4 sets 0.
float softcap_inv = (attn_softcap > 0.0f) ? (1.0f / attn_softcap) : 0.0f;
uint q_offset = seq_k - seq_q;
// 1. Compute scores[qi, ki] = scale * (Q[bi, qi, hi*dh:] · K[bi, ki, hi*dh:]) + mask.
uint total = seq_q * seq_k;
for (uint idx = tid; idx < total; idx += tsize) {
uint qi = idx / seq_k;
uint ki = idx % seq_k;
float dot = 0.0;
uint q_base = qkv_q_offset(bi, hi, qi, heads, seq_q, head_dim, q_stride, bhsd);
uint k_base = qkv_kv_offset(bi, hi, ki, heads, seq_k, head_dim, k_stride, bhsd);
for (uint d = 0; d < head_dim; ++d) {
dot += Q[q_base + d] * K[k_base + d];
}
float s = dot * scale;
if (softcap_inv > 0.0f) {
s = precise::tanh(s * softcap_inv) * attn_softcap;
}
if (mask_kind == 1u) {
if (ki > q_offset + qi) s = -1e9;
} else if (mask_kind == 2u) {
if (M[bi * k_stride + ki] < 0.5) s = -1e9;
} else if (mask_kind == 4u) {
uint abs_q = q_offset + qi;
uint lo = abs_q > window ? abs_q - window : 0u;
if (ki < lo || ki > abs_q) s = -1e9;
}
scores[qi * seq_k + ki] = s;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// 2. Softmax row-by-row over scores[seq_q, seq_k]. `precise::exp`
// matches CPU `f32::exp` to within 1 ULP; the default fast-math
// `exp` accumulates several ULPs of error per token, which the
// softcap + LM head amplify into visible logit drift.
for (uint qi = 0; qi < seq_q; ++qi) {
if (tid == 0) {
float mx = -1e30;
for (uint ki = 0; ki < seq_k; ++ki) {
mx = max(mx, scores[qi * seq_k + ki]);
}
row_max = mx;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
if (tid == 0) {
float sum = 0.0;
for (uint ki = 0; ki < seq_k; ++ki) {
float e = precise::exp(scores[qi * seq_k + ki] - row_max);
scores[qi * seq_k + ki] = e;
sum += e;
}
row_sum = sum;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint ki = tid; ki < seq_k; ki += tsize) {
scores[qi * seq_k + ki] /= row_sum;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
// 3. Output[qi, d] = sum_ki scores[qi, ki] * V[bi, ki, hi*dh + d]
uint out_total = seq_q * head_dim;
for (uint idx = tid; idx < out_total; idx += tsize) {
uint qi = idx / head_dim;
uint d = idx % head_dim;
float acc = 0.0;
for (uint ki = 0; ki < seq_k; ++ki) {
uint v_base = qkv_kv_offset(bi, hi, ki, heads, seq_k, head_dim, k_stride, bhsd);
acc += scores[qi * seq_k + ki] * V[v_base + d];
}
uint o_base = qkv_q_offset(bi, hi, qi, heads, seq_q, head_dim, q_stride, bhsd);
OUT[o_base + d] = acc;
}
}
// Online-softmax SDPA (FlashAttention v1 inner-row form). Same algorithm
// as `wgpu/src/kernels/attention.wgsl` and `cpu/src/thunk.rs` Attention.
// One thread per (batch, head, q_row); each thread walks the K dimension
// exactly once, maintaining a running (m, l, O[D]) tuple — no scores
// matrix in threadgroup memory, so it scales to arbitrary seq length.
//
// The plain `sdpa` kernel above uses `threadgroup float scores[64*64]`;
// for vision (seq=257) that overflows. This kernel handles seq > 64.
//
// Mask layout (vision constant all-ones is `[batch, seq_stride]`):
// reads M[bi * seq_stride + ki] just like `sdpa`.
//
// MAX_HEAD_DIM = 128 covers BERT/Nomic/Vision (head_dim ≤ 128); larger
// head dims would need a per-thread spill buffer.
kernel void sdpa_long(
device const float* arena_q [[buffer(0)]],
device const float* arena_k [[buffer(1)]],
device const float* arena_v [[buffer(2)]],
device const float* arena_m [[buffer(3)]],
device float* arena_o [[buffer(4)]],
constant uint& batch [[buffer(5)]],
constant uint& seq_q [[buffer(6)]], // query length Lq
constant uint& heads [[buffer(7)]],
constant uint& head_dim [[buffer(8)]],
constant uint& q_stride [[buffer(9)]], // per-batch Q row stride (= Lq for dense)
constant uint& mask_kind [[buffer(10)]],
constant uint& seq_k [[buffer(11)]], // key/value length Lk
constant uint& k_stride [[buffer(12)]], // per-batch K/V row stride (= Lk for dense)
constant uint& bhsd [[buffer(13)]], // 1 = [B,H,S,D]
constant uint& window [[buffer(14)]], // SlidingWindow lookback (0 otherwise)
constant float& score_scale [[buffer(15)]],
constant float& attn_softcap [[buffer(16)]],
// Task #50: > 4 GB activations need offsets in inline constants.
constant SdpaOffsets& byte_offs [[buffer(17)]],
uint tid_x [[thread_position_in_grid]]
) {
device const float* Q = (device const float*)((device const char*)arena_q + byte_offs.q);
device const float* K = (device const float*)((device const char*)arena_k + byte_offs.k);
device const float* V = (device const float*)((device const char*)arena_v + byte_offs.v);
device const float* M = (device const float*)((device const char*)arena_m + byte_offs.m);
device float* OUT = (device float*)((device char*)arena_o + byte_offs.o);
// mask_kind:
// 0 = None
// 1 = Causal (prefill — Lq == Lk required)
// 2 = Custom (binary key-padding mask M[B, Lk])
// 3 = Bias (additive per-head bias M[B, H, Lq, Lk])
// 4 = SlidingWindow (visible range [abs_q - window, abs_q])
//
// Gemma 4 12B's SWA layers use head_dim=256 and the FULL layers use
// head_dim=512. The previous 128 cap silently overflowed q_reg/o_acc
// and produced all-NaN logits when decode picked this kernel.
constexpr uint MAX_HEAD_DIM = 512u;
uint total = batch * heads * seq_q;
if (tid_x >= total) return;
uint qi = tid_x % seq_q;
uint bh = tid_x / seq_q;
uint hi = bh % heads;
uint bi = bh / heads;
// Precise scale; honour caller's `score_scale` (Gemma 4 = 1.0).
float scale = (score_scale > 0.0f) ? score_scale : 1.0f / precise::sqrt(float(head_dim));
float softcap_inv = (attn_softcap > 0.0f) ? (1.0f / attn_softcap) : 0.0f;
// Cache Q[qi, hi*dh : (hi+1)*dh] in registers — read seq_k times below.
float q_reg[MAX_HEAD_DIM];
uint q_base = qkv_q_offset(bi, hi, qi, heads, seq_q, head_dim, q_stride, bhsd);
for (uint d = 0; d < head_dim; ++d) q_reg[d] = Q[q_base + d];
// Bias base offset (only read when mask_kind == 3).
uint bias_row_base = ((bi * heads + hi) * seq_q + qi) * seq_k;
uint q_offset = seq_k - seq_q;
// Online softmax accumulators.
float m_acc = -1e30;
float l_acc = 0.0;
float o_acc[MAX_HEAD_DIM];
for (uint d = 0; d < head_dim; ++d) o_acc[d] = 0.0;
for (uint ki = 0; ki < seq_k; ++ki) {
// Score: scale * (Q · K[ki]) + mask
uint k_base = qkv_kv_offset(bi, hi, ki, heads, seq_k, head_dim, k_stride, bhsd);
float dot = 0.0;
for (uint d = 0; d < head_dim; ++d) dot += q_reg[d] * K[k_base + d];
float s = dot * scale;
if (softcap_inv > 0.0f) {
s = precise::tanh(s * softcap_inv) * attn_softcap;
}
if (mask_kind == 1u) {
if (ki > q_offset + qi) s = -1e9;
} else if (mask_kind == 2u) {
if (M[bi * k_stride + ki] < 0.5) s = -1e9;
} else if (mask_kind == 3u) {
s += M[bias_row_base + ki];
} else if (mask_kind == 4u) {
uint abs_q = q_offset + qi;
uint lo = abs_q > window ? abs_q - window : 0u;
if (ki < lo || ki > abs_q) s = -1e9;
}
// Online softmax update with precise exp.
float m_new = max(m_acc, s);
float e_old = precise::exp(m_acc - m_new);
float e_cur = precise::exp(s - m_new);
l_acc = e_old * l_acc + e_cur;
uint v_base = qkv_kv_offset(bi, hi, ki, heads, seq_k, head_dim, k_stride, bhsd);
for (uint d = 0; d < head_dim; ++d) {
o_acc[d] = e_old * o_acc[d] + e_cur * V[v_base + d];
}
m_acc = m_new;
}
// Normalize and emit.
float inv_l = 1.0 / l_acc;
uint o_base = qkv_q_offset(bi, hi, qi, heads, seq_q, head_dim, q_stride, bhsd);
for (uint d = 0; d < head_dim; ++d) {
OUT[o_base + d] = o_acc[d] * inv_l;
}
}
// Flash-attention tile kernel with optional additive bias mask.
//
// Targets the SAM3 detector decoder image cross-attention where the
// scalar `sdpa_long` is bandwidth-bound (each query thread re-reads K
// and V for all 5184 positions). This kernel processes Br=8 query
// rows per threadgroup with K, V, and bias tiles loaded cooperatively
// into threadgroup memory — each K/V/bias element is read once per
// row tile instead of once per query.
//
// Layout matches `sdpa_long`: Q/K/V are [B, Lq_or_Lk, heads*head_dim],
// bias is [B, H, Lq, Lk]. head_dim is dynamic but capped at 128 for
// the per-thread output accumulator.
kernel void sdpa_fa_f32(
device const float* Q [[buffer(0)]],
device const float* K [[buffer(1)]],
device const float* V [[buffer(2)]],
device const float* M [[buffer(3)]],
device float* OUT [[buffer(4)]],
constant uint& batch [[buffer(5)]],
constant uint& seq_q [[buffer(6)]],
constant uint& heads [[buffer(7)]],
constant uint& head_dim [[buffer(8)]],
constant uint& q_stride [[buffer(9)]],
constant uint& mask_kind [[buffer(10)]],
constant uint& seq_k [[buffer(11)]],
constant uint& k_stride [[buffer(12)]],
constant uint& bhsd [[buffer(13)]],
constant uint& window [[buffer(14)]], // reserved; not yet wired in FA tile path
uint3 tgid [[threadgroup_position_in_grid]],
uint tid_in_tg [[thread_index_in_threadgroup]]
) {
(void)window; // SlidingWindow falls through to sdpa_long today
// Tile sizes — tuned for SAM3 image CA (dh=16) but kernel is
// generic. With Br=8, Bc=64, the per-TG threadgroup memory is
// 8*128 (Q) + 64*128 (K) + 64*128 (V) + 8*64 (S/M) ≈ 71KB at
// dh=128; well under the 32–64KB per-TG hard limit at dh=16
// (where it's ~10KB).
// Tile sizes — the threadgroup-memory cap on Apple7/8 (32KB) and
// Apple9 (64KB) bounds `MAX_DH`. At MAX_DH=32 we use ~20KB,
// leaving headroom for larger Bc later. dh up to 32 covers SAM
// family models (dh=16) and DETR-style detectors. Larger dh
// (LLM 64–128) falls back to scalar sdpa_long via the dispatch
// guard in `encode_sdpa`.
constexpr uint Br = 8u;
constexpr uint Bc = 64u;
constexpr uint MAX_DH = 32u;
constexpr uint THREADS = 64u;
threadgroup float Q_tg[Br * MAX_DH]; // 1 KB
threadgroup float K_tg[Bc * MAX_DH]; // 8 KB
threadgroup float V_tg[Bc * MAX_DH]; // 8 KB
threadgroup float S_tg[Br * Bc]; // 2 KB
// Per-row online softmax state.
threadgroup float m_row[Br];
threadgroup float l_row[Br];
threadgroup float o_row[Br * MAX_DH]; // 1 KB
uint q_tile = tgid.x; // index over Lq / Br
uint hi = tgid.y; // head
uint bi = tgid.z; // batch
uint q_start = q_tile * Br;
float scale = rsqrt(float(head_dim));
// ── Load Q tile cooperatively ────────────────────────────────────
for (uint i = tid_in_tg; i < Br * head_dim; i += THREADS) {
uint qi = i / head_dim;
uint di = i % head_dim;
uint pos = q_start + qi;
Q_tg[qi * MAX_DH + di] = (pos < seq_q)
? Q[qkv_q_offset(bi, hi, pos, heads, seq_q, head_dim, q_stride, bhsd) + di]
: 0.0f;
}
// Initialize per-row state.
if (tid_in_tg < Br) {
m_row[tid_in_tg] = -1e30f;
l_row[tid_in_tg] = 0.0f;
}
for (uint i = tid_in_tg; i < Br * head_dim; i += THREADS) {
o_row[(i / head_dim) * MAX_DH + (i % head_dim)] = 0.0f;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// ── Iterate K/V tiles ─────────────────────────────────────────────
uint bias_row_base = (bi * heads + hi) * seq_q * seq_k;
for (uint kt = 0; kt < seq_k; kt += Bc) {
// Load K and V tiles (Bc * head_dim elements each).
for (uint i = tid_in_tg; i < Bc * head_dim; i += THREADS) {
uint ki = i / head_dim;
uint di = i % head_dim;
uint pos = kt + ki;
uint kv_off = qkv_kv_offset(bi, hi, pos, heads, seq_k, head_dim, k_stride, bhsd);
bool in_range = pos < seq_k;
K_tg[ki * MAX_DH + di] = in_range ? K[kv_off + di] : 0.0f;
V_tg[ki * MAX_DH + di] = in_range ? V[kv_off + di] : 0.0f;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// Compute scores S[Br, Bc] = Q_tg @ K_tg^T, scaled, +bias, +pad-mask.
// Each thread covers Br*Bc/THREADS = 8*64/64 = 8 cells.
for (uint c = tid_in_tg; c < Br * Bc; c += THREADS) {
uint qi = c / Bc;
uint ki = c % Bc;
uint pos = kt + ki;
bool valid = (q_start + qi) < seq_q && pos < seq_k;
float s = 0.0f;
if (valid) {
for (uint di = 0; di < head_dim; ++di) {
s += Q_tg[qi * MAX_DH + di] * K_tg[ki * MAX_DH + di];
}
s *= scale;
if (mask_kind == 1u) {
uint q_offset = seq_k - seq_q;
if (pos > q_offset + q_start + qi) s = -1e9f;
} else if (mask_kind == 2u) {
if (M[bi * k_stride + pos] < 0.5f) s = -1e9f;
} else if (mask_kind == 3u) {
s += M[bias_row_base + (q_start + qi) * seq_k + pos];
}
} else {
s = -1e9f;
}
S_tg[c] = s;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
// Online softmax update — one thread per row (Br threads).
if (tid_in_tg < Br) {
uint qi = tid_in_tg;
float m_new = m_row[qi];
for (uint ki = 0; ki < Bc; ++ki) {
m_new = max(m_new, S_tg[qi * Bc + ki]);
}
float e_old = exp(m_row[qi] - m_new);
float l_new = e_old * l_row[qi];
for (uint ki = 0; ki < Bc; ++ki) {
float p = exp(S_tg[qi * Bc + ki] - m_new);
S_tg[qi * Bc + ki] = p;
l_new += p;
}
// O ← e_old * O + P @ V
for (uint di = 0; di < head_dim; ++di) {
float o = o_row[qi * MAX_DH + di] * e_old;
for (uint ki = 0; ki < Bc; ++ki) {
o += S_tg[qi * Bc + ki] * V_tg[ki * MAX_DH + di];
}
o_row[qi * MAX_DH + di] = o;
}
m_row[qi] = m_new;
l_row[qi] = l_new;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
// ── Normalize + emit ─────────────────────────────────────────────
for (uint i = tid_in_tg; i < Br * head_dim; i += THREADS) {
uint qi = i / head_dim;
uint di = i % head_dim;
uint pos = q_start + qi;
if (pos < seq_q) {
float o = o_row[qi * MAX_DH + di] / l_row[qi];
OUT[qkv_q_offset(bi, hi, pos, heads, seq_q, head_dim, q_stride, bhsd) + di] = o;
}
}
}
// RoPE: apply rotary position embeddings to one tensor (Q or K).
// x: [batch, seq, hidden], hidden = num_heads * head_dim
// cos/sin: [max_pos, head_dim/2]
// Out-of-place into out (or in-place via aliasing).
kernel void rope(
device const float* x [[buffer(0)]],
device const float* cos [[buffer(1)]],
device const float* sin [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& batch [[buffer(4)]],
constant uint& seq [[buffer(5)]],
constant uint& hidden [[buffer(6)]],
constant uint& head_dim [[buffer(7)]],
constant uint& src_row_stride [[buffer(8)]],
constant uint& seq_stride [[buffer(9)]],
constant uint& n_rot [[buffer(10)]],
uint3 gid [[thread_position_in_grid]]
) {
// gid.x = dim index within head (0..head_dim)
// gid.y = head index
// gid.z = batch * seq + seq pos (linearized)
uint half_dh = head_dim / 2;
uint rot_half = n_rot / 2;
if (gid.x >= head_dim) return;
uint bs = gid.z;
uint bi = bs / seq;
uint si = bs % seq;
if (bi >= batch || si >= seq) return;
uint nh = hidden / head_dim;
uint hi = gid.y;
if (hi >= nh) return;
// PLAN L1 — `seq_stride` is the compile-time full extent for buffer
// offsets; `seq` is the (possibly scaled) iteration bound. This
// separation lets active-extent dispatch shrink the loop without
// corrupting per-batch strides.
uint src_base = bi * seq_stride * src_row_stride + si * src_row_stride + hi * head_dim;
uint dst_base = bi * seq_stride * hidden + si * hidden + hi * head_dim;
uint d = gid.x;
if (d < rot_half) {
float x1 = x[src_base + d];
float x2 = x[src_base + rot_half + d];
float c = cos[si * half_dh + d];
float s = sin[si * half_dh + d];
out[dst_base + d] = x1 * c - x2 * s;
out[dst_base + rot_half + d] = x2 * c + x1 * s;
} else if (d >= n_rot) {
out[dst_base + d] = x[src_base + d];
}
}
// in-place SiLU: x * sigmoid(x)
kernel void silu_inplace(
device float* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float x = data[gid];
data[gid] = x / (1.0 + exp(-x));
}
// Fused SwiGLU: input is concat'd [outer, 2N] (per-row up || gate).
// Output: [outer, N] where out[r,i] = up[r,i] * silu(gate[r,i]).
// One thread per output element. Each thread reads exactly two source
// values from the same row (up + gate) and writes one — no inter-thread
// communication, no shared memory, no reductions.
//
// Grid: total output elements (outer * N). The thread maps to (row, col)
// via the n_half stride. Up and gate live at offsets [row*2N + col] and
// [row*2N + N + col] respectively.
kernel void fused_swiglu(
device const float* x [[buffer(0)]], // [outer, 2*n_half]
device float* out [[buffer(1)]], // [outer, n_half]
constant uint& n_half [[buffer(2)]],
constant uint& total [[buffer(3)]], // outer * n_half
constant uint& gate_first [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= total) return;
uint row = gid / n_half;
uint col = gid % n_half;
uint base = row * (2u * n_half);
float up;
float gate;
if (gate_first != 0u) {
gate = x[base + col];
up = x[base + n_half + col];
} else {
up = x[base + col];
gate = x[base + n_half + col];
}
out[gid] = up * (gate / (1.0f + exp(-gate)));
}
// Half-precision variant: f16 in/out. Computation in f32 (silu's exp can
// underflow at half precision). Same dispatch as fused_swiglu.
kernel void fused_swiglu_h(
device const half* x [[buffer(0)]],
device half* out [[buffer(1)]],
constant uint& n_half [[buffer(2)]],
constant uint& total [[buffer(3)]],
constant uint& gate_first [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= total) return;
uint row = gid / n_half;
uint col = gid % n_half;
uint base = row * (2u * n_half);
float up;
float gate;
if (gate_first != 0u) {
gate = float(x[base + col]);
up = float(x[base + n_half + col]);
} else {
up = float(x[base + col]);
gate = float(x[base + n_half + col]);
}
out[gid] = half(up * (gate / (1.0f + exp(-gate))));
}
// SwiGLU + cast: f32 input, f16 output. Saves a separate cast pass when
// the next consumer wants half precision. Reserved for paths where the
// AutoMixedPrecision boundary lands right after SwiGLU.
kernel void fused_swiglu_cast_f32_to_f16(
device const float* x [[buffer(0)]],
device half* out [[buffer(1)]],
constant uint& n_half [[buffer(2)]],
constant uint& total [[buffer(3)]],
constant uint& gate_first [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= total) return;
uint row = gid / n_half;
uint col = gid % n_half;
uint base = row * (2u * n_half);
float up;
float gate;
if (gate_first != 0u) {
gate = x[base + col];
up = x[base + n_half + col];
} else {
up = x[base + col];
gate = x[base + n_half + col];
}
out[gid] = half(up * (gate / (1.0f + exp(-gate))));
}
// SwiGLU + cast: f16 input, f32 output. Symmetric to the above.
kernel void fused_swiglu_cast_f16_to_f32(
device const half* x [[buffer(0)]],
device float* out [[buffer(1)]],
constant uint& n_half [[buffer(2)]],
constant uint& total [[buffer(3)]],
constant uint& gate_first [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= total) return;
uint row = gid / n_half;
uint col = gid % n_half;
uint base = row * (2u * n_half);
float up;
float gate;
if (gate_first != 0u) {
gate = float(x[base + col]);
up = float(x[base + n_half + col]);
} else {
up = float(x[base + col]);
gate = float(x[base + n_half + col]);
}
out[gid] = up * (gate / (1.0f + exp(-gate)));
}
// LayerNorm: out = (x - mean) * inv_std * gamma + beta, per row
// One threadgroup per row; reductions via threadgroup memory.
kernel void layer_norm(
device const float* input [[buffer(0)]],
device const float* gamma [[buffer(1)]],
device const float* beta [[buffer(2)]],
device float* output [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant float& eps [[buffer(5)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
// Pass 1: compute mean + variance via reduction
float local_sum = 0.0;
float local_sumsq = 0.0;
for (uint i = tid; i < h; i += tsize) {
float v = input[row * h + i];
local_sum += v;
local_sumsq += v * v;
}
partial_sum[tid] = local_sum;
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
// Reduction within threadgroup
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sum[tid] += partial_sum[tid + stride];
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float mean = partial_sum[0] / float(h);
float var = partial_sumsq[0] / float(h) - mean * mean;
float inv_std = rsqrt(var + eps);
// Pass 2: normalize
for (uint i = tid; i < h; i += tsize) {
float v = input[row * h + i];
output[row * h + i] = (v - mean) * inv_std * gamma[i] + beta[i];
}
}
// RMSNorm: out = (x / sqrt(mean(x^2) + eps)) * gamma + beta. No mean
// subtraction. Same dispatch shape as layer_norm (one threadgroup per row,
// power-of-2 reduction within the group).
kernel void rms_norm(
device const char* arena [[buffer(0)]],
constant ulong& in_off [[buffer(1)]],
constant ulong& g_off [[buffer(2)]],
constant ulong& b_off [[buffer(3)]],
constant ulong& out_off [[buffer(4)]],
constant uint& h [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
device const float* input = (device const float*)(arena + in_off);
device const float* gamma = (device const float*)(arena + g_off);
device const float* beta = (device const float*)(arena + b_off);
device float* output = (device float*)(arena + out_off);
threadgroup float partial_sumsq[256];
float local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = input[row * h + i];
local_sumsq += v * v;
}
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_rms = rsqrt(partial_sumsq[0] / float(h) + eps);
for (uint i = tid; i < h; i += tsize) {
output[row * h + i] = input[row * h + i] * inv_rms * gamma[i] + beta[i];
}
}
// f16 RMSNorm: half I/O, float accumulation.
kernel void rms_norm_h(
device const char* arena [[buffer(0)]],
constant ulong& in_off [[buffer(1)]],
constant ulong& g_off [[buffer(2)]],
constant ulong& b_off [[buffer(3)]],
constant ulong& out_off [[buffer(4)]],
constant uint& h [[buffer(5)]],
constant float& eps [[buffer(6)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
device const half* input = (device const half*)(arena + in_off);
device const half* gamma = (device const half*)(arena + g_off);
device const half* beta = (device const half*)(arena + b_off);
device half* output = (device half*)(arena + out_off);
threadgroup float partial_sumsq[256];
float local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = float(input[row * h + i]);
local_sumsq += v * v;
}
partial_sumsq[tid] = local_sumsq;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial_sumsq[tid] += partial_sumsq[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_rms = rsqrt(partial_sumsq[0] / float(h) + eps);
for (uint i = tid; i < h; i += tsize) {
float v = float(input[row * h + i]);
output[row * h + i] = half(v * inv_rms * float(gamma[i]) + float(beta[i]));
}
}
// f16 standalone softmax along the last axis. Half I/O, float accumulation
// for max + exp-sum (matters: f16 sum overflows above ~65k summands and
// exp() loses precision for moderate negatives).
kernel void softmax_lastax_h(
device half* data [[buffer(0)]],
constant uint& cols [[buffer(1)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float partial[256];
uint base = row * cols;
float local_max = -INFINITY;
for (uint i = tid; i < cols; i += tsize) {
local_max = max(local_max, float(data[base + i]));
}
partial[tid] = local_max;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial[tid] = max(partial[tid], partial[tid + stride]);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float row_max = partial[0];
float local_sum = 0.0f;
for (uint i = tid; i < cols; i += tsize) {
float e = exp(float(data[base + i]) - row_max);
data[base + i] = half(e);
local_sum += e;
}
partial[tid] = local_sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial[tid] += partial[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_sum = 1.0f / partial[0];
for (uint i = tid; i < cols; i += tsize) {
data[base + i] = half(float(data[base + i]) * inv_sum);
}
}
// f16 multi-axis reduce. Same op_kind encoding as reduce_axes; accumulate
// in float so 1e-2 .. 1e+4 f16 values don't lose precision summing across
// the reduced axis.
kernel void reduce_axes_h(
device const half* src [[buffer(0)]],
device half* dst [[buffer(1)]],
constant uint& reduced [[buffer(2)]],
constant uint& inner [[buffer(3)]],
constant uint& op_kind [[buffer(4)]],
uint2 gid [[thread_position_in_grid]]
) {
uint i = gid.x;
uint o = gid.y;
if (i >= inner) return;
float acc;
if (op_kind == 2) acc = -INFINITY;
else if (op_kind == 3) acc = INFINITY;
else if (op_kind == 4) acc = 1.0f;
else acc = 0.0f;
uint base = o * reduced * inner + i;
for (uint r = 0; r < reduced; ++r) {
float v = float(src[base + r * inner]);
if (op_kind == 0 || op_kind == 1) acc += v;
else if (op_kind == 2) acc = max(acc, v);
else if (op_kind == 3) acc = min(acc, v);
else acc *= v;
}
if (op_kind == 1) acc /= float(reduced);
dst[o * inner + i] = half(acc);
}
// PLAN L2 — interpreted N-ary element-wise chain kernel.
// One thread per output element. Walks the chain encoding (4 u32s
// per step: op_kind, op_sub, lhs_enc, rhs_enc) into a private
// scratch register array. Operand encoding: bit 31 = src kind
// (0=Input, 1=Step), bits 0..30 = index. Caps: 32 steps, 16 inputs.
inline uint region_input_row_resize2x_nchw(
uint gid,
uint out_n,
uint out_c,
uint out_h,
uint out_w
) {
uint plane = out_c * out_h * out_w;
uint local = gid % plane;
uint batch = gid / plane;
uint w_pos = local % out_w;
uint tmp = local / out_w;
uint h_pos = tmp % out_h;
uint c_pos = tmp / out_h;
uint in_w = out_w / 2u;
uint in_h = out_h / 2u;
uint in_plane = out_c * in_h * in_w;
return batch * in_plane + c_pos * in_h * in_w + (h_pos / 2u) * in_w + (w_pos / 2u);
}
inline uint region_resolve_row(
uint gid,
uint kind,
uint idx,
uint prologue_row0,
uint has_prologue_row0,
uint prologue_input,
uint scalar_input_mask,
device const uint* input_modulus
) {
if (kind != 0u) { return 0u; }
if (has_prologue_row0 != 0u && idx == prologue_input) {
return prologue_row0;
}
if ((scalar_input_mask & (1u << idx)) != 0u) { return 0u; }
if (input_modulus[idx] != 0u) { return gid % input_modulus[idx]; }
return gid;
}
kernel void elementwise_region(
device float* arena [[buffer(0)]],
constant uint& len [[buffer(1)]],
constant uint& num_inputs [[buffer(2)]],
constant uint& num_steps [[buffer(3)]],
constant uint& dst_off [[buffer(4)]],
device const uint* input_offs [[buffer(5)]], // 16 entries
device const uint* chain [[buffer(6)]], // 128 entries (32 steps * 4)
constant uint& scalar_input_mask [[buffer(7)]],
device const uint* input_modulus [[buffer(8)]], // 16 entries
constant uint& prologue [[buffer(9)]],
constant uint& out_n [[buffer(10)]],
constant uint& out_c [[buffer(11)]],
constant uint& out_h [[buffer(12)]],
constant uint& out_w [[buffer(13)]],
constant uint& prologue_input [[buffer(14)]],
uint3 gpos [[thread_position_in_grid]]
) {
uint gid;
if (prologue == 1u) {
uint nc = gpos.z;
uint ho = gpos.y;
uint wo = gpos.x;
if (nc >= out_n * out_c || ho >= out_h || wo >= out_w) { return; }
gid = nc * out_h * out_w + ho * out_w + wo;
} else {
gid = gpos.x;
if (gid >= len) { return; }
}
uint prologue_row0 = 0u;
uint has_prologue_row0 = 0u;
if (prologue == 1u) {
prologue_row0 = region_input_row_resize2x_nchw(gid, out_n, out_c, out_h, out_w);
has_prologue_row0 = 1u;
}
float scratch[32];
uint last_idx = 0;
for (uint k = 0; k < num_steps; ++k) {
uint base = k * 4;
uint op_kind = chain[base + 0];
uint op_sub = chain[base + 1];
uint lhs_enc = chain[base + 2];
uint rhs_enc = chain[base + 3];
// resolve_operand inline. Scalar-broadcast inputs read element
// 0 regardless of gid (fast path); trailing-shape broadcast
// reads `gid % input_modulus[idx]`. `input_modulus[idx]==0`
// means "no broadcast" and the kernel reads gid directly.
float lhs;
{
uint kind = lhs_enc >> 31;
uint idx = lhs_enc & 0x7FFFFFFFu;
uint row = region_resolve_row(
gid, kind, idx, prologue_row0, has_prologue_row0, prologue_input,
scalar_input_mask, input_modulus);
lhs = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
float result;
if (op_kind == 4u) {
// Where (3-operand select). op_sub carries cond_enc; lhs_enc
// / rhs_enc carry on_true / on_false. lhs already resolved
// above is on_true; resolve cond from op_sub and on_false
// from rhs_enc here.
float cond;
{
uint kind = op_sub >> 31;
uint idx = op_sub & 0x7FFFFFFFu;
uint row = region_resolve_row(
gid, kind, idx, prologue_row0, has_prologue_row0, prologue_input,
scalar_input_mask, input_modulus);
cond = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
float on_false;
{
uint kind = rhs_enc >> 31;
uint idx = rhs_enc & 0x7FFFFFFFu;
uint row = region_resolve_row(
gid, kind, idx, prologue_row0, has_prologue_row0, prologue_input,
scalar_input_mask, input_modulus);
on_false = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
result = (cond != 0.0f) ? lhs : on_false;
} else if (op_kind == 0u) {
// Activation
if (op_sub == 3u) result = max(lhs, 0.0f); // Relu
else if (op_sub == 0u || op_sub == 1u) {
float c = 0.7978845608f;
float inner = c * (lhs + 0.044715f * lhs * lhs * lhs);
result = 0.5f * lhs * (1.0f + tanh(inner)); // Gelu
}
else if (op_sub == 2u) result = lhs / (1.0f + exp(-lhs)); // Silu
else if (op_sub == 4u) result = 1.0f / (1.0f + exp(-lhs)); // Sigmoid
else if (op_sub == 5u) result = tanh(lhs);
else if (op_sub == 6u) result = exp(lhs);
else if (op_sub == 7u) result = log(lhs);
else if (op_sub == 8u) result = sqrt(lhs);
else if (op_sub == 9u) result = 1.0f / sqrt(lhs);
else if (op_sub == 10u) result = -lhs;
else if (op_sub == 11u) result = fabs(lhs);
else if (op_sub == 12u) result = round(lhs);
else if (op_sub == 13u) result = sin(lhs);
else if (op_sub == 14u) result = cos(lhs);
else if (op_sub == 15u) result = tan(lhs);
else if (op_sub == 16u) result = atan(lhs);
else result = lhs;
} else if (op_kind == 1u) {
// Cast at f32-arena layer is identity
result = lhs;
} else {
float rhs;
{
uint kind = rhs_enc >> 31;
uint idx = rhs_enc & 0x7FFFFFFFu;
uint row = region_resolve_row(
gid, kind, idx, prologue_row0, has_prologue_row0, prologue_input,
scalar_input_mask, input_modulus);
rhs = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
if (op_kind == 2u) {
if (op_sub == 0u) result = lhs + rhs;
else if (op_sub == 1u) result = lhs - rhs;
else if (op_sub == 2u) result = lhs * rhs;
else if (op_sub == 3u) result = lhs / rhs;
else if (op_sub == 4u) result = max(lhs, rhs);
else if (op_sub == 5u) result = min(lhs, rhs);
else result = pow(lhs, rhs);
} else {
bool b;
if (op_sub == 0u) b = (lhs == rhs);
else if (op_sub == 1u) b = (lhs != rhs);
else if (op_sub == 2u) b = (lhs < rhs);
else if (op_sub == 3u) b = (lhs <= rhs);
else if (op_sub == 4u) b = (lhs > rhs);
else b = (lhs >= rhs);
result = b ? 1.0f : 0.0f;
}
}
scratch[k] = result;
last_idx = k;
}
arena[dst_off + gid] = scratch[last_idx];
}
inline uint batch_region_resolve_row(
uint gid,
uint kind,
uint idx,
uint scalar_input_mask,
constant uint* input_modulus
) {
if (kind != 0u) { return 0u; }
if ((scalar_input_mask & (1u << idx)) != 0u) { return 0u; }
if (input_modulus[idx] != 0u) { return gid % input_modulus[idx]; }
return gid;
}
// FKL batch horizontal fusion: one dispatch, thread_position_in_grid.z = slice index.
// Requires prologue == 0 (no resize prologue on batch slices).
kernel void batch_elementwise_region(
device float* arena [[buffer(0)]],
constant uint& slice_len [[buffer(1)]],
constant uint& num_batch [[buffer(2)]],
constant uint& num_steps [[buffer(3)]],
constant uint& base_dst_off [[buffer(4)]],
constant uint& slice_elems [[buffer(5)]],
constant uint* batch_input_offs [[buffer(6)]], // 64 entries
constant uint* chain [[buffer(7)]], // 128 entries
constant uint& scalar_input_mask [[buffer(8)]],
constant uint* input_modulus [[buffer(9)]], // 16 entries
uint3 gpos [[thread_position_in_grid]]
) {
uint batch_idx = gpos.z;
if (batch_idx >= num_batch) { return; }
uint i = gpos.x;
if (i >= slice_len) { return; }
uint input_offs[16];
for (uint k = 0; k < 16u; ++k) { input_offs[k] = 0u; }
input_offs[0] = batch_input_offs[batch_idx];
uint dst_off = base_dst_off + batch_idx * slice_elems;
float scratch[32];
uint last_idx = 0;
for (uint k = 0; k < num_steps; ++k) {
uint base = k * 4;
uint op_kind = chain[base + 0];
uint op_sub = chain[base + 1];
uint lhs_enc = chain[base + 2];
uint rhs_enc = chain[base + 3];
float lhs;
{
uint kind = lhs_enc >> 31;
uint idx = lhs_enc & 0x7FFFFFFFu;
uint row = batch_region_resolve_row(
i, kind, idx, scalar_input_mask, input_modulus);
lhs = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
float result;
if (op_kind == 4u) {
float cond;
{
uint kind = op_sub >> 31;
uint idx = op_sub & 0x7FFFFFFFu;
uint row = batch_region_resolve_row(
i, kind, idx, scalar_input_mask, input_modulus);
cond = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
float on_false;
{
uint kind = rhs_enc >> 31;
uint idx = rhs_enc & 0x7FFFFFFFu;
uint row = batch_region_resolve_row(
i, kind, idx, scalar_input_mask, input_modulus);
on_false = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
result = (cond != 0.0f) ? lhs : on_false;
} else if (op_kind == 0u) {
if (op_sub == 3u) result = max(lhs, 0.0f);
else if (op_sub == 0u || op_sub == 1u) {
float c = 0.7978845608f;
float inner = c * (lhs + 0.044715f * lhs * lhs * lhs);
result = 0.5f * lhs * (1.0f + tanh(inner));
}
else if (op_sub == 2u) result = lhs / (1.0f + exp(-lhs));
else if (op_sub == 4u) result = 1.0f / (1.0f + exp(-lhs));
else if (op_sub == 5u) result = tanh(lhs);
else if (op_sub == 6u) result = exp(lhs);
else if (op_sub == 7u) result = log(lhs);
else if (op_sub == 8u) result = sqrt(lhs);
else if (op_sub == 9u) result = 1.0f / sqrt(lhs);
else if (op_sub == 10u) result = -lhs;
else if (op_sub == 11u) result = fabs(lhs);
else if (op_sub == 12u) result = round(lhs);
else if (op_sub == 13u) result = sin(lhs);
else if (op_sub == 14u) result = cos(lhs);
else if (op_sub == 15u) result = tan(lhs);
else if (op_sub == 16u) result = atan(lhs);
else result = lhs;
} else if (op_kind == 1u) {
result = lhs;
} else {
float rhs;
{
uint kind = rhs_enc >> 31;
uint idx = rhs_enc & 0x7FFFFFFFu;
uint row = batch_region_resolve_row(
i, kind, idx, scalar_input_mask, input_modulus);
rhs = (kind == 0u) ? arena[input_offs[idx] + row] : scratch[idx];
}
if (op_kind == 2u) {
if (op_sub == 0u) result = lhs + rhs;
else if (op_sub == 1u) result = lhs - rhs;
else if (op_sub == 2u) result = lhs * rhs;
else if (op_sub == 3u) result = lhs / rhs;
else if (op_sub == 4u) result = max(lhs, rhs);
else if (op_sub == 5u) result = min(lhs, rhs);
else result = pow(lhs, rhs);
} else {
bool b;
if (op_sub == 0u) b = (lhs == rhs);
else if (op_sub == 1u) b = (lhs != rhs);
else if (op_sub == 2u) b = (lhs < rhs);
else if (op_sub == 3u) b = (lhs <= rhs);
else if (op_sub == 4u) b = (lhs > rhs);
else b = (lhs >= rhs);
result = b ? 1.0f : 0.0f;
}
}
scratch[k] = result;
last_idx = k;
}
arena[dst_off + i] = scratch[last_idx];
}
// ── Gated DeltaNet scan (f32) ───────────────────────────────────────
// One threadgroup per (batch, head), `n` threads parallelize the state
// dimension (n ≤ 128). Matches `execute_gated_delta_net_f32` on CPU.
#define GDN_MAX_N 128u
kernel void gated_delta_net(
device float* arena [[buffer(0)]],
constant uint& q_off [[buffer(1)]],
constant uint& k_off [[buffer(2)]],
constant uint& v_off [[buffer(3)]],
constant uint& g_off [[buffer(4)]],
constant uint& beta_off [[buffer(5)]],
constant uint& state_off [[buffer(6)]],
constant uint& dst_off [[buffer(7)]],
constant uint4& dims [[buffer(8)]], // batch, seq, heads, n
constant uint& use_carry [[buffer(9)]],
uint gid [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]
) {
uint b = dims.x, s = dims.y, h = dims.z, n = dims.w;
if (n > GDN_MAX_N || gid >= b * h || tid >= n) return;
uint bi = gid / h;
uint hi = gid % h;
uint j = tid;
float scale = rsqrt(float(n));
uint s_base = state_off + (bi * h + hi) * n * n;
device float* s_mat = arena + s_base;
if (use_carry == 0u && tid == 0u) {
for (uint i = 0; i < n * n; ++i) {
s_mat[i] = 0.0f;
}
}
threadgroup float sk_sh[GDN_MAX_N];
threadgroup_barrier(mem_flags::mem_threadgroup);
uint hs_n = h * n;
for (uint ti = 0; ti < s; ++ti) {
uint qkv_step = bi * s * hs_n + ti * hs_n + hi * n;
uint gb_step = bi * s * h + ti * h + hi;
uint q_row = q_off + qkv_step;
uint k_row = k_off + qkv_step;
uint v_row = v_off + qkv_step;
float g_t = arena[g_off + gb_step];
float beta_t = arena[beta_off + gb_step];
float g_exp = exp(g_t);
if (tid == 0u) {
for (uint idx = 0; idx < n * n; ++idx) {
s_mat[idx] *= g_exp;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
float acc = 0.0f;
for (uint i = 0; i < n; ++i) {
acc += s_mat[i * n + j] * arena[k_row + i];
}
sk_sh[j] = acc;
threadgroup_barrier(mem_flags::mem_threadgroup);
sk_sh[j] = (arena[v_row + j] - sk_sh[j]) * beta_t;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint i = 0; i < n; ++i) {
float ki = arena[k_row + i];
s_mat[i * n + j] += ki * sk_sh[j];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
uint out_row = dst_off + qkv_step;
acc = 0.0f;
for (uint i = 0; i < n; ++i) {
acc += s_mat[i * n + j] * arena[q_row + i];
}
arena[out_row + j] = acc * scale;
}
}
// Single-layer unidirectional LSTM (gate order i, f, g, o; h0 = c0 = 0).
// One threadgroup per batch item; thread `k` owns hidden unit `k` and
// keeps c[k] in a register; h_prev lives in threadgroup memory so every
// thread can read it for the w_hh matvec. Requires hidden <= LSTM_MAX_H.
// Matches `execute_lstm_f32` on CPU.
#define LSTM_MAX_H 1024u
kernel void lstm(
device float* arena [[buffer(0)]],
constant uint& x_off [[buffer(1)]],
constant uint& wih_off [[buffer(2)]],
constant uint& whh_off [[buffer(3)]],
constant uint& bias_off [[buffer(4)]],
constant uint& dst_off [[buffer(5)]],
constant uint4& dims [[buffer(6)]], // batch, seq, input, hidden
uint gid [[threadgroup_position_in_grid]],
uint tid [[thread_index_in_threadgroup]]
) {
uint b = dims.x, s = dims.y, in_sz = dims.z, h = dims.w;
if (h > LSTM_MAX_H || gid >= b || tid >= h) return;
uint bi = gid;
uint k = tid;
threadgroup float h_sh[LSTM_MAX_H];
h_sh[k] = 0.0f;
float c_k = 0.0f;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint t = 0; t < s; ++t) {
uint x_base = x_off + (bi * s + t) * in_sz;
// Gate pre-activations for hidden unit k: rows i=k, f=h+k, g=2h+k, o=3h+k.
float z[4];
for (uint gate = 0u; gate < 4u; ++gate) {
uint r = gate * h + k;
float acc = arena[bias_off + r];
uint wih_row = wih_off + r * in_sz;
for (uint j = 0u; j < in_sz; ++j) {
acc += arena[wih_row + j] * arena[x_base + j];
}
uint whh_row = whh_off + r * h;
for (uint j = 0u; j < h; ++j) {
acc += arena[whh_row + j] * h_sh[j];
}
z[gate] = acc;
}
float i_g = 1.0f / (1.0f + exp(-z[0]));
float f_g = 1.0f / (1.0f + exp(-z[1]));
float g_g = tanh(z[2]);
float o_g = 1.0f / (1.0f + exp(-z[3]));
c_k = f_g * c_k + i_g * g_g;
float h_k = o_g * tanh(c_k);
// Finish reading h_prev across all threads before overwriting it.
threadgroup_barrier(mem_flags::mem_threadgroup);
h_sh[k] = h_k;
threadgroup_barrier(mem_flags::mem_threadgroup);
arena[dst_off + (bi * s + t) * h + k] = h_k;
}
}
// RMSNorm backward (wrt: 0=dx, 1=dgamma, 2=dbeta). One threadgroup per row.
kernel void rms_norm_bwd(
device const float* x [[buffer(0)]],
device const float* gamma [[buffer(1)]],
device const float* beta [[buffer(2)]],
device const float* dy [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& inner [[buffer(5)]],
constant float& eps [[buffer(6)]],
constant uint& wrt [[buffer(7)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
if (wrt != 0u) return;
threadgroup float partial[256];
float local_dot = 0.0f;
for (uint i = tid; i < inner; i += tsize) {
float xv = x[row * inner + i];
float gv = gamma[i];
float dyv = dy[row * inner + i];
local_dot += dyv * gv * xv;
}
partial[tid] = local_dot;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) partial[tid] += partial[tid + stride];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float dot = partial[0] / float(inner);
float local_ss = 0.0f;
for (uint i = tid; i < inner; i += tsize) {
float xv = x[row * inner + i];
local_ss += xv * xv;
}
partial[tid] = local_ss;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) partial[tid] += partial[tid + stride];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float inv_r = rsqrt(partial[0] / float(inner) + eps);
float inv_r3 = inv_r * inv_r * inv_r;
for (uint i = tid; i < inner; i += tsize) {
float xv = x[row * inner + i];
float gv = gamma[i];
float dyv = dy[row * inner + i];
float term = gv * dyv - xv * dot * inv_r3;
out[row * inner + i] = term * inv_r;
}
}
kernel void rms_norm_bwd_param(
device const float* x [[buffer(0)]],
device const float* gamma [[buffer(1)]],
device const float* beta [[buffer(2)]],
device const float* dy [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& rows [[buffer(5)]],
constant uint& inner [[buffer(6)]],
constant float& eps [[buffer(7)]],
constant uint& wrt [[buffer(8)]],
uint tid [[thread_position_in_threadgroup]]
) {
if (tid != 0u) return;
for (uint i = 0; i < inner; ++i) out[i] = 0.0f;
for (uint row = 0; row < rows; ++row) {
float sumsq = 0.0f;
for (uint i = 0; i < inner; ++i) {
float xv = x[row * inner + i];
sumsq += xv * xv;
}
float inv_r = rsqrt(sumsq / float(inner) + eps);
if (wrt == 1u) {
for (uint i = 0; i < inner; ++i) {
out[i] += dy[row * inner + i] * x[row * inner + i] * inv_r;
}
} else {
for (uint i = 0; i < inner; ++i) {
out[i] += dy[row * inner + i];
}
}
}
}
// Per-row RMS inverse scale — scratch for parallel dgamma/dbeta.
kernel void rms_norm_bwd_inv_r_f32(
device const float* x [[buffer(0)]],
device float* inv_r [[buffer(1)]],
constant uint& inner [[buffer(2)]],
constant float& eps [[buffer(3)]],
uint row [[thread_position_in_grid]]
) {
float sumsq = 0.0f;
for (uint i = 0; i < inner; ++i) {
float xv = x[row * inner + i];
sumsq += xv * xv;
}
inv_r[row] = rsqrt(sumsq / float(inner) + eps);
}
// Reduce dy·x·inv_r (gamma) or dy (beta) across rows — one thread per param.
kernel void rms_norm_bwd_param_reduce_f32(
device const float* x [[buffer(0)]],
device const float* dy [[buffer(1)]],
device const float* inv_r [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& rows [[buffer(4)]],
constant uint& inner [[buffer(5)]],
constant uint& wrt [[buffer(6)]],
uint i [[thread_position_in_grid]]
) {
if (i >= inner) return;
float acc = 0.0f;
for (uint row = 0; row < rows; ++row) {
if (wrt == 1u) {
acc += dy[row * inner + i] * x[row * inner + i] * inv_r[row];
} else {
acc += dy[row * inner + i];
}
}
out[i] = acc;
}
kernel void rope_bwd(
device const float* dy [[buffer(0)]],
device const float* cos [[buffer(1)]],
device const float* sin [[buffer(2)]],
device float* dx [[buffer(3)]],
constant uint& batch [[buffer(4)]],
constant uint& seq [[buffer(5)]],
constant uint& hidden [[buffer(6)]],
constant uint& head_dim [[buffer(7)]],
constant uint& n_rot [[buffer(8)]],
constant uint& cos_len [[buffer(9)]],
uint3 gid [[thread_position_in_grid]]
) {
uint d = gid.x;
uint hi = gid.y;
uint bs = gid.z;
if (d >= head_dim) return;
uint nh = hidden / head_dim;
if (hi >= nh) return;
if (bs >= batch * seq) return;
uint bi = bs / seq;
uint si = bs % seq;
uint rot_half = n_rot / 2u;
uint half_dh = head_dim / 2u;
uint tab_off = (si * half_dh) % max(cos_len, 1u);
uint dy_base = bi * seq * hidden + si * hidden + hi * head_dim;
uint dx_base = dy_base;
if (d < rot_half) {
float y1 = dy[dy_base + d];
float y2 = dy[dy_base + rot_half + d];
float c = cos[tab_off + d];
float s = sin[tab_off + d];
dx[dx_base + d] = y1 * c + y2 * s;
dx[dx_base + rot_half + d] = -y1 * s + y2 * c;
} else if (d >= n_rot) {
dx[dx_base + d] = dy[dy_base + d];
}
}
kernel void cumsum_fwd(
device const float* src [[buffer(0)]],
device float* dst [[buffer(1)]],
constant uint& inner [[buffer(2)]],
constant uint& exclusive [[buffer(3)]],
uint row [[threadgroup_position_in_grid]]
) {
float acc = 0.0f;
for (uint i = 0; i < inner; ++i) {
if (exclusive != 0u) {
dst[row * inner + i] = acc;
acc += src[row * inner + i];
} else {
acc += src[row * inner + i];
dst[row * inner + i] = acc;
}
}
}
kernel void cumsum_bwd(
device const float* dy [[buffer(0)]],
device float* dx [[buffer(1)]],
constant uint& inner [[buffer(2)]],
constant uint& exclusive [[buffer(3)]],
uint row [[threadgroup_position_in_grid]]
) {
float suffix = 0.0f;
for (int i = int(inner) - 1; i >= 0; --i) {
uint ui = uint(i);
if (exclusive != 0u) {
dx[row * inner + ui] = suffix;
suffix += dy[row * inner + ui];
} else {
suffix += dy[row * inner + ui];
dx[row * inner + ui] = suffix;
}
}
}
// Single im2col element for conv weight backward GEMM (B[k_idx, n_col]).
inline float conv_bwd_im2col_elem(
device const float* x,
uint k_idx,
uint n_col,
uint c_in,
uint h,
uint w_in,
uint h_out,
uint w_out,
uint kh,
uint kw,
uint sh,
uint sw,
uint ph,
uint pw,
uint dh,
uint dw_dil
) {
uint ho = k_idx / w_out;
uint wo = k_idx % w_out;
uint rem = n_col;
uint ci = rem / (kh * kw);
rem = rem % (kh * kw);
uint ki = rem / kw;
uint kj = rem % kw;
int hi = (int)(ho * sh + ki * dh) - (int)ph;
int wi = (int)(wo * sw + kj * dw_dil) - (int)pw;
if (hi < 0 || wi < 0 || hi >= (int)h || wi >= (int)w_in) {
return 0.0f;
}
return x[(ci * h + (uint)hi) * w_in + (uint)wi];
}
// dw = dy @ im2col(x) — 8×8 simdgroup tiles, B generated on the fly (no scratch).
kernel void conv2d_bwd_weight_gemm(
device const float* dy [[buffer(0)]],
device const float* x [[buffer(1)]],
device float* dw [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
constant uint4& nchw [[buffer(6)]],
constant uint4& out_dims [[buffer(7)]],
constant uint4& kshape [[buffer(8)]],
constant uint4& padd [[buffer(9)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint slid [[thread_index_in_threadgroup]]
) {
uint row_base = tgid.y * 8;
uint col_base = tgid.x * 8;
if (row_base >= M || col_base >= N) return;
uint c_in = nchw.x;
uint h = nchw.y;
uint w_in = nchw.z;
uint h_out = out_dims.y;
uint w_out = out_dims.z;
uint kh = kshape.x;
uint kw = kshape.y;
uint sh = kshape.z;
uint sw = kshape.w;
uint ph = padd.x;
uint pw = padd.y;
uint dh = padd.z;
uint dw_dil = padd.w;
threadgroup float B_tg[64];
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint k0 = 0; k0 < K; k0 += 8) {
for (uint i = 0; i < 2; ++i) {
uint idx = i * 32 + slid;
if (idx < 64) {
uint br = idx / 8;
uint bc = idx % 8;
uint k_idx = k0 + br;
uint n_col = col_base + bc;
B_tg[idx] = (k_idx < K && n_col < N)
? conv_bwd_im2col_elem(
x, k_idx, n_col, c_in, h, w_in, h_out, w_out, kh, kw, sh, sw, ph, pw,
dh, dw_dil)
: 0.0f;
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
simdgroup_load(a, dy + row_base * K + k0, K);
simdgroup_load(b, B_tg, 8);
simdgroup_multiply_accumulate(c, a, b, c);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
simdgroup_store(c, dw + row_base * N + col_base, N);
}
// dw = dy @ im2col(x) — 32×32 threadgroup tiles (requires M,K,N % 32 == 0).
kernel void conv2d_bwd_weight_gemm_4x4(
device const float* dy [[buffer(0)]],
device const float* x [[buffer(1)]],
device float* dw [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
constant uint4& nchw [[buffer(6)]],
constant uint4& out_dims [[buffer(7)]],
constant uint4& kshape [[buffer(8)]],
constant uint4& padd [[buffer(9)]],
uint2 tgid [[threadgroup_position_in_grid]],
uint sgid [[simdgroup_index_in_threadgroup]],
uint slid [[thread_index_in_simdgroup]]
) {
uint sg_row = sgid / 4;
uint sg_col = sgid % 4;
uint tg_row_base = tgid.y * 32;
uint tg_col_base = tgid.x * 32;
uint c_in = nchw.x;
uint h = nchw.y;
uint w_in = nchw.z;
uint h_out = out_dims.y;
uint w_out = out_dims.z;
uint kh = kshape.x;
uint kw = kshape.y;
uint sh = kshape.z;
uint sw = kshape.w;
uint ph = padd.x;
uint pw = padd.y;
uint dh = padd.z;
uint dw_dil = padd.w;
threadgroup float A_tg[32 * 32];
threadgroup float B_tg[32 * 32];
simdgroup_float8x8 a, b, c;
c = simdgroup_float8x8(0.0f);
for (uint kk = 0; kk < K; kk += 32) {
uint linear = sgid * 32 + slid;
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint ar = idx / 32;
uint ac = idx % 32;
A_tg[idx] = dy[(tg_row_base + ar) * K + (kk + ac)];
}
for (uint i = 0; i < 2; ++i) {
uint idx = i * 512 + linear;
uint br = idx / 32;
uint bc = idx % 32;
uint k_idx = kk + br;
uint n_col = tg_col_base + bc;
B_tg[idx] = conv_bwd_im2col_elem(
x, k_idx, n_col, c_in, h, w_in, h_out, w_out, kh, kw, sh, sw, ph, pw, dh,
dw_dil);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint k_inner = 0; k_inner < 32; k_inner += 8) {
simdgroup_load(a, &A_tg[sg_row * 8 * 32 + k_inner], 32);
simdgroup_load(b, &B_tg[k_inner * 32 + sg_col * 8], 32);
simdgroup_multiply_accumulate(c, a, b, c);
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
uint out_row = tg_row_base + sg_row * 8;
uint out_col = tg_col_base + sg_col * 8;
simdgroup_store(c, &dw[out_row * N + out_col], N);
}
// Fast im2col when W_in = W_out = 1 (Voxtral codec `[1,C,T,1]` slices).
kernel void im2col_group_w1(
device const float* x [[buffer(0)]],
device float* col [[buffer(1)]],
constant uint4& nchw [[buffer(2)]], // [C_in/g, H, 1, unused]
constant uint4& out_dims [[buffer(3)]], // [unused, H_out, 1, unused]
constant uint4& kshape [[buffer(4)]],
constant uint4& padd [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
uint h_out = out_dims.y;
uint c_in = nchw.x;
uint h = nchw.y;
uint kh = kshape.x;
uint kw = kshape.y;
uint sh = kshape.z;
uint ph = padd.x;
uint dh = padd.z;
uint n_dim = c_in * kh * kw;
uint k_dim = h_out;
uint idx = gid;
if (idx >= n_dim * k_dim) return;
uint row = idx / k_dim;
uint ho = idx % k_dim;
uint rem = row;
uint ci = rem / (kh * kw);
rem = rem % (kh * kw);
uint ki = rem / kw;
int hi = (int)(ho * sh + ki * dh) - (int)ph;
col[ho * n_dim + row] = (hi < 0 || hi >= (int)h)
? 0.0f
: x[ci * h + (uint)hi];
}
// im2col for one (batch, group) slice — layout matches `rlx_cpu::conv_bwd` /
// `[n_dim, k_dim]` row-major with `n_dim = C_in/g · kH · kW`, `k_dim = H_out · W_out`.
kernel void im2col_group(
device const float* x [[buffer(0)]],
device float* col [[buffer(1)]],
constant uint4& nchw [[buffer(2)]], // [C_in/g, H, W, unused] (group slice)
constant uint4& out_dims [[buffer(3)]], // [unused, H_out, W_out, unused]
constant uint4& kshape [[buffer(4)]],
constant uint4& padd [[buffer(5)]],
uint2 gid [[thread_position_in_grid]]
) {
uint h_out = out_dims.y;
uint w_out = out_dims.z;
uint c_in = nchw.x;
uint h = nchw.y;
uint w = nchw.z;
uint kh = kshape.x;
uint kw = kshape.y;
uint sh = kshape.z;
uint sw = kshape.w;
uint ph = padd.x;
uint pw = padd.y;
uint dh = padd.z;
uint dw_dil = padd.w;
uint n_dim = c_in * kh * kw;
uint k_dim = h_out * w_out;
uint idx = gid.x;
if (idx >= n_dim * k_dim) return;
uint row = idx / k_dim;
uint k_idx = idx % k_dim;
uint ho = k_idx / w_out;
uint wo = k_idx % w_out;
uint rem = row;
uint ci = rem / (kh * kw);
rem = rem % (kh * kw);
uint ki = rem / kw;
uint kj = rem % kw;
int hi = (int)(ho * sh + ki * dh) - (int)ph;
int wi = (int)(wo * sw + kj * dw_dil) - (int)pw;
col[k_idx * n_dim + row] = (hi < 0 || wi < 0 || hi >= (int)h || wi >= (int)w)
? 0.0f
: x[(ci * h + (uint)hi) * w + (uint)wi];
}
// ── Attention backward (recompute scores + softmax) ─────────────────────
kernel void attn_bwd_scores_f32(
device const float* q [[buffer(0)]],
device const float* k [[buffer(1)]],
device float* scores [[buffer(2)]],
constant uint& sq [[buffer(3)]],
constant uint& sk [[buffer(4)]],
constant uint& hs [[buffer(5)]],
constant uint& head_dim [[buffer(6)]],
constant float& scale [[buffer(7)]],
constant uint& mask_kind [[buffer(8)]],
constant uint& window [[buffer(9)]],
uint2 gid [[thread_position_in_grid]]
) {
uint qi = gid.y;
uint ki = gid.x;
if (qi >= sq || ki >= sk) return;
float dot = 0.0f;
for (uint d = 0; d < head_dim; ++d) {
dot += q[qi * hs + d] * k[ki * hs + d];
}
float s = dot * scale;
if (mask_kind == 1u) {
if (ki > qi) s = -1e4f;
} else if (mask_kind == 3u) {
uint lo = qi > window ? qi - window : 0u;
if (ki < lo || ki > qi) s = -1e4f;
}
scores[qi * sk + ki] = s;
}
kernel void attn_bwd_dp_f32(
device const float* dy [[buffer(0)]],
device const float* v [[buffer(1)]],
device float* dp [[buffer(2)]],
constant uint& sq [[buffer(3)]],
constant uint& sk [[buffer(4)]],
constant uint& hs [[buffer(5)]],
constant uint& head_dim [[buffer(6)]],
uint2 gid [[thread_position_in_grid]]
) {
uint qi = gid.y;
uint ki = gid.x;
if (qi >= sq || ki >= sk) return;
float acc = 0.0f;
for (uint d = 0; d < head_dim; ++d) {
acc += dy[qi * hs + d] * v[ki * hs + d];
}
dp[qi * sk + ki] = acc;
}
kernel void attn_bwd_ds_f32(
device const float* scores [[buffer(0)]],
device const float* dp [[buffer(1)]],
device float* ds [[buffer(2)]],
constant uint& sq [[buffer(3)]],
constant uint& sk [[buffer(4)]],
constant float& scale [[buffer(5)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
if (row >= sq) return;
float row_sum = 0.0f;
for (uint ki = tid; ki < sk; ki += tsize) {
row_sum += scores[row * sk + ki] * dp[row * sk + ki];
}
threadgroup float partial[256];
partial[tid] = row_sum;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = tsize / 2; stride > 0; stride /= 2) {
if (tid < stride) {
partial[tid] += partial[tid + stride];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
float sum = partial[0];
for (uint ki = tid; ki < sk; ki += tsize) {
uint idx = row * sk + ki;
float p = scores[idx];
ds[idx] = p * (dp[idx] - sum) * scale;
}
}
kernel void attn_bwd_dv_f32(
device const float* scores [[buffer(0)]],
device const float* dy [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& sq [[buffer(3)]],
constant uint& sk [[buffer(4)]],
constant uint& hs [[buffer(5)]],
constant uint& head_dim [[buffer(6)]],
uint2 gid [[thread_position_in_grid]]
) {
uint ki = gid.y;
uint d = gid.x;
if (ki >= sk || d >= head_dim) return;
float acc = 0.0f;
for (uint qi = 0; qi < sq; ++qi) {
acc += scores[qi * sk + ki] * dy[qi * hs + d];
}
out[ki * hs + d] = acc;
}
kernel void attn_bwd_dq_f32(
device const float* ds [[buffer(0)]],
device const float* k [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& sq [[buffer(3)]],
constant uint& sk [[buffer(4)]],
constant uint& hs [[buffer(5)]],
constant uint& head_dim [[buffer(6)]],
uint2 gid [[thread_position_in_grid]]
) {
uint qi = gid.y;
uint d = gid.x;
if (qi >= sq || d >= head_dim) return;
float acc = 0.0f;
for (uint ki = 0; ki < sk; ++ki) {
acc += ds[qi * sk + ki] * k[ki * hs + d];
}
out[qi * hs + d] = acc;
}
kernel void attn_bwd_dk_f32(
device const float* ds [[buffer(0)]],
device const float* q [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& sq [[buffer(3)]],
constant uint& sk [[buffer(4)]],
constant uint& hs [[buffer(5)]],
constant uint& head_dim [[buffer(6)]],
uint2 gid [[thread_position_in_grid]]
) {
uint ki = gid.y;
uint d = gid.x;
if (ki >= sk || d >= head_dim) return;
float acc = 0.0f;
for (uint qi = 0; qi < sq; ++qi) {
acc += ds[qi * sk + ki] * q[qi * hs + d];
}
out[ki * hs + d] = acc;
}
kernel void gather_bwd_zero(
device float* dst [[buffer(0)]],
constant uint& n [[buffer(1)]],
uint i [[thread_position_in_grid]]
) {
if (i < n) dst[i] = 0.0f;
}
kernel void gather_bwd_acc(
device const float* dy [[buffer(0)]],
device const float* idx [[buffer(1)]],
device float* dst [[buffer(2)]],
constant uint& outer [[buffer(3)]],
constant uint& axis_dim [[buffer(4)]],
constant uint& num_idx [[buffer(5)]],
constant uint& trailing [[buffer(6)]],
uint o [[threadgroup_position_in_grid]]
) {
if (o >= outer) return;
for (uint k = 0; k < num_idx; ++k) {
uint row = uint(idx[k]);
if (row >= axis_dim) continue;
for (uint j = 0; j < trailing; ++j) {
float v = dy[(o * num_idx + k) * trailing + j];
dst[(o * axis_dim + row) * trailing + j] += v;
}
}
}
"#;
const RLX_KERNELS_MSL_DEQUANT: &str = include_str!("dequant_gguf.msl");
const RLX_KERNELS_MSL_FFT_GPU: &str = include_str!("fft_gpu.msl");
const RLX_KERNELS_MSL_SPLAT: &str = include_str!("splat.msl");
const RLX_KERNELS_MSL_SPLAT_CONIC: &str = include_str!("splat_conic_bin.msl");
fn msl_source() -> String {
format!(
"{RLX_KERNELS_MSL}\n{RLX_KERNELS_MSL_DEQUANT}\n{RLX_KERNELS_MSL_FFT_GPU}\n{RLX_KERNELS_MSL_SPLAT}\n{RLX_KERNELS_MSL_SPLAT_CONIC}"
)
}
pub struct Kernels {
pub library: Library,
pub sgemm: ComputePipelineState,
pub sgemm_simd: ComputePipelineState,
pub sgemm_simd_bias: ComputePipelineState,
pub sgemm_simd_4x4: ComputePipelineState,
pub sgemm_simd_4x4_bias: ComputePipelineState,
pub hgemm_simd_4x4: ComputePipelineState,
pub hgemm_simd_4x4_bias: ComputePipelineState,
pub bias_add_h: ComputePipelineState,
pub gelu_inplace_h: ComputePipelineState,
pub gelu_approx_inplace_h: ComputePipelineState,
pub silu_inplace_h: ComputePipelineState,
pub layer_norm_h: ComputePipelineState,
pub fused_residual_ln_h: ComputePipelineState,
pub fused_residual_rms_norm_h: ComputePipelineState,
pub rms_norm_h: ComputePipelineState,
pub softmax_lastax_h: ComputePipelineState,
pub reduce_axes_h: ComputePipelineState,
pub elem_add_h: ComputePipelineState,
pub elem_mul_h: ComputePipelineState,
pub gather_axis0_h: ComputePipelineState,
pub narrow_lastax_h: ComputePipelineState,
pub sdpa_h: ComputePipelineState,
pub rope_h: ComputePipelineState,
pub cast_f32_to_f16: ComputePipelineState,
pub cast_f16_to_f32: ComputePipelineState,
pub copy_f32: ComputePipelineState,
pub copy4: ComputePipelineState,
pub sgemm_simd_padded: ComputePipelineState,
pub sgemm_simd_padded_bias: ComputePipelineState,
pub sgemm_tiled: ComputePipelineState,
pub bias_add: ComputePipelineState,
pub gelu_inplace: ComputePipelineState,
pub gelu_inplace4: ComputePipelineState,
pub gelu_approx_inplace: ComputePipelineState,
pub gelu_approx_inplace4: ComputePipelineState,
pub gelu_approx_out4: ComputePipelineState,
pub silu_inplace: ComputePipelineState,
pub silu_inplace4: ComputePipelineState,
pub binary_broadcast_rhs_col_f32: ComputePipelineState,
pub binary_broadcast_rhs_col4: ComputePipelineState,
pub binary_broadcast_rhs_row_f32: ComputePipelineState,
pub binary_broadcast_rhs_row4: ComputePipelineState,
pub binary_broadcast_rhs_scalar_f32: ComputePipelineState,
pub binary_broadcast_rhs_scalar4: ComputePipelineState,
pub binary_broadcast_1ax_f32: ComputePipelineState,
pub binary_broadcast_1ax4: ComputePipelineState,
pub fused_binary_activation_f32: ComputePipelineState,
pub fused_binary_activation4: ComputePipelineState,
pub fused_ternary_activation_f32: ComputePipelineState,
pub fused_ternary_activation4: ComputePipelineState,
pub layer_norm: ComputePipelineState,
pub rms_norm: ComputePipelineState,
pub elem_add: ComputePipelineState,
pub elem_add4: ComputePipelineState,
pub elem_sub4: ComputePipelineState,
pub binary_broadcast_f32: ComputePipelineState,
pub binary_broadcast_rank2_f32: ComputePipelineState,
pub binary_broadcast_rank24: ComputePipelineState,
pub elem_mul: ComputePipelineState,
pub elem_mul4: ComputePipelineState,
pub elem_div4: ComputePipelineState,
pub gather_axis0: ComputePipelineState,
pub narrow_lastax: ComputePipelineState,
pub narrow_lastax4: ComputePipelineState,
pub split_lastax: ComputePipelineState,
pub split_lastax4: ComputePipelineState,
pub fused_residual_ln: ComputePipelineState,
pub fused_residual_rms_norm: ComputePipelineState,
pub sdpa: ComputePipelineState,
pub sdpa_long: ComputePipelineState,
pub sdpa_fa_f32: ComputePipelineState,
pub rope: ComputePipelineState,
pub fused_swiglu: ComputePipelineState,
pub fused_swiglu_h: ComputePipelineState,
pub elementwise_region: ComputePipelineState,
pub batch_elementwise_region: ComputePipelineState,
pub fused_swiglu_cast_f32_to_f16: ComputePipelineState,
pub fused_swiglu_cast_f16_to_f32: ComputePipelineState,
pub concat_segment_lastax: ComputePipelineState,
pub concat_segment_lastax4: ComputePipelineState,
pub concat_lastax_multi: ComputePipelineState,
pub concat_lastax_multi4: ComputePipelineState,
pub concat_segment_lastax_h: ComputePipelineState,
pub elem_sub: ComputePipelineState,
pub elem_div: ComputePipelineState,
pub elem_max: ComputePipelineState,
pub elem_min: ComputePipelineState,
pub elem_pow: ComputePipelineState,
pub elem_compare: ComputePipelineState,
pub elem_where: ComputePipelineState,
pub reduce_axes: ComputePipelineState,
pub topk_lastax: ComputePipelineState,
pub grouped_matmul: ComputePipelineState,
pub scatter_add_zero: ComputePipelineState,
pub scatter_add_accumulate: ComputePipelineState,
pub transpose_nd: ComputePipelineState,
pub transpose_2d_f32: ComputePipelineState,
pub transpose_2d_tiled_f32: ComputePipelineState,
pub transpose_last2_batched_f32: ComputePipelineState,
pub transpose_last2_batched_tiled_f32: ComputePipelineState,
pub transpose_swap12_batched_trail_f32: ComputePipelineState,
pub transpose_swap12_batched_trail_tiled_f32: ComputePipelineState,
pub gather_axis: ComputePipelineState,
pub pool2d: ComputePipelineState,
pub conv2d: ComputePipelineState,
pub conv2d_w1: ComputePipelineState,
pub layer_norm2d: ComputePipelineState,
pub group_norm: ComputePipelineState,
pub resize_nearest_2x: ComputePipelineState,
pub conv_transpose2d: ComputePipelineState,
pub relu_inplace: ComputePipelineState,
pub sigmoid_inplace: ComputePipelineState,
pub tanh_inplace: ComputePipelineState,
pub exp_inplace: ComputePipelineState,
pub log_inplace: ComputePipelineState,
pub sqrt_inplace: ComputePipelineState,
pub rsqrt_inplace: ComputePipelineState,
pub neg_inplace: ComputePipelineState,
pub abs_inplace: ComputePipelineState,
pub round_inplace: ComputePipelineState,
pub sin_inplace: ComputePipelineState,
pub cos_inplace: ComputePipelineState,
pub tan_inplace: ComputePipelineState,
pub atan_inplace: ComputePipelineState,
pub softmax_lastax: ComputePipelineState,
pub fft_radix2_full_f32: ComputePipelineState,
pub fft_bit_reverse_f32: ComputePipelineState,
pub fft_inner_f32: ComputePipelineState,
pub fft_outer_r4_f32: ComputePipelineState,
pub fft_outer_r2_f32: ComputePipelineState,
pub gated_delta_net: ComputePipelineState,
pub lstm: ComputePipelineState,
pub dequant_gguf: ComputePipelineState,
pub q4k_mv_f32: ComputePipelineState,
pub q4k_mv_f32_sg: ComputePipelineState,
iq_grid_lut: Buffer,
pub rms_norm_bwd: ComputePipelineState,
pub rms_norm_bwd_param: ComputePipelineState,
pub rms_norm_bwd_inv_r_f32: ComputePipelineState,
pub rms_norm_bwd_param_reduce_f32: ComputePipelineState,
pub rope_bwd: ComputePipelineState,
pub cumsum_fwd: ComputePipelineState,
pub cumsum_bwd: ComputePipelineState,
pub im2col_group: ComputePipelineState,
pub im2col_group_w1: ComputePipelineState,
pub conv2d_bwd_weight_gemm: ComputePipelineState,
pub conv2d_bwd_weight_gemm_4x4: ComputePipelineState,
pub attn_bwd_scores_f32: ComputePipelineState,
pub attn_bwd_dp_f32: ComputePipelineState,
pub attn_bwd_ds_f32: ComputePipelineState,
pub attn_bwd_dv_f32: ComputePipelineState,
pub attn_bwd_dq_f32: ComputePipelineState,
pub attn_bwd_dk_f32: ComputePipelineState,
pub gather_bwd_zero: ComputePipelineState,
pub gather_bwd_acc: ComputePipelineState,
pub gaussian_splat_rasterize: ComputePipelineState,
pub gaussian_splat_rasterize_linear: ComputePipelineState,
pub gaussian_splat_rasterize_linear_traced: ComputePipelineState,
pub gaussian_splat_rasterize_backward_linear: ComputePipelineState,
pub gaussian_splat_adam_step: ComputePipelineState,
pub gaussian_splat_mse_loss_grad: ComputePipelineState,
pub gaussian_splat_ssim_stats: ComputePipelineState,
pub gaussian_splat_blended_loss_grad: ComputePipelineState,
pub gaussian_splat_project_training: ComputePipelineState,
pub gaussian_splat_geometry_backward: ComputePipelineState,
pub gaussian_splat_scene_grad_projection: ComputePipelineState,
pub gaussian_splat_splat_color_backward: ComputePipelineState,
pub gaussian_splat_emit_tile_keys: ComputePipelineState,
pub gaussian_splat_project_screen_ellipse: ComputePipelineState,
pub gaussian_splat_emit_tile_keys_conic: ComputePipelineState,
pub gaussian_splat_bin_histogram: ComputePipelineState,
pub gaussian_splat_bin_copy_counts: ComputePipelineState,
pub gaussian_splat_bin_prefix_sum: ComputePipelineState,
pub gaussian_splat_bin_scatter: ComputePipelineState,
pub gaussian_splat_build_tile_ranges: ComputePipelineState,
pub gaussian_splat_pack_grads: ComputePipelineState,
}
unsafe impl Send for Kernels {}
unsafe impl Sync for Kernels {}
impl Kernels {
fn new() -> Self {
let dev = metal_device().expect("Metal device required");
let library = crate::pipeline_cache::load_or_compile_library(&dev.device, &msl_source());
let pipeline = |name: &str| -> ComputePipelineState {
let f = library.get_function(name, None).expect(name);
dev.device
.new_compute_pipeline_state_with_function(&f)
.unwrap_or_else(|_| panic!("pipeline {name}"))
};
Self {
sgemm: pipeline("sgemm"),
sgemm_simd: pipeline("sgemm_simd"),
sgemm_simd_bias: pipeline("sgemm_simd_bias"),
sgemm_simd_4x4: pipeline("sgemm_simd_4x4"),
sgemm_simd_4x4_bias: pipeline("sgemm_simd_4x4_bias"),
hgemm_simd_4x4: pipeline("hgemm_simd_4x4"),
hgemm_simd_4x4_bias: pipeline("hgemm_simd_4x4_bias"),
bias_add_h: pipeline("bias_add_h"),
gelu_inplace_h: pipeline("gelu_inplace_h"),
gelu_approx_inplace_h: pipeline("gelu_approx_inplace_h"),
silu_inplace_h: pipeline("silu_inplace_h"),
layer_norm_h: pipeline("layer_norm_h"),
fused_residual_ln_h: pipeline("fused_residual_ln_h"),
fused_residual_rms_norm_h: pipeline("fused_residual_rms_norm_h"),
rms_norm_h: pipeline("rms_norm_h"),
softmax_lastax_h: pipeline("softmax_lastax_h"),
reduce_axes_h: pipeline("reduce_axes_h"),
elem_add_h: pipeline("elem_add_h"),
elem_mul_h: pipeline("elem_mul_h"),
gather_axis0_h: pipeline("gather_axis0_h"),
narrow_lastax_h: pipeline("narrow_lastax_h"),
sdpa_h: pipeline("sdpa_h"),
rope_h: pipeline("rope_h"),
cast_f32_to_f16: pipeline("cast_f32_to_f16"),
cast_f16_to_f32: pipeline("cast_f16_to_f32"),
copy_f32: pipeline("copy_f32"),
copy4: pipeline("copy4"),
sgemm_simd_padded: pipeline("sgemm_simd_padded"),
sgemm_simd_padded_bias: pipeline("sgemm_simd_padded_bias"),
sgemm_tiled: pipeline("sgemm_tiled"),
bias_add: pipeline("bias_add"),
gelu_inplace: pipeline("gelu_inplace"),
gelu_inplace4: pipeline("gelu_inplace4"),
gelu_approx_inplace: pipeline("gelu_approx_inplace"),
gelu_approx_inplace4: pipeline("gelu_approx_inplace4"),
gelu_approx_out4: pipeline("gelu_approx_out4"),
silu_inplace: pipeline("silu_inplace"),
silu_inplace4: pipeline("silu_inplace4"),
binary_broadcast_rhs_col_f32: pipeline("binary_broadcast_rhs_col_f32"),
binary_broadcast_rhs_col4: pipeline("binary_broadcast_rhs_col4"),
binary_broadcast_rhs_row_f32: pipeline("binary_broadcast_rhs_row_f32"),
binary_broadcast_rhs_row4: pipeline("binary_broadcast_rhs_row4"),
binary_broadcast_rhs_scalar_f32: pipeline("binary_broadcast_rhs_scalar_f32"),
binary_broadcast_rhs_scalar4: pipeline("binary_broadcast_rhs_scalar4"),
binary_broadcast_1ax_f32: pipeline("binary_broadcast_1ax_f32"),
binary_broadcast_1ax4: pipeline("binary_broadcast_1ax4"),
fused_binary_activation_f32: pipeline("fused_binary_activation_f32"),
fused_binary_activation4: pipeline("fused_binary_activation4"),
fused_ternary_activation_f32: pipeline("fused_ternary_activation_f32"),
fused_ternary_activation4: pipeline("fused_ternary_activation4"),
layer_norm: pipeline("layer_norm"),
rms_norm: pipeline("rms_norm"),
elem_add: pipeline("elem_add"),
elem_add4: pipeline("elem_add4"),
elem_sub4: pipeline("elem_sub4"),
binary_broadcast_f32: pipeline("binary_broadcast_f32"),
binary_broadcast_rank2_f32: pipeline("binary_broadcast_rank2_f32"),
binary_broadcast_rank24: pipeline("binary_broadcast_rank24"),
elem_mul: pipeline("elem_mul"),
elem_mul4: pipeline("elem_mul4"),
elem_div4: pipeline("elem_div4"),
gather_axis0: pipeline("gather_axis0"),
narrow_lastax: pipeline("narrow_lastax"),
narrow_lastax4: pipeline("narrow_lastax4"),
split_lastax: pipeline("split_lastax"),
split_lastax4: pipeline("split_lastax4"),
fused_residual_ln: pipeline("fused_residual_ln"),
fused_residual_rms_norm: pipeline("fused_residual_rms_norm"),
sdpa: pipeline("sdpa"),
sdpa_long: pipeline("sdpa_long"),
sdpa_fa_f32: pipeline("sdpa_fa_f32"),
rope: pipeline("rope"),
fused_swiglu: pipeline("fused_swiglu"),
fused_swiglu_h: pipeline("fused_swiglu_h"),
elementwise_region: pipeline("elementwise_region"),
batch_elementwise_region: pipeline("batch_elementwise_region"),
fused_swiglu_cast_f32_to_f16: pipeline("fused_swiglu_cast_f32_to_f16"),
fused_swiglu_cast_f16_to_f32: pipeline("fused_swiglu_cast_f16_to_f32"),
concat_segment_lastax: pipeline("concat_segment_lastax"),
concat_segment_lastax4: pipeline("concat_segment_lastax4"),
concat_lastax_multi: pipeline("concat_lastax_multi"),
concat_lastax_multi4: pipeline("concat_lastax_multi4"),
concat_segment_lastax_h: pipeline("concat_segment_lastax_h"),
elem_sub: pipeline("elem_sub"),
elem_div: pipeline("elem_div"),
elem_max: pipeline("elem_max"),
elem_min: pipeline("elem_min"),
elem_pow: pipeline("elem_pow"),
elem_compare: pipeline("elem_compare"),
elem_where: pipeline("elem_where"),
reduce_axes: pipeline("reduce_axes"),
topk_lastax: pipeline("topk_lastax"),
grouped_matmul: pipeline("grouped_matmul"),
scatter_add_zero: pipeline("scatter_add_zero"),
scatter_add_accumulate: pipeline("scatter_add_accumulate"),
transpose_nd: pipeline("transpose_nd"),
transpose_2d_f32: pipeline("transpose_2d_f32"),
transpose_2d_tiled_f32: pipeline("transpose_2d_tiled_f32"),
transpose_last2_batched_f32: pipeline("transpose_last2_batched_f32"),
transpose_last2_batched_tiled_f32: pipeline("transpose_last2_batched_tiled_f32"),
transpose_swap12_batched_trail_f32: pipeline("transpose_swap12_batched_trail_f32"),
transpose_swap12_batched_trail_tiled_f32: pipeline(
"transpose_swap12_batched_trail_tiled_f32",
),
gather_axis: pipeline("gather_axis"),
pool2d: pipeline("pool2d"),
conv2d: pipeline("conv2d"),
conv2d_w1: pipeline("conv2d_w1"),
layer_norm2d: pipeline("layer_norm2d"),
group_norm: pipeline("group_norm"),
resize_nearest_2x: pipeline("resize_nearest_2x"),
conv_transpose2d: pipeline("conv_transpose2d"),
relu_inplace: pipeline("relu_inplace"),
sigmoid_inplace: pipeline("sigmoid_inplace"),
tanh_inplace: pipeline("tanh_inplace"),
exp_inplace: pipeline("exp_inplace"),
log_inplace: pipeline("log_inplace"),
sqrt_inplace: pipeline("sqrt_inplace"),
rsqrt_inplace: pipeline("rsqrt_inplace"),
neg_inplace: pipeline("neg_inplace"),
abs_inplace: pipeline("abs_inplace"),
round_inplace: pipeline("round_inplace"),
sin_inplace: pipeline("sin_inplace"),
cos_inplace: pipeline("cos_inplace"),
tan_inplace: pipeline("tan_inplace"),
atan_inplace: pipeline("atan_inplace"),
softmax_lastax: pipeline("softmax_lastax"),
fft_radix2_full_f32: pipeline("fft_radix2_full_f32"),
fft_bit_reverse_f32: pipeline("fft_bit_reverse_f32"),
fft_inner_f32: pipeline("fft_inner_f32"),
fft_outer_r4_f32: pipeline("fft_outer_r4_f32"),
fft_outer_r2_f32: pipeline("fft_outer_r2_f32"),
gated_delta_net: pipeline("gated_delta_net"),
lstm: pipeline("lstm"),
dequant_gguf: pipeline("dequant_gguf"),
q4k_mv_f32: pipeline("q4k_mv_f32"),
q4k_mv_f32_sg: pipeline("q4k_mv_f32_sg"),
rms_norm_bwd: pipeline("rms_norm_bwd"),
rms_norm_bwd_param: pipeline("rms_norm_bwd_param"),
rms_norm_bwd_inv_r_f32: pipeline("rms_norm_bwd_inv_r_f32"),
rms_norm_bwd_param_reduce_f32: pipeline("rms_norm_bwd_param_reduce_f32"),
rope_bwd: pipeline("rope_bwd"),
cumsum_fwd: pipeline("cumsum_fwd"),
cumsum_bwd: pipeline("cumsum_bwd"),
im2col_group: pipeline("im2col_group"),
im2col_group_w1: pipeline("im2col_group_w1"),
conv2d_bwd_weight_gemm: pipeline("conv2d_bwd_weight_gemm"),
conv2d_bwd_weight_gemm_4x4: pipeline("conv2d_bwd_weight_gemm_4x4"),
attn_bwd_scores_f32: pipeline("attn_bwd_scores_f32"),
attn_bwd_dp_f32: pipeline("attn_bwd_dp_f32"),
attn_bwd_ds_f32: pipeline("attn_bwd_ds_f32"),
attn_bwd_dv_f32: pipeline("attn_bwd_dv_f32"),
attn_bwd_dq_f32: pipeline("attn_bwd_dq_f32"),
attn_bwd_dk_f32: pipeline("attn_bwd_dk_f32"),
gather_bwd_zero: pipeline("gather_bwd_zero"),
gather_bwd_acc: pipeline("gather_bwd_acc"),
gaussian_splat_rasterize: pipeline("gaussian_splat_rasterize"),
gaussian_splat_rasterize_linear: pipeline("gaussian_splat_rasterize_linear"),
gaussian_splat_rasterize_linear_traced: pipeline(
"gaussian_splat_rasterize_linear_traced",
),
gaussian_splat_rasterize_backward_linear: pipeline(
"gaussian_splat_rasterize_backward_linear",
),
gaussian_splat_adam_step: pipeline("gaussian_splat_adam_step"),
gaussian_splat_mse_loss_grad: pipeline("gaussian_splat_mse_loss_grad"),
gaussian_splat_ssim_stats: pipeline("gaussian_splat_ssim_stats"),
gaussian_splat_blended_loss_grad: pipeline("gaussian_splat_blended_loss_grad"),
gaussian_splat_project_training: pipeline("gaussian_splat_project_training"),
gaussian_splat_geometry_backward: pipeline("gaussian_splat_geometry_backward"),
gaussian_splat_scene_grad_projection: pipeline("gaussian_splat_scene_grad_projection"),
gaussian_splat_splat_color_backward: pipeline("gaussian_splat_splat_color_backward"),
gaussian_splat_emit_tile_keys: pipeline("gaussian_splat_emit_tile_keys"),
gaussian_splat_project_screen_ellipse: pipeline(
"gaussian_splat_project_screen_ellipse",
),
gaussian_splat_emit_tile_keys_conic: pipeline("gaussian_splat_emit_tile_keys_conic"),
gaussian_splat_bin_histogram: pipeline("gaussian_splat_bin_histogram"),
gaussian_splat_bin_copy_counts: pipeline("gaussian_splat_bin_copy_counts"),
gaussian_splat_bin_prefix_sum: pipeline("gaussian_splat_bin_prefix_sum"),
gaussian_splat_bin_scatter: pipeline("gaussian_splat_bin_scatter"),
gaussian_splat_build_tile_ranges: pipeline("gaussian_splat_build_tile_ranges"),
gaussian_splat_pack_grads: pipeline("gaussian_splat_pack_grads"),
iq_grid_lut: build_iq_grid_lut(
&metal_device()
.expect("rlx-metal: no Metal device for IQ grid LUT staging")
.device,
),
library,
}
}
pub fn iq_grid_buffer(&self) -> &Buffer {
&self.iq_grid_lut
}
}
fn build_iq_grid_lut(device: &metal::DeviceRef) -> Buffer {
use rlx_gguf::iq_grids::{
IQ1S_GRID, IQ2S_GRID, IQ2XS_GRID, IQ2XXS_GRID, IQ3S_GRID, IQ3XXS_GRID, KMASK_IQ2XS,
KSIGNS_IQ2XS,
};
let mut bytes = Vec::with_capacity(33_944);
bytes.extend_from_slice(&KMASK_IQ2XS);
bytes.extend_from_slice(&KSIGNS_IQ2XS);
for v in IQ2XXS_GRID.iter() {
bytes.extend_from_slice(&v.to_le_bytes());
}
for v in IQ2XS_GRID.iter() {
bytes.extend_from_slice(&v.to_le_bytes());
}
for v in IQ2S_GRID.iter() {
bytes.extend_from_slice(&v.to_le_bytes());
}
for v in IQ3XXS_GRID.iter() {
bytes.extend_from_slice(&v.to_le_bytes());
}
for v in IQ3S_GRID.iter() {
bytes.extend_from_slice(&v.to_le_bytes());
}
for v in IQ1S_GRID.iter() {
bytes.extend_from_slice(&v.to_le_bytes());
}
device.new_buffer_with_data(
bytes.as_ptr() as *const _,
bytes.len() as u64,
MTLResourceOptions::StorageModeShared,
)
}
pub fn kernels() -> &'static Kernels {
static K: OnceLock<Kernels> = OnceLock::new();
K.get_or_init(Kernels::new)
}
pub fn prewarm() -> &'static Kernels {
kernels()
}