Constant BATCH_MSL
Source pub(crate) const BATCH_MSL: &str = "#include <metal_stdlib>\nusing namespace metal;\n\n#line 1 \"qwen4_exp/types.metal\"\n// Host ABI types shared by qwen4-exp row-batched subsystems.\n\nstruct FnQmvParams {\n uint out_dim;\n uint in_dim;\n};\n\nstruct GroupParams {\n uint n; // group width (hidden)\n uint groups; // streams\n float eps;\n float shift; // added to the norm weight (1.0 for raw HF norms)\n};\n\n\n#line 1 \"qwen4_exp/quantized_rows.metal\"\n// Q4 row helpers shared by hyper-connections and experts; compiled before both.\n\n// Affine-Q4 dot of one weight row against NB half-stream rows (the\n// qmv_multi_h inner loop). Row r reads xe/xo + r*bstride4 (half4 units)\n// and xsum + r*gstride. Leaves the simd-reduced sums in acc[].\ntemplate <uint NB>\nstatic inline void fn_q4_rows_h(\n device const uint4* wr,\n device const bfloat* sr,\n device const bfloat* br,\n device const half4* xe,\n device const half4* xo,\n device const float* xsum,\n uint halves_per_row,\n uint bstride4,\n uint gstride,\n uint lane,\n thread float* acc)\n{\n for (uint hg = lane; hg < halves_per_row; hg += 32) {\n const uint4 w4 = wr[hg];\n const float s = (float)sr[hg >> 1];\n const float bb = (float)br[hg >> 1];\n half4 qd[NB];\n for (uint r = 0; r < NB; r++) qd[r] = 0.0h;\n uint word;\n half4 lo4;\n half4 hi4;\n #define FN_ROWS_WORD(W, J) \\\n word = (W); \\\n lo4 = half4(as_type<uchar4>(word & 0x0F0F0F0Fu)); \\\n hi4 = half4(as_type<uchar4>((word >> 4) & 0x0F0F0F0Fu)); \\\n for (uint r = 0; r < NB; r++) { \\\n qd[r] = fma(lo4, xe[r * bstride4 + hg * 4 + (J)], qd[r]); \\\n qd[r] = fma(hi4, xo[r * bstride4 + hg * 4 + (J)], qd[r]); \\\n }\n FN_ROWS_WORD(w4.x, 0) FN_ROWS_WORD(w4.y, 1)\n FN_ROWS_WORD(w4.z, 2) FN_ROWS_WORD(w4.w, 3)\n #undef FN_ROWS_WORD\n for (uint r = 0; r < NB; r++) {\n const float qdf = (float)qd[r].x + (float)qd[r].y + (float)qd[r].z + (float)qd[r].w;\n acc[r] = fma(s, qdf, fma(bb, xsum[r * gstride + hg], acc[r]));\n }\n }\n for (uint r = 0; r < NB; r++) acc[r] = simd_sum(acc[r]);\n}\n\n// Affine-Q4 dot of TWO weight rows against NB half-stream rows, sharing\n// every x load. Same per-row lane assignment and reduction order as\n// fn_q4_rows_h, so results are bit-identical; the point is that a short\n// row (the experts\' 640-wide down projection is 20 half-groups over 32\n// lanes, one iteration each) amortizes its setup and x loads over two\n// output rows. Measured 26 percent faster at that shape.\ntemplate <uint NB>\nstatic inline void fn_q4_2rows_h(\n device const uint4* wr0,\n device const uint4* wr1,\n device const bfloat* sr0,\n device const bfloat* br0,\n device const bfloat* sr1,\n device const bfloat* br1,\n device const half4* xe,\n device const half4* xo,\n device const float* xsum,\n uint halves_per_row,\n uint bstride4,\n uint gstride,\n uint lane,\n thread float* acc0,\n thread float* acc1)\n{\n for (uint hg = lane; hg < halves_per_row; hg += 32) {\n const uint4 w0 = wr0[hg];\n const uint4 w1 = wr1[hg];\n const float s0 = (float)sr0[hg >> 1];\n const float b0 = (float)br0[hg >> 1];\n const float s1 = (float)sr1[hg >> 1];\n const float b1 = (float)br1[hg >> 1];\n half4 qd0[NB];\n half4 qd1[NB];\n for (uint r = 0; r < NB; r++) { qd0[r] = 0.0h; qd1[r] = 0.0h; }\n #define FN_2ROWS_WORD(W0, W1, J) \\\n { \\\n const uint word0 = (W0); \\\n const uint word1 = (W1); \\\n const half4 lo0 = half4(as_type<uchar4>(word0 & 0x0F0F0F0Fu)); \\\n const half4 hi0 = half4(as_type<uchar4>((word0 >> 4) & 0x0F0F0F0Fu)); \\\n const half4 lo1 = half4(as_type<uchar4>(word1 & 0x0F0F0F0Fu)); \\\n const half4 hi1 = half4(as_type<uchar4>((word1 >> 4) & 0x0F0F0F0Fu)); \\\n for (uint r = 0; r < NB; r++) { \\\n const half4 xa = xe[r * bstride4 + hg * 4 + (J)]; \\\n const half4 xb = xo[r * bstride4 + hg * 4 + (J)]; \\\n qd0[r] = fma(lo0, xa, qd0[r]); \\\n qd0[r] = fma(hi0, xb, qd0[r]); \\\n qd1[r] = fma(lo1, xa, qd1[r]); \\\n qd1[r] = fma(hi1, xb, qd1[r]); \\\n } \\\n }\n FN_2ROWS_WORD(w0.x, w1.x, 0) FN_2ROWS_WORD(w0.y, w1.y, 1)\n FN_2ROWS_WORD(w0.z, w1.z, 2) FN_2ROWS_WORD(w0.w, w1.w, 3)\n #undef FN_2ROWS_WORD\n for (uint r = 0; r < NB; r++) {\n const float f0 = (float)qd0[r].x + (float)qd0[r].y + (float)qd0[r].z + (float)qd0[r].w;\n const float f1 = (float)qd1[r].x + (float)qd1[r].y + (float)qd1[r].z + (float)qd1[r].w;\n const float xs = xsum[r * gstride + hg];\n acc0[r] = fma(s0, f0, fma(b0, xs, acc0[r]));\n acc1[r] = fma(s1, f1, fma(b1, xs, acc1[r]));\n }\n }\n for (uint r = 0; r < NB; r++) {\n acc0[r] = simd_sum(acc0[r]);\n acc1[r] = simd_sum(acc1[r]);\n }\n}\n\n\n#line 1 \"qwen4_exp/hyperconnection.metal\"\n// Hyper-connection replication, normalization, bottleneck projection, mixing, and injection.\n\n// hyper[b][g][i] = e[b][i]\nkernel void fn_replicate_b(\n device const float* e [[buffer(0)]],\n device float* hyper [[buffer(1)]],\n constant GroupParams& p [[buffer(2)]],\n constant uint& nb [[buffer(3)]],\n uint gi [[thread_position_in_grid]])\n{\n const uint hh = p.n * p.groups;\n if (gi >= nb * hh) return;\n const uint b = gi / hh;\n const uint i = (gi % hh) % p.n;\n hyper[gi] = e[b * p.n + i];\n}\n\n// One threadgroup per (row, group):\n// y[b][g*n+i] = rmsnorm(x[b][g*n..])[i] * (w[g*n+i] + shift)\nkernel void fn_group_norm_b(\n device const float* x [[buffer(0)]],\n device const bfloat* w [[buffer(1)]],\n device float* y [[buffer(2)]],\n constant GroupParams& p [[buffer(3)]],\n uint tg [[threadgroup_position_in_grid]],\n uint tid [[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 float partial[32];\n const uint g = tg % p.groups;\n const ulong base = (ulong)tg * p.n;\n const ulong wbase = (ulong)g * p.n;\n float ss = 0.0f;\n for (uint i = tid; i < p.n; i += tpg) {\n float v = x[base + i];\n ss += v * v;\n }\n ss = simd_sum(ss);\n if (lane == 0) partial[sgid] = ss;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (sgid == 0) {\n float total = (lane < (tpg + 31) / 32) ? partial[lane] : 0.0f;\n total = simd_sum(total);\n if (lane == 0) partial[0] = rsqrt(total / (float)p.n + p.eps);\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n const float inv = partial[0];\n for (uint i = tid; i < p.n; i += tpg) {\n y[base + i] = x[base + i] * inv * ((float)w[wbase + i] + p.shift);\n }\n}\n\nstruct NormPrepParams {\n uint n; // group width (hidden)\n uint groups; // streams\n float eps;\n uint inject; // 1: first add out[b][i] * 2 sigmoid(r[b][g]/groups) into x\n};\n\n// Fused gated-residual read, stage 1, one 256-thread threadgroup per\n// (row, group): optional injection of `out` into x (in place), grouped\n// RMSNorm to y, and y\'s half even/odd copy plus per-32 group sums for the\n// Q4 matvecs (deinterleave_bh layout). Requires n % 256 == 0.\nkernel void fn_norm_prep_b(\n device float* x [[buffer(0)]], // [nb][hh]\n device const bfloat* w [[buffer(1)]], // [hh]\n device float* y [[buffer(2)]], // [nb][hh]\n device half* xe [[buffer(3)]], // [nb][hh/2]\n device half* xo [[buffer(4)]],\n device float* xsum [[buffer(5)]], // [nb][hh/32]\n device const float* out [[buffer(6)]], // [nb][n]\n device const float* r [[buffer(7)]], // [nb][groups]\n constant NormPrepParams& p [[buffer(8)]],\n uint tg [[threadgroup_position_in_grid]],\n uint tid [[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 float partial[32];\n const uint b = tg / p.groups;\n const uint g = tg % p.groups;\n const uint hh = p.n * p.groups;\n const ulong base = (ulong)tg * p.n;\n const ulong wbase = (ulong)g * p.n;\n float wgt = 0.0f;\n if (p.inject != 0) wgt = 2.0f / (1.0f + exp(-r[b * p.groups + g] / (float)p.groups));\n float ss = 0.0f;\n for (uint i = tid; i < p.n; i += tpg) {\n float v = x[base + i];\n if (p.inject != 0) {\n v += out[(ulong)b * p.n + i] * wgt;\n x[base + i] = v;\n }\n ss += v * v;\n }\n ss = simd_sum(ss);\n if (lane == 0) partial[sgid] = ss;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (sgid == 0) {\n float total = (lane < (tpg + 31) / 32) ? partial[lane] : 0.0f;\n total = simd_sum(total);\n if (lane == 0) partial[0] = rsqrt(total / (float)p.n + p.eps);\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n const float inv = partial[0];\n // With 256 threads, element i = tid + 256k sits in 32-chunk i/32 =\n // sgid + 8k with lane i % 32, so a simd_sum is the chunk sum.\n device float* xs = xsum + (ulong)b * (hh / 32) + wbase / 32;\n for (uint i = tid; i < p.n; i += tpg) {\n const ulong j = base + i;\n float v = x[j] * inv * (float)w[wbase + i];\n y[j] = v;\n if ((j & 1) == 0) xe[j >> 1] = (half)v; else xo[j >> 1] = (half)v;\n float s = simd_sum(v);\n if (lane == 0) xs[i >> 5] = s;\n }\n}\n\n// Q4 matvec over NB rows whose outputs go through silu(y / div): the\n// gated-residual bottleneck. y is [NB][out_dim] f32.\n#define FN_QMV_SILU_B_ARGS \\\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 const float* xsum [[buffer(5)]], \\\n device float* y [[buffer(6)]], \\\n constant FnQmvParams& p [[buffer(7)]], \\\n constant float& div [[buffer(8)]], \\\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\ntemplate <uint NB>\n[[kernel]] void fn_qmv_silu_b(FN_QMV_SILU_B_ARGS)\n{\n const uint row = tgpos * spt + sgid;\n if (row >= p.out_dim) return;\n const uint halves = p.in_dim / 32;\n device const uint4* wr = (device const uint4*)w + (ulong)row * halves;\n device const bfloat* sr = scales + (ulong)row * (halves / 2);\n device const bfloat* br = biases + (ulong)row * (halves / 2);\n float acc[NB];\n for (uint r = 0; r < NB; r++) acc[r] = 0.0f;\n fn_q4_rows_h<NB>(wr, sr, br, (device const half4*)xe, (device const half4*)xo, xsum,\n halves, p.in_dim / 8, halves, lane, acc);\n if (lane == 0) {\n for (uint r = 0; r < NB; r++) {\n float v = acc[r] / div;\n y[(ulong)r * p.out_dim + row] = v / (1.0f + exp(-v));\n }\n }\n}\n#define FN_INST_QMV_SILU(N) \\\n template [[host_name(\"fn_qmv_silu_b\" #N)]] [[kernel]] void fn_qmv_silu_b<N>(FN_QMV_SILU_B_ARGS);\nFN_INST_QMV_SILU(1)\nFN_INST_QMV_SILU(2)\nFN_INST_QMV_SILU(3)\nFN_INST_QMV_SILU(4)\n\n// mixed[b][i] = mean over groups of sigmoid(u[b][g*n+i]) * normed[b][g*n+i]\nkernel void fn_hc_mix_b(\n device const float* u [[buffer(0)]],\n device const float* normed [[buffer(1)]],\n device float* mixed [[buffer(2)]],\n constant GroupParams& p [[buffer(3)]],\n constant uint& nb [[buffer(4)]],\n uint gi [[thread_position_in_grid]])\n{\n if (gi >= nb * p.n) return;\n const uint b = gi / p.n;\n const uint i = gi % p.n;\n const ulong base = (ulong)b * p.n * p.groups + i;\n float acc = 0.0f;\n for (uint g = 0; g < p.groups; g++) {\n const ulong idx = base + (ulong)g * p.n;\n acc += normed[idx] / (1.0f + exp(-u[idx]));\n }\n mixed[gi] = acc / (float)p.groups;\n}\n\n// hyper[b][g][i] += out[b][i] * 2 sigmoid(r[b][g] / groups)\nkernel void fn_inject_b(\n device float* hyper [[buffer(0)]],\n device const float* out [[buffer(1)]],\n device const float* r [[buffer(2)]],\n constant GroupParams& p [[buffer(3)]],\n constant uint& nb [[buffer(4)]],\n uint gi [[thread_position_in_grid]])\n{\n const uint hh = p.n * p.groups;\n if (gi >= nb * hh) return;\n const uint b = gi / hh;\n const uint rem = gi % hh;\n const uint g = rem / p.n;\n const uint i = rem % p.n;\n const float wgt = 2.0f / (1.0f + exp(-r[b * p.groups + g] / (float)p.groups));\n hyper[gi] += out[(ulong)b * p.n + i] * wgt;\n}\n\n\n#line 1 \"qwen4_exp/experts.metal\"\n// Expert routing, address tables, Q2/Q3 projections, and resident/fetched MoE execution.\n\n#define FN_MAX_NB 4u\n#define FN_SLOT_STRIDE 64u\n\n// y[b][r] = sum_c w[r][c] * x[b][c], bf16 weights. One simdgroup per (b, r).\nkernel void fn_bf16_matvec_b(\n device const bfloat* w [[buffer(0)]],\n device const float* x [[buffer(1)]],\n device float* y [[buffer(2)]],\n constant uint& rows [[buffer(3)]],\n constant uint& cols [[buffer(4)]],\n constant uint& nb [[buffer(5)]],\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 >= nb * rows) return;\n const uint b = gsg / rows;\n const uint r = gsg % rows;\n device const bfloat* wr = w + (ulong)r * cols;\n device const float* xb = x + (ulong)b * cols;\n float acc = 0.0f;\n for (uint c = lane; c < cols; c += 32) acc += (float)wr[c] * xb[c];\n acc = simd_sum(acc);\n if (lane == 0) y[gsg] = acc;\n}\n\n// Softmax over n router logits per row (n <= 1024, one threadgroup of n\n// threads per row), then top-k by repeated argmax; writes indices and the\n// renormalized (or raw) probabilities.\nkernel void fn_topk_softmax_b(\n device const float* logits_all [[buffer(0)]], // [nb][n]\n device uint* idx_all [[buffer(1)]], // [nb][k]\n device float* wts_all [[buffer(2)]], // [nb][k]\n constant uint& n [[buffer(3)]],\n constant uint& k [[buffer(4)]],\n constant uint& renorm [[buffer(5)]],\n uint b [[threadgroup_position_in_grid]],\n uint tid [[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 float red[32];\n threadgroup uint redi[32];\n threadgroup float chosen_sum;\n device const float* logits = logits_all + (ulong)b * n;\n device uint* idx = idx_all + (ulong)b * k;\n device float* wts = wts_all + (ulong)b * k;\n const bool live = tid < n;\n float v = live ? logits[tid] : -INFINITY;\n float m = simd_max(v);\n if (lane == 0) red[sgid] = m;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n m = -INFINITY;\n for (uint s = 0; s < (tpg + 31) / 32; s++) m = max(m, red[s]);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float e = live ? exp(v - m) : 0.0f;\n float sum = simd_sum(e);\n if (lane == 0) red[sgid] = sum;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n sum = 0.0f;\n for (uint s = 0; s < (tpg + 31) / 32; s++) sum += red[s];\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float prob = e / sum;\n float cand = live ? prob : -1.0f;\n if (tid == 0) chosen_sum = 0.0f;\n for (uint round = 0; round < k; round++) {\n float bv = cand;\n uint bi = tid;\n for (uint off = 16; off > 0; off >>= 1) {\n float ov = simd_shuffle_down(bv, off);\n uint oi = simd_shuffle_down(bi, off);\n if (ov > bv || (ov == bv && oi < bi)) { bv = ov; bi = oi; }\n }\n if (lane == 0) { red[sgid] = bv; redi[sgid] = bi; }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (tid == 0) {\n float best = red[0];\n uint besti = redi[0];\n for (uint s = 1; s < (tpg + 31) / 32; s++) {\n if (red[s] > best || (red[s] == best && redi[s] < besti)) {\n best = red[s];\n besti = redi[s];\n }\n }\n idx[round] = besti;\n wts[round] = best;\n chosen_sum += best;\n redi[31] = besti;\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (tid == redi[31]) cand = -1.0f;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n }\n if (renorm != 0 && tid == 0) {\n for (uint j = 0; j < k; j++) wts[j] = wts[j] / chosen_sum;\n }\n}\n\n// ---- grouped experts over the union of the rows\' routed experts ----\n// The union is published by the CPU after the router: tab[layer][u] holds\n// the GPU address of unique expert u\'s record (a residency-set member;\n// records already resident first, then the ones being fetched),\n// tab[layer][62] the resident count, tab[layer][63] the union size, and\n// wmap[layer][b][u] each row\'s routing weight for it (0 when the row did\n// not pick it). The dispatch runs in two parts: part 0 covers the\n// resident experts plus the shared expert (stored at index p.n_u), part 1\n// the late arrivals, which the combine accumulates. Grids are sized for\n// the largest union (p.n_u = nb * k); jobs outside their part exit.\nstruct MoeBParams {\n uint inter;\n uint hidden;\n uint n_u; // grid bound on unique routed experts (nb * k)\n uint gate_w; // byte offsets inside a record\n uint up_w;\n uint down_w;\n uint gate_s;\n uint gate_b;\n uint up_s;\n uint up_b;\n uint down_s;\n uint down_b;\n uint layer; // row of slot_tab / wmap\n uint shared; // 1: expert n_u is the shared expert, weights in `dense`\n uint nb;\n uint sh_gate_w;\n uint sh_gate_s;\n uint sh_gate_b;\n uint sh_up_w;\n uint sh_up_s;\n uint sh_up_b;\n uint sh_down_w;\n uint sh_down_s;\n uint sh_down_b;\n uint sh_gate_vec;\n uint part; // 0: resident experts + shared; 1: late arrivals\n // Byte offsets inside a low-bit record (3-bit: 32 codes per 12 bytes;\n // 2-bit: 16 codes per u32; scales and biases as in the 4-bit record).\n // A table entry with bit 63 set points at a 3-bit record, bit 62 at\n // a 2-bit one.\n uint low_gate_w;\n uint low_up_w;\n uint low_down_w;\n uint low_gate_s;\n uint low_gate_b;\n uint low_up_s;\n uint low_up_b;\n uint low_down_s;\n uint low_down_b;\n};\n\n// Which unique-expert jobs belong to this part; `u == p.n_u` is the shared\n// expert\'s job in part 0.\nstatic inline bool fn_moe_live(constant MoeBParams& p, device const ulong* tab, uint u, thread bool& is_shared)\n{\n const uint n_res = (uint)tab[p.layer * FN_SLOT_STRIDE + FN_SLOT_STRIDE - 2];\n const uint nu = (uint)tab[p.layer * FN_SLOT_STRIDE + FN_SLOT_STRIDE - 1];\n is_shared = (u == p.n_u) && (p.shared != 0);\n if (p.part == 0) return u < n_res || is_shared;\n return u >= n_res && u < nu;\n}\n\n// The record of unique expert u, by GPU address (bits 63 and 62 flag\n// 3-bit and 2-bit records).\nstatic inline device const uchar* fn_moe_rec(constant MoeBParams& p, device const ulong* tab, uint u)\n{\n return (device const uchar*)(tab[p.layer * FN_SLOT_STRIDE + u] & 0x3FFFFFFFFFFFFFFFul);\n}\n\n// 0: 4-bit record, 1: 3-bit, 2: 2-bit.\nstatic inline uint fn_moe_kind(constant MoeBParams& p, device const ulong* tab, uint u)\n{\n const ulong e = tab[p.layer * FN_SLOT_STRIDE + u];\n return (e >> 63) ? 1u : ((e >> 62) & 1ul ? 2u : 0u);\n}\n\n// 2-bit dot of one weight row against NB half-stream rows. Codes are\n// q4 >> 2, reconstructed at the quad midpoint: w = (4q + 1.5) s + b, so a\n// chunk contributes s (4 sum(q x) + 1.5 sum(x)) + b sum(x).\n//\n// The store in src/qwen4_exp/lowbit.rs places a word\'s 16 codes so one\n// masked cast yields four consecutive even codes (casts 0 and 1, paired with xe) or\n// four consecutive odd ones (casts 2 and 3, paired with xo), the same\n// shape as the 4-bit nibble trick. That is eight masked casts and eight\n// fma pairs per 32 codes, exactly the 4-bit kernel\'s instruction count,\n// over half the bytes: measured 26 percent faster than 4-bit at nb = 2\n// where the code-at-a-time version was 133 percent slower.\ntemplate <uint NB>\nstatic inline void fn_q2_rows_h(\n device const uint* wr,\n device const bfloat* sr,\n device const bfloat* br,\n device const half* xe,\n device const half* xo,\n device const float* xsum,\n uint halves,\n uint hstride,\n uint gstride,\n uint lane,\n thread float* acc)\n{\n device const half4* xe4 = (device const half4*)xe;\n device const half4* xo4 = (device const half4*)xo;\n const uint bs4 = hstride / 4;\n device const uint2* w2r = (device const uint2*)wr;\n for (uint hg = lane; hg < halves; hg += 32) {\n const uint2 w2 = w2r[hg];\n const float s = (float)sr[hg >> 1];\n const float bb = (float)br[hg >> 1];\n half4 qd[NB];\n for (uint r = 0; r < NB; r++) qd[r] = 0.0h;\n for (uint t = 0; t < 2; t++) {\n const uint word = (t == 0) ? w2.x : w2.y;\n const half4 e0 = half4(as_type<uchar4>(word & 0x03030303u));\n const half4 e1 = half4(as_type<uchar4>((word >> 2) & 0x03030303u));\n const half4 o0 = half4(as_type<uchar4>((word >> 4) & 0x03030303u));\n const half4 o1 = half4(as_type<uchar4>((word >> 6) & 0x03030303u));\n for (uint r = 0; r < NB; r++) {\n qd[r] = fma(e0, xe4[r * bs4 + hg * 4 + 2 * t], qd[r]);\n qd[r] = fma(e1, xe4[r * bs4 + hg * 4 + 2 * t + 1], qd[r]);\n qd[r] = fma(o0, xo4[r * bs4 + hg * 4 + 2 * t], qd[r]);\n qd[r] = fma(o1, xo4[r * bs4 + hg * 4 + 2 * t + 1], qd[r]);\n }\n }\n for (uint r = 0; r < NB; r++) {\n const float f = (float)qd[r].x + (float)qd[r].y + (float)qd[r].z + (float)qd[r].w;\n const float xs = xsum[r * gstride + hg];\n // s (4 sum(q x) + 1.5 sum(x)) + b sum(x)\n acc[r] = fma(4.0f * s, f, fma(fma(1.5f, s, bb), xs, acc[r]));\n }\n }\n for (uint r = 0; r < NB; r++) acc[r] = simd_sum(acc[r]);\n}\n\n// 3-bit dot of one weight row against NB half-stream rows. A 32-code\n// chunk is three u32 words (packed by src/qwen4_exp/lowbit.rs): the first\n// two hold each code\'s upper two bits in the 2-bit kernel\'s order. The\n// third holds the lowest bits, positioned so each masked cast yields\n// the four bits belonging to the corresponding upper-bit cast.\n// Both planes come out four codes at a time: code = 2 * upper + lowest.\n// This costs 16 to 22 percent more than the 4-bit kernel\n// instead of the 133 percent the code-at-a-time version cost.\n//\n// Codes are q4 >> 1 and w = (2q + 0.5) s + b, so a chunk contributes\n// s (2 sum(q x) + 0.5 sum(x)) + b sum(x).\ntemplate <uint NB>\nstatic inline void fn_q3_rows_h(\n device const uint* wr,\n device const bfloat* sr,\n device const bfloat* br,\n device const half* xe,\n device const half* xo,\n device const float* xsum,\n uint halves,\n uint hstride,\n uint gstride,\n uint lane,\n thread float* acc)\n{\n device const half4* xe4 = (device const half4*)xe;\n device const half4* xo4 = (device const half4*)xo;\n const uint bs4 = hstride / 4;\n for (uint hg = lane; hg < halves; hg += 32) {\n const uint upper0 = wr[3 * hg];\n const uint upper1 = wr[3 * hg + 1];\n const uint lowest_bits = wr[3 * hg + 2];\n const float s = (float)sr[hg >> 1];\n const float bb = (float)br[hg >> 1];\n half4 qd[NB];\n for (uint r = 0; r < NB; r++) qd[r] = 0.0h;\n for (uint t = 0; t < 2; t++) {\n const uint word = (t == 0) ? upper0 : upper1;\n for (uint j = 0; j < 4; j++) {\n // Packing keeps the upper two bits and lowest bit in\n // separate planes: code = 2*upper + lowest.\n const half4 upper = half4(as_type<uchar4>((word >> (2 * j)) & 0x03030303u));\n const half4 lowest = half4(as_type<uchar4>((lowest_bits >> (4 * t + j)) & 0x01010101u));\n const half4 code = fma((half)2.0h, upper, lowest);\n const uint idx = hg * 4 + 2 * t + (j & 1);\n for (uint r = 0; r < NB; r++) {\n qd[r] = fma(code, (j < 2) ? xe4[r * bs4 + idx] : xo4[r * bs4 + idx], qd[r]);\n }\n }\n }\n for (uint r = 0; r < NB; r++) {\n const float f = (float)qd[r].x + (float)qd[r].y + (float)qd[r].z + (float)qd[r].w;\n const float xs = xsum[r * gstride + hg];\n acc[r] = fma(2.0f * s, f, fma(fma(0.5f, s, bb), xs, acc[r]));\n }\n }\n for (uint r = 0; r < NB; r++) acc[r] = simd_sum(acc[r]);\n}\n\n// gate_up[u*NB+b][0..inter] = gate_u(x_b), [inter..2*inter] = up_u(x_b).\n// One simdgroup per (expert, output row), all NB rows at once.\n#define FN_MOE_GATE_UP_B_ARGS \\\n device const half* xe [[buffer(0)]], \\\n device const half* xo [[buffer(1)]], \\\n device const float* xsum [[buffer(2)]], \\\n device float* out [[buffer(3)]], \\\n constant MoeBParams& p [[buffer(4)]], \\\n device const ulong* tab [[buffer(5)]], \\\n device const uchar* dense [[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\ntemplate <uint NB>\n[[kernel]] void fn_moe_gate_up_b(FN_MOE_GATE_UP_B_ARGS)\n{\n const uint g = tgpos * spt + sgid;\n const uint rows2 = 2 * p.inter;\n const uint u = g / rows2;\n bool is_shared;\n if (u > p.n_u || !fn_moe_live(p, tab, u, is_shared)) return;\n const uint r = g % rows2;\n const bool up = r >= p.inter;\n const uint row = up ? r - p.inter : r;\n device const uchar* wb;\n device const uchar* sb;\n device const uchar* bb;\n const uint kind = is_shared ? 0u : fn_moe_kind(p, tab, u);\n if (!is_shared) {\n device const uchar* rec = fn_moe_rec(p, tab, u);\n if (kind != 0) {\n wb = rec + (up ? p.low_up_w : p.low_gate_w);\n sb = rec + (up ? p.low_up_s : p.low_gate_s);\n bb = rec + (up ? p.low_up_b : p.low_gate_b);\n } else {\n wb = rec + (up ? p.up_w : p.gate_w);\n sb = rec + (up ? p.up_s : p.gate_s);\n bb = rec + (up ? p.up_b : p.gate_b);\n }\n } else {\n wb = dense + (up ? p.sh_up_w : p.sh_gate_w);\n sb = dense + (up ? p.sh_up_s : p.sh_gate_s);\n bb = dense + (up ? p.sh_up_b : p.sh_gate_b);\n }\n const uint halves = p.hidden / 32;\n device const bfloat* sr = (device const bfloat*)sb + (ulong)row * (halves / 2);\n device const bfloat* br = (device const bfloat*)bb + (ulong)row * (halves / 2);\n float acc[NB];\n for (uint i = 0; i < NB; i++) acc[i] = 0.0f;\n if (kind == 1) {\n device const uint* wr = (device const uint*)wb + (ulong)row * halves * 3;\n fn_q3_rows_h<NB>(wr, sr, br, xe, xo, xsum, halves, p.hidden / 2, halves, lane, acc);\n } else if (kind == 2) {\n device const uint* wr = (device const uint*)wb + (ulong)row * halves * 2;\n fn_q2_rows_h<NB>(wr, sr, br, xe, xo, xsum, halves, p.hidden / 2, halves, lane, acc);\n } else {\n device const uint4* wr = (device const uint4*)wb + (ulong)row * halves;\n fn_q4_rows_h<NB>(wr, sr, br, (device const half4*)xe, (device const half4*)xo, xsum,\n halves, p.hidden / 8, halves, lane, acc);\n }\n if (lane == 0) {\n for (uint b = 0; b < NB; b++) out[((ulong)u * NB + b) * rows2 + r] = acc[b];\n }\n}\n#define FN_INST_MOE_GATE_UP(N) \\\n template [[host_name(\"fn_moe_gate_up_b\" #N)]] [[kernel]] void fn_moe_gate_up_b<N>(FN_MOE_GATE_UP_B_ARGS);\nFN_INST_MOE_GATE_UP(1)\nFN_INST_MOE_GATE_UP(2)\nFN_INST_MOE_GATE_UP(3)\nFN_INST_MOE_GATE_UP(4)\n\n// h[u*nb+b][j] = silu(gate) * up, written as half even/odd streams plus\n// per-32 group sums for the down projection. Grid = (n_u+shared)*nb*inter/2\n// threads (a multiple of 16: inter/2 % 16 == 0).\nkernel void fn_moe_act_b(\n device const float* gate_up [[buffer(0)]],\n device half* xe2 [[buffer(1)]],\n device half* xo2 [[buffer(2)]],\n device float* xsum2 [[buffer(3)]],\n constant MoeBParams& p [[buffer(4)]],\n device const ulong* tab [[buffer(5)]],\n uint gi [[thread_position_in_grid]],\n uint lane [[thread_index_in_simdgroup]])\n{\n const uint half_i = p.inter / 2;\n const uint ub = gi / half_i;\n const uint i = gi % half_i;\n const uint u = ub / p.nb;\n bool is_shared;\n const bool live = u <= p.n_u && fn_moe_live(p, tab, u, is_shared);\n float e = 0.0f;\n float o = 0.0f;\n if (live) {\n device const float* gu = gate_up + (ulong)ub * 2 * p.inter;\n float g0 = gu[2 * i], g1 = gu[2 * i + 1];\n float u0 = gu[p.inter + 2 * i], u1 = gu[p.inter + 2 * i + 1];\n e = (g0 / (1.0f + exp(-g0))) * u0;\n o = (g1 / (1.0f + exp(-g1))) * u1;\n xe2[gi] = (half)e;\n xo2[gi] = (half)o;\n }\n float s = (float)(half)e + (float)(half)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 (live && (lane & 15) == 0) xsum2[ub * (p.inter / 32) + i / 16] = s;\n}\n\n// y[u*NB+b][row] = down_u(h[u*NB+b]). One simdgroup per (expert, row).\n#define FN_MOE_DOWN_B_ARGS \\\n device const half* xe2 [[buffer(0)]], \\\n device const half* xo2 [[buffer(1)]], \\\n device const float* xsum2 [[buffer(2)]], \\\n device float* y [[buffer(3)]], \\\n constant MoeBParams& p [[buffer(4)]], \\\n device const ulong* tab [[buffer(5)]], \\\n device const uchar* dense [[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// One simdgroup computes TWO consecutive output rows of one expert\n// (p.hidden is even), sharing the x loads: the down projection\'s rows are\n// only 20 half-groups long, so pairing them is 26 percent faster at the\n// same numerics.\ntemplate <uint NB>\n[[kernel]] void fn_moe_down_b(FN_MOE_DOWN_B_ARGS)\n{\n const uint half_h = p.hidden / 2;\n const uint g = tgpos * spt + sgid;\n const uint u = g / half_h;\n bool is_shared;\n if (u > p.n_u || !fn_moe_live(p, tab, u, is_shared)) return;\n const uint row = 2 * (g % half_h);\n device const uchar* wb;\n device const uchar* sb;\n device const uchar* bb;\n const uint kind = is_shared ? 0u : fn_moe_kind(p, tab, u);\n if (!is_shared) {\n device const uchar* rec = fn_moe_rec(p, tab, u);\n if (kind != 0) {\n wb = rec + p.low_down_w;\n sb = rec + p.low_down_s;\n bb = rec + p.low_down_b;\n } else {\n wb = rec + p.down_w;\n sb = rec + p.down_s;\n bb = rec + p.down_b;\n }\n } else {\n wb = dense + p.sh_down_w;\n sb = dense + p.sh_down_s;\n bb = dense + p.sh_down_b;\n }\n const uint halves = p.inter / 32;\n const ulong xrow = (ulong)u * NB;\n float acc[NB];\n float acc1[NB];\n for (uint i = 0; i < NB; i++) { acc[i] = 0.0f; acc1[i] = 0.0f; }\n for (uint j = 0; j < 2; j++) {\n const uint rj = row + j;\n device const bfloat* sr = (device const bfloat*)sb + (ulong)rj * (halves / 2);\n device const bfloat* br = (device const bfloat*)bb + (ulong)rj * (halves / 2);\n thread float* a = (j == 0) ? acc : acc1;\n if (kind == 1) {\n device const uint* wr = (device const uint*)wb + (ulong)rj * halves * 3;\n fn_q3_rows_h<NB>(wr, sr, br,\n xe2 + xrow * (p.inter / 2),\n xo2 + xrow * (p.inter / 2),\n xsum2 + xrow * halves,\n halves, p.inter / 2, halves, lane, a);\n } else if (kind == 2) {\n device const uint* wr = (device const uint*)wb + (ulong)rj * halves * 2;\n fn_q2_rows_h<NB>(wr, sr, br,\n xe2 + xrow * (p.inter / 2),\n xo2 + xrow * (p.inter / 2),\n xsum2 + xrow * halves,\n halves, p.inter / 2, halves, lane, a);\n } else if (kind == 0 && j == 0) {\n // Both rows at once, sharing the x loads.\n device const uint4* wr0 = (device const uint4*)wb + (ulong)row * halves;\n device const uint4* wr1 = wr0 + halves;\n device const bfloat* sr1 = sr + (halves / 2);\n device const bfloat* br1 = br + (halves / 2);\n fn_q4_2rows_h<NB>(wr0, wr1, sr, br, sr1, br1,\n (device const half4*)xe2 + xrow * (p.inter / 8),\n (device const half4*)xo2 + xrow * (p.inter / 8),\n xsum2 + xrow * halves,\n halves, p.inter / 8, halves, lane, acc, acc1);\n }\n if (kind == 0) break;\n }\n if (lane == 0) {\n for (uint b = 0; b < NB; b++) {\n y[(xrow + b) * p.hidden + row] = acc[b];\n y[(xrow + b) * p.hidden + row + 1] = acc1[b];\n }\n }\n}\n#define FN_INST_MOE_DOWN(N) \\\n template [[host_name(\"fn_moe_down_b\" #N)]] [[kernel]] void fn_moe_down_b<N>(FN_MOE_DOWN_B_ARGS);\nFN_INST_MOE_DOWN(1)\nFN_INST_MOE_DOWN(2)\nFN_INST_MOE_DOWN(3)\nFN_INST_MOE_DOWN(4)\n\n// Part 0: out[b][i] = sum_{u < n_res} wmap[layer][b][u] * y[u*nb+b][i]\n// + sigmoid(gate_vec . x[b]) * y[n_u*nb+b][i]\n// Part 1: out[b][i] += sum_{n_res <= u < nu} wmap[layer][b][u] * y[u*nb+b][i]\n// One threadgroup per (row, 256 outputs); the gate dot is recomputed per\n// threadgroup.\nkernel void fn_moe_combine_b(\n device const float* y [[buffer(0)]],\n device const float* wmap [[buffer(1)]], // [layer][FN_MAX_NB][FN_SLOT_STRIDE]\n device float* out [[buffer(2)]],\n constant MoeBParams& p [[buffer(3)]],\n device const uchar* dense [[buffer(4)]],\n device const float* x [[buffer(5)]], // [nb][hidden]\n device const ulong* tab [[buffer(6)]],\n uint tg [[threadgroup_position_in_grid]],\n uint tid [[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 float red[32];\n const uint nchunk = (p.hidden + 255) / 256;\n const uint b = tg / nchunk;\n const uint i = (tg % nchunk) * 256 + tid;\n const uint n_res = (uint)tab[p.layer * FN_SLOT_STRIDE + FN_SLOT_STRIDE - 2];\n const uint nu = (uint)tab[p.layer * FN_SLOT_STRIDE + FN_SLOT_STRIDE - 1];\n if (p.part != 0) {\n if (i >= p.hidden) return;\n device const float* wm = wmap + ((ulong)p.layer * FN_MAX_NB + b) * FN_SLOT_STRIDE;\n float acc = 0.0f;\n for (uint u = n_res; u < nu; u++) {\n const float wv = wm[u];\n if (wv != 0.0f) acc = fma(wv, y[((ulong)u * p.nb + b) * p.hidden + i], acc);\n }\n out[(ulong)b * p.hidden + i] += acc;\n return;\n }\n float gs = 0.0f;\n if (p.shared != 0) {\n device const bfloat* gv = (device const bfloat*)(dense + p.sh_gate_vec);\n device const float* xb = x + (ulong)b * p.hidden;\n float acc = 0.0f;\n for (uint j = tid; j < p.hidden; j += tpg) acc += (float)gv[j] * xb[j];\n acc = simd_sum(acc);\n if (lane == 0) red[sgid] = acc;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float dot = 0.0f;\n for (uint s = 0; s < (tpg + 31) / 32; s++) dot += red[s];\n gs = 1.0f / (1.0f + exp(-dot));\n }\n if (i >= p.hidden) return;\n device const float* wm = wmap + ((ulong)p.layer * FN_MAX_NB + b) * FN_SLOT_STRIDE;\n float acc = 0.0f;\n for (uint u = 0; u < n_res; u++) {\n const float wv = wm[u];\n if (wv != 0.0f) acc = fma(wv, y[((ulong)u * p.nb + b) * p.hidden + i], acc);\n }\n if (p.shared != 0) acc = fma(gs, y[((ulong)p.n_u * p.nb + b) * p.hidden + i], acc);\n out[(ulong)b * p.hidden + i] = acc;\n}\n\n\n#line 1 \"qwen4_exp/expert_gemm.metal\"\n// Low-bit routed-expert GEMMs: 128 output channels x 8/16/32 token rows.\n// Uses the decode stores directly; shared/dense projections remain Q4.\n// K is a multiple of 64, output channels a multiple of 8. The caller pads\n// token buffers to the tile width, as for the existing Q4 GEMMs. Matrix\n// staging/reduction follows those GEMMs; decode\'s half-dot math is separate.\n\ntemplate <uint BITS, uint TOKENS>\n[[kernel]] void fn_expert_qmm(\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 FnQmvParams& 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[TOKENS * 32];\n\n const uint r1 = (tgpos % ntt) * TOKENS;\n const uint r0 = (tgpos / ntt) * 128;\n const uint words_per_row = p.in_dim / 32 * BITS;\n constexpr uint TOKEN_FRAGMENTS = TOKENS / 8;\n const uint groups_per_row = p.in_dim / 64;\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[TOKEN_FRAGMENTS];\n simdgroup_float8x8 mc[4 * TOKEN_FRAGMENTS];\n for (short i = 0; i < 4 * TOKEN_FRAGMENTS; 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 // Decode the same 32-code records used by the decode kernels.\n // Each staging unit owns 16 consecutive K values, unpacked four\n // even/odd codes at a time. Reconstruct Q4-code midpoints before\n // applying the original bf16 scale and bias, then stage as half.\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 device const uint* chunk = w + (ulong)arow * words_per_row\n + (loop_k / 32) * BITS;\n uint upper = chunk[il0];\n uint lowest = 0;\n if (BITS == 3) lowest = chunk[2] >> (4 * il0);\n float4 s = (float)scales[(ulong)arow * groups_per_row + loop_k / 64];\n float4 b = (float)biases[(ulong)arow * groups_per_row + loop_k / 64];\n for (uint h = 0; h < 2; h++) {\n float4 even = float4(as_type<uchar4>((upper >> (2 * h)) & 0x03030303u));\n float4 odd = float4(as_type<uchar4>((upper >> (2 * (h + 2))) & 0x03030303u));\n if (BITS == 3) {\n even = 2.0f * even + float4(as_type<uchar4>((lowest >> h) & 0x01010101u));\n odd = 2.0f * odd + float4(as_type<uchar4>((lowest >> (h + 2)) & 0x01010101u));\n even = 2.0f * even + 0.5f;\n odd = 2.0f * odd + 0.5f;\n } else {\n even = 4.0f * even + 1.5f;\n odd = 4.0f * odd + 1.5f;\n }\n de[uu][h] = half4(fma(s, even, b));\n dodd[uu][h] = half4(fma(s, odd, b));\n }\n }\n if (btok < TOKENS) {\n xv[0] = *(device const float4*)(yb);\n xv[1] = *(device const float4*)(yb + 4);\n }\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 // codes 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 if (btok < TOKENS) {\n const uint ib = TOKEN_FRAGMENTS * (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 owns 32 output rows and all token fragments.\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 < TOKEN_FRAGMENTS; 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 < 4 * TOKEN_FRAGMENTS; i++) {\n simdgroup_multiply_accumulate(mc[i], mb[i / 4], ma[i % 4], mc[i]);\n }\n lsma += 16 * 64;\n lsmb += TOKEN_FRAGMENTS * 64;\n }\n }\n\n const uint crow = r0 + 32 * sgitg;\n for (short i = 0; i < 4 * TOKEN_FRAGMENTS; 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#define FN_EXPERT_QMM_ARGS \\\n device const uint*, device const bfloat*, device const bfloat*, \\\n device const float*, device float*, constant FnQmvParams&, constant uint&, \\\n uint, uint, uint\n#define FN_EXPERT_QMM(B, T) \\\n template [[host_name(\"fn_expert_qmm_q\" #B \"_n\" #T)]] [[kernel]] \\\n void fn_expert_qmm<B, T>(FN_EXPERT_QMM_ARGS);\nFN_EXPERT_QMM(2, 8)\nFN_EXPERT_QMM(2, 16)\nFN_EXPERT_QMM(2, 32)\nFN_EXPERT_QMM(3, 8)\nFN_EXPERT_QMM(3, 16)\nFN_EXPERT_QMM(3, 32)\n#undef FN_EXPERT_QMM\n#undef FN_EXPERT_QMM_ARGS\n\n\n#line 1 \"qwen4_exp/ple.metal\"\n// PLE n-gram gating and dilated convolution.\n\n// ---- PLE (n-gram) block, batched ----\n\n// gated[b][g*n+i] = sigmoid(sq(dot_bg)) * value[b][i] with\n// dot_bg = key[b][g] . query[b][g] / sqrt(n), sq(x) = sign(x) sqrt(max(|x|,1e-6)).\n// One threadgroup per (row, group).\nkernel void fn_ple_gate_b(\n device const float* key [[buffer(0)]],\n device const float* query [[buffer(1)]],\n device const float* value [[buffer(2)]],\n device float* gated [[buffer(3)]],\n constant GroupParams& p [[buffer(4)]],\n uint tg [[threadgroup_position_in_grid]],\n uint tid [[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 float red[32];\n const uint b = tg / p.groups;\n const ulong base = (ulong)tg * p.n;\n float acc = 0.0f;\n for (uint i = tid; i < p.n; i += tpg) acc += key[base + i] * query[base + i];\n acc = simd_sum(acc);\n if (lane == 0) red[sgid] = acc;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float dot = 0.0f;\n for (uint s = 0; s < (tpg + 31) / 32; s++) dot += red[s];\n dot /= sqrt((float)p.n);\n float gate = sqrt(max(fabs(dot), 1e-6f)) * (dot < 0.0f ? -1.0f : 1.0f);\n const float s = 1.0f / (1.0f + exp(-gate));\n device const float* vb = value + (ulong)b * p.n;\n for (uint i = tid; i < p.n; i += tpg) gated[base + i] = s * vb[i];\n}\n\nstruct PleConvParams {\n uint channels; // hc_hidden\n uint ksize; // 4\n uint dilation; // 3\n uint span; // (kernel-1)*dilation history slots\n uint filled; // positions already recorded (= position of row 0)\n uint nb;\n};\n\n// Dilated causal depthwise conv over time with SiLU, added to `gated`,\n// for nb consecutive positions; each channel walks its rows in order and\n// records gvn into the ring, so in-batch taps read what earlier rows wrote.\nkernel void fn_ple_conv_b(\n device const float* gated [[buffer(0)]],\n device const float* gvn [[buffer(1)]],\n device const bfloat* w [[buffer(2)]],\n device float* hist [[buffer(3)]],\n device float* out [[buffer(4)]],\n constant PleConvParams& p [[buffer(5)]],\n uint c [[thread_position_in_grid]])\n{\n if (c >= p.channels) return;\n device const bfloat* wr = w + (ulong)c * p.ksize;\n for (uint b = 0; b < p.nb; b++) {\n const uint pos = p.filled + b;\n const ulong bc = (ulong)b * p.channels + c;\n float acc = (float)wr[p.ksize - 1] * gvn[bc];\n for (uint k = 0; k + 1 < p.ksize; k++) {\n const uint back = p.dilation * (p.ksize - 1 - k);\n if (back <= pos) {\n const uint slot = (pos - back) % p.span;\n acc += (float)wr[k] * hist[(ulong)slot * p.channels + c];\n }\n }\n out[bc] = gated[bc] + acc / (1.0f + exp(-acc));\n hist[(ulong)(pos % p.span) * p.channels + c] = gvn[bc];\n }\n}\n\n\n#line 1 \"qwen4_exp/mtp.metal\"\n// MTP input folding into the hyper-connection streams.\n\n// hyper[b][g][i] = fe[b][i] + fh[b*groups+g][i] (MTP input fold)\nkernel void fn_mtp_fold(\n device const float* fe [[buffer(0)]],\n device const float* fh [[buffer(1)]],\n device float* hyper [[buffer(2)]],\n constant GroupParams& p [[buffer(3)]],\n constant uint& nb [[buffer(4)]],\n uint gi [[thread_position_in_grid]])\n{\n const uint hh = p.n * p.groups;\n if (gi >= nb * hh) return;\n const uint b = gi / hh;\n const uint rem = gi % hh;\n const uint g = rem / p.n;\n const uint i = rem % p.n;\n hyper[gi] = fe[(ulong)b * p.n + i] + fh[((ulong)b * p.groups + g) * p.n + i];\n}\n\n\n#line 1 \"qwen4_exp/deltanet.metal\"\n// DeltaNet output normalization and sigmoid gating.\n\n// The host binds DeltaPrepParams here too; keep the full layout identical\n// to that Rust struct and the DeltaPrepParams in common/deltanet.metal.\nstruct GateNormParams {\n uint n_k;\n uint n_v;\n uint d_k;\n uint d_v;\n float eps;\n uint nb;\n uint snap_after;\n};\n\n// DeltaNet output: (norm(y) * w) * sigmoid(z), per (row, head).\nkernel void fn_delta_gate_norm_sigmoid_b(\n device float* y [[buffer(0)]],\n device const float* z [[buffer(1)]],\n device const bfloat* nw [[buffer(2)]],\n constant GateNormParams& p [[buffer(3)]],\n uint tgpos [[threadgroup_position_in_grid]],\n uint iv [[thread_position_in_threadgroup]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n threadgroup float shm[32];\n const uint dv = p.d_v;\n const uint b = tgpos / p.n_v;\n const uint head = tgpos % p.n_v;\n const ulong off = (ulong)b * p.n_v * dv + (ulong)head * dv + iv;\n float yv = y[off];\n float ss = simd_sum(yv * yv);\n if (lane == 0) shm[sgid] = ss;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float tot = 0.0f;\n for (uint i = 0; i < (dv + 31) / 32; i++) tot += shm[i];\n float inv = rsqrt(tot / (float)dv + p.eps);\n float zv = z[off];\n y[off] = (yv * inv * (float)nw[iv]) * (1.0f / (1.0f + exp(-zv)));\n}\n\n\n#line 1 \"qwen4_exp/qsa.metal\"\n// QSA block indexing, selection, and masked prefill/decode attention.\n\n// ---- QSA indexer: which 4-token blocks a query past the budget sees ----\n// Raw index keys (one per position) are cached; a block\'s key is the mean\n// of its `ratio` raw keys, RMS-normed and roped at the block\'s first\n// position. A query\'s `inh` index heads are normed and roped at its\n// position; block score = sum over heads of relu(q_h . key) / sqrt(ihd).\n// The top `k` blocks (ties: lowest index) plus the incomplete tail are\n// visible. Rows with at most k complete blocks attend densely.\nstruct IndexParams {\n uint ihd; // index head dim\n uint inh; // index query heads\n uint qk_dim; // index_qk projection width (inh*ihd + ihd)\n uint ratio; // tokens per block\n uint rot; // rotary dims\n float theta;\n float eps;\n uint base_pos;\n uint nb;\n uint b0; // first block to (re)compute\n uint b1; // one past the last\n uint k; // blocks kept (budget / ratio)\n uint max_blocks; // score row stride\n uint vis_stride; // visible token list row stride (k*ratio + ratio)\n uint mask_words; // block bitmask row stride (u32 words)\n};\n\n// ikc[base_pos + b][d] = iqk[b][inh*ihd + d]\nkernel void fn_index_append(\n device const float* iqk [[buffer(0)]],\n device float* ikc [[buffer(1)]],\n constant IndexParams& p [[buffer(2)]],\n uint gi [[thread_position_in_grid]])\n{\n if (gi >= p.nb * p.ihd) return;\n const uint b = gi / p.ihd;\n const uint d = gi % p.ihd;\n ikc[(ulong)(p.base_pos + b) * p.ihd + d] = iqk[(ulong)b * p.qk_dim + p.inh * p.ihd + d];\n}\n\n// RMSNorm of the `ihd` values held one per thread, times w, then partial\n// rope (rotate_half pairing over the first `rot` dims) at `pos`; result\n// written to dst. Threadgroup = ihd threads.\nstatic inline void fn_index_norm_rope(\n float v, device const bfloat* w, device float* dst, constant IndexParams& p, uint pos,\n threadgroup float* red, threadgroup float* vals, uint d, uint sgid, uint lane)\n{\n float ss = simd_sum(v * v);\n if (lane == 0) red[sgid] = ss;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float total = 0.0f;\n for (uint i = 0; i < (p.ihd + 31) / 32; i++) total += red[i];\n const float inv = rsqrt(total / (float)p.ihd + p.eps);\n const float val = v * inv * (float)w[d];\n vals[d] = val;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n const uint hr = p.rot / 2;\n if (d < hr) {\n const float inv_freq = pow(p.theta, -2.0f * (float)d / (float)p.rot);\n const float angle = (float)pos * inv_freq;\n const float c = cos(angle);\n const float s = sin(angle);\n const float a = vals[d];\n const float bb = vals[d + hr];\n dst[d] = a * c - bb * s;\n dst[d + hr] = bb * c + a * s;\n } else if (d >= p.rot) {\n dst[d] = val;\n }\n}\n\n// One threadgroup (ihd threads) per block in [b0, b1): its key from the\n// cached raw keys.\nkernel void fn_index_blocks(\n device const float* ikc [[buffer(0)]],\n device float* blk [[buffer(1)]],\n device const bfloat* w [[buffer(2)]],\n constant IndexParams& p [[buffer(3)]],\n uint tg [[threadgroup_position_in_grid]],\n uint d [[thread_position_in_threadgroup]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n threadgroup float red[32];\n threadgroup float vals[256];\n const uint b = p.b0 + tg;\n if (b >= p.b1) return;\n float acc = 0.0f;\n for (uint t = 0; t < p.ratio; t++) acc += ikc[(ulong)(b * p.ratio + t) * p.ihd + d];\n acc /= (float)p.ratio;\n fn_index_norm_rope(acc, w, blk + (ulong)b * p.ihd, p, b * p.ratio, red, vals, d, sgid, lane);\n}\n\n// One threadgroup (ihd threads) per (row, index head): the roped query.\nkernel void fn_index_q(\n device const float* iqk [[buffer(0)]],\n device float* iq [[buffer(1)]],\n device const bfloat* w [[buffer(2)]],\n constant IndexParams& p [[buffer(3)]],\n uint tg [[threadgroup_position_in_grid]],\n uint d [[thread_position_in_threadgroup]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n threadgroup float red[32];\n threadgroup float vals[256];\n const uint b = tg / p.inh;\n const uint h = tg % p.inh;\n const float v = iqk[(ulong)b * p.qk_dim + h * p.ihd + d];\n fn_index_norm_rope(v, w, iq + ((ulong)b * p.inh + h) * p.ihd, p, p.base_pos + b, red, vals, d, sgid, lane);\n}\n\n// score[b][j] = sum_h relu(iq[b][h] . blk[j]) / sqrt(ihd) for the complete\n// blocks of row b. Threadgroups of 256 blocks per row.\nkernel void fn_index_score(\n device const float* iq [[buffer(0)]],\n device const float* blk [[buffer(1)]],\n device float* score [[buffer(2)]],\n constant IndexParams& p [[buffer(3)]],\n uint tg [[threadgroup_position_in_grid]],\n uint tid [[thread_position_in_threadgroup]])\n{\n threadgroup float q[1024];\n const uint per = (p.max_blocks + 255) / 256;\n const uint b = tg / per;\n const uint j = (tg % per) * 256 + tid;\n const uint qn = p.inh * p.ihd;\n for (uint i = tid; i < qn; i += 256) q[i] = iq[(ulong)b * qn + i];\n threadgroup_barrier(mem_flags::mem_threadgroup);\n const uint blocks = (p.base_pos + b + 1) / p.ratio;\n if (j >= blocks) return;\n device const float* key = blk + (ulong)j * p.ihd;\n float acc = 0.0f;\n for (uint h = 0; h < p.inh; h++) {\n float dot = 0.0f;\n for (uint d = 0; d < p.ihd; d++) dot += q[h * p.ihd + d] * key[d];\n acc += max(dot, 0.0f);\n }\n score[(ulong)b * p.max_blocks + j] = acc / sqrt((float)p.ihd);\n}\n\n// Top-k blocks of a row (ties: lowest index) plus the incomplete tail, as\n// an ascending token list vis[b][..nvis[b]] and as a block bitmask\n// vmask[b] (for the prefill softmax). One 1024-thread threadgroup per\n// row; rows within the budget are left alone (dense attention). Radix\n// select over the score bits (scores are >= 0, so their bit patterns\n// order like the values), 4 bits per pass from the top.\nkernel void fn_index_select(\n device const float* score [[buffer(0)]],\n device uint* vis [[buffer(1)]],\n device uint* nvis [[buffer(2)]],\n constant IndexParams& p [[buffer(3)]],\n device uint* vmask [[buffer(4)]],\n uint b [[threadgroup_position_in_grid]],\n uint tid [[thread_position_in_threadgroup]],\n uint sgid [[simdgroup_index_in_threadgroup]],\n uint lane [[thread_index_in_simdgroup]])\n{\n threadgroup atomic_uint hist[16];\n threadgroup uint sh_digit;\n threadgroup uint sh_above;\n threadgroup uint sh_sum[32];\n const uint t_len = p.base_pos + b + 1;\n const uint n = t_len / p.ratio;\n if (n <= p.k) return;\n device const float* sc = score + (ulong)b * p.max_blocks;\n device uint* mrow = vmask + (ulong)b * p.mask_words;\n for (uint w = tid; w < p.mask_words; w += 1024) mrow[w] = 0u;\n threadgroup_barrier(mem_flags::mem_device);\n uint prefix = 0;\n uint mask = 0;\n uint remaining = p.k;\n for (int shift = 28; shift >= 0; shift -= 4) {\n if (tid < 16) atomic_store_explicit(&hist[tid], 0u, memory_order_relaxed);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n for (uint i = tid; i < n; i += 1024) {\n const uint key = as_type<uint>(sc[i]);\n if ((key & mask) == prefix) {\n atomic_fetch_add_explicit(&hist[(key >> shift) & 15u], 1u, memory_order_relaxed);\n }\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (tid == 0) {\n uint cum = 0;\n uint digit = 0;\n uint above = 0;\n for (int dg = 15; dg >= 0; dg--) {\n const uint c = atomic_load_explicit(&hist[dg], memory_order_relaxed);\n if (cum + c >= remaining) { digit = (uint)dg; above = cum; break; }\n cum += c;\n }\n sh_digit = digit;\n sh_above = above;\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n prefix |= sh_digit << shift;\n mask |= 0xFu << shift;\n remaining -= sh_above;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n }\n const uint T = prefix;\n // Contiguous index chunks per thread keep ascending order.\n const uint chunk = (n + 1023) / 1024;\n const uint i0 = min(tid * chunk, n);\n const uint i1 = min(i0 + chunk, n);\n uint local_gt = 0;\n uint local_eq = 0;\n for (uint i = i0; i < i1; i++) {\n const uint key = as_type<uint>(sc[i]);\n local_gt += key > T;\n local_eq += key == T;\n }\n // Exclusive prefix of local_eq across the threadgroup, and the totals.\n uint eq_pre = simd_prefix_exclusive_sum(local_eq);\n uint gt_sum = simd_sum(local_gt);\n if (lane == 31) sh_sum[sgid] = eq_pre + local_eq;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n uint eq_base = 0;\n for (uint s = 0; s < sgid; s++) eq_base += sh_sum[s];\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (lane == 0) sh_sum[sgid] = gt_sum;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n uint count_gt = 0;\n for (uint s = 0; s < 32; s++) count_gt += sh_sum[s];\n const uint need = p.k - count_gt; // keys == T taken lowest-index first\n uint eq_rank = eq_base + eq_pre;\n uint local_sel = 0;\n for (uint i = i0; i < i1; i++) {\n const uint key = as_type<uint>(sc[i]);\n if (key > T) local_sel++;\n else if (key == T) { if (eq_rank < need) local_sel++; eq_rank++; }\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n uint sel_pre = simd_prefix_exclusive_sum(local_sel);\n if (lane == 31) sh_sum[sgid] = sel_pre + local_sel;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n uint out = sel_pre;\n for (uint s = 0; s < sgid; s++) out += sh_sum[s];\n device uint* row = vis + (ulong)b * p.vis_stride;\n eq_rank = eq_base + eq_pre;\n for (uint i = i0; i < i1; i++) {\n const uint key = as_type<uint>(sc[i]);\n bool take = key > T;\n if (key == T) { take = eq_rank < need; eq_rank++; }\n if (take) {\n for (uint t = 0; t < p.ratio; t++) row[out * p.ratio + t] = i * p.ratio + t;\n atomic_fetch_or_explicit((device atomic_uint*)(mrow + (i >> 5)), 1u << (i & 31u), memory_order_relaxed);\n out++;\n }\n }\n if (tid == 0) {\n const uint tail0 = n * p.ratio;\n for (uint t = tail0; t < t_len; t++) row[p.k * p.ratio + (t - tail0)] = t;\n nvis[b] = p.k * p.ratio + (t_len - tail0);\n }\n}\n\nstruct FnAttnPartParams {\n uint n_heads;\n uint n_kv;\n uint head_dim;\n uint t_len; // visible tokens\n uint q_stride;\n uint q_off;\n uint max_blk;\n float scale;\n};\n\n// attn_part2_q8 (common/attention.metal) over a visible token list instead of the\n// contiguous prefix: split-T flash-decode partials over the q8 KV cache.\nkernel void fn_attn_part2_q8_sel(\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 FnAttnPartParams& p [[buffer(4)]],\n constant uint& n_wg [[buffer(5)]],\n constant uint& scale_off [[buffer(6)]],\n device const uint* vis [[buffer(7)]],\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 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 = vis[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\nstruct FnAttnGemmParams {\n uint kv_row;\n uint blocks_per_row;\n uint head;\n uint kl;\n uint kl_pad;\n uint scale_off;\n uint base; // context length before this query sub-chunk\n uint nb; // queries in the sub-chunk\n uint n_rep;\n float scale;\n};\n\nstruct SelMaskParams {\n uint row0; // first query\'s row in the block bitmask\n uint mask_words;\n uint ratio;\n uint k;\n};\n\n// Prefill row softmax with the QSA visibility mask: a query\n// with more than k complete blocks sees only its selected blocks (bit\n// set in vmask) and the incomplete tail.\nkernel void fn_attn_softmax_sel(\n device const float* S [[buffer(0)]], // [m][kl_pad]\n device half* P [[buffer(1)]], // [m][kl_pad]\n constant FnAttnGemmParams& p [[buffer(2)]],\n device const uint* vmask [[buffer(3)]],\n constant SelMaskParams& sm [[buffer(4)]],\n uint row [[threadgroup_position_in_grid]],\n uint tid [[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 float partial[32];\n device const float* r = S + (ulong)row * p.kl_pad;\n device half* o = P + (ulong)row * p.kl_pad;\n const uint t = row % p.nb;\n const uint e = p.base + t + 1; // causal extent\n const uint blocks = e / sm.ratio;\n const bool selective = blocks > sm.k;\n const uint tail0 = blocks * sm.ratio;\n device const uint* mrow = vmask + (ulong)(sm.row0 + t) * sm.mask_words;\n #define FN_VIS(i) (!selective || (i) >= tail0 || ((mrow[((i) / sm.ratio) >> 5] >> (((i) / sm.ratio) & 31u)) & 1u))\n float m = -INFINITY;\n for (uint i = tid; i < e; i += tpg) if (FN_VIS(i)) m = max(m, r[i]);\n m = simd_max(m);\n if (lane == 0) partial[sgid] = m;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (sgid == 0) {\n float v = (lane < (tpg + 31) / 32) ? partial[lane] : -INFINITY;\n v = simd_max(v);\n if (lane == 0) partial[0] = v;\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n m = partial[0];\n float s = 0.0f;\n for (uint i = tid; i < e; i += tpg) if (FN_VIS(i)) s += exp2(r[i] - m);\n s = simd_sum(s);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (lane == 0) partial[sgid] = s;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n if (sgid == 0) {\n float v = (lane < (tpg + 31) / 32) ? partial[lane] : 0.0f;\n v = simd_sum(v);\n if (lane == 0) partial[0] = v;\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n const float inv = 1.0f / partial[0];\n for (uint i = tid; i < e; i += tpg) o[i] = FN_VIS(i) ? (half)(exp2(r[i] - m) * inv) : 0.0h;\n for (uint i = e + tid; i < p.kl_pad; i += tpg) o[i] = 0.0h;\n #undef FN_VIS\n}\n\n\n#line 1 \"qwen4_exp/rows.metal\"\n// General row utilities and narrow projections used by batched prefill.\n\nkernel void fn_zero(\n device float* x [[buffer(0)]],\n constant uint& n [[buffer(1)]],\n uint i [[thread_position_in_grid]])\n{\n if (i < n) x[i] = 0.0f;\n}\n\n// ---- prefill engine helpers (row batches of any size) ----\n\n// Q4 projection with few output rows over any number of input rows,\n// straight from f32 x: one simdgroup per (row, output). For the narrow\n// weights (injection logits, DeltaNet gates) the GEMM tiles waste or\n// corrupt more than they compute.\nkernel void fn_qmv_small_b(\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 FnQmvParams& p [[buffer(5)]],\n constant uint& nb [[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 g = tgpos * spt + sgid;\n if (g >= nb * p.out_dim) return;\n const uint b = g / p.out_dim;\n const uint row = g % p.out_dim;\n const uint words = p.in_dim / 8;\n device const uint* wr = w + (ulong)row * words;\n device const bfloat* sr = scales + (ulong)row * (p.in_dim / 64);\n device const bfloat* br = biases + (ulong)row * (p.in_dim / 64);\n device const float* xb = x + (ulong)b * p.in_dim;\n float acc = 0.0f;\n for (uint wi = lane; wi < words; wi += 32) {\n const uint word = wr[wi];\n const uint grp = wi >> 3;\n float qd = 0.0f;\n float xs = 0.0f;\n for (uint i = 0; i < 8; i++) {\n const float xv = xb[wi * 8 + i];\n qd = fma((float)((word >> (4 * i)) & 0xFu), xv, qd);\n xs += xv;\n }\n acc = fma((float)sr[grp], qd, fma((float)br[grp], xs, acc));\n }\n acc = simd_sum(acc);\n if (lane == 0) y[(ulong)b * p.out_dim + row] = acc;\n}\n\n// x[i] = silu(x[i] / div)\nkernel void fn_silu_rows(\n device float* x [[buffer(0)]],\n constant uint& n [[buffer(1)]],\n constant float& div [[buffer(2)]],\n uint i [[thread_position_in_grid]])\n{\n if (i >= n) return;\n float v = x[i] / div;\n x[i] = v / (1.0f + exp(-v));\n}\n\n// out[r] = x[idx[r]] for n_rows rows of `width`\nkernel void fn_gather_rows(\n device const float* x [[buffer(0)]],\n device const uint* idx [[buffer(1)]],\n device float* out [[buffer(2)]],\n constant uint& n_rows [[buffer(3)]],\n constant uint& width [[buffer(4)]],\n uint gi [[thread_position_in_grid]])\n{\n if (gi >= n_rows * width) return;\n const uint r = gi / width;\n const uint d = gi % width;\n out[gi] = x[(ulong)idx[r] * width + d];\n}\n\n// out[idx[r]] += w[r] * y[r]; rows of one expert are distinct, so no\n// two threads touch the same element.\nkernel void fn_scatter_add_rows(\n device const float* y [[buffer(0)]],\n device const uint* idx [[buffer(1)]],\n device const float* w [[buffer(2)]],\n device float* out [[buffer(3)]],\n constant uint& n_rows [[buffer(4)]],\n constant uint& width [[buffer(5)]],\n uint gi [[thread_position_in_grid]])\n{\n if (gi >= n_rows * width) return;\n const uint r = gi / width;\n const uint d = gi % width;\n out[(ulong)idx[r] * width + d] += w[r] * y[gi];\n}\n\n// out[t] += sigmoid(gate . x[t]) * y[t]; one threadgroup per row.\nkernel void fn_shared_add_rows(\n device float* out [[buffer(0)]],\n device const float* y [[buffer(1)]],\n device const bfloat* gate [[buffer(2)]],\n device const float* x [[buffer(3)]],\n constant uint& width [[buffer(4)]],\n uint t [[threadgroup_position_in_grid]],\n uint tid [[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 float red[32];\n device const float* xr = x + (ulong)t * width;\n float acc = 0.0f;\n for (uint i = tid; i < width; i += tpg) acc += (float)gate[i] * xr[i];\n acc = simd_sum(acc);\n if (lane == 0) red[sgid] = acc;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n float dot = 0.0f;\n for (uint s = 0; s < (tpg + 31) / 32; s++) dot += red[s];\n const float g = 1.0f / (1.0f + exp(-dot));\n device float* o = out + (ulong)t * width;\n device const float* yr = y + (ulong)t * width;\n for (uint i = tid; i < width; i += tpg) o[i] += g * yr[i];\n}\n\n";