Constant FORWARD_MSL
Source pub(crate) const FORWARD_MSL: &str = "#include <metal_stdlib>\nusing namespace metal;\n\n#line 1 \"common/quantized.metal\"\n// Affine 4-bit dense projections: prefill GEMMs, verify matvecs, and input staging.\n\nconstant constexpr uint GROUP_SIZE = 64;\n\nstruct QmvParams {\n uint out_dim;\n uint in_dim;\n};\n\n// Affine-Q4 GEMM: 128 output rows x 32 token rows x 32 K per\n// threadgroup (128 threads). Four simdgroups each own a 32x32 output\n// tile. Device loads stage in registers before reusing the shared half\n// tiles; barriers separate staging from matrix work. Output edges are\n// clamp-loaded and guard-stored. qmm_from passes ntt = ceil(nb/32) and\n// pads input rows so unused tile lanes cannot read beyond the buffer.\nkernel void qmm_af4_w(\n device const uint* w [[buffer(0)]],\n device const bfloat* scales [[buffer(1)]],\n device const bfloat* biases [[buffer(2)]],\n device const float* x [[buffer(3)]], // [NB][in_dim]\n device float* y [[buffer(4)]], // [NB][out_dim]\n constant QmvParams& p [[buffer(5)]],\n constant uint& ntt [[buffer(6)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint tiitg [[thread_position_in_threadgroup]],\n uint sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half sa[128 * 32];\n threadgroup half sb[32 * 32];\n\n const uint r1 = (tgpos % ntt) * 32;\n const uint r0 = (tgpos / ntt) * 128;\n const uint words_per_row = p.in_dim / 8;\n const uint groups_per_row = p.in_dim / GROUP_SIZE;\n // A-staging: two (row, 16-K-chunk) units per thread. Unit u covers\n // row = u / 2 and K-chunk il = u % 2 of the current 32-K slice.\n const uint u0 = tiitg * 2;\n // B-staging assignment: token + which 8-element K-chunk.\n const uint btok = tiitg / 4;\n const uint bky = 8 * (tiitg % 4);\n device const float* yb = x + (ulong)(r1 + btok) * p.in_dim + bky;\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb[4];\n simdgroup_float8x8 mc[16];\n for (short i = 0; i < 16; i++) {\n mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n }\n\n for (uint loop_k = 0; loop_k < p.in_dim; loop_k += 32) {\n // Dequant this thread\'s two staging units into registers before the\n // barrier so device loads overlap the previous MMA phase. Bytewise\n // nibble split: uchar4 of (w & 0x0F0F0F0F) yields K positions\n // 0,2,4,6 (even) and the >>4 half yields 1,3,5,7 (odd); the\n // scattered swizzle store remaps them to sequential K.\n half4 de[2][2];\n half4 dodd[2][2];\n float4 xv[2];\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint arow = min(r0 + ar, p.out_dim - 1);\n uint2 w2 = *((device const uint2*)(w + (ulong)arow * words_per_row\n + loop_k / 8) + il0);\n float s = (float)scales[(ulong)arow * groups_per_row + loop_k / GROUP_SIZE];\n float b = (float)biases[(ulong)arow * groups_per_row + loop_k / GROUP_SIZE];\n float4 s4 = s;\n float4 b4 = b;\n de[uu][0] = half4(fma(s4, float4(as_type<uchar4>(w2.x & 0x0F0F0F0Fu)), b4));\n dodd[uu][0] =\n half4(fma(s4, float4(as_type<uchar4>((w2.x >> 4) & 0x0F0F0F0Fu)), b4));\n de[uu][1] = half4(fma(s4, float4(as_type<uchar4>(w2.y & 0x0F0F0F0Fu)), b4));\n dodd[uu][1] =\n half4(fma(s4, float4(as_type<uchar4>((w2.y >> 4) & 0x0F0F0F0Fu)), b4));\n }\n xv[0] = *(device const float4*)(yb);\n xv[1] = *(device const float4*)(yb + 4);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n // Swizzled A store: 8x8 blocks, transposed within block (sa laid\n // out as 16 row-blocks x 4 K-blocks of 64 halves each). Even\n // nibbles land at K = 2i, odd at K = 2i+1 within each 8-half.\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint sy = ar / 8;\n const uint lx = ar % 8;\n for (uint h = 0; h < 2; h++) {\n threadgroup half* base = sa + 64 * (16 * (2 * il0 + h) + sy) + lx;\n for (uint i = 0; i < 4; i++) {\n base[8 * (2 * i)] = de[uu][h][i];\n base[8 * (2 * i + 1)] = dodd[uu][h][i];\n }\n }\n }\n // B store: token-major 8x8 blocks.\n {\n const uint ib = 4 * (tiitg % 4) + btok / 8;\n const uint ly = btok % 8;\n for (uint i = 0; i < 4; i++) {\n *(sb + 64 * ib + 8 * ly + i) = (half)xv[0][i];\n *(sb + 64 * ib + 8 * ly + 4 + i) = (half)xv[1][i];\n }\n }\n yb += 32;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n // Each simdgroup: rows sgitg*32..+32, all 32 tokens.\n threadgroup const half* lsma = sa + 4 * 64 * sgitg;\n threadgroup const half* lsmb = sb;\n for (short ik = 0; ik < 4; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_load(ma[i], lsma + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_load(mb[i], lsmb + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 16; i++) {\n simdgroup_multiply_accumulate(mc[i], mb[i / 4], ma[i % 4], mc[i]);\n }\n lsma += 16 * 64;\n lsmb += 4 * 64;\n }\n }\n\n const uint crow = r0 + 32 * sgitg;\n for (short i = 0; i < 16; i++) {\n if (crow + 8 * (i % 4) + 8 <= p.out_dim) {\n device float* c = y + crow + (ulong)(r1 + 8 * (i / 4)) * p.out_dim;\n simdgroup_store(mc[i], c + 8 * (i % 4), p.out_dim, 0, false);\n }\n }\n}\n\n// Small-batch qmm (nb <= 8, one 8-token tile): avoids computing the\n// padding of a full 32-token tile. Same A-staging/dequant as qmm_af4_w;\n// B shrinks to one 8x8 block per K-chunk and each simdgroup keeps four\n// accumulator fragments (32 output rows x 8 tokens).\nkernel void qmm_af4_n8(\n device const uint* w [[buffer(0)]],\n device const bfloat* scales [[buffer(1)]],\n device const bfloat* biases [[buffer(2)]],\n device const float* x [[buffer(3)]], // [NB][in_dim]\n device float* y [[buffer(4)]], // [NB][out_dim]\n constant QmvParams& p [[buffer(5)]],\n constant uint& ntt [[buffer(6)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint tiitg [[thread_position_in_threadgroup]],\n uint sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half sa[128 * 32];\n threadgroup half sb[8 * 32];\n\n const uint r1 = (tgpos % ntt) * 8;\n const uint r0 = (tgpos / ntt) * 128;\n const uint words_per_row = p.in_dim / 8;\n const uint groups_per_row = p.in_dim / GROUP_SIZE;\n const uint u0 = tiitg * 2;\n // B-staging: threads 0..31 own (token btok, 8-K-chunk bky).\n const uint btok = tiitg / 4;\n const uint bky = 8 * (tiitg % 4);\n device const float* yb = x + (ulong)(r1 + btok) * p.in_dim + bky;\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb;\n simdgroup_float8x8 mc[4];\n for (short i = 0; i < 4; i++) {\n mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n }\n\n for (uint loop_k = 0; loop_k < p.in_dim; loop_k += 32) {\n half4 de[2][2];\n half4 dodd[2][2];\n float4 xv[2];\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint arow = min(r0 + ar, p.out_dim - 1);\n uint2 w2 = *((device const uint2*)(w + (ulong)arow * words_per_row\n + loop_k / 8) + il0);\n float s = (float)scales[(ulong)arow * groups_per_row + loop_k / GROUP_SIZE];\n float b = (float)biases[(ulong)arow * groups_per_row + loop_k / GROUP_SIZE];\n float4 s4 = s;\n float4 b4 = b;\n de[uu][0] = half4(fma(s4, float4(as_type<uchar4>(w2.x & 0x0F0F0F0Fu)), b4));\n dodd[uu][0] =\n half4(fma(s4, float4(as_type<uchar4>((w2.x >> 4) & 0x0F0F0F0Fu)), b4));\n de[uu][1] = half4(fma(s4, float4(as_type<uchar4>(w2.y & 0x0F0F0F0Fu)), b4));\n dodd[uu][1] =\n half4(fma(s4, float4(as_type<uchar4>((w2.y >> 4) & 0x0F0F0F0Fu)), b4));\n }\n if (btok < 8) {\n xv[0] = *(device const float4*)(yb);\n xv[1] = *(device const float4*)(yb + 4);\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint sy = ar / 8;\n const uint lx = ar % 8;\n for (uint h = 0; h < 2; h++) {\n threadgroup half* base = sa + 64 * (16 * (2 * il0 + h) + sy) + lx;\n for (uint i = 0; i < 4; i++) {\n base[8 * (2 * i)] = de[uu][h][i];\n base[8 * (2 * i + 1)] = dodd[uu][h][i];\n }\n }\n }\n if (btok < 8) {\n const uint ib = tiitg % 4;\n const uint ly = btok;\n for (uint i = 0; i < 4; i++) {\n *(sb + 64 * ib + 8 * ly + i) = (half)xv[0][i];\n *(sb + 64 * ib + 8 * ly + 4 + i) = (half)xv[1][i];\n }\n }\n yb += 32;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n threadgroup const half* lsma = sa + 4 * 64 * sgitg;\n for (short ik = 0; ik < 4; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_load(ma[i], lsma + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n simdgroup_load(mb, sb + 64 * ik, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_multiply_accumulate(mc[i], mb, ma[i], mc[i]);\n }\n lsma += 16 * 64;\n }\n }\n\n const uint crow = r0 + 32 * sgitg;\n for (short i = 0; i < 4; i++) {\n if (crow + 8 * i + 8 <= p.out_dim) {\n device float* c = y + crow + (ulong)r1 * p.out_dim;\n simdgroup_store(mc[i], c + 8 * i, p.out_dim, 0, false);\n }\n }\n}\n\n// 16-token sibling of qmm_af4_n8 (one 16-token tile, 2 B fragments per\n// K-chunk, 8 accumulator fragments per simdgroup): streams the weights\n// once for 9\u{2013}16 rows instead of twice via two 8-token tiles.\nkernel void qmm_af4_n16(\n device const uint* w [[buffer(0)]],\n device const bfloat* scales [[buffer(1)]],\n device const bfloat* biases [[buffer(2)]],\n device const float* x [[buffer(3)]], // [NB][in_dim]\n device float* y [[buffer(4)]], // [NB][out_dim]\n constant QmvParams& p [[buffer(5)]],\n constant uint& ntt [[buffer(6)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint tiitg [[thread_position_in_threadgroup]],\n uint sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half sa[128 * 32];\n threadgroup half sb[16 * 32];\n\n const uint r1 = (tgpos % ntt) * 16;\n const uint r0 = (tgpos / ntt) * 128;\n const uint words_per_row = p.in_dim / 8;\n const uint groups_per_row = p.in_dim / GROUP_SIZE;\n const uint u0 = tiitg * 2;\n const uint btok = tiitg / 4;\n const uint bky = 8 * (tiitg % 4);\n device const float* yb = x + (ulong)(r1 + btok) * p.in_dim + bky;\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb[2];\n simdgroup_float8x8 mc[8];\n for (short i = 0; i < 8; i++) {\n mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n }\n\n for (uint loop_k = 0; loop_k < p.in_dim; loop_k += 32) {\n half4 de[2][2];\n half4 dodd[2][2];\n float4 xv[2];\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint arow = min(r0 + ar, p.out_dim - 1);\n uint2 w2 = *((device const uint2*)(w + (ulong)arow * words_per_row\n + loop_k / 8) + il0);\n float s = (float)scales[(ulong)arow * groups_per_row + loop_k / GROUP_SIZE];\n float b = (float)biases[(ulong)arow * groups_per_row + loop_k / GROUP_SIZE];\n float4 s4 = s;\n float4 b4 = b;\n de[uu][0] = half4(fma(s4, float4(as_type<uchar4>(w2.x & 0x0F0F0F0Fu)), b4));\n dodd[uu][0] =\n half4(fma(s4, float4(as_type<uchar4>((w2.x >> 4) & 0x0F0F0F0Fu)), b4));\n de[uu][1] = half4(fma(s4, float4(as_type<uchar4>(w2.y & 0x0F0F0F0Fu)), b4));\n dodd[uu][1] =\n half4(fma(s4, float4(as_type<uchar4>((w2.y >> 4) & 0x0F0F0F0Fu)), b4));\n }\n if (btok < 16) {\n xv[0] = *(device const float4*)(yb);\n xv[1] = *(device const float4*)(yb + 4);\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint sy = ar / 8;\n const uint lx = ar % 8;\n for (uint h = 0; h < 2; h++) {\n threadgroup half* base = sa + 64 * (16 * (2 * il0 + h) + sy) + lx;\n for (uint i = 0; i < 4; i++) {\n base[8 * (2 * i)] = de[uu][h][i];\n base[8 * (2 * i + 1)] = dodd[uu][h][i];\n }\n }\n }\n if (btok < 16) {\n // Blocks: ib = 2*kc + tokgroup(btok/8); [token%8][K elem].\n const uint ib = 2 * (tiitg % 4) + btok / 8;\n const uint ly = btok % 8;\n for (uint i = 0; i < 4; i++) {\n *(sb + 64 * ib + 8 * ly + i) = (half)xv[0][i];\n *(sb + 64 * ib + 8 * ly + 4 + i) = (half)xv[1][i];\n }\n }\n yb += 32;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n threadgroup const half* lsma = sa + 4 * 64 * sgitg;\n threadgroup const half* lsmb = sb;\n for (short ik = 0; ik < 4; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_load(ma[i], lsma + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 2; i++) {\n simdgroup_load(mb[i], lsmb + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 8; i++) {\n simdgroup_multiply_accumulate(mc[i], mb[i / 4], ma[i % 4], mc[i]);\n }\n lsma += 16 * 64;\n lsmb += 2 * 64;\n }\n }\n\n const uint crow = r0 + 32 * sgitg;\n for (short i = 0; i < 8; i++) {\n if (crow + 8 * (i % 4) + 8 <= p.out_dim) {\n device float* c = y + crow + (ulong)(r1 + 8 * (i / 4)) * p.out_dim;\n simdgroup_store(mc[i], c + 8 * (i % 4), p.out_dim, 0, false);\n }\n }\n}\n\n// Deinterleave x into HALF even/odd streams for the half-math verify qmv,\n// and emit per-32-element-group sums (in float) so the qmv bias term reads\n// one value per group instead of re-summing x per output row. Group g of\n// stream b covers pair-indices [g*16, g*16+16); n2 is always a multiple of\n// 32, so each 16-lane simd half maps to exactly one group.\nkernel void deinterleave_bh(\n device const float* x [[buffer(0)]],\n device half* xe [[buffer(1)]],\n device half* xo [[buffer(2)]],\n constant uint& n2 [[buffer(3)]],\n constant uint& nb [[buffer(4)]],\n device float* xsum [[buffer(5)]],\n uint gi [[thread_position_in_grid]],\n uint lane [[thread_index_in_simdgroup]])\n{\n if (gi >= nb * n2) return;\n uint b = gi / n2;\n uint i = gi % n2;\n device const float* xb = x + (ulong)b * n2 * 2;\n float e = xb[2 * i];\n float o = xb[2 * i + 1];\n xe[gi] = (half)e;\n xo[gi] = (half)o;\n float s = e + o;\n s += simd_shuffle_xor(s, 1);\n s += simd_shuffle_xor(s, 2);\n s += simd_shuffle_xor(s, 4);\n s += simd_shuffle_xor(s, 8);\n if ((lane & 15) == 0) {\n xsum[b * (n2 / 16) + i / 16] = s;\n }\n}\n\n// One-to-three-row verify qmv with half dot math: per\n// half-group, half4 accumulation promoted into f32 accumulators.\nkernel void qmv_multi_h(\n device const uint* w [[buffer(0)]],\n device const bfloat* scales [[buffer(1)]],\n device const bfloat* biases [[buffer(2)]],\n device const half* xe [[buffer(3)]],\n device const half* xo [[buffer(4)]],\n device float* y [[buffer(5)]], // [b][out_dim]\n constant QmvParams& p [[buffer(6)]],\n constant uint& nb [[buffer(7)]],\n device const float* xsum [[buffer(8)]], // [b][in_dim/32] group sums\n uint tgpos [[threadgroup_position_in_grid]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint spt [[simdgroups_per_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint row = tgpos * spt + sgid;\n if (row >= p.out_dim) return;\n const uint words_per_row = p.in_dim / 8;\n const uint halves_per_row = p.in_dim / 32;\n const uint bstride4 = p.in_dim / 8; // half4s per token stream\n device const uint4* wr = (device const uint4*)(w + (ulong)row * words_per_row);\n device const bfloat* sr = scales + (ulong)row * (halves_per_row / 2);\n device const bfloat* br = biases + (ulong)row * (halves_per_row / 2);\n device const half4* xe0 = (device const half4*)(xe);\n device const half4* xo0 = (device const half4*)(xo);\n device const half4* xe1 = xe0 + bstride4;\n device const half4* xo1 = xo0 + bstride4;\n device const half4* xe2 = xe1 + bstride4;\n device const half4* xo2 = xo1 + bstride4;\n device const float* xg0 = xsum;\n device const float* xg1 = xsum + halves_per_row;\n device const float* xg2 = xsum + 2 * halves_per_row;\n const bool two = nb > 1;\n const bool three = nb > 2;\n float acc0 = 0.0f;\n float acc1 = 0.0f;\n float acc2 = 0.0f;\n for (uint hg = lane; hg < halves_per_row; hg += 32) {\n uint4 w4 = wr[hg];\n float s = (float)sr[hg >> 1];\n float b = (float)br[hg >> 1];\n half4 qd0 = 0.0h;\n half4 qd1 = 0.0h;\n half4 qd2 = 0.0h;\n uint word;\n half4 lo4;\n half4 hi4;\n half4 xa;\n half4 xb;\n #define QMH_WORD(W, J) \\\n word = (W); \\\n lo4 = half4(as_type<uchar4>(word & 0x0F0F0F0Fu)); \\\n hi4 = half4(as_type<uchar4>((word >> 4) & 0x0F0F0F0Fu)); \\\n xa = xe0[hg * 4 + (J)]; \\\n xb = xo0[hg * 4 + (J)]; \\\n qd0 = fma(lo4, xa, qd0); \\\n qd0 = fma(hi4, xb, qd0); \\\n if (two) { \\\n xa = xe1[hg * 4 + (J)]; \\\n xb = xo1[hg * 4 + (J)]; \\\n qd1 = fma(lo4, xa, qd1); \\\n qd1 = fma(hi4, xb, qd1); \\\n } \\\n if (three) { \\\n xa = xe2[hg * 4 + (J)]; \\\n xb = xo2[hg * 4 + (J)]; \\\n qd2 = fma(lo4, xa, qd2); \\\n qd2 = fma(hi4, xb, qd2); \\\n }\n QMH_WORD(w4.x, 0) QMH_WORD(w4.y, 1)\n QMH_WORD(w4.z, 2) QMH_WORD(w4.w, 3)\n #undef QMH_WORD\n float qdf0 = (float)qd0.x + (float)qd0.y + (float)qd0.z + (float)qd0.w;\n acc0 = fma(s, qdf0, fma(b, xg0[hg], acc0));\n if (two) {\n float qdf1 = (float)qd1.x + (float)qd1.y + (float)qd1.z + (float)qd1.w;\n acc1 = fma(s, qdf1, fma(b, xg1[hg], acc1));\n }\n if (three) {\n float qdf2 = (float)qd2.x + (float)qd2.y + (float)qd2.z + (float)qd2.w;\n acc2 = fma(s, qdf2, fma(b, xg2[hg], acc2));\n }\n }\n acc0 = simd_sum(acc0);\n if (lane == 0) y[row] = acc0;\n if (two) {\n acc1 = simd_sum(acc1);\n if (lane == 0) y[(ulong)p.out_dim + row] = acc1;\n }\n if (three) {\n acc2 = simd_sum(acc2);\n if (lane == 0) y[2 * (ulong)p.out_dim + row] = acc2;\n }\n}\n\n// N-stream verify qmv (nb <= 8): one weight stream shared across up to\n// 8 token rows. Per-row reduction order is identical to qmv_multi_h /\n// qmv (same hg walk, same fma order), so verify numerics match the\n// historically gated path exactly. Runtime row guards; ~40 GPRs.\nkernel void qmv_multi_hn(\n device const uint* w [[buffer(0)]],\n device const bfloat* scales [[buffer(1)]],\n device const bfloat* biases [[buffer(2)]],\n device const half* xe [[buffer(3)]],\n device const half* xo [[buffer(4)]],\n device float* y [[buffer(5)]], // [b][out_dim]\n constant QmvParams& p [[buffer(6)]],\n constant uint& nb [[buffer(7)]],\n device const float* xsum [[buffer(8)]], // [b][in_dim/32] group sums\n uint tgpos [[threadgroup_position_in_grid]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint spt [[simdgroups_per_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint row = tgpos * spt + sgid;\n if (row >= p.out_dim) return;\n const uint words_per_row = p.in_dim / 8;\n const uint halves_per_row = p.in_dim / 32;\n const uint bstride4 = p.in_dim / 8; // half4s per token stream\n device const uint4* wr = (device const uint4*)(w + (ulong)row * words_per_row);\n device const bfloat* sr = scales + (ulong)row * (halves_per_row / 2);\n device const bfloat* br = biases + (ulong)row * (halves_per_row / 2);\n device const half4* xe0 = (device const half4*)(xe);\n device const half4* xo0 = (device const half4*)(xo);\n const bool r1 = nb > 1;\n const bool r2 = nb > 2;\n const bool r3 = nb > 3;\n const bool r4 = nb > 4;\n const bool r5 = nb > 5;\n const bool r6 = nb > 6;\n const bool r7 = nb > 7;\n float acc0 = 0.0f;\n float acc1 = 0.0f;\n float acc2 = 0.0f;\n float acc3 = 0.0f;\n float acc4 = 0.0f;\n float acc5 = 0.0f;\n float acc6 = 0.0f;\n float acc7 = 0.0f;\n for (uint hg = lane; hg < halves_per_row; hg += 32) {\n uint4 w4 = wr[hg];\n float s = (float)sr[hg >> 1];\n float b = (float)br[hg >> 1];\n half4 qd0 = 0.0h;\n half4 qd1 = 0.0h;\n half4 qd2 = 0.0h;\n half4 qd3 = 0.0h;\n half4 qd4 = 0.0h;\n half4 qd5 = 0.0h;\n half4 qd6 = 0.0h;\n half4 qd7 = 0.0h;\n uint word;\n half4 lo4;\n half4 hi4;\n half4 xa;\n half4 xb;\n #define QMN_WORD(W, J) \\\n word = (W); \\\n lo4 = half4(as_type<uchar4>(word & 0x0F0F0F0Fu)); \\\n hi4 = half4(as_type<uchar4>((word >> 4) & 0x0F0F0F0Fu)); \\\n xa = xe0[hg * 4 + (J)]; \\\n xb = xo0[hg * 4 + (J)]; \\\n qd0 = fma(lo4, xa, qd0); \\\n qd0 = fma(hi4, xb, qd0); \\\n if (r1) { \\\n xa = xe0[1u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[1u * bstride4 + hg * 4 + (J)]; \\\n qd1 = fma(lo4, xa, qd1); \\\n qd1 = fma(hi4, xb, qd1); \\\n } \\\n if (r2) { \\\n xa = xe0[2u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[2u * bstride4 + hg * 4 + (J)]; \\\n qd2 = fma(lo4, xa, qd2); \\\n qd2 = fma(hi4, xb, qd2); \\\n } \\\n if (r3) { \\\n xa = xe0[3u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[3u * bstride4 + hg * 4 + (J)]; \\\n qd3 = fma(lo4, xa, qd3); \\\n qd3 = fma(hi4, xb, qd3); \\\n } \\\n if (r4) { \\\n xa = xe0[4u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[4u * bstride4 + hg * 4 + (J)]; \\\n qd4 = fma(lo4, xa, qd4); \\\n qd4 = fma(hi4, xb, qd4); \\\n } \\\n if (r5) { \\\n xa = xe0[5u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[5u * bstride4 + hg * 4 + (J)]; \\\n qd5 = fma(lo4, xa, qd5); \\\n qd5 = fma(hi4, xb, qd5); \\\n } \\\n if (r6) { \\\n xa = xe0[6u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[6u * bstride4 + hg * 4 + (J)]; \\\n qd6 = fma(lo4, xa, qd6); \\\n qd6 = fma(hi4, xb, qd6); \\\n } \\\n if (r7) { \\\n xa = xe0[7u * bstride4 + hg * 4 + (J)]; \\\n xb = xo0[7u * bstride4 + hg * 4 + (J)]; \\\n qd7 = fma(lo4, xa, qd7); \\\n qd7 = fma(hi4, xb, qd7); \\\n }\n QMN_WORD(w4.x, 0) QMN_WORD(w4.y, 1)\n QMN_WORD(w4.z, 2) QMN_WORD(w4.w, 3)\n #undef QMN_WORD\n float qdf0 = (float)qd0.x + (float)qd0.y + (float)qd0.z + (float)qd0.w;\n acc0 = fma(s, qdf0, fma(b, xsum[hg], acc0));\n if (r1) {\n float qdf1 = (float)qd1.x + (float)qd1.y + (float)qd1.z + (float)qd1.w;\n acc1 = fma(s, qdf1, fma(b, xsum[1u * halves_per_row + hg], acc1));\n }\n if (r2) {\n float qdf2 = (float)qd2.x + (float)qd2.y + (float)qd2.z + (float)qd2.w;\n acc2 = fma(s, qdf2, fma(b, xsum[2u * halves_per_row + hg], acc2));\n }\n if (r3) {\n float qdf3 = (float)qd3.x + (float)qd3.y + (float)qd3.z + (float)qd3.w;\n acc3 = fma(s, qdf3, fma(b, xsum[3u * halves_per_row + hg], acc3));\n }\n if (r4) {\n float qdf4 = (float)qd4.x + (float)qd4.y + (float)qd4.z + (float)qd4.w;\n acc4 = fma(s, qdf4, fma(b, xsum[4u * halves_per_row + hg], acc4));\n }\n if (r5) {\n float qdf5 = (float)qd5.x + (float)qd5.y + (float)qd5.z + (float)qd5.w;\n acc5 = fma(s, qdf5, fma(b, xsum[5u * halves_per_row + hg], acc5));\n }\n if (r6) {\n float qdf6 = (float)qd6.x + (float)qd6.y + (float)qd6.z + (float)qd6.w;\n acc6 = fma(s, qdf6, fma(b, xsum[6u * halves_per_row + hg], acc6));\n }\n if (r7) {\n float qdf7 = (float)qd7.x + (float)qd7.y + (float)qd7.z + (float)qd7.w;\n acc7 = fma(s, qdf7, fma(b, xsum[7u * halves_per_row + hg], acc7));\n }\n }\n acc0 = simd_sum(acc0);\n if (lane == 0) y[row] = acc0;\n if (r1) {\n acc1 = simd_sum(acc1);\n if (lane == 0) y[1u * (ulong)p.out_dim + row] = acc1;\n }\n if (r2) {\n acc2 = simd_sum(acc2);\n if (lane == 0) y[2u * (ulong)p.out_dim + row] = acc2;\n }\n if (r3) {\n acc3 = simd_sum(acc3);\n if (lane == 0) y[3u * (ulong)p.out_dim + row] = acc3;\n }\n if (r4) {\n acc4 = simd_sum(acc4);\n if (lane == 0) y[4u * (ulong)p.out_dim + row] = acc4;\n }\n if (r5) {\n acc5 = simd_sum(acc5);\n if (lane == 0) y[5u * (ulong)p.out_dim + row] = acc5;\n }\n if (r6) {\n acc6 = simd_sum(acc6);\n if (lane == 0) y[6u * (ulong)p.out_dim + row] = acc6;\n }\n if (r7) {\n acc7 = simd_sum(acc7);\n if (lane == 0) y[7u * (ulong)p.out_dim + row] = acc7;\n }\n}\n\n\n#line 1 \"common/gemm.metal\"\n// Half-input, float-output GEMM used by prefill attention.\n\n// ---------------- GEMM-attention (prefill) ----------------\n// Attention as two dense GEMMs: S = Q K^T, masked row-softmax, O = P V.\n// Causality and QSA visibility are applied by fn_attn_softmax_sel, so\n// both GEMMs can operate on rectangular tiles.\n\nstruct GemmHParams {\n uint m; // rows of A/C\n uint n; // cols of C (rounded up to 32; garbage cols never read)\n uint k; // contraction length\n uint lda; // row stride of A (elements)\n uint ldb; // row stride of B (elements)\n uint ldc; // row stride of C (elements)\n};\n\n// C[m][n] = sum_k A[m][k] * B[n][k]; A,B half, C float. 128x32 tile per\n// threadgroup (4 simdgroups of 32x32), K-step 32 - the qmm_af4_w\n// geometry with direct half staging instead of dequant.\nkernel void gemm_hh(\n device const half* A [[buffer(0)]],\n device const half* B [[buffer(1)]],\n device float* C [[buffer(2)]],\n constant GemmHParams& p [[buffer(3)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint tiitg [[thread_position_in_threadgroup]],\n uint sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half sa[128 * 32];\n threadgroup half sb[32 * 32];\n const uint ntn = p.n / 32;\n const uint rn = (tgpos % ntn) * 32;\n const uint r0 = (tgpos / ntn) * 128;\n const uint u0 = tiitg * 2;\n const uint btok = tiitg / 4;\n const uint bky = 8 * (tiitg % 4);\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb[4];\n simdgroup_float8x8 mc[16];\n for (short i = 0; i < 16; i++) {\n mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n }\n\n for (uint loop_k = 0; loop_k < p.k; loop_k += 32) {\n threadgroup_barrier(mem_flags::mem_threadgroup);\n for (uint uu = 0; uu < 2; uu++) {\n const uint unit = u0 + uu;\n const uint ar = unit / 2;\n const uint il0 = unit % 2;\n const uint arow = min(r0 + ar, p.m - 1);\n device const half* src = A + (ulong)arow * p.lda + loop_k + il0 * 16;\n const uint sy = ar / 8;\n const uint lx = ar % 8;\n for (uint i = 0; i < 16; i++) {\n const uint sx = 2 * il0 + i / 8;\n *(sa + 64 * (16 * sx + sy) + 8 * (i % 8) + lx) = src[i];\n }\n }\n {\n const uint ib = 4 * (tiitg % 4) + btok / 8;\n const uint ly = btok % 8;\n device const half* src = B + (ulong)(rn + btok) * p.ldb + loop_k + bky;\n for (uint i = 0; i < 8; i++) {\n *(sb + 64 * ib + 8 * ly + i) = src[i];\n }\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n threadgroup const half* lsma = sa + 4 * 64 * sgitg;\n threadgroup const half* lsmb = sb;\n for (short ik = 0; ik < 4; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_load(ma[i], lsma + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 4; i++) {\n simdgroup_load(mb[i], lsmb + 64 * i, 8, 0, false);\n }\n simdgroup_barrier(mem_flags::mem_none);\n for (short i = 0; i < 16; i++) {\n simdgroup_multiply_accumulate(mc[i], mb[i / 4], ma[i % 4], mc[i]);\n }\n lsma += 16 * 64;\n lsmb += 4 * 64;\n }\n }\n\n // mc[i] covers m-block (i % 4) within this sg\'s 32 rows and n-block\n // (i / 4); fragment (row, col) = (n_local, m_local), so store\n // transposed to land C[m][n] row-major.\n // Store whenever the fragment STARTS in range: C rows are allocated\n // padded past m, and rows >= m are never read downstream, so partial\n // tail fragments write harmless garbage instead of being dropped.\n const uint crow = r0 + 32 * sgitg;\n for (short i = 0; i < 16; i++) {\n const uint m0 = crow + 8 * (i % 4);\n if (m0 < p.m) {\n simdgroup_store(mc[i], C + (ulong)m0 * p.ldc + rn + 8 * (i / 4),\n p.ldc, 0, true);\n }\n }\n}\n\n\n#line 1 \"common/attention.metal\"\n// QK normalization/RoPE, prefill staging, q8 KV storage, and decode attention.\n\nstruct QkRopeParams {\n uint n_heads;\n uint head_dim; // 256\n uint stride; // elements between consecutive heads\' query starts\n uint rot; // 64\n uint pos;\n float theta;\n float eps;\n};\n\nstruct AttnGemmParams {\n uint kv_row; // n_kv * head_dim\n uint blocks_per_row; // kv_row / 32\n uint head; // kv head index\n uint kl; // real cache length (base + nb)\n uint kl_pad; // kl rounded up to 32\n uint scale_off; // byte offset of scales plane in kc/vc\n uint base; // context length before this chunk\n uint nb; // tokens in this chunk\n uint n_rep; // q heads per kv head\n float scale; // 1/sqrt(head_dim), folded into staged q\n};\n\n// Dequantize one kv head\'s K rows ([t][256] half) and V columns\n// (transposed: [d][kl_pad] half) from the q8 cache. Rows/cols beyond kl\n// are zeroed so downstream GEMMs never see uninitialized memory.\nkernel void attn_kv_stage(\n device const char* kc [[buffer(0)]],\n device const char* vc [[buffer(1)]],\n device half* kh [[buffer(2)]], // [kl_pad][256]\n device half* vt [[buffer(3)]], // [256][kl_pad]\n constant AttnGemmParams& p [[buffer(4)]],\n uint gid [[thread_position_in_grid]])\n{\n const uint t = gid / 256;\n const uint d = gid % 256;\n if (t >= p.kl_pad) return;\n if (t >= p.kl) {\n kh[(ulong)t * 256 + d] = 0.0h;\n vt[(ulong)d * p.kl_pad + t] = 0.0h;\n return;\n }\n device const half* ks = (device const half*)(kc + p.scale_off);\n device const half* vs = (device const half*)(vc + p.scale_off);\n const uint e = p.head * 256 + d;\n const uint blk = e / 32;\n float dk = (float)ks[(ulong)t * p.blocks_per_row + blk];\n float dv = (float)vs[(ulong)t * p.blocks_per_row + blk];\n kh[(ulong)t * 256 + d] =\n (half)((float)kc[(ulong)t * p.kv_row + e] * dk);\n vt[(ulong)d * p.kl_pad + t] =\n (half)((float)vc[(ulong)t * p.kv_row + e] * dv);\n}\n\n// Stage q rows for ALL heads as half with the softmax scale folded in:\n// qh[(head*nb + t)][256] from qg [t][n_heads][2*256] (q half of each\n// head\'s [query|gate] pair).\nkernel void attn_q_stage(\n device const float* qg [[buffer(0)]],\n device half* qh [[buffer(1)]],\n constant AttnGemmParams& p [[buffer(2)]],\n constant uint& n_heads [[buffer(3)]],\n uint gid [[thread_position_in_grid]])\n{\n const uint d = gid % 256;\n const uint t = (gid / 256) % p.nb;\n const uint h = gid / (256 * p.nb);\n if (h >= n_heads) return;\n float q = qg[((ulong)t * n_heads + h) * 512 + d];\n qh[((ulong)h * p.nb + t) * 256 + d] = (half)(q * p.scale);\n}\n\n// Scatter one kv head\'s O rows into attn_out with the sigmoid output\n// gate applied: attn_out[t][(head*n_rep+qh)*256 + d].\nkernel void attn_o_scatter(\n device const float* O [[buffer(0)]], // [(qh*nb + t)][256]\n device const float* qg [[buffer(1)]],\n device float* out [[buffer(2)]],\n constant AttnGemmParams& p [[buffer(3)]],\n constant uint& n_heads [[buffer(4)]],\n uint gid [[thread_position_in_grid]])\n{\n const uint d = gid % 256;\n const uint t = (gid / 256) % p.nb;\n const uint qh = gid / (256 * p.nb);\n if (qh >= p.n_rep) return;\n const uint gh = p.head * p.n_rep + qh;\n float gate = qg[((ulong)t * n_heads + gh) * 512 + 256 + d];\n out[((ulong)t * n_heads + gh) * 256 + d] =\n O[((ulong)qh * p.nb + t) * 256 + d] / (1.0f + exp(-gate));\n}\n\n// Batched per-head QK-norm + RoPE. One simdgroup per (token, head);\n// x layout [b][n_heads*stride]; position = base_pos + b.\nkernel void qk_norm_rope_b(\n device float* x [[buffer(0)]],\n device const bfloat* w [[buffer(1)]],\n constant QkRopeParams& p [[buffer(2)]],\n constant uint& nb [[buffer(3)]],\n uint sg [[simdgroup_index_in_threadgroup]],\n uint spt [[simdgroups_per_threadgroup]],\n uint gid [[threadgroup_position_in_grid]],\n uint lane [[thread_index_in_simdgroup]])\n{\n uint gsg = gid * spt + sg;\n if (gsg >= nb * p.n_heads) return;\n uint b = gsg / p.n_heads;\n uint head = gsg % p.n_heads;\n device float* h = x + (ulong)b * p.n_heads * p.stride + (ulong)head * p.stride;\n const uint per_lane = p.head_dim / 32;\n float ss = 0.0f;\n for (uint j = 0; j < per_lane; j++) {\n float v = h[lane * per_lane + j];\n ss += v * v;\n }\n float total = simd_sum(ss);\n float inv = rsqrt(total / (float)p.head_dim + p.eps);\n for (uint j = 0; j < per_lane; j++) {\n uint idx = lane * per_lane + j;\n h[idx] = h[idx] * inv * (float)w[idx];\n }\n simdgroup_barrier(mem_flags::mem_device);\n uint half_rot = p.rot / 2;\n if (lane < half_rot) {\n float inv_freq = pow(p.theta, -2.0f * (float)lane / (float)p.rot);\n float angle = (float)(p.pos + b) * inv_freq;\n float c = cos(angle);\n float s = sin(angle);\n float a = h[lane];\n float bb = h[lane + half_rot];\n h[lane] = a * c - bb * s;\n h[lane + half_rot] = bb * c + a * s;\n }\n}\n\n// ---------------- Split-T attention (flash-decode style) ----------------\n\nconstant constexpr uint ATTN_TB = 512; // positions per partial block\n\nstruct AttnPartParams {\n uint n_heads;\n uint n_kv;\n uint head_dim; // 256\n uint t_len;\n uint q_stride; // 512 (per-head [query|gate])\n uint q_off; // element offset of this token\'s q block\n uint max_blk; // partials stride per head\n float scale;\n};\n\n// One threadgroup per head: merge partials (log-sum-exp), apply the\n// sigmoid gate, write out[head][head_dim].\nkernel void attn_combine(\n device const float* q [[buffer(0)]],\n device const float* part [[buffer(1)]],\n device float* out [[buffer(2)]],\n constant AttnPartParams& p [[buffer(3)]],\n constant uint& nblk [[buffer(4)]],\n constant uint& out_off [[buffer(5)]],\n uint head [[threadgroup_position_in_grid]],\n uint tid [[thread_position_in_threadgroup]],\n uint tpg [[threads_per_threadgroup]])\n{\n const uint hd = p.head_dim;\n device const float* ph = part + (ulong)head * p.max_blk * (2 + hd);\n float m = -INFINITY;\n for (uint b = 0; b < nblk; b++) m = max(m, ph[b * (2 + hd)]);\n float s = 0.0f;\n for (uint b = 0; b < nblk; b++) {\n s += exp2(ph[b * (2 + hd)] - m) * ph[b * (2 + hd) + 1];\n }\n float inv_s = 1.0f / s;\n device const float* qh = q + p.q_off + (ulong)head * p.q_stride;\n for (uint d = tid; d < hd; d += tpg) {\n float acc = 0.0f;\n for (uint b = 0; b < nblk; b++) {\n acc += exp2(ph[b * (2 + hd)] - m) * ph[b * (2 + hd) + 2 + d];\n }\n float gate = qh[hd + d];\n out[out_off + (ulong)head * hd + d] =\n acc * inv_s * (1.0f / (1.0f + exp(-gate)));\n }\n}\n\n// q8 KV cache: separate quantized-value and scale planes.\n// Layout per cache side: int8 quants plane [max_t][1024] then f16 scales\n// plane [max_t][32] (one symmetric scale per 32-element block, d = amax/127).\n// scale_off = byte offset of the scales plane within the buffer.\n\nstruct KvQParams {\n uint row; // kv_row elements (1024)\n uint t0; // first position\n uint nb; // tokens\n uint scale_off; // byte offset of the scales plane\n};\n\n// Quantize-on-append: one simdgroup per 32-element block; lane = element.\n// Handles K and V in one dispatch (side = z index of the simdgroup).\nkernel void kv_append_q8(\n device const float* k [[buffer(0)]], // [nb][row] f32 (post norm+rope)\n device const float* v [[buffer(1)]],\n device char* kc [[buffer(2)]],\n device char* vc [[buffer(3)]],\n constant KvQParams& p [[buffer(4)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint spt [[simdgroups_per_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint blocks_per_row = p.row / 32;\n const uint gsg = tgpos * spt + sgid;\n const uint total = p.nb * blocks_per_row * 2; // both sides\n if (gsg >= total) return;\n const uint side = gsg / (p.nb * blocks_per_row);\n const uint rem = gsg % (p.nb * blocks_per_row);\n const uint b = rem / blocks_per_row;\n const uint blk = rem % blocks_per_row;\n\n device const float* src = (side == 0 ? k : v) + (ulong)b * p.row + blk * 32;\n device char* dstq = (side == 0 ? kc : vc);\n device half* dsts = (device half*)((side == 0 ? kc : vc) + p.scale_off);\n\n float val = src[lane];\n float amax = simd_max(fabs(val));\n float d = amax / 127.0f;\n float q = d > 0.0f ? rint(val / d) : 0.0f;\n dstq[(ulong)(p.t0 + b) * p.row + blk * 32 + lane] =\n (char)clamp(q, -127.0f, 127.0f);\n if (lane == 0) {\n dsts[(ulong)(p.t0 + b) * blocks_per_row + blk] = (half)d;\n }\n}\n\n// GQA lock-step flash-decode partials over the q8 cache: same structure as\n// attn_part2, dequant fused into the dot (char4 -> float4 times the block\n// scale; lane l covers dims l*4 (block l/8) and (l+32)*4 (block 4+l/8)).\nkernel void attn_part2_q8(\n device const float* q [[buffer(0)]],\n device const char* kc [[buffer(1)]],\n device const char* vc [[buffer(2)]],\n device float* part [[buffer(3)]],\n constant AttnPartParams& p [[buffer(4)]],\n constant uint& n_wg [[buffer(5)]],\n constant uint& scale_off [[buffer(6)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint hd = p.head_dim;\n const uint kv_row = p.n_kv * hd;\n const uint blocks_per_row = kv_row / 32;\n const uint n_rep = p.n_heads / p.n_kv;\n const uint hk = tgpos % p.n_kv;\n const uint iwg = tgpos / p.n_kv;\n const uint head = hk * n_rep + sgid;\n if (sgid >= n_rep) return;\n\n device const half* ks = (device const half*)(kc + scale_off);\n device const half* vs = (device const half*)(vc + scale_off);\n // This lane\'s element offsets within the row, and their block indices.\n const uint e0 = hk * hd + lane * 4;\n const uint e1 = hk * hd + (lane + 32) * 4;\n const uint b0 = e0 / 32;\n const uint b1 = e1 / 32;\n\n device const float* qh = q + p.q_off + (ulong)head * p.q_stride;\n const float4 qa = *(device const float4*)(qh + lane * 4);\n const float4 qb = *(device const float4*)(qh + (lane + 32) * 4);\n\n float m = -INFINITY;\n float s = 0.0f;\n float4 oa = 0.0f;\n float4 ob = 0.0f;\n\n const uint n_chunks = (p.t_len + 31) / 32;\n for (uint c = iwg; c < n_chunks; c += n_wg) {\n const uint c0 = c * 32;\n const uint cn = min(p.t_len - c0, 32u);\n for (uint i = 0; i < cn; i++) {\n const ulong t = c0 + i;\n float ds0 = (float)ks[t * blocks_per_row + b0];\n float ds1 = (float)ks[t * blocks_per_row + b1];\n float4 ka = float4(*(device const char4*)(kc + t * kv_row + e0)) * ds0;\n float4 kb = float4(*(device const char4*)(kc + t * kv_row + e1)) * ds1;\n float partial = dot(qa, ka) + dot(qb, kb);\n float sc = simd_sum(partial) * p.scale;\n float mnew = max(m, sc);\n float factor = exp2(m - mnew);\n float e = exp2(sc - mnew);\n float dv0 = (float)vs[t * blocks_per_row + b0];\n float dv1 = (float)vs[t * blocks_per_row + b1];\n float4 va = float4(*(device const char4*)(vc + t * kv_row + e0)) * dv0;\n float4 vb = float4(*(device const char4*)(vc + t * kv_row + e1)) * dv1;\n oa = oa * factor + e * va;\n ob = ob * factor + e * vb;\n s = s * factor + e;\n m = mnew;\n }\n }\n\n device float* pb = part + ((ulong)head * p.max_blk + iwg) * (2 + hd);\n if (lane == 0) {\n pb[0] = m;\n pb[1] = s;\n }\n *(device float4*)(pb + 2 + lane * 4) = oa;\n *(device float4*)(pb + 2 + (lane + 32) * 4) = ob;\n}\n\n\n#line 1 \"common/deltanet.metal\"\n// DeltaNet causal convolution, gate preparation, and recurrent state scan.\n\nstruct ConvParams {\n uint channels; // 10240\n uint ksize; // 4\n};\n\n// Batched conv step: each thread owns a channel, iterating nb tokens.\nkernel void conv_b(\n device float* qkv [[buffer(0)]], // [b][channels]\n device const bfloat* w [[buffer(1)]],\n device float* hist [[buffer(2)]],\n constant ConvParams& p [[buffer(3)]],\n constant uint& nb [[buffer(4)]],\n constant uint& snap_after [[buffer(5)]],\n device float* mid [[buffer(6)]],\n uint c [[thread_position_in_grid]])\n{\n if (c >= p.channels) return;\n const uint k = p.ksize;\n const uint km1 = k - 1;\n device const bfloat* wr = w + (ulong)c * k;\n device float* h = hist + (ulong)c * km1;\n float h0 = h[0];\n float h1 = h[1];\n float h2 = h[2];\n float w0 = (float)wr[0];\n float w1 = (float)wr[1];\n float w2 = (float)wr[2];\n float w3 = (float)wr[3];\n for (uint b = 0; b < nb; b++) {\n device float* xc = qkv + (ulong)b * p.channels + c;\n float cur = *xc;\n float acc = w0 * h0 + w1 * h1 + w2 * h2 + w3 * cur;\n h0 = h1; h1 = h2; h2 = cur;\n *xc = acc / (1.0f + exp(-acc));\n if (b < snap_after) {\n device float* plane = mid + (ulong)b * p.channels * km1;\n plane[(ulong)c * km1 + 0] = h0;\n plane[(ulong)c * km1 + 1] = h1;\n plane[(ulong)c * km1 + 2] = h2;\n }\n }\n h[0] = h0;\n h[1] = h1;\n h[2] = h2;\n}\n\n// DeltaNet recurrent scan; state is transposed for contiguous simdgroup loads.\n\nstruct DeltaPrepParams {\n uint n_k; // 16\n uint n_v; // 48\n uint d_k; // 128\n uint d_v; // 128\n float eps;\n uint nb;\n uint snap_after; // save planes after rows [0, snap_after); 0 disables snapshots\n};\n\n// Per (token, k-head): L2-normalize k and q (query gets the 1/sqrt(dk)\n// scale) out of the post-conv qkv into kqn = [b][ k(2048) | q(2048) ].\nkernel void delta_norms(\n device const float* qkv [[buffer(0)]],\n device float* kqn [[buffer(1)]],\n constant DeltaPrepParams& p [[buffer(2)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint spt [[simdgroups_per_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint gsg = tgpos * spt + sgid;\n const uint total = p.nb * p.n_k;\n if (gsg >= total) return;\n const uint b = gsg / p.n_k;\n const uint kh = gsg % p.n_k;\n const uint dk = p.d_k;\n const uint qk = p.n_k * dk;\n const uint conv_dim = 2 * qk + p.n_v * p.d_v;\n device const float* qh = qkv + (ulong)b * conv_dim + kh * dk;\n device const float* kh_p = qkv + (ulong)b * conv_dim + qk + kh * dk;\n float4 q4 = *(device const float4*)(qh + lane * 4);\n float4 k4 = *(device const float4*)(kh_p + lane * 4);\n float qss = simd_sum(dot(q4, q4));\n float kss = simd_sum(dot(k4, k4));\n float qinv = rsqrt(qss + p.eps) * rsqrt((float)dk);\n float kinv = rsqrt(kss + p.eps);\n device float* out = kqn + (ulong)b * 2 * qk;\n *(device float4*)(out + kh * dk + lane * 4) = k4 * kinv;\n *(device float4*)(out + qk + kh * dk + lane * 4) = q4 * qinv;\n}\n\n// Per (token, v-head): decay and beta scalars into gb = [b][n_v][2].\nkernel void delta_gates(\n device const float* a [[buffer(0)]],\n device const float* bb [[buffer(1)]],\n device const bfloat* a_log [[buffer(2)]],\n device const bfloat* dt_b [[buffer(3)]],\n device float* gb [[buffer(4)]],\n constant DeltaPrepParams& p [[buffer(5)]],\n uint gi [[thread_position_in_grid]])\n{\n if (gi >= p.nb * p.n_v) return;\n const uint h = gi % p.n_v;\n float av = a[gi] + (float)dt_b[h];\n float sp = av > 20.0f ? av : log(1.0f + exp(av));\n gb[2 * gi] = exp(-exp((float)a_log[h]) * sp);\n gb[2 * gi + 1] = 1.0f / (1.0f + exp(-bb[gi]));\n}\n\n// The scan. One simdgroup per (v-head, dv-row); state row (dk=128 floats,\n// TRANSPOSED layout [head][dv][dk]) lives in 4 registers per lane across\n// the whole scan. Two simd_sums per token; no threadgroup memory, no\n// barriers. Emits raw readout y; fn_delta_gate_norm_sigmoid_b applies\n// normalization and the output gate afterwards.\nkernel void delta_scan2(\n device const float* qkv [[buffer(0)]], // [b][conv_dim] post-conv\n device const float* kqn [[buffer(1)]], // [b][k|q normalized]\n device const float* gb [[buffer(2)]], // [b][n_v][decay, beta]\n device float* state [[buffer(3)]], // [head][d_v][d_k]\n device float* y [[buffer(4)]], // [b][n_v*d_v] raw\n device float* mid [[buffer(5)]], // snapshot plane, same layout\n constant DeltaPrepParams& p [[buffer(6)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint spt [[simdgroups_per_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint gsg = tgpos * spt + sgid;\n if (gsg >= p.n_v * p.d_v) return;\n const uint head = gsg / p.d_v;\n const uint iv = gsg % p.d_v;\n const uint dk = p.d_k;\n const uint qk = p.n_k * dk;\n const uint conv_dim = 2 * qk + p.n_v * p.d_v;\n const uint hk = head / (p.n_v / p.n_k);\n\n device float* srow = state + ((ulong)head * p.d_v + iv) * dk;\n float4 s4 = *(device const float4*)(srow + lane * 4);\n\n for (uint b = 0; b < p.nb; b++) {\n device const float* kq = kqn + (ulong)b * 2 * qk;\n float4 k4 = *(device const float4*)(kq + hk * dk + lane * 4);\n float4 q4 = *(device const float4*)(kq + qk + hk * dk + lane * 4);\n float decay = gb[2 * (b * p.n_v + head)];\n float beta = gb[2 * (b * p.n_v + head) + 1];\n float vv = qkv[(ulong)b * conv_dim + 2 * qk + head * p.d_v + iv];\n\n // Decay the state, correct its key readout, then read with q.\n // Keep the float4/FMA/reduction order: it defines verify numerics.\n s4 *= decay;\n float kv_mem = simd_sum(dot(s4, k4));\n float delta = (vv - kv_mem) * beta;\n s4 = fma(k4, delta, s4);\n float yv = simd_sum(dot(s4, q4));\n if (lane == 0) {\n y[(ulong)b * p.n_v * p.d_v + (ulong)head * p.d_v + iv] = yv;\n }\n // Plane b is the state AFTER row b, restored by commit(b + 1)\n // when the following speculative rows are rejected.\n if (b < p.snap_after) {\n device float* plane = mid + (ulong)b * p.n_v * p.d_v * dk;\n *(device float4*)(plane + ((ulong)head * p.d_v + iv) * dk + lane * 4) = s4;\n }\n }\n *(device float4*)(srow + lane * 4) = s4;\n}\n\n\n#line 1 \"common/elementwise.metal\"\n// Elementwise activation, residual addition, and buffer copying.\n\n// y[i] = silu(a[i]) * b[i]\nkernel void silu_mul(\n device const float* a [[buffer(0)]],\n device const float* b [[buffer(1)]],\n device float* y [[buffer(2)]],\n uint i [[thread_position_in_grid]])\n{\n float v = a[i];\n y[i] = (v / (1.0f + exp(-v))) * b[i];\n}\n\n// x[i] += r[i]\nkernel void add_inplace(\n device float* x [[buffer(0)]],\n device const float* r [[buffer(1)]],\n uint i [[thread_position_in_grid]])\n{\n x[i] += r[i];\n}\n\n// Plain device-to-device f32 copy.\nkernel void copy_f32(\n device const float* src [[buffer(0)]],\n device float* dst [[buffer(1)]],\n constant uint& n [[buffer(2)]],\n uint i [[thread_position_in_grid]])\n{\n if (i < n) dst[i] = src[i];\n}\n\n\n#line 1 \"common/sampling.metal\"\n// Token embeddings and greedy argmax reduction.\n\nstruct ArgPair {\n float v;\n uint i;\n};\n\n// Stage 1: per-threadgroup argmax over a grid-strided slice of logits.\n// Lowest index wins ties.\nkernel void argmax_partial(\n device const float* logits [[buffer(0)]],\n device ArgPair* partials [[buffer(1)]],\n constant uint& n [[buffer(2)]],\n uint tid [[thread_position_in_grid]],\n uint grid [[threads_per_grid]],\n uint tgid [[threadgroup_position_in_grid]],\n uint ltid [[thread_position_in_threadgroup]],\n uint tpg [[threads_per_threadgroup]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n threadgroup ArgPair shm[32];\n float best = -INFINITY;\n uint besti = 0;\n for (uint i = tid; i < n; i += grid) {\n float v = logits[i];\n if (v > best) { best = v; besti = i; }\n }\n for (uint off = 16; off > 0; off >>= 1) {\n float ov = simd_shuffle_down(best, off);\n uint oi = simd_shuffle_down(besti, off);\n if (ov > best || (ov == best && oi < besti)) { best = ov; besti = oi; }\n }\n if (lane == 0) { shm[sgid].v = best; shm[sgid].i = besti; }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (ltid == 0) {\n for (uint s = 1; s < (tpg + 31) / 32; s++) {\n if (shm[s].v > best || (shm[s].v == best && shm[s].i < besti)) {\n best = shm[s].v;\n besti = shm[s].i;\n }\n }\n partials[tgid].v = best;\n partials[tgid].i = besti;\n }\n}\n\n// Stage 2: reduce partials, write the winner into ids[step + 1].\nkernel void argmax_final(\n device const ArgPair* partials [[buffer(0)]],\n device uint* ids [[buffer(1)]],\n constant uint& np [[buffer(2)]],\n constant uint& step [[buffer(3)]],\n uint ltid [[thread_position_in_threadgroup]],\n uint tpg [[threads_per_threadgroup]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n threadgroup ArgPair shm[32];\n float best = -INFINITY;\n uint besti = 0;\n for (uint i = ltid; i < np; i += tpg) {\n if (partials[i].v > best || (partials[i].v == best && partials[i].i < besti)) {\n best = partials[i].v;\n besti = partials[i].i;\n }\n }\n for (uint off = 16; off > 0; off >>= 1) {\n float ov = simd_shuffle_down(best, off);\n uint oi = simd_shuffle_down(besti, off);\n if (ov > best || (ov == best && oi < besti)) { best = ov; besti = oi; }\n }\n if (lane == 0) { shm[sgid].v = best; shm[sgid].i = besti; }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (ltid == 0) {\n for (uint s = 1; s < (tpg + 31) / 32; s++) {\n if (shm[s].v > best || (shm[s].v == best && shm[s].i < besti)) {\n best = shm[s].v;\n besti = shm[s].i;\n }\n }\n ids[step + 1] = besti;\n }\n}\n\n// ---------------- Batched prefill path (NB tokens per pass) ----------------\n\n// Embed nb rows: ids[i0..i0+nb] -> hidden[b][hsize].\nkernel void embed_rows(\n device const uint* ids [[buffer(0)]],\n device const uint* w [[buffer(1)]],\n device const bfloat* scales [[buffer(2)]],\n device const bfloat* biases [[buffer(3)]],\n device float* hidden [[buffer(4)]],\n constant uint& hsize [[buffer(5)]],\n constant uint& i0 [[buffer(6)]],\n constant uint& nb [[buffer(7)]],\n uint gi [[thread_position_in_grid]])\n{\n if (gi >= nb * hsize) return;\n uint b = gi / hsize;\n uint i = gi % hsize;\n ulong row = ids[i0 + b];\n uint word = w[row * (hsize / 8) + i / 8];\n float q = float((word >> (4 * (i % 8))) & 0xFu);\n float s = scales[row * (hsize / 64) + i / 64];\n float bb = biases[row * (hsize / 64) + i / 64];\n hidden[gi] = fma(s, q, bb);\n}\n\n";