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;
}
// F32 activations × F16 weights → F32 (Zonos Metal Linear). Avoids casting
// the full weight matrix to F32 on every decode step.
kernel void sgemm_f16w(
device const float* A [[buffer(0)]],
device const half* 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] * float(B[k * N + col]);
}
C[row * N + col] = sum;
}
// Small-M F16-weight GEMM (CFG decode: M=2). One thread per output column;
// all M rows share each B load — padded simdgroup wastes 6/8 A rows when M=2.
kernel void sgemm_f16w_small_m(
device const float* A [[buffer(0)]],
device const half* B [[buffer(1)]],
device float* C [[buffer(2)]],
constant uint& M [[buffer(3)]],
constant uint& K [[buffer(4)]],
constant uint& N [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
uint col = gid;
if (col >= N || M == 0u || M > 4u) return;
// Hot path: M==2 (Zonos CFG) — no per-k row branches.
if (M == 2u) {
float acc0 = 0.0f;
float acc1 = 0.0f;
uint k = 0;
for (; k + 3u < K; k += 4u) {
float4 bv = float4(
float(B[(k + 0u) * N + col]),
float(B[(k + 1u) * N + col]),
float(B[(k + 2u) * N + col]),
float(B[(k + 3u) * N + col]));
float4 a0 = *(device const float4*)(A + k);
float4 a1 = *(device const float4*)(A + K + k);
acc0 += a0.x * bv.x + a0.y * bv.y + a0.z * bv.z + a0.w * bv.w;
acc1 += a1.x * bv.x + a1.y * bv.y + a1.z * bv.z + a1.w * bv.w;
}
for (; k < K; ++k) {
float b = float(B[k * N + col]);
acc0 += A[k] * b;
acc1 += A[K + k] * b;
}
C[col] = acc0;
C[N + col] = acc1;
return;
}
float acc0 = 0.0f;
float acc1 = 0.0f;
float acc2 = 0.0f;
float acc3 = 0.0f;
uint k = 0;
for (; k + 3u < K; k += 4u) {
float b0 = float(B[(k + 0u) * N + col]);
float b1 = float(B[(k + 1u) * N + col]);
float b2 = float(B[(k + 2u) * N + col]);
float b3 = float(B[(k + 3u) * N + col]);
float4 bv = float4(b0, b1, b2, b3);
{
float4 a = *(device const float4*)(A + k);
acc0 += a.x * bv.x + a.y * bv.y + a.z * bv.z + a.w * bv.w;
}
if (M > 1u) {
float4 a = *(device const float4*)(A + K + k);
acc1 += a.x * bv.x + a.y * bv.y + a.z * bv.z + a.w * bv.w;
}
if (M > 2u) {
float4 a = *(device const float4*)(A + 2u * K + k);
acc2 += a.x * bv.x + a.y * bv.y + a.z * bv.z + a.w * bv.w;
}
if (M > 3u) {
float4 a = *(device const float4*)(A + 3u * K + k);
acc3 += a.x * bv.x + a.y * bv.y + a.z * bv.z + a.w * bv.w;
}
}
for (; k < K; ++k) {
float b = float(B[k * N + col]);
acc0 += A[k] * b;
if (M > 1u) acc1 += A[K + k] * b;
if (M > 2u) acc2 += A[2u * K + k] * b;
if (M > 3u) acc3 += A[3u * K + k] * b;
}
C[col] = acc0;
if (M > 1u) C[N + col] = acc1;
if (M > 2u) C[2u * N + col] = acc2;
if (M > 3u) C[3u * N + col] = acc3;
}
// Padded simdgroup with F16 B (same staging as sgemm_simd_padded).
kernel void sgemm_simd_padded_f16w(
device const float* A [[buffer(0)]],
device const half* 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;
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;
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;
}
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) ? float(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);
}
// Write back with bounds (same as sgemm_simd_padded).
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 out_row = row_base + cr;
uint out_col = col_base + cc;
if (out_row < M && out_col < N) {
C[out_row * N + out_col] = C_pad[idx];
}
}
}
// ── 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(-min(xa * xa, 80.0f)));
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(-min(xa * xa, 80.0f)));
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]);
// Clamp to the tanh saturation range: tanh(±15)≈±1 to f32 precision, but a
// huge argument (packed-QAT gate outliers → large x³, or +inf) makes Metal's
// fast-math tanh return NaN. clamp() also folds ±inf to ±15.
float inner = clamp(0.7978845608f * (x + 0.044715f * x * x * x), -15.0f, 15.0f);
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)));
}
kernel void relu_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
data[gid] = max(data[gid], half(0.0h));
}
kernel void sigmoid_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(1.0f / (1.0f + exp(-x)));
}
kernel void tanh_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(tanh(clamp(x, -15.0f, 15.0f)));
}
kernel void exp_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(exp(float(data[gid]))); }
kernel void log_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(log(float(data[gid]))); }
kernel void sqrt_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(sqrt(float(data[gid]))); }
kernel void rsqrt_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(rsqrt(float(data[gid]))); }
kernel void neg_inplace_h(
device half* 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_h(
device half* 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 sin_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(sin(float(data[gid]))); }
kernel void cos_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(cos(float(data[gid]))); }
kernel void tan_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(tan(float(data[gid]))); }
kernel void atan_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(atan(float(data[gid]))); }
kernel void round_inplace_h(
device half* data [[buffer(0)]],
constant uint& len [[buffer(1)]],
uint gid [[thread_position_in_grid]]
) { if (gid >= len) return; data[gid] = half(round(float(data[gid]))); }
// 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 = fmax(0.0f, 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 = fmax(0.0f, 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 char* arena [[buffer(0)]],
constant ulong& x_off [[buffer(1)]],
constant ulong& res_off [[buffer(2)]],
constant ulong& g_off [[buffer(3)]],
constant ulong& b_off [[buffer(4)]],
constant ulong& out_off [[buffer(5)]],
constant uint& h [[buffer(6)]],
constant float& eps [[buffer(7)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
device const half* x = (device const half*)(arena + x_off);
device const half* res = (device const half*)(arena + res_off);
device const half* gamma = (device const half*)(arena + g_off);
device const half* beta = (device const half*)(arena + b_off);
device half* out = (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(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 elem_sub_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_div_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_max_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] = max(a[gid], b[gid]); }
kernel void elem_min_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] = min(a[gid], b[gid]); }
kernel void elem_pow_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] = half(pow(float(a[gid]), float(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)]],
constant uint& cos_per_token [[buffer(11)]],
constant uint& interleaved [[buffer(12)]],
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;
// Per-seq-position table by default; per global token for ragged decode.
uint cos_row = (cos_per_token != 0u) ? bs : si;
// 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 (interleaved != 0u) {
// GPT-J / llama.cpp-NORM: pairs are adjacent (2d, 2d+1). cos/sin
// row index is the freq d (0..rot_half).
if (d < rot_half) {
uint a = 2u * d;
uint b = 2u * d + 1u;
float x1 = float(x[src_base + a]);
float x2 = float(x[src_base + b]);
float c = float(cos[cos_row * half_dh + d]);
float s = float(sin[cos_row * half_dh + d]);
out[dst_base + a] = half(x1 * c - x2 * s);
out[dst_base + b] = half(x2 * c + x1 * s);
} else if (d >= n_rot) {
out[dst_base + d] = x[src_base + d];
}
} else if (d < rot_half) {
float x1 = float(x[src_base + d]);
float x2 = float(x[src_base + rot_half + d]);
float c = float(cos[cos_row * half_dh + d]);
float s = float(sin[cos_row * 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];
}
}
// Native f32 fused-attention core for `Op::FusedAttentionBlock` (no-bias
// path). Reads the PACKED QKV projection `[B,S,3*inner]` (per token
// `[Q(inner)|K(inner)|V(inner)]`, heads interleaved), applies optional NeoX
// RoPE to Q/K inline, runs softmax SDPA with the score matrix resident in
// threadgroup memory (one threadgroup per batch·head), and writes the
// attention output `[B,S,inner]`. Collapses narrow×3 + transpose×3 + rope×2
// + attention into a single dispatch; the QKV / out projections stay GEMMs.
// mask_kind: 0=None, 1=Causal, 2=Custom (binary [B,S], <0.5 ⇒ drop).
kernel void fused_attn_block(
device const float* QKV [[buffer(0)]],
device const float* M [[buffer(1)]],
device const float* COS [[buffer(2)]],
device const float* SIN [[buffer(3)]],
device float* 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& mask_kind [[buffer(9)]],
constant uint& scale_bits[[buffer(10)]],
constant uint& has_rope [[buffer(11)]],
uint tgid [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
threadgroup float scores[64 * 64]; // seq ≤ 64 (gated in rlx-metal)
uint inner = heads * head_dim;
uint bi = tgid / heads;
uint hi = tgid % heads;
if (bi >= batch) return;
float scale = as_type<float>(scale_bits);
uint half_d = head_dim / 2;
uint tok = 3u * inner; // per-token stride in QKV
uint total = seq * seq;
for (uint idx = tid; idx < total; idx += tsize) {
uint qi = idx / seq;
uint ki = idx % seq;
uint qb = (bi * seq + qi) * tok + hi * head_dim; // Q
uint kb = (bi * seq + ki) * tok + inner + hi * head_dim; // K
float dot = 0.0f;
if (has_rope != 0u) {
uint qc = qi * half_d;
uint kc = ki * half_d;
for (uint i = 0; i < half_d; ++i) {
float q1 = QKV[qb + i], q2 = QKV[qb + half_d + i];
float k1 = QKV[kb + i], k2 = QKV[kb + half_d + i];
float cq = COS[qc + i], sq = SIN[qc + i];
float ck = COS[kc + i], sk = SIN[kc + i];
float qr1 = q1 * cq - q2 * sq, qr2 = q2 * cq + q1 * sq;
float kr1 = k1 * ck - k2 * sk, kr2 = k2 * ck + k1 * sk;
dot += qr1 * kr1 + qr2 * kr2;
}
} else {
for (uint d = 0; d < head_dim; ++d) dot += QKV[qb + d] * QKV[kb + d];
}
float s = dot * scale;
if (mask_kind == 1u) {
if (ki > qi) s = -1e9f;
} else if (mask_kind == 2u) {
if (M[bi * seq + ki] < 0.5f) s = -1e9f;
}
scores[qi * seq + ki] = s;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint qi = tid; qi < seq; qi += tsize) {
float mx = -1e30f;
for (uint ki = 0; ki < seq; ++ki) mx = max(mx, scores[qi * seq + ki]);
float sum = 0.0f;
for (uint ki = 0; ki < seq; ++ki) {
float e = precise::exp(scores[qi * seq + ki] - mx);
scores[qi * seq + ki] = e;
sum += e;
}
float inv = (sum > 0.0f) ? (1.0f / sum) : 0.0f;
for (uint ki = 0; ki < seq; ++ki) scores[qi * seq + ki] *= inv;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
uint otot = seq * head_dim;
for (uint idx = tid; idx < otot; idx += tsize) {
uint qi = idx / head_dim;
uint d = idx % head_dim;
float acc = 0.0f;
for (uint ki = 0; ki < seq; ++ki) {
uint vb = (bi * seq + ki) * tok + 2u * inner + hi * head_dim;
acc += scores[qi * seq + ki] * QKV[vb + d];
}
OUT[(bi * seq + qi) * inner + hi * head_dim + d] = acc;
}
}
// 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(-min(xa * xa, 80.0f)));
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(-min(xa * xa, 80.0f)));
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(-min(xa * xa, 80.0f)));
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(-min(xa * xa, 80.0f)));
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];
// Clamp to the tanh saturation range: tanh(±15)≈±1 to f32 precision, but a
// huge argument (packed-QAT gate outliers → large x³, or +inf) makes Metal's
// fast-math tanh return NaN. clamp() also folds ±inf to ±15.
float inner = clamp(0.7978845608f * (x + 0.044715f * x * x * x), -15.0f, 15.0f);
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];
// Clamp to the tanh saturation range: tanh(±15)≈±1 to f32 precision, but a
// huge argument (packed-QAT gate outliers → large x³, or +inf) makes Metal's
// fast-math tanh return NaN. clamp() also folds ±inf to ±15.
float inner = clamp(0.7978845608f * (x + 0.044715f * x * x * x), -15.0f, 15.0f);
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];
// Clamp to the tanh saturation range: tanh(±15)≈±1 to f32 precision, but a
// huge argument (packed-QAT gate outliers → large x³, or +inf) makes Metal's
// fast-math tanh return NaN. clamp() also folds ±inf to ±15.
float inner = clamp(0.7978845608f * (x + 0.044715f * x * x * x), -15.0f, 15.0f);
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(-min(xa * xa, 80.0f)));
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;
}
// Out-of-place SiLU: dst = src * sigmoid(src) — skips a prior Copy.
kernel void silu_out4(
device 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);
packed_float4 px = src[gid];
packed_float4 out;
for (uint c = 0; c < 4; ++c) {
float xv = px[c];
out[c] = xv / (1.0 + exp(-xv));
}
dst[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(-min(xa * xa, 80.0f)));
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(clamp(x, -15.0f, 15.0f)); // fast-math tanh NaNs on large |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;
}
/// Compare with optional scalar broadcast (`flags`: bit0=lhs scalar, bit1=rhs).
kernel void elem_compare_bcast(
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)]],
constant uint& flags [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float x = (flags & 1u) ? a[0] : a[gid];
float y = (flags & 2u) ? b[0] : 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;
}
// Depthwise causal 1-D conv on BSC layout `[B, W, C]` → `[B, out_seq, C]`.
// Weight is NCHW-packed `[C, 1, 1, K]` (same as `Op::Conv` depthwise). Optional
// fused SiLU. Replaces Transpose→Copy→Conv2D→Copy→Transpose(+Silu) for GDN.
kernel void depthwise_conv1d_bsc(
device float* arena [[buffer(0)]],
constant ulong& src_off [[buffer(1)]],
constant ulong& wt_off [[buffer(2)]],
constant ulong& dst_off [[buffer(3)]],
constant uint4& dims [[buffer(4)]], // batch, width, out_seq, channels
constant uint2& k_silu [[buffer(5)]], // k, silu!=0
device const float* wt_buf [[buffer(7)]],
uint gid [[thread_position_in_grid]]
) {
const uint batch = dims.x;
const uint width = dims.y;
const uint out_seq = dims.z;
const uint channels = dims.w;
const uint k = k_silu.x;
const bool do_silu = k_silu.y != 0u;
const uint total = batch * out_seq * channels;
if (gid >= total || k == 0u) return;
const uint c = gid % channels;
const uint tmp = gid / channels;
const uint t = tmp % out_seq;
const uint b = tmp / out_seq;
device const float* src = (device const float*)((device char*)arena + src_off);
device float* dst = (device float*)((device char*)arena + dst_off);
device const float* wt = (device const float*)((device const char*)wt_buf + wt_off);
const uint src_base = b * width * channels + t * channels + c;
const uint wt_base = c * k;
float acc = 0.f;
if (k == 4u) {
acc += src[src_base + 0u * channels] * wt[wt_base + 0u];
acc += src[src_base + 1u * channels] * wt[wt_base + 1u];
acc += src[src_base + 2u * channels] * wt[wt_base + 2u];
acc += src[src_base + 3u * channels] * wt[wt_base + 3u];
} else {
for (uint ki = 0u; ki < k; ++ki) {
acc += src[src_base + ki * channels] * wt[wt_base + ki];
}
}
if (do_silu) {
acc = acc / (1.0f + exp(-acc));
}
dst[b * out_seq * channels + t * channels + c] = 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 = fmax(0.0f, 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;
}
// ──────────────────────────────────────────────────────────────────────
// Training backward kernels. All three are OUTPUT-PARALLEL: each thread owns
// exactly one output element and writes it once — no atomics, no pre-zeroing,
// no scratch buffers, and (critically) no GPU→CPU sync. They mirror the CPU
// reference (crates/rlx-cpu/src/{conv_bwd,training_bwd}.rs) bit-for-bit,
// including the max-pool strict-`>` first-in-scan arg-max tie-break.
// ──────────────────────────────────────────────────────────────────────
// Max-pool backward. One thread per INPUT element (n,c,ih,iw); it accumulates
// dy from every output window in which it is the arg-max. Handles overlapping
// windows and padding. Race-free: distinct threads write distinct dx.
kernel void maxpool2d_backward(
device const float* x [[buffer(0)]],
device const float* dy [[buffer(1)]],
device float* dx [[buffer(2)]],
constant uint4& p0 [[buffer(3)]], // [N, C, H, W]
constant uint4& p1 [[buffer(4)]], // [H_out, W_out, kh, kw]
constant uint4& p2 [[buffer(5)]], // [sh, sw, ph, pw]
uint3 gid [[thread_position_in_grid]]
) {
uint N = p0.x, C = p0.y, H = p0.z, W = p0.w;
uint h_out = p1.x, w_out = p1.y, kh = p1.z, kw = p1.w;
uint sh = p2.x, sw = p2.y, ph = p2.z, pw = p2.w;
uint iw = gid.x, ih = gid.y, nc = gid.z;
if (nc >= N * C || ih >= H || iw >= W) return;
int p_h = (int)ih + (int)ph;
int p_w = (int)iw + (int)pw;
int oh_max = p_h / (int)sh;
int ow_max = p_w / (int)sw;
if (oh_max >= (int)h_out) oh_max = (int)h_out - 1;
if (ow_max >= (int)w_out) ow_max = (int)w_out - 1;
// floor((p - k)/s)+1, clamped: integer div is floor only for non-neg args.
int oh_min = (p_h - (int)kh < 0) ? 0 : (p_h - (int)kh) / (int)sh + 1;
int ow_min = (p_w - (int)kw < 0) ? 0 : (p_w - (int)kw) / (int)sw + 1;
uint in_chan = nc * H * W;
uint out_chan = nc * h_out * w_out;
float acc = 0.0f;
for (int oh = oh_min; oh <= oh_max; ++oh) {
for (int ow = ow_min; ow <= ow_max; ++ow) {
float best_v = -INFINITY;
int best_h = -1, best_w = -1;
for (uint ki = 0; ki < kh; ++ki) {
int hh = oh * (int)sh + (int)ki - (int)ph;
if (hh < 0 || hh >= (int)H) continue;
for (uint kj = 0; kj < kw; ++kj) {
int ww = ow * (int)sw + (int)kj - (int)pw;
if (ww < 0 || ww >= (int)W) continue;
float v = x[in_chan + (uint)hh * W + (uint)ww];
if (v > best_v) { best_v = v; best_h = hh; best_w = ww; }
}
}
if (best_h == (int)ih && best_w == (int)iw)
acc += dy[out_chan + (uint)oh * w_out + (uint)ow];
}
}
dx[in_chan + ih * W + iw] = acc;
}
// Conv2d backward-input (transposed-conv gather). One thread per dx element
// (n, ci, ih, iw); gathers from every (co,ki,kj) whose forward map lands here.
kernel void conv2d_backward_input(
device const float* dy [[buffer(0)]],
device const float* wt [[buffer(1)]],
device float* dx [[buffer(2)]],
constant uint4& a [[buffer(3)]], // [N, C_in, H, W_in]
constant uint4& b [[buffer(4)]], // [C_out, H_out, W_out, kh]
constant uint4& cc [[buffer(5)]], // [kw, sh, sw, ph]
constant uint4& d [[buffer(6)]], // [pw, dh, dw, groups]
uint3 gid [[thread_position_in_grid]]
) {
uint N=a.x, C_in=a.y, H=a.z, W_in=a.w;
uint C_out=b.x, H_out=b.y, W_out=b.z, kh=b.w;
uint kw=cc.x, sh=cc.y, sw=cc.z, ph=cc.w;
uint pw=d.x, dh=d.y, dw=d.z, groups=d.w;
uint iw = gid.x, ih = gid.y, nci = gid.z;
if (nci >= N * C_in || ih >= H || iw >= W_in) return;
uint n = nci / C_in;
uint ci = nci % C_in;
uint c_in_per_g = C_in / groups;
uint c_out_per_g = C_out / groups;
uint g = ci / c_in_per_g;
uint ci_local = ci % c_in_per_g;
float acc = 0.0f;
for (uint ki = 0; ki < kh; ++ki) {
int num_h = (int)ih + (int)ph - (int)(ki * dh);
if (num_h < 0 || (num_h % (int)sh) != 0) continue;
int ho = num_h / (int)sh;
if (ho >= (int)H_out) continue;
for (uint kj = 0; kj < kw; ++kj) {
int num_w = (int)iw + (int)pw - (int)(kj * dw);
if (num_w < 0 || (num_w % (int)sw) != 0) continue;
int wo = num_w / (int)sw;
if (wo >= (int)W_out) continue;
for (uint col = 0; col < c_out_per_g; ++col) {
uint co = g * c_out_per_g + col;
uint w_idx = ((co * c_in_per_g + ci_local) * kh + ki) * kw + kj;
uint dy_idx = ((n * C_out + co) * H_out + (uint)ho) * W_out + (uint)wo;
acc += wt[w_idx] * dy[dy_idx];
}
}
}
dx[((n * C_in + ci) * H + ih) * W_in + iw] = acc;
}
// Conv2d backward-weight (direct, batch-reduced). One thread per dw element
// (co, ci_local, ki, kj); sums dy*x over (n, ho, wo).
kernel void conv2d_backward_weight(
device const float* x [[buffer(0)]],
device const float* dy [[buffer(1)]],
device float* dw [[buffer(2)]],
constant uint4& a [[buffer(3)]], // [N, C_in, H, W]
constant uint4& b [[buffer(4)]], // [C_out, H_out, W_out, kh]
constant uint4& cc [[buffer(5)]], // [kw, sh, sw, ph]
constant uint4& d [[buffer(6)]], // [pw, dh, dw_dil, groups]
uint3 gid [[thread_position_in_grid]]
) {
uint N=a.x, C_in=a.y, H=a.z, W=a.w;
uint C_out=b.x, H_out=b.y, W_out=b.z, kh=b.w;
uint kw=cc.x, sh=cc.y, sw=cc.z, ph=cc.w;
uint pw=d.x, dh=d.y, dwd=d.z, groups=d.w;
uint kj = gid.x, ki = gid.y, coci = gid.z;
uint c_in_per_g = C_in / groups;
uint c_out_per_g = C_out / groups;
if (kj >= kw || ki >= kh || coci >= C_out * c_in_per_g) return;
uint co = coci / c_in_per_g;
uint ci_local = coci % c_in_per_g;
uint g = co / c_out_per_g;
uint ci = g * c_in_per_g + ci_local;
float acc = 0.0f;
for (uint n = 0; n < N; ++n) {
for (uint ho = 0; ho < H_out; ++ho) {
int hh = (int)(ho * sh + ki * dh) - (int)ph;
if (hh < 0 || hh >= (int)H) continue;
for (uint wo = 0; wo < W_out; ++wo) {
int ww = (int)(wo * sw + kj * dwd) - (int)pw;
if (ww < 0 || ww >= (int)W) continue;
uint x_idx = ((n * C_in + ci) * H + (uint)hh) * W + (uint)ww;
uint dy_idx = ((n * C_out + co) * H_out + ho) * W_out + wo;
acc += x[x_idx] * dy[dy_idx];
}
}
}
dw[((co * c_in_per_g + ci_local) * kh + ki) * kw + kj] = acc;
}
// Conv2d backward-weight, pass 1 (batch-parallel). One thread per
// (n, co, ci_local, ki, kj) writes a per-sample partial sum into `part`,
// laid out [N, C_out, c_in_per_g, kh, kw]. Threads scale with N, so small
// kernels (conv1: 288 dw elems) no longer starve the GPU.
kernel void conv2d_backward_weight_partial(
device const float* x [[buffer(0)]],
device const float* dy [[buffer(1)]],
device float* part [[buffer(2)]],
constant uint4& a [[buffer(3)]], // [N, C_in, H, W]
constant uint4& b [[buffer(4)]], // [C_out, H_out, W_out, kh]
constant uint4& cc [[buffer(5)]], // [kw, sh, sw, ph]
constant uint4& d [[buffer(6)]], // [pw, dh, dw_dil, groups]
uint3 gid [[thread_position_in_grid]]
) {
uint N=a.x, C_in=a.y, H=a.z, W=a.w;
uint C_out=b.x, H_out=b.y, W_out=b.z, kh=b.w;
uint kw=cc.x, sh=cc.y, sw=cc.z, ph=cc.w;
uint pw=d.x, dh=d.y, dwd=d.z, groups=d.w;
uint c_in_per_g = C_in / groups;
uint c_out_per_g = C_out / groups;
uint wsz = c_in_per_g * kh * kw; // per-(n,co) weight slab
uint j = gid.x, co = gid.y, n = gid.z;
if (j >= wsz || co >= C_out || n >= N) return;
uint kj = j % kw;
uint ki = (j / kw) % kh;
uint ci_local = j / (kw * kh);
uint g = co / c_out_per_g;
uint ci = g * c_in_per_g + ci_local;
float acc = 0.0f;
for (uint ho = 0; ho < H_out; ++ho) {
int hh = (int)(ho * sh + ki * dh) - (int)ph;
if (hh < 0 || hh >= (int)H) continue;
for (uint wo = 0; wo < W_out; ++wo) {
int ww = (int)(wo * sw + kj * dwd) - (int)pw;
if (ww < 0 || ww >= (int)W) continue;
uint x_idx = ((n * C_in + ci) * H + (uint)hh) * W + (uint)ww;
uint dy_idx = ((n * C_out + co) * H_out + ho) * W_out + wo;
acc += x[x_idx] * dy[dy_idx];
}
}
part[(n * C_out + co) * wsz + j] = acc;
}
// Conv2d backward-weight, pass 2: sum the per-sample partials over the batch.
// One thread per dw element. `wslab` = C_out * c_in_per_g * kh * kw.
kernel void conv2d_backward_weight_reduce(
device const float* part [[buffer(0)]],
device float* dw [[buffer(1)]],
constant uint2& dims [[buffer(2)]], // [N, wslab]
uint gid [[thread_position_in_grid]]
) {
uint N = dims.x, wslab = dims.y;
if (gid >= wslab) return;
float acc = 0.0f;
for (uint n = 0; n < N; ++n) acc += part[n * wslab + gid];
dw[gid] = 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];
}
/// Where with optional scalar broadcast on any operand.
/// `flags`: bit0=cond scalar, bit1=a scalar, bit2=b scalar.
kernel void elem_where_bcast(
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)]],
constant uint& flags [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
float c = (flags & 1u) ? cond[0] : cond[gid];
float x = (flags & 2u) ? a[0] : a[gid];
float y = (flags & 4u) ? b[0] : b[gid];
out[gid] = c != 0.0f ? x : y;
}
// Single-rounded fused multiply-add: out = fma(a, b, c). MSL `fma` is a true
// fused op (one rounding) — required for compensated / error-free-transform math.
kernel void elem_fma(
device const float* a [[buffer(0)]],
device const float* b [[buffer(1)]],
device const float* c [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& len [[buffer(4)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= len) return;
out[gid] = fma(a[gid], b[gid], c[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;
// Clamp to the tanh saturation range: Metal's fast-math `tanh` returns NaN for
// large |x| (it evaluates via exp, which overflows) — e.g. VITS posterior
// WaveNet pre-activations reach ~300. tanh(±15)≈±1 to f32 precision, matching
// CPU's stable tanh; clamp also folds ±inf to ±15.
data[gid] = tanh(clamp(data[gid], -15.0f, 15.0f));
}
// 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;
}
}
// Fused dense / soft-label softmax cross-entropy along the last axis.
// One threadgroup per row computes, numerically stably,
// loss[n] = logsumexp(logits[n]) - Σ_c targets[n,c]·logits[n,c]
// via three threadgroup reductions (row max, Σexp, Σtargets·logits).
// `cols` is the class count C; output is one scalar per row.
kernel void softmax_cross_entropy_dense(
device const float* logits [[buffer(0)]],
device const float* targets [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& cols [[buffer(3)]],
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: row max for numerical stability.
float local_max = -INFINITY;
for (uint i = tid; i < cols; i += tsize) {
local_max = max(local_max, logits[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];
threadgroup_barrier(mem_flags::mem_threadgroup);
// Pass 2: Σ exp(x - max) and Σ targets·logits in one sweep.
float local_sum = 0.0f;
float local_dot = 0.0f;
for (uint i = tid; i < cols; i += tsize) {
float v = logits[base + i];
local_sum += exp(v - row_max);
local_dot += targets[base + i] * v;
}
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 sum_exp = partial[0];
threadgroup_barrier(mem_flags::mem_threadgroup);
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];
if (tid == 0) {
out[row] = (row_max + log(sum_exp)) - dot;
}
}
// Softmax cross-entropy with integer labels (forward). One threadgroup per row.
// loss[row] = logsumexp(logits[row]) - logits[row, label]. Replaces the
// softmax + one-hot(compare/where) + gather decomposition on Metal.
kernel void softmax_cross_entropy_with_logits(
device const float* logits [[buffer(0)]],
device const float* labels [[buffer(1)]],
device float* out [[buffer(2)]],
constant uint& cols [[buffer(3)]],
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, logits[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];
threadgroup_barrier(mem_flags::mem_threadgroup);
float local_sum = 0.0f;
for (uint i = tid; i < cols; i += tsize) local_sum += exp(logits[base + i] - row_max);
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);
}
if (tid == 0) {
uint label = (uint)labels[row];
out[row] = (row_max + log(partial[0])) - logits[base + label];
}
}
// Softmax cross-entropy backward (integer labels). One threadgroup per row.
// dlogits[row,k] = (softmax(logits[row])[k] - [k==label]) * d_loss[row].
kernel void softmax_cross_entropy_backward(
device const float* logits [[buffer(0)]],
device const float* labels [[buffer(1)]],
device const float* d_loss [[buffer(2)]],
device float* dlogits [[buffer(3)]],
constant uint& cols [[buffer(4)]],
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, logits[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];
threadgroup_barrier(mem_flags::mem_threadgroup);
float local_sum = 0.0f;
for (uint i = tid; i < cols; i += tsize) local_sum += exp(logits[base + i] - row_max);
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];
float scale = d_loss[row];
uint label = (uint)labels[row];
for (uint k = tid; k < cols; k += tsize) {
float p = exp(logits[base + k] - row_max) * inv_sum;
dlogits[base + k] = (p - (k == label ? 1.0f : 0.0f)) * scale;
}
}
// 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];
}
// Mid-axis concat (inner > 1): copy one segment src[outer][src_axis][inner]
// into dst[outer][dst_axis][inner] starting at axis offset `dst_col`. One 1D
// dispatch per segment, encoded into the live command buffer (NO commit/wait)
// — the host-copy fallback used to commit+wait per concat, serializing a
// decode step into ~100 tiny GPU submissions (the dominant cost on GGUF KV
// caches). `arena` bound at 0 + ulong byte offsets so it is correct on
// >4 GiB arenas (task #50). Generic: subsumes last-axis concat when inner==1.
kernel void concat_midaxis_seg(
device char* arena [[buffer(0)]],
constant ulong& dst_byte [[buffer(1)]],
constant ulong& src_byte [[buffer(2)]],
constant uint& outer [[buffer(3)]],
constant uint& dst_axis [[buffer(4)]],
constant uint& src_axis [[buffer(5)]],
constant uint& inner [[buffer(6)]],
constant uint& dst_col [[buffer(7)]],
uint gid [[thread_position_in_grid]]
) {
uint total = outer * src_axis * inner;
if (gid >= total) return;
uint ii = gid % inner;
uint tmp = gid / inner;
uint a = tmp % src_axis;
uint o = tmp / src_axis;
device const float* src = (device const float*)(arena + src_byte);
device float* dst = (device float*)(arena + dst_byte);
dst[(o * dst_axis + dst_col + a) * inner + ii] = src[(o * src_axis + a) * inner + ii];
}
kernel void concat_midaxis_seg_h(
device char* arena [[buffer(0)]],
constant ulong& dst_byte [[buffer(1)]],
constant ulong& src_byte [[buffer(2)]],
constant uint& outer [[buffer(3)]],
constant uint& dst_axis [[buffer(4)]],
constant uint& src_axis [[buffer(5)]],
constant uint& inner [[buffer(6)]],
constant uint& dst_col [[buffer(7)]],
uint gid [[thread_position_in_grid]]
) {
uint total = outer * src_axis * inner;
if (gid >= total) return;
uint ii = gid % inner;
uint tmp = gid / inner;
uint a = tmp % src_axis;
uint o = tmp / src_axis;
device const half* src = (device const half*)(arena + src_byte);
device half* dst = (device half*)(arena + dst_byte);
dst[(o * dst_axis + dst_col + a) * inner + ii] = src[(o * src_axis + a) * inner + ii];
}
// 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 = fmax(0.0f, 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];
}
}
// DiT adaLN-Zero: out = norm(x) * (1 + scale) + shift
// scale/shift broadcast over leading dims (typically [B,1,D] over [B,S,D]).
// `layer_norm != 0` → mean-subtract; else RMS only.
// Packed lead dims: [lead_rank, x_lead[8], mod_lead[8]] as 17 uints.
kernel void ada_layer_norm(
device const float* x [[buffer(0)]],
device const float* scale [[buffer(1)]],
device const float* shift [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant float& eps [[buffer(5)]],
constant uint& layer_norm [[buffer(6)]],
constant uint* lead_pack [[buffer(7)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
uint lead_rank = lead_pack[0];
// Decode `row` into a multi-index over x_lead; fold into mod base offset.
uint rem = row;
uint mod_base = 0;
uint mod_stride = h;
for (int j = int(lead_rank) - 1; j >= 0; --j) {
uint xd = lead_pack[1 + j];
if (xd == 0u) { xd = 1u; }
uint xi = rem % xd;
rem /= xd;
uint md = lead_pack[9 + j];
if (md == 0u) { md = 1u; }
if (md != 1u) {
mod_base += xi * mod_stride;
}
mod_stride *= md;
}
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
float local_sum = 0.0;
float local_sumsq = 0.0;
for (uint i = tid; i < h; i += tsize) {
float v = x[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 = 0.0;
float inv;
if (layer_norm != 0u) {
mean = partial_sum[0] / float(h);
float var = fmax(0.0f, partial_sumsq[0] / float(h) - mean * mean);
inv = rsqrt(var + eps);
} else {
float ms = partial_sumsq[0] / float(h);
inv = rsqrt(ms + eps);
}
for (uint i = tid; i < h; i += tsize) {
float n = (x[row * h + i] - mean) * inv;
out[row * h + i] = n * (1.0f + scale[mod_base + i]) + shift[mod_base + i];
}
}
// f16 DiT adaLN — accumulate in f32, store half.
kernel void ada_layer_norm_h(
device const half* x [[buffer(0)]],
device const half* scale [[buffer(1)]],
device const half* shift [[buffer(2)]],
device half* out [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant float& eps [[buffer(5)]],
constant uint& layer_norm [[buffer(6)]],
constant uint* lead_pack [[buffer(7)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
uint lead_rank = lead_pack[0];
uint rem = row;
uint mod_base = 0;
uint mod_stride = h;
for (int j = int(lead_rank) - 1; j >= 0; --j) {
uint xd = lead_pack[1 + j];
if (xd == 0u) { xd = 1u; }
uint xi = rem % xd;
rem /= xd;
uint md = lead_pack[9 + j];
if (md == 0u) { md = 1u; }
if (md != 1u) {
mod_base += xi * mod_stride;
}
mod_stride *= md;
}
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
float local_sum = 0.0;
float local_sumsq = 0.0;
for (uint i = tid; i < h; i += tsize) {
float v = float(x[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 = 0.0;
float inv;
if (layer_norm != 0u) {
mean = partial_sum[0] / float(h);
float var = fmax(0.0f, partial_sumsq[0] / float(h) - mean * mean);
inv = rsqrt(var + eps);
} else {
float ms = partial_sumsq[0] / float(h);
inv = rsqrt(ms + eps);
}
for (uint i = tid; i < h; i += tsize) {
float n = (float(x[row * h + i]) - mean) * inv;
out[row * h + i] = half(n * (1.0f + float(scale[mod_base + i])) + float(shift[mod_base + i]));
}
}
// DiT gated residual: out = x + gate * y (gate broadcasts like adaLN scale).
// lead_pack: [lead_rank, x_lead[8], gate_lead[8]].
kernel void gated_residual(
device const float* x [[buffer(0)]],
device const float* y [[buffer(1)]],
device const float* gate [[buffer(2)]],
device float* out [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant uint* lead_pack [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
uint lead_rank = lead_pack[0];
uint row = gid / h;
uint col = gid % h;
uint rem = row;
uint gate_base = 0;
uint gate_stride = h;
for (int j = int(lead_rank) - 1; j >= 0; --j) {
uint xd = lead_pack[1 + j];
if (xd == 0u) { xd = 1u; }
uint xi = rem % xd;
rem /= xd;
uint gd = lead_pack[9 + j];
if (gd == 0u) { gd = 1u; }
if (gd != 1u) {
gate_base += xi * gate_stride;
}
gate_stride *= gd;
}
uint i = row * h + col;
out[i] = x[i] + gate[gate_base + col] * y[i];
}
kernel void gated_residual_h(
device const half* x [[buffer(0)]],
device const half* y [[buffer(1)]],
device const half* gate [[buffer(2)]],
device half* out [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant uint* lead_pack [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
uint lead_rank = lead_pack[0];
uint row = gid / h;
uint col = gid % h;
uint rem = row;
uint gate_base = 0;
uint gate_stride = h;
for (int j = int(lead_rank) - 1; j >= 0; --j) {
uint xd = lead_pack[1 + j];
if (xd == 0u) { xd = 1u; }
uint xi = rem % xd;
rem /= xd;
uint gd = lead_pack[9 + j];
if (gd == 0u) { gd = 1u; }
if (gd != 1u) {
gate_base += xi * gate_stride;
}
gate_stride *= gd;
}
uint i = row * h + col;
out[i] = half(float(x[i]) + float(gate[gate_base + col]) * float(y[i]));
}
// Packed DiT gated residual backward: out = [dx ∥ dy ∥ dgate] (1-D).
// Threadgroups = mod_rows (unique gate rows); each TG owns one gate slice
// and loops `seq_per_mod` x-rows that share it (DiT [B,S,D] / [B,1,D]).
kernel void gated_residual_backward(
device const float* y [[buffer(0)]],
device const float* gate [[buffer(1)]],
device const float* dy [[buffer(2)]],
device float* packed [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant uint& seq_per_mod [[buffer(5)]],
constant uint& mod_rows [[buffer(6)]],
uint m [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
if (m >= mod_rows) return;
uint nx = mod_rows * seq_per_mod * h;
uint gate_base = m * h;
device float* dx = packed;
device float* dy_out = packed + nx;
device float* dgate = packed + 2 * nx;
for (uint i = tid; i < h; i += tsize) {
float acc = 0.0f;
for (uint s = 0; s < seq_per_mod; s++) {
uint row = m * seq_per_mod + s;
uint idx = row * h + i;
float g = gate[gate_base + i];
float d = dy[idx];
dx[idx] = d;
dy_out[idx] = d * g;
acc += d * y[idx];
}
dgate[gate_base + i] = acc;
}
}
kernel void gated_residual_backward_h(
device const half* y [[buffer(0)]],
device const half* gate [[buffer(1)]],
device const half* dy [[buffer(2)]],
device half* packed [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant uint& seq_per_mod [[buffer(5)]],
constant uint& mod_rows [[buffer(6)]],
uint m [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
if (m >= mod_rows) return;
uint nx = mod_rows * seq_per_mod * h;
uint gate_base = m * h;
device half* dx = packed;
device half* dy_out = packed + nx;
device half* dgate = packed + 2 * nx;
for (uint i = tid; i < h; i += tsize) {
float acc = 0.0f;
for (uint s = 0; s < seq_per_mod; s++) {
uint row = m * seq_per_mod + s;
uint idx = row * h + i;
float g = float(gate[gate_base + i]);
float d = float(dy[idx]);
dx[idx] = half(d);
dy_out[idx] = half(d * g);
acc += d * float(y[idx]);
}
dgate[gate_base + i] = half(acc);
}
}
// Packed AdaLayerNorm backward: out = [dx ∥ dscale ∥ dshift] (1-D).
// Same launch geometry as gated_residual_backward.
kernel void ada_layer_norm_backward(
device const float* x [[buffer(0)]],
device const float* scale [[buffer(1)]],
device const float* dy [[buffer(2)]],
device float* packed [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant float& eps [[buffer(5)]],
constant uint& layer_norm [[buffer(6)]],
constant uint& seq_per_mod [[buffer(7)]],
constant uint& mod_rows [[buffer(8)]],
uint m [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
if (m >= mod_rows) return;
uint nx = mod_rows * seq_per_mod * h;
uint mod_len = mod_rows * h;
uint mod_base = m * h;
device float* dx = packed;
device float* dscale = packed + nx;
device float* dshift = packed + nx + mod_len;
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
// Zero modulation grads for this slice.
for (uint i = tid; i < h; i += tsize) {
dscale[mod_base + i] = 0.0f;
dshift[mod_base + i] = 0.0f;
}
threadgroup_barrier(mem_flags::mem_device);
float inv_h = 1.0f / float(h);
for (uint s = 0; s < seq_per_mod; s++) {
uint row = m * seq_per_mod + s;
float local_sum = 0.0f;
float local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = x[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 = 0.0f;
float inv;
if (layer_norm != 0u) {
mean = partial_sum[0] * inv_h;
float var = fmax(0.0f, partial_sumsq[0] * inv_h - mean * mean);
inv = rsqrt(var + eps);
} else {
inv = rsqrt(partial_sumsq[0] * inv_h + eps);
}
// First pass: accumulate dscale/dshift and reduction stats for dx.
float local_sy = 0.0f;
float local_sxh = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float n = (x[row * h + i] - mean) * inv;
float d = dy[row * h + i];
float sc = scale[mod_base + i];
float sy = d * (1.0f + sc);
dscale[mod_base + i] += d * n;
dshift[mod_base + i] += d;
local_sy += sy;
local_sxh += sy * n;
}
partial_sum[tid] = local_sy;
partial_sumsq[tid] = local_sxh;
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 m_sy = partial_sum[0] * inv_h;
float m_sxh = partial_sumsq[0] * inv_h;
for (uint i = tid; i < h; i += tsize) {
float n = (x[row * h + i] - mean) * inv;
float d = dy[row * h + i];
float sc = scale[mod_base + i];
float sy = d * (1.0f + sc);
if (layer_norm != 0u) {
dx[row * h + i] = inv * (sy - m_sy - n * m_sxh);
} else {
float n_rms = x[row * h + i] * inv;
dx[row * h + i] = inv * (sy - n_rms * m_sxh);
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}
kernel void ada_layer_norm_backward_h(
device const half* x [[buffer(0)]],
device const half* scale [[buffer(1)]],
device const half* dy [[buffer(2)]],
device half* packed [[buffer(3)]],
constant uint& h [[buffer(4)]],
constant float& eps [[buffer(5)]],
constant uint& layer_norm [[buffer(6)]],
constant uint& seq_per_mod [[buffer(7)]],
constant uint& mod_rows [[buffer(8)]],
uint m [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
if (m >= mod_rows) return;
uint nx = mod_rows * seq_per_mod * h;
uint mod_len = mod_rows * h;
uint mod_base = m * h;
device half* dx = packed;
device half* dscale = packed + nx;
device half* dshift = packed + nx + mod_len;
threadgroup float partial_sum[256];
threadgroup float partial_sumsq[256];
for (uint i = tid; i < h; i += tsize) {
dscale[mod_base + i] = half(0.0h);
dshift[mod_base + i] = half(0.0h);
}
threadgroup_barrier(mem_flags::mem_device);
float inv_h = 1.0f / float(h);
for (uint s = 0; s < seq_per_mod; s++) {
uint row = m * seq_per_mod + s;
float local_sum = 0.0f;
float local_sumsq = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float v = float(x[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 = 0.0f;
float inv;
if (layer_norm != 0u) {
mean = partial_sum[0] * inv_h;
float var = fmax(0.0f, partial_sumsq[0] * inv_h - mean * mean);
inv = rsqrt(var + eps);
} else {
inv = rsqrt(partial_sumsq[0] * inv_h + eps);
}
float local_sy = 0.0f;
float local_sxh = 0.0f;
for (uint i = tid; i < h; i += tsize) {
float n = (float(x[row * h + i]) - mean) * inv;
float d = float(dy[row * h + i]);
float sc = float(scale[mod_base + i]);
float sy = d * (1.0f + sc);
dscale[mod_base + i] = half(float(dscale[mod_base + i]) + d * n);
dshift[mod_base + i] = half(float(dshift[mod_base + i]) + d);
local_sy += sy;
local_sxh += sy * n;
}
partial_sum[tid] = local_sy;
partial_sumsq[tid] = local_sxh;
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 m_sy = partial_sum[0] * inv_h;
float m_sxh = partial_sumsq[0] * inv_h;
for (uint i = tid; i < h; i += tsize) {
float n = (float(x[row * h + i]) - mean) * inv;
float d = float(dy[row * h + i]);
float sc = float(scale[mod_base + i]);
float sy = d * (1.0f + sc);
if (layer_norm != 0u) {
dx[row * h + i] = half(inv * (sy - m_sy - n * m_sxh));
} else {
float n_rms = float(x[row * h + i]) * inv;
dx[row * h + i] = half(inv * (sy - n_rms * m_sxh));
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
}
// Fused residual + RMSNorm: out = RmsNorm(x + residual, gamma, beta).
// Arena + byte offsets (task #50) — `set_buffer(..., large_offset)` silently
// drops writes on M-series for big Qwen3.5 / Bonsai arenas.
kernel void fused_residual_rms_norm(
device const char* arena [[buffer(0)]],
constant ulong& x_off [[buffer(1)]],
constant ulong& res_off [[buffer(2)]],
constant ulong& g_off [[buffer(3)]],
constant ulong& b_off [[buffer(4)]],
constant ulong& out_off [[buffer(5)]],
constant uint& h [[buffer(6)]],
constant float& eps [[buffer(7)]],
uint row [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint tsize [[threads_per_threadgroup]]
) {
device const float* x = (device const float*)(arena + x_off);
device const float* res = (device const float*)(arena + res_off);
device const float* gamma = (device const float*)(arena + g_off);
device const float* beta = (device const float*)(arena + b_off);
device float* out = (device float*)(arena + out_off);
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;
// GQA/MQA: K/V head count (equals `heads` for MHA). Packed here so we
// don't need a 7th set_bytes slot past the Metal arg-table quirk.
uint kv_heads;
uint _pad;
};
// 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 kv_heads, uint seq_k, uint head_dim, uint k_stride, uint bhsd
) {
// Map query head → shared KV head (GQA/MQA). MHA: kv_heads == heads.
uint nkv = (kv_heads == 0u) ? heads : kv_heads;
uint group = heads / nkv;
uint hi_kv = (group > 1u) ? (hi / group) : hi;
if (bhsd != 0u) {
return bi * nkv * seq_k * head_dim + hi_kv * seq_k * head_dim + ki * head_dim;
}
uint hs = nkv * head_dim;
return bi * k_stride * hs + ki * hs + hi_kv * 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, byte_offs.kv_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, byte_offs.kv_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, byte_offs.kv_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, byte_offs.kv_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;
}
}
// Decode-step SDPA fast path (Lq == 1, Lk arbitrary).
//
// One threadgroup per (batch, head); threads split the K axis and merge
// online-softmax states. The old 1-thread-per-head launch left Apple GPUs
// nearly idle on GQA decode (Zonos: batch=2 × 16 heads = 32 threads).
kernel void sdpa_decode_m1(
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& heads [[buffer(6)]],
constant uint& head_dim [[buffer(7)]],
constant uint& q_stride [[buffer(8)]],
constant uint& mask_kind [[buffer(9)]],
constant uint& seq_k [[buffer(10)]],
constant uint& k_stride [[buffer(11)]],
constant uint& bhsd [[buffer(12)]],
constant uint& window [[buffer(13)]],
constant float& score_scale [[buffer(14)]],
constant float& attn_softcap [[buffer(15)]],
constant SdpaOffsets& byte_offs [[buffer(16)]],
uint tid [[thread_position_in_threadgroup]],
uint tgid [[threadgroup_position_in_grid]],
uint tsize [[threads_per_threadgroup]]
) {
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);
// TG=32 × dh≤128 → 16 KiB for partial O; fits Apple7/8 32 KiB tg limit.
constexpr uint MAX_DH = 128u;
constexpr uint TG = 32u;
threadgroup float tg_m[TG];
threadgroup float tg_l[TG];
threadgroup float tg_o[TG * MAX_DH];
uint bi = tgid / heads;
uint hi = tgid % heads;
if (bi >= batch) return;
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;
uint q_offset = seq_k - 1u;
uint q_base = qkv_q_offset(bi, hi, 0u, heads, 1u, head_dim, q_stride, bhsd);
// Large head_dim: one thread walks all of K (rare for LLM decode; Gemma-4).
if (head_dim > MAX_DH || tsize < 2u) {
if (tid != 0u) return;
float q_reg[512];
for (uint d = 0; d < head_dim; ++d) q_reg[d] = Q[q_base + d];
float m_acc = -1e30f;
float l_acc = 0.0f;
float o_acc[512];
for (uint d = 0; d < head_dim; ++d) o_acc[d] = 0.0f;
for (uint ki = 0; ki < seq_k; ++ki) {
uint k_base = qkv_kv_offset(bi, hi, ki, heads, byte_offs.kv_heads, seq_k, head_dim, k_stride, bhsd);
float dot = 0.0f;
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) s = -1e9f;
} else if (mask_kind == 2u) {
if (M[bi * k_stride + ki] < 0.5f) s = -1e9f;
} else if (mask_kind == 4u) {
uint lo = q_offset > window ? q_offset - window : 0u;
if (ki < lo || ki > q_offset) s = -1e9f;
}
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, byte_offs.kv_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;
}
float inv_l = 1.0f / l_acc;
uint o_base = qkv_q_offset(bi, hi, 0u, heads, 1u, head_dim, q_stride, bhsd);
for (uint d = 0; d < head_dim; ++d) OUT[o_base + d] = o_acc[d] * inv_l;
return;
}
float q_reg[MAX_DH];
for (uint d = tid; d < head_dim; d += tsize) tg_o[d] = Q[q_base + d];
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint d = 0; d < head_dim; ++d) q_reg[d] = tg_o[d];
threadgroup_barrier(mem_flags::mem_threadgroup);
float m_acc = -1e30f;
float l_acc = 0.0f;
float o_acc[MAX_DH];
for (uint d = 0; d < head_dim; ++d) o_acc[d] = 0.0f;
for (uint ki = tid; ki < seq_k; ki += tsize) {
uint k_base = qkv_kv_offset(bi, hi, ki, heads, byte_offs.kv_heads, seq_k, head_dim, k_stride, bhsd);
float dot = 0.0f;
for (uint d = 0; d < head_dim; d += 4u) {
if (d + 3u < head_dim) {
float4 qv = float4(q_reg[d], q_reg[d + 1u], q_reg[d + 2u], q_reg[d + 3u]);
float4 kv = *(device const float4*)(K + k_base + d);
dot += qv.x * kv.x + qv.y * kv.y + qv.z * kv.z + qv.w * kv.w;
} else {
for (uint dd = d; dd < head_dim; ++dd) dot += q_reg[dd] * K[k_base + dd];
}
}
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) s = -1e9f;
} else if (mask_kind == 2u) {
if (M[bi * k_stride + ki] < 0.5f) s = -1e9f;
} else if (mask_kind == 4u) {
uint lo = q_offset > window ? q_offset - window : 0u;
if (ki < lo || ki > q_offset) s = -1e9f;
}
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, byte_offs.kv_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;
}
// Idle lanes beyond TG keep identity merge state (already set if tid>=TG).
if (tid < TG) {
tg_m[tid] = m_acc;
tg_l[tid] = l_acc;
for (uint d = 0; d < head_dim; ++d) tg_o[tid * MAX_DH + d] = o_acc[d];
}
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint stride = TG / 2u; stride > 0u; stride /= 2u) {
if (tid < stride) {
float m1 = tg_m[tid];
float m2 = tg_m[tid + stride];
float m_new = max(m1, m2);
float e1 = precise::exp(m1 - m_new);
float e2 = precise::exp(m2 - m_new);
tg_l[tid] = e1 * tg_l[tid] + e2 * tg_l[tid + stride];
tg_m[tid] = m_new;
for (uint d = 0; d < head_dim; ++d) {
tg_o[tid * MAX_DH + d] =
e1 * tg_o[tid * MAX_DH + d] + e2 * tg_o[(tid + stride) * MAX_DH + d];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0u) {
float inv_l = 1.0f / tg_l[0];
uint o_base = qkv_q_offset(bi, hi, 0u, heads, 1u, head_dim, q_stride, bhsd);
for (uint d = 0; d < head_dim; ++d) OUT[o_base + d] = tg_o[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, 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)]],
constant uint& cos_per_token [[buffer(11)]],
constant uint& interleaved [[buffer(12)]],
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;
// RoPE table row: per-seq-position by default; per global (batch·seq)
// token for ragged batched decode, where each sequence sits at its own
// absolute position.
uint cos_row = (cos_per_token != 0u) ? bs : si;
// 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 (interleaved != 0u) {
// GPT-J / llama.cpp-NORM: rotated pairs are adjacent (2d, 2d+1);
// cos/sin indexed by freq d. GGUF Llama weights need this flavor.
if (d < rot_half) {
uint a = 2u * d;
uint b = 2u * d + 1u;
float x1 = x[src_base + a];
float x2 = x[src_base + b];
float c = cos[cos_row * half_dh + d];
float s = sin[cos_row * half_dh + d];
out[dst_base + a] = x1 * c - x2 * s;
out[dst_base + b] = x2 * c + x1 * s;
} else if (d >= n_rot) {
out[dst_base + d] = x[src_base + d];
}
} else if (d < rot_half) {
float x1 = x[src_base + d];
float x2 = x[src_base + rot_half + d];
float c = cos[cos_row * half_dh + d];
float s = sin[cos_row * 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];
}
}
// ArgMax / ArgMin along the middle axis of [outer, reduced, inner], emitting
// the winning index (f32). One thread per (outer, inner) output element. Strict
// comparison with first-best tie-break — matches rlx-cpu execute_argreduce_f32.
kernel void argreduce(
device const float* src [[buffer(0)]],
device float* out [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& reduced [[buffer(3)]],
constant uint& inner [[buffer(4)]],
constant uint& is_max [[buffer(5)]],
uint gid [[thread_position_in_grid]]
) {
uint total = outer * inner;
if (gid >= total) return;
uint o = gid / inner;
uint i = gid % inner;
uint base = o * reduced * inner + i;
float best = src[base];
uint best_idx = 0u;
for (uint r = 1u; r < reduced; ++r) {
float v = src[base + r * inner];
bool better = (is_max != 0u) ? (v > best) : (v < best);
if (better) { best = v; best_idx = r; }
}
out[o * inner + i] = float(best_idx);
}
// Cooperative last-axis ArgMax/ArgMin: one threadgroup reduces one `outer`
// row over the `reduced` axis (inner == 1 — the decode logits case, where the
// naive one-thread-per-output `argreduce` would loop a 128k-vocab row on a
// single GPU lane). Tie-break = lowest index wins, matching the strict `>`/`<`
// in rlx-cpu execute_argreduce_f32. Threadgroup size must be a power of two.
kernel void argreduce_lastaxis(
device const float* src [[buffer(0)]],
device float* out [[buffer(1)]],
constant uint& outer [[buffer(2)]],
constant uint& reduced [[buffer(3)]],
constant uint& is_max [[buffer(4)]],
uint tg [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint nthreads [[threads_per_threadgroup]]
) {
if (tg >= outer) return;
device const float* row = src + (ulong)tg * (ulong)reduced;
threadgroup float sval[256];
threadgroup uint sidx[256];
float best = (is_max != 0u) ? -INFINITY : INFINITY;
uint bidx = 0u;
// Strided scan: within a lane, strict comparison keeps the lowest index.
for (uint r = tid; r < reduced; r += nthreads) {
float v = row[r];
bool better = (is_max != 0u) ? (v > best) : (v < best);
if (better) { best = v; bidx = r; }
}
sval[tid] = best;
sidx[tid] = bidx;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) {
float v = sval[tid + s];
uint i = sidx[tid + s];
float cv = sval[tid];
uint ci = sidx[tid];
bool better = (is_max != 0u) ? (v > cv) : (v < cv);
// Equal value → keep the lower source index (CPU first-best).
if (better || (v == cv && i < ci)) { sval[tid] = v; sidx[tid] = i; }
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0u) out[tg] = float(sidx[0]);
}
// Fused vector-quantization: one threadgroup per input row cooperatively finds
// the nearest codebook entry over the reduced K axis — no [N,K] materialization,
// no D2H copy (the host-callback custom-op path was ~100× slower on Metal).
// metric 0 = L2 (argmin ‖x−C‖²), 1 = cosine (argmax x·C/‖C‖). Output is the
// f32-encoded code index (ready for gather). Lowest index wins ties.
kernel void vq_assign(
device const float* x [[buffer(0)]], // [N, D]
device const float* cb [[buffer(1)]], // [K, D]
device float* out [[buffer(2)]], // [N]
constant uint& n [[buffer(3)]],
constant uint& d [[buffer(4)]],
constant uint& k [[buffer(5)]],
constant uint& metric [[buffer(6)]],
uint tg [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint nthreads [[threads_per_threadgroup]]
) {
if (tg >= n) return;
device const float* xi = x + (ulong)tg * (ulong)d;
threadgroup float sval[256];
threadgroup uint sidx[256];
// float4 fast path when D is a multiple of 4 (the common case); the tail
// loop handles any remainder. Vectorizing the dot closes most of the gap to
// MPS's matrix units.
const uint d4 = d >> 2;
const uint drem = d & 3u;
device const float4* xi4 = (device const float4*)xi;
float best = (metric == 0u) ? INFINITY : -INFINITY;
uint bidx = 0u;
for (uint j = tid; j < k; j += nthreads) {
device const float* cj = cb + (ulong)j * (ulong)d;
device const float4* cj4 = (device const float4*)cj;
if (metric == 0u) {
float4 acc = 0.0f;
for (uint t = 0u; t < d4; ++t) { float4 df = xi4[t] - cj4[t]; acc += df * df; }
float dist = acc.x + acc.y + acc.z + acc.w;
for (uint t = d - drem; t < d; ++t) { float df = xi[t] - cj[t]; dist += df * df; }
if (dist < best) { best = dist; bidx = j; }
} else {
float4 dacc = 0.0f, nacc = 0.0f;
for (uint t = 0u; t < d4; ++t) { dacc += xi4[t] * cj4[t]; nacc += cj4[t] * cj4[t]; }
float dot = dacc.x + dacc.y + dacc.z + dacc.w;
float nc = nacc.x + nacc.y + nacc.z + nacc.w;
for (uint t = d - drem; t < d; ++t) { dot += xi[t] * cj[t]; nc += cj[t] * cj[t]; }
float sim = (nc > 0.0f) ? dot * rsqrt(nc) : 0.0f;
if (sim > best) { best = sim; bidx = j; }
}
}
sval[tid] = best;
sidx[tid] = bidx;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) {
bool take = (metric == 0u) ? (sval[tid + s] < sval[tid]) : (sval[tid + s] > sval[tid]);
if (take || (sval[tid + s] == sval[tid] && sidx[tid + s] < sidx[tid])) {
sval[tid] = sval[tid + s];
sidx[tid] = sidx[tid + s];
}
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0u) out[tg] = float(sidx[0]);
}
// On-GPU logit sampling: temperature -> top-k -> softmax -> top-p -> Philox
// inverse-CDF draw. One threadgroup per batch row. Mirrors the CPU algorithm
// in rlx-cpu `sample_row` / `execute_sample_f32` exactly, including the
// Philox4x32 stream (rlx_ir::rng), so a fixed seed is bit-comparable.
//
// top-k / top-p cutoffs are found by parallel bisection on the value/prob
// range rather than a full sort: for distinct logits this selects the
// identical kept set (the cutoff lands strictly between the two order
// statistics that bracket it). The final inverse-CDF walk (thread 0) recomputes
// each token's filtered probability in original index order so the sequential
// float accumulation matches the CPU reference element-for-element.
//
// Threadgroup size must be a power of two (dispatched at 256).
kernel void sample_logits(
device float* arena [[buffer(0)]],
constant ulong& logits_off [[buffer(1)]],
constant ulong& dst_off [[buffer(2)]],
constant uint& batch [[buffer(3)]],
constant uint& vocab [[buffer(4)]],
constant uint& top_k [[buffer(5)]],
constant float& top_p [[buffer(6)]],
constant float& temperature [[buffer(7)]],
constant ulong& seed [[buffer(8)]],
uint tg [[threadgroup_position_in_grid]],
uint tid [[thread_position_in_threadgroup]],
uint nthreads [[threads_per_threadgroup]]
) {
if (tg >= batch) return;
device const float* logits =
(device const float*)((device char*)arena + logits_off);
device float* dst = (device float*)((device char*)arena + dst_off);
device const float* row = logits + (ulong)tg * (ulong)vocab;
const uint v = vocab;
const float MIN_POS = 1.1754944e-38f; // f32::MIN_POSITIVE
const float temp = max(temperature, 1e-6f);
const uint kk = min(top_k, v);
const bool use_topk = (kk > 0u) && (kk < v);
const bool use_topp = (top_p < 1.0f);
if (v == 0u) { if (tid == 0u) dst[tg] = 0.0f; return; }
// `red` is the reduction scratch; `bounds[0..1]` carries the bisection
// lo/hi (an array, so per-element-init warnings don't fire). All cross-lane
// values are read back from `red[0]` after the reduction barrier.
threadgroup float red[256];
threadgroup float bounds[2];
// ── max(scaled) ────────────────────────────────────────────────
float lmax = -INFINITY;
for (uint i = tid; i < v; i += nthreads) lmax = max(lmax, row[i]);
red[tid] = lmax;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) red[tid] = max(red[tid], red[tid + s]);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const float max_l = red[0] / temp;
threadgroup_barrier(mem_flags::mem_threadgroup);
// ── min(scaled) (bisection lower bound) ────────────────────────
float lmin = INFINITY;
for (uint i = tid; i < v; i += nthreads) lmin = min(lmin, row[i]);
red[tid] = lmin;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) red[tid] = min(red[tid], red[tid + s]);
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const float min_l = red[0] / temp;
threadgroup_barrier(mem_flags::mem_threadgroup);
// ── top-k cutoff = kk-th largest scaled value (bisection) ──────
float cutoff = -INFINITY;
if (use_topk) {
if (tid == 0u) { bounds[0] = min_l; bounds[1] = max_l; }
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint it = 0u; it < 50u; ++it) {
float mid = 0.5f * (bounds[0] + bounds[1]);
float cnt = 0.0f;
for (uint i = tid; i < v; i += nthreads)
if (row[i] / temp >= mid) cnt += 1.0f;
red[tid] = cnt;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) red[tid] += red[tid + s];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0u) {
if (red[0] >= float(kk)) bounds[0] = mid; else bounds[1] = mid;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
cutoff = bounds[0];
}
// ── softmax denom over the top-k set ───────────────────────────
float s1 = 0.0f;
for (uint i = tid; i < v; i += nthreads) {
float sc = row[i] / temp;
if (!use_topk || sc >= cutoff) s1 += exp(sc - max_l);
}
red[tid] = s1;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) red[tid] += red[tid + s];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
const float sum1 = red[0];
const float inv1 = 1.0f / max(sum1, MIN_POS);
threadgroup_barrier(mem_flags::mem_threadgroup);
// ── top-p prob cutoff (bisection over [0,1]) ───────────────────
float pcut = 0.0f;
if (use_topp) {
if (tid == 0u) { bounds[0] = 0.0f; bounds[1] = 1.0f; }
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint it = 0u; it < 60u; ++it) {
float mid = 0.5f * (bounds[0] + bounds[1]);
float psum = 0.0f;
for (uint i = tid; i < v; i += nthreads) {
float sc = row[i] / temp;
if (use_topk && sc < cutoff) continue;
float p = exp(sc - max_l) * inv1;
if (p >= mid) psum += p;
}
red[tid] = psum;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) red[tid] += red[tid + s];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
if (tid == 0u) {
if (red[0] >= top_p) bounds[0] = mid; else bounds[1] = mid;
}
threadgroup_barrier(mem_flags::mem_threadgroup);
}
pcut = bounds[0];
}
// ── renorm denom over the top-p set ────────────────────────────
float sum2 = 1.0f;
if (use_topp) {
float s2 = 0.0f;
for (uint i = tid; i < v; i += nthreads) {
float sc = row[i] / temp;
if (use_topk && sc < cutoff) continue;
float p = exp(sc - max_l) * inv1;
if (p >= pcut) s2 += p;
}
red[tid] = s2;
threadgroup_barrier(mem_flags::mem_threadgroup);
for (uint s = nthreads >> 1; s > 0u; s >>= 1) {
if (tid < s) red[tid] += red[tid + s];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
sum2 = red[0];
threadgroup_barrier(mem_flags::mem_threadgroup);
}
// ── thread 0: Philox draw + sequential inverse-CDF in index order ─
if (tid == 0u) {
ulong sd = (seed == 0ul) ? 0xDEADBEEFul : seed;
uint k0 = (uint)(sd & 0xFFFFFFFFul);
uint k1 = (uint)(sd >> 32);
uint st0 = (uint)(tg / 4u), st1 = 0u, st2 = 0u, st3 = 0u;
for (uint rr = 0u; rr < 10u; ++rr) {
ulong p0 = (ulong)st0 * 0xD2561A75ul;
ulong p1 = (ulong)st2 * 0xCD9E8D57ul;
uint hi0 = (uint)(p0 >> 32), lo0 = (uint)(p0 & 0xFFFFFFFFul);
uint hi1 = (uint)(p1 >> 32), lo1 = (uint)(p1 & 0xFFFFFFFFul);
uint n0 = hi1 ^ st1 ^ k0;
uint n1 = lo1;
uint n2 = hi0 ^ st3 ^ k1;
uint n3 = lo0;
st0 = n0; st1 = n1; st2 = n2; st3 = n3;
k0 += 0x9E3779B9u; k1 += 0xBB67AE85u;
}
uint lane = tg % 4u;
uint bits = (lane == 0u ? st0 : (lane == 1u ? st1 : (lane == 2u ? st2 : st3))) >> 8;
float r = (float)bits / 16777216.0f;
float inv2 = use_topp ? (1.0f / max(sum2, MIN_POS)) : 1.0f;
float acc = 0.0f;
uint chosen = v - 1u;
for (uint i = 0u; i < v; ++i) {
float sc = row[i] / temp;
float p;
if (use_topk && sc < cutoff) {
p = 0.0f;
} else {
p = exp(sc - max_l) * inv1;
if (use_topp) { p = (p >= pcut) ? (p * inv2) : 0.0f; }
}
acc += p;
if (r <= acc) { chosen = i; break; }
}
dst[tg] = float(chosen);
}
}
// Block-quantized int8 weight matmul: out[m,n] = x[m,k] @ dequant(wq[k,n]).
// Per-(block-of-k, n) scale (+ optional zero-point). One thread per output
// element. Matches rlx-cpu dequant_matmul_int8.
kernel void dequant_matmul_int8(
device const float* x [[buffer(0)]],
device const char* wq [[buffer(1)]],
device const float* scales [[buffer(2)]],
device const float* zps [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& m [[buffer(5)]],
constant uint& k [[buffer(6)]],
constant uint& n [[buffer(7)]],
constant uint& block_size [[buffer(8)]],
constant uint& asym [[buffer(9)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= m * n) return;
uint i = gid / n;
uint j = gid % n;
float acc = 0.0f;
for (uint p = 0; p < k; ++p) {
uint block = p / block_size;
float s = scales[block * n + j];
float z = (asym != 0u) ? zps[block * n + j] : 0.0f;
float q = float(wq[p * n + j]);
acc += x[i * k + p] * ((q - z) * s);
}
out[i * n + j] = acc;
}
// Block-quantized int4 (two nibbles per byte) weight matmul. Low nibble first.
kernel void dequant_matmul_int4(
device const float* x [[buffer(0)]],
device const uchar* wq [[buffer(1)]],
device const float* scales [[buffer(2)]],
device const float* zps [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& m [[buffer(5)]],
constant uint& k [[buffer(6)]],
constant uint& n [[buffer(7)]],
constant uint& block_size [[buffer(8)]],
constant uint& asym [[buffer(9)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= m * n) return;
uint i = gid / n;
uint j = gid % n;
float acc = 0.0f;
for (uint p = 0; p < k; ++p) {
uint block = p / block_size;
float s = scales[block * n + j];
float z = (asym != 0u) ? zps[block * n + j] : 0.0f;
uint idx = p * n + j;
uchar byte = wq[idx >> 1];
uint nib = ((idx & 1u) == 0u) ? (uint(byte) & 0x0Fu) : (uint(byte) >> 4);
acc += x[i * k + p] * ((float(nib) - z) * s);
}
out[i * n + j] = acc;
}
inline float dq_fp8_e4m3(uchar byte) {
uint sign = (uint(byte) >> 7) & 1u;
uint exp_v = (uint(byte) >> 3) & 0x0Fu;
uint mant = uint(byte) & 0x7u;
float v;
if (exp_v == 0u) {
v = (mant == 0u) ? 0.0f : (float(mant) / 8.0f) * exp2(-6.0f);
} else if (exp_v == 0x0Fu && mant == 0x7u) {
v = 0.0f;
} else {
v = (1.0f + float(mant) / 8.0f) * exp2(float(int(exp_v) - 7));
}
return (sign != 0u) ? -v : v;
}
inline float dq_fp8_e5m2(uchar byte) {
uint sign = (uint(byte) >> 7) & 1u;
uint exp_v = (uint(byte) >> 2) & 0x1Fu;
uint mant = uint(byte) & 0x3u;
float v;
if (exp_v == 0u) {
v = (mant == 0u) ? 0.0f : (float(mant) / 4.0f) * exp2(-14.0f);
} else if (exp_v == 0x1Fu) {
v = 0.0f;
} else {
v = (1.0f + float(mant) / 4.0f) * exp2(float(int(exp_v) - 15));
}
return (sign != 0u) ? -v : v;
}
constant float DQ_FP4_E2M1[16] = {
0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f,
-0.0f, -0.5f, -1.0f, -1.5f, -2.0f, -3.0f, -4.0f, -6.0f
};
kernel void dequant_matmul_fp8(
device const float* x [[buffer(0)]],
device const uchar* wq [[buffer(1)]],
device const float* scales [[buffer(2)]],
device float* out [[buffer(4)]],
constant uint& m [[buffer(5)]],
constant uint& k [[buffer(6)]],
constant uint& n [[buffer(7)]],
constant uint& e5m2 [[buffer(8)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= m * n) return;
uint i = gid / n;
uint j = gid % n;
float acc = 0.0f;
float col_scale = scales[j];
for (uint p = 0u; p < k; ++p) {
uchar byte = wq[p * n + j];
float w = (e5m2 != 0u) ? dq_fp8_e5m2(byte) : dq_fp8_e4m3(byte);
acc += x[i * k + p] * w * col_scale;
}
out[i * n + j] = acc;
}
kernel void dequant_matmul_nvfp4(
device const float* x [[buffer(0)]],
device const uchar* wq [[buffer(1)]],
device const uchar* scales [[buffer(2)]],
device const float* gs_ptr [[buffer(3)]],
device float* out [[buffer(4)]],
constant uint& m [[buffer(5)]],
constant uint& k [[buffer(6)]],
constant uint& n [[buffer(7)]],
constant uint& group_size [[buffer(8)]],
uint gid [[thread_position_in_grid]]
) {
if (gid >= m * n) return;
uint i = gid / n;
uint j = gid % n;
float gs = gs_ptr[0];
float acc = 0.0f;
for (uint p = 0u; p < k; ++p) {
uint idx = p * n + j;
uint byte_idx = idx >> 1;
uint nib = ((idx & 1u) == 0u) ? (uint(wq[byte_idx]) & 0x0Fu) : (uint(wq[byte_idx]) >> 4);
uint block = p / group_size;
float s = dq_fp8_e4m3(scales[block * n + j]);
acc += x[i * k + p] * DQ_FP4_E2M1[nib] * s * gs;
}
out[i * n + j] = acc;
}
// 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 = fmax(0.0f, 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];
}
}
// GDN gated norm: out = rms_norm(x) * silu(z). Same dispatch as rms_norm.
kernel void rms_norm_mul_silu(
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& z_off [[buffer(4)]],
constant ulong& out_off [[buffer(5)]],
constant uint& h [[buffer(6)]],
constant float& eps [[buffer(7)]],
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 const float* z = (device const float*)(arena + z_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) {
float zv = z[row * h + i];
float silu = zv / (1.0f + exp(-zv));
output[row * h + i] =
(input[row * h + i] * inv_rms * gamma[i] + beta[i]) * silu;
}
}
// 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 = clamp(c * (lhs + 0.044715f * lhs * lhs * lhs), -15.0f, 15.0f);
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(clamp(lhs, -15.0f, 15.0f));
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 = clamp(c * (lhs + 0.044715f * lhs * lhs * lhs), -15.0f, 15.0f);
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(clamp(lhs, -15.0f, 15.0f));
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;
}
}
// Simdgroup-cooperative GDN for n == 128 (Bonsai / Qwen3.5). llama.cpp-style:
// one simdgroup owns one state column; 32 threads × 4 floats keep that column
// in registers and use simd_sum — no threadgroup barriers inside the scan.
// Grid: (n/NSG, heads, batch) threadgroups, threads (32, NSG, 1).
constant uint GDN_SG_N = 128u;
constant uint GDN_SG_NSG = 4u;
kernel void gated_delta_net_sg(
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)]],
uint3 tgpig [[threadgroup_position_in_grid]],
uint3 tpitg [[thread_position_in_threadgroup]]
) {
const uint b = dims.x, s = dims.y, h = dims.z, n = dims.w;
if (n != GDN_SG_N) return;
const uint tx = tpitg.x; // 0..31 within simdgroup
const uint ty = tpitg.y; // simdgroup index in threadgroup
const uint bi = tgpig.z;
const uint hi = tgpig.y;
const uint col = tgpig.x * GDN_SG_NSG + ty; // state column j
if (bi >= b || hi >= h || col >= n || tx >= 32u) return;
const float scale = rsqrt(float(n));
const uint s_base = state_off + (bi * h + hi) * n * n;
device float* s_mat = arena + s_base;
// Register tile of column `col`: S[is][col] for is = tx*4 .. tx*4+3.
float ls[GDN_SG_NSG];
#pragma unroll
for (uint j = 0u; j < GDN_SG_NSG; ++j) {
const uint is = tx * GDN_SG_NSG + j;
ls[j] = (use_carry != 0u) ? s_mat[is * n + col] : 0.0f;
}
const uint hs_n = h * n;
for (uint ti = 0u; ti < s; ++ti) {
const uint qkv_step = bi * s * hs_n + ti * hs_n + hi * n;
const uint gb_step = bi * s * h + ti * h + hi;
device const float* q_ptr = arena + q_off + qkv_step;
device const float* k_ptr = arena + k_off + qkv_step;
device const float* v_ptr = arena + v_off + qkv_step;
const float g_exp = exp(arena[g_off + gb_step]);
const float beta_t = arena[beta_off + gb_step];
float s_k = 0.0f;
#pragma unroll
for (uint j = 0u; j < GDN_SG_NSG; ++j) {
const uint is = tx * GDN_SG_NSG + j;
ls[j] *= g_exp;
s_k += ls[j] * k_ptr[is];
}
s_k = simd_sum(s_k);
const float d = (v_ptr[col] - s_k) * beta_t;
float y = 0.0f;
#pragma unroll
for (uint j = 0u; j < GDN_SG_NSG; ++j) {
const uint is = tx * GDN_SG_NSG + j;
ls[j] += k_ptr[is] * d;
y += ls[j] * q_ptr[is];
}
y = simd_sum(y);
if (tx == 0u) {
arena[dst_off + qkv_step + col] = y * scale;
}
}
// Write column back (carry / next step).
#pragma unroll
for (uint j = 0u; j < GDN_SG_NSG; ++j) {
const uint is = tx * GDN_SG_NSG + j;
s_mat[is * n + col] = ls[j];
}
}
// ── Selective scan (Mamba SSM, f32) ─────────────────────────────────
// One thread per (batch, channel); each thread owns a private state
// vector of size n (n ≤ SSM_MAX_N) and scans sequentially over seq.
// Inputs (float indices): x, delta [b,s,h]; a [h,n]; b, c [b,s,n].
// Output [b,s,h]. Matches `execute_selective_scan_f32` on CPU:
// h[t] = exp(Δ·A)·h[t-1] + Δ·B·x; y[t] = Σ_n C·h[t]
#define SSM_MAX_N 128u
kernel void selective_scan(
device float* arena [[buffer(0)]],
constant uint& x_off [[buffer(1)]],
constant uint& delta_off [[buffer(2)]],
constant uint& a_off [[buffer(3)]],
constant uint& b_off [[buffer(4)]],
constant uint& c_off [[buffer(5)]],
constant uint& dst_off [[buffer(6)]],
constant uint4& dims [[buffer(7)]], // batch, seq, hidden, n
uint gid [[thread_position_in_grid]]
) {
uint b = dims.x, s = dims.y, h = dims.z, n = dims.w;
if (n > SSM_MAX_N || gid >= b * h) return;
uint bi = gid / h;
uint ci = gid % h;
float state[SSM_MAX_N];
for (uint i = 0; i < n; ++i) {
state[i] = 0.0f;
}
// a[ci, :] is constant across the sequence for this channel.
uint a_base = a_off + ci * n;
for (uint si = 0; si < s; ++si) {
uint bsh = bi * s * h + si * h + ci; // x/delta/out element offset
uint bsn = (bi * s + si) * n; // b/c row base
float d = arena[delta_off + bsh];
float xv = arena[x_off + bsh];
float acc = 0.0f;
for (uint ni = 0; ni < n; ++ni) {
float da = exp(d * arena[a_base + ni]);
float st = da * state[ni] + d * arena[b_off + bsn + ni] * xv;
state[ni] = st;
acc += arena[c_off + bsn + ni] * st;
}
arena[dst_off + bsh] = acc;
}
}
// 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(clamp(z[2], -15.0f, 15.0f));
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;
}
}
// Single-layer unidirectional GRU (gate order r, z, n; linear_before_reset=1;
// separate b_ih/b_hh; h0 = 0). One threadgroup per batch item; thread `k` owns
// hidden unit `k`. Matches `execute_gru_f32` on CPU.
#define GRU_MAX_H 1024u
kernel void gru(
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& bih_off [[buffer(4)]],
constant uint& bhh_off [[buffer(5)]],
constant uint& dst_off [[buffer(6)]],
constant uint4& dims [[buffer(7)]], // 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 > GRU_MAX_H || gid >= b || tid >= h) return;
uint bi = gid, k = tid;
threadgroup float h_sh[GRU_MAX_H];
h_sh[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 rows r=k, z=h+k, n=2h+k. Input and hidden parts kept separate
// because the reset gate multiplies the hidden term after its bias.
float xi[3], hi[3];
for (uint g = 0u; g < 3u; ++g) {
uint r = g * h + k;
float ax = arena[bih_off + r];
uint wih_row = wih_off + r * in_sz;
for (uint j = 0u; j < in_sz; ++j) {
ax += arena[wih_row + j] * arena[x_base + j];
}
float ah = arena[bhh_off + r];
uint whh_row = whh_off + r * h;
for (uint j = 0u; j < h; ++j) {
ah += arena[whh_row + j] * h_sh[j];
}
xi[g] = ax;
hi[g] = ah;
}
float rg = 1.0f / (1.0f + exp(-(xi[0] + hi[0])));
float zg = 1.0f / (1.0f + exp(-(xi[1] + hi[1])));
float ng = tanh(clamp(xi[2] + rg * hi[2], -15.0f, 15.0f));
float h_k = (1.0f - zg) * ng + zg * h_sh[k];
// Finish reading h_prev across all threads before overwriting.
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;
}
}
// Single-layer unidirectional Elman RNN (`relu_flag` ? relu : tanh; h0 = 0).
// One threadgroup per batch item; thread `k` owns hidden unit `k`. Matches
// `execute_rnn_f32` on CPU.
#define RNN_MAX_H 1024u
kernel void rnn(
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
constant uint& relu_flag [[buffer(7)]],
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 > RNN_MAX_H || gid >= b || tid >= h) return;
uint bi = gid, k = tid;
threadgroup float h_sh[RNN_MAX_H];
h_sh[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;
float acc = arena[bias_off + k];
uint wih_row = wih_off + k * in_sz;
for (uint j = 0u; j < in_sz; ++j) {
acc += arena[wih_row + j] * arena[x_base + j];
}
uint whh_row = whh_off + k * h;
for (uint j = 0u; j < h; ++j) {
acc += arena[whh_row + j] * h_sh[j];
}
float h_k = relu_flag != 0u ? fmax(acc, 0.0f) : tanh(acc);
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;
}
}
// Mamba-2 / SSD scalar-decay scan. One thread per (batch, head, head_dim_pos);
// each owns a private N-state vector and scans the sequence. Matches
// `execute_mamba2_f32` on CPU. n ≤ MAMBA2_MAX_N.
#define MAMBA2_MAX_N 128u
kernel void mamba2(
device float* arena [[buffer(0)]],
constant uint& x_off [[buffer(1)]],
constant uint& dt_off [[buffer(2)]],
constant uint& a_off [[buffer(3)]],
constant uint& b_off [[buffer(4)]],
constant uint& c_off [[buffer(5)]],
constant uint& dst_off [[buffer(6)]],
constant uint4& dims [[buffer(7)]], // batch, seq, heads, (head_dim<<16 | state_size)
uint gid [[thread_position_in_grid]]
) {
uint bn = dims.x, s = dims.y, hh = dims.z;
uint p = dims.w >> 16, n = dims.w & 0xffffu;
if (n > MAMBA2_MAX_N || gid >= bn * hh * p) return;
uint pi = gid % p;
uint hi = (gid / p) % hh;
uint bi = gid / (p * hh);
float state[MAMBA2_MAX_N];
for (uint i = 0u; i < n; ++i) {
state[i] = 0.0f;
}
float ah = arena[a_off + hi];
for (uint t = 0u; t < s; ++t) {
uint bsh = (bi * s + t) * hh + hi;
float dt_t = arena[dt_off + bsh];
float da = exp(dt_t * ah);
float dtx = dt_t * arena[x_off + bsh * p + pi];
uint bc = bsh * n;
float acc = 0.0f;
for (uint ni = 0u; ni < n; ++ni) {
float st = da * state[ni] + dtx * arena[b_off + bc + ni];
state[ni] = st;
acc += st * arena[c_off + bc + ni];
}
arena[dst_off + bsh * p + pi] = acc;
}
}
// 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);
// Cross term is inv_r³ (= inv_r2·inv_r below), NOT inv_r⁴: outer ·inv_r already supplies
// one factor, so the inner term uses inv_r2. The prior inv_r3-then-·inv_r was a stray 1/r.
float inv_r2 = 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_r2;
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_f16w: ComputePipelineState,
pub sgemm_f16w_small_m: 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 relu_inplace_h: ComputePipelineState,
pub sigmoid_inplace_h: ComputePipelineState,
pub tanh_inplace_h: ComputePipelineState,
pub exp_inplace_h: ComputePipelineState,
pub log_inplace_h: ComputePipelineState,
pub sqrt_inplace_h: ComputePipelineState,
pub rsqrt_inplace_h: ComputePipelineState,
pub neg_inplace_h: ComputePipelineState,
pub abs_inplace_h: ComputePipelineState,
pub sin_inplace_h: ComputePipelineState,
pub cos_inplace_h: ComputePipelineState,
pub tan_inplace_h: ComputePipelineState,
pub atan_inplace_h: ComputePipelineState,
pub round_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 elem_sub_h: ComputePipelineState,
pub elem_div_h: ComputePipelineState,
pub elem_max_h: ComputePipelineState,
pub elem_min_h: ComputePipelineState,
pub elem_pow_h: ComputePipelineState,
pub gather_axis0_h: ComputePipelineState,
pub narrow_lastax_h: ComputePipelineState,
pub sdpa_h: ComputePipelineState,
pub fused_attn_block: 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_f16w: 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 silu_out4: 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 rms_norm_mul_silu: 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 ada_layer_norm: ComputePipelineState,
pub ada_layer_norm_h: ComputePipelineState,
pub gated_residual: ComputePipelineState,
pub gated_residual_h: ComputePipelineState,
pub ada_layer_norm_backward: ComputePipelineState,
pub ada_layer_norm_backward_h: ComputePipelineState,
pub gated_residual_backward: ComputePipelineState,
pub gated_residual_backward_h: ComputePipelineState,
pub sdpa: ComputePipelineState,
pub sdpa_long: ComputePipelineState,
pub sdpa_decode_m1: ComputePipelineState,
pub sdpa_fa_f32: ComputePipelineState,
pub argreduce: ComputePipelineState,
pub argreduce_lastaxis: ComputePipelineState,
pub vq_assign: ComputePipelineState,
pub sample_logits: ComputePipelineState,
pub dequant_matmul_int8: ComputePipelineState,
pub dequant_matmul_int4: ComputePipelineState,
pub dequant_matmul_fp8: ComputePipelineState,
pub dequant_matmul_nvfp4: 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 concat_midaxis_seg: ComputePipelineState,
pub concat_midaxis_seg_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_compare_bcast: ComputePipelineState,
pub elem_where: ComputePipelineState,
pub elem_where_bcast: ComputePipelineState,
pub elem_fma: 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 maxpool2d_backward: ComputePipelineState,
pub conv2d_backward_input: ComputePipelineState,
pub conv2d_backward_weight: ComputePipelineState,
pub conv2d_backward_weight_partial: ComputePipelineState,
pub conv2d_backward_weight_reduce: ComputePipelineState,
pub conv2d: ComputePipelineState,
pub depthwise_conv1d_bsc: 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 softmax_cross_entropy_dense: ComputePipelineState,
pub softmax_cross_entropy_with_logits: ComputePipelineState,
pub softmax_cross_entropy_backward: ComputePipelineState,
pub fft_radix2_full_f32: ComputePipelineState,
#[cfg(feature = "native-gpu-fft")]
pub fft_radix2_full_big_f32: ComputePipelineState,
#[cfg(feature = "native-gpu-fft")]
pub fft_radix4_full_f32: ComputePipelineState,
#[cfg(feature = "native-gpu-fft")]
pub fft_radix8_full_f32: ComputePipelineState,
#[cfg(feature = "native-gpu-fft")]
pub fft_radix16_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 gated_delta_net_sg: ComputePipelineState,
pub selective_scan: ComputePipelineState,
pub lstm: ComputePipelineState,
pub gru: ComputePipelineState,
pub rnn: ComputePipelineState,
pub mamba2: ComputePipelineState,
pub dequant_gguf: ComputePipelineState,
pub q4k_mv_f32: ComputePipelineState,
pub q4_0_mv_f32: ComputePipelineState,
pub q4_1_mv_f32: ComputePipelineState,
pub q8_0_mv_f32: ComputePipelineState,
pub q1_0_mv_f32: ComputePipelineState,
pub q1_0_mv_f32_sg: ComputePipelineState,
pub q1_0_dual_mv_f32_sg: ComputePipelineState,
pub q1_0_mm_f32: ComputePipelineState,
pub q2_0_mv_f32: ComputePipelineState,
pub q2_0_mv_f32_sg: ComputePipelineState,
pub q2_0_dual_mv_f32_sg: ComputePipelineState,
pub q2_0_mm_f32: ComputePipelineState,
pub iq4_nl_mv_f32: ComputePipelineState,
pub iq2_xxs_mv_f32: ComputePipelineState,
pub iq2_xs_mv_f32: ComputePipelineState,
pub iq2_s_mv_f32: ComputePipelineState,
pub iq3_xxs_mv_f32: ComputePipelineState,
pub iq3_s_mv_f32: ComputePipelineState,
pub iq1_s_mv_f32: ComputePipelineState,
pub iq1_m_mv_f32: ComputePipelineState,
pub q4k_mv_f32_sg: ComputePipelineState,
pub q4k_mm_f32: ComputePipelineState,
pub q6k_mm_f32: ComputePipelineState,
pub q4k_swiglu_mv_f32: ComputePipelineState,
pub q4k_gelu_mv_f32: ComputePipelineState,
pub q5_0_swiglu_mv_f32: ComputePipelineState,
pub q5_0_gelu_mv_f32: ComputePipelineState,
pub q4k_mv_residual_f32: ComputePipelineState,
pub q6k_mv_residual_f32: ComputePipelineState,
pub q5_0_mv_residual_f32: ComputePipelineState,
pub q1_0_swiglu_mv_f32: ComputePipelineState,
pub q1_0_swiglu_mv_f32_sg: ComputePipelineState,
pub q1_0_mv_residual_f32: ComputePipelineState,
pub q1_0_mv_residual_f32_sg: ComputePipelineState,
pub q2_0_swiglu_mv_f32: ComputePipelineState,
pub q2_0_swiglu_mv_f32_sg: ComputePipelineState,
pub q2_0_mv_residual_f32: ComputePipelineState,
pub q2_0_mv_residual_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_f16w: pipeline("sgemm_f16w"),
sgemm_f16w_small_m: pipeline("sgemm_f16w_small_m"),
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"),
relu_inplace_h: pipeline("relu_inplace_h"),
sigmoid_inplace_h: pipeline("sigmoid_inplace_h"),
tanh_inplace_h: pipeline("tanh_inplace_h"),
exp_inplace_h: pipeline("exp_inplace_h"),
log_inplace_h: pipeline("log_inplace_h"),
sqrt_inplace_h: pipeline("sqrt_inplace_h"),
rsqrt_inplace_h: pipeline("rsqrt_inplace_h"),
neg_inplace_h: pipeline("neg_inplace_h"),
abs_inplace_h: pipeline("abs_inplace_h"),
sin_inplace_h: pipeline("sin_inplace_h"),
cos_inplace_h: pipeline("cos_inplace_h"),
tan_inplace_h: pipeline("tan_inplace_h"),
atan_inplace_h: pipeline("atan_inplace_h"),
round_inplace_h: pipeline("round_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"),
elem_sub_h: pipeline("elem_sub_h"),
elem_div_h: pipeline("elem_div_h"),
elem_max_h: pipeline("elem_max_h"),
elem_min_h: pipeline("elem_min_h"),
elem_pow_h: pipeline("elem_pow_h"),
gather_axis0_h: pipeline("gather_axis0_h"),
narrow_lastax_h: pipeline("narrow_lastax_h"),
sdpa_h: pipeline("sdpa_h"),
fused_attn_block: pipeline("fused_attn_block"),
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_f16w: pipeline("sgemm_simd_padded_f16w"),
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"),
silu_out4: pipeline("silu_out4"),
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"),
rms_norm_mul_silu: pipeline("rms_norm_mul_silu"),
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"),
ada_layer_norm: pipeline("ada_layer_norm"),
ada_layer_norm_h: pipeline("ada_layer_norm_h"),
gated_residual: pipeline("gated_residual"),
gated_residual_h: pipeline("gated_residual_h"),
ada_layer_norm_backward: pipeline("ada_layer_norm_backward"),
ada_layer_norm_backward_h: pipeline("ada_layer_norm_backward_h"),
gated_residual_backward: pipeline("gated_residual_backward"),
gated_residual_backward_h: pipeline("gated_residual_backward_h"),
sdpa: pipeline("sdpa"),
sdpa_long: pipeline("sdpa_long"),
sdpa_decode_m1: pipeline("sdpa_decode_m1"),
sdpa_fa_f32: pipeline("sdpa_fa_f32"),
argreduce: pipeline("argreduce"),
argreduce_lastaxis: pipeline("argreduce_lastaxis"),
vq_assign: pipeline("vq_assign"),
sample_logits: pipeline("sample_logits"),
dequant_matmul_int8: pipeline("dequant_matmul_int8"),
dequant_matmul_int4: pipeline("dequant_matmul_int4"),
dequant_matmul_fp8: pipeline("dequant_matmul_fp8"),
dequant_matmul_nvfp4: pipeline("dequant_matmul_nvfp4"),
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"),
concat_midaxis_seg: pipeline("concat_midaxis_seg"),
concat_midaxis_seg_h: pipeline("concat_midaxis_seg_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_compare_bcast: pipeline("elem_compare_bcast"),
elem_where: pipeline("elem_where"),
elem_where_bcast: pipeline("elem_where_bcast"),
elem_fma: pipeline("elem_fma"),
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"),
maxpool2d_backward: pipeline("maxpool2d_backward"),
conv2d_backward_input: pipeline("conv2d_backward_input"),
conv2d_backward_weight: pipeline("conv2d_backward_weight"),
conv2d_backward_weight_partial: pipeline("conv2d_backward_weight_partial"),
conv2d_backward_weight_reduce: pipeline("conv2d_backward_weight_reduce"),
conv2d: pipeline("conv2d"),
depthwise_conv1d_bsc: pipeline("depthwise_conv1d_bsc"),
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"),
softmax_cross_entropy_dense: pipeline("softmax_cross_entropy_dense"),
softmax_cross_entropy_with_logits: pipeline("softmax_cross_entropy_with_logits"),
softmax_cross_entropy_backward: pipeline("softmax_cross_entropy_backward"),
fft_radix2_full_f32: pipeline("fft_radix2_full_f32"),
#[cfg(feature = "native-gpu-fft")]
fft_radix2_full_big_f32: pipeline("fft_radix2_full_big_f32"),
#[cfg(feature = "native-gpu-fft")]
fft_radix4_full_f32: pipeline("fft_radix4_full_f32"),
#[cfg(feature = "native-gpu-fft")]
fft_radix8_full_f32: pipeline("fft_radix8_full_f32"),
#[cfg(feature = "native-gpu-fft")]
fft_radix16_full_f32: pipeline("fft_radix16_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"),
gated_delta_net_sg: pipeline("gated_delta_net_sg"),
selective_scan: pipeline("selective_scan"),
lstm: pipeline("lstm"),
gru: pipeline("gru"),
rnn: pipeline("rnn"),
mamba2: pipeline("mamba2"),
dequant_gguf: pipeline("dequant_gguf"),
q4k_mv_f32: pipeline("q4k_mv_f32"),
q1_0_mv_f32: pipeline("q1_0_mv_f32"),
q1_0_mv_f32_sg: pipeline("q1_0_mv_f32_sg"),
q1_0_dual_mv_f32_sg: pipeline("q1_0_dual_mv_f32_sg"),
q1_0_mm_f32: pipeline("q1_0_mm_f32"),
q2_0_mv_f32: pipeline("q2_0_mv_f32"),
q2_0_mv_f32_sg: pipeline("q2_0_mv_f32_sg"),
q2_0_dual_mv_f32_sg: pipeline("q2_0_dual_mv_f32_sg"),
q2_0_mm_f32: pipeline("q2_0_mm_f32"),
q4_0_mv_f32: pipeline("q4_0_mv_f32"),
q4_1_mv_f32: pipeline("q4_1_mv_f32"),
q8_0_mv_f32: pipeline("q8_0_mv_f32"),
iq4_nl_mv_f32: pipeline("iq4_nl_mv_f32"),
iq2_xxs_mv_f32: pipeline("iq2_xxs_mv_f32"),
iq2_xs_mv_f32: pipeline("iq2_xs_mv_f32"),
iq2_s_mv_f32: pipeline("iq2_s_mv_f32"),
iq3_xxs_mv_f32: pipeline("iq3_xxs_mv_f32"),
iq3_s_mv_f32: pipeline("iq3_s_mv_f32"),
iq1_s_mv_f32: pipeline("iq1_s_mv_f32"),
iq1_m_mv_f32: pipeline("iq1_m_mv_f32"),
q4k_mv_f32_sg: pipeline("q4k_mv_f32_sg"),
q4k_mm_f32: pipeline("q4k_mm_f32"),
q6k_mm_f32: pipeline("q6k_mm_f32"),
q4k_swiglu_mv_f32: pipeline("q4k_swiglu_mv_f32"),
q4k_gelu_mv_f32: pipeline("q4k_gelu_mv_f32"),
q5_0_swiglu_mv_f32: pipeline("q5_0_swiglu_mv_f32"),
q5_0_gelu_mv_f32: pipeline("q5_0_gelu_mv_f32"),
q4k_mv_residual_f32: pipeline("q4k_mv_residual_f32"),
q6k_mv_residual_f32: pipeline("q6k_mv_residual_f32"),
q5_0_mv_residual_f32: pipeline("q5_0_mv_residual_f32"),
q1_0_swiglu_mv_f32: pipeline("q1_0_swiglu_mv_f32"),
q1_0_swiglu_mv_f32_sg: pipeline("q1_0_swiglu_mv_f32_sg"),
q1_0_mv_residual_f32: pipeline("q1_0_mv_residual_f32"),
q1_0_mv_residual_f32_sg: pipeline("q1_0_mv_residual_f32_sg"),
q2_0_swiglu_mv_f32: pipeline("q2_0_swiglu_mv_f32"),
q2_0_swiglu_mv_f32_sg: pipeline("q2_0_swiglu_mv_f32_sg"),
q2_0_mv_residual_f32: pipeline("q2_0_mv_residual_f32"),
q2_0_mv_residual_f32_sg: pipeline("q2_0_mv_residual_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()
}