ML::GGUF

ALIGNMENT = 32

BERT_FUSED_SOURCE = "// FP16 intermediate kernels for BERT transformer\n// All intermediate buffers (hidden, QKV, attn, FFN) use half precision.\n// Accumulation in F32, final conversion to half on output.\n\n#include <metal_stdlib>\nusing namespace metal;\n\n#define FOR_UNROLL _Pragma(\"clang loop unroll(full)\")\n\n// ============================================================================\n// QKV split: [seq, 3*dim] half → Q, K half [n_heads, seq, head_dim]\n// V_t half [n_heads, head_dim, seq]\n// ============================================================================\nkernel void qkv_split(\n device const half* qkv [[buffer(0)]],\n device half* Q [[buffer(1)]],\n device half* K [[buffer(2)]],\n device half* V [[buffer(3)]], // original layout\n device half* V_t [[buffer(4)]], // transposed layout\n constant uint& seq_len [[buffer(5)]],\n constant uint& dim [[buffer(6)]],\n constant uint& n_heads [[buffer(7)]],\n constant uint& head_dim [[buffer(8)]],\n uint tid [[thread_position_in_grid]])\n{\n if (tid >= seq_len * dim) return;\n const uint pos = tid / dim;\n const uint d = tid % dim;\n const uint h = d / head_dim;\n const uint hd = d % head_dim;\n const uint src = pos * 3 * dim;\n const uint dst = h * seq_len * head_dim + pos * head_dim + hd;\n const uint vt = h * head_dim * seq_len + hd * seq_len + pos;\n half v_val = qkv[src + 2 * dim + d];\n Q[dst] = qkv[src + d];\n K[dst] = qkv[src + dim + d];\n V[dst] = v_val;\n V_t[vt] = v_val;\n}\n\n// ============================================================================\n// Fused QKV split + RoPE: [seq, 3*dim] → Q(rope'd), K(rope'd), V, V_t\n// Eliminates 3 dispatches (split + 2×rope) + 2 barriers per layer\n// Dispatch: dispatch_1d(seq_len * dim, 256)\n// ============================================================================\nkernel void qkv_split_rope(\n device const half* qkv [[buffer(0)]],\n device half* Q [[buffer(1)]],\n device half* K [[buffer(2)]],\n device half* V [[buffer(3)]],\n device half* V_t [[buffer(4)]],\n device const float* cos_t [[buffer(5)]],\n device const float* sin_t [[buffer(6)]],\n constant uint& seq_len [[buffer(7)]],\n constant uint& dim [[buffer(8)]],\n constant uint& n_heads [[buffer(9)]],\n constant uint& head_dim [[buffer(10)]],\n uint tid [[thread_position_in_grid]])\n{\n if (tid >= seq_len * dim) return;\n const uint pos = tid / dim;\n const uint d = tid % dim;\n const uint h = d / head_dim;\n const uint hd = d % head_dim;\n const uint hd2 = head_dim / 2;\n const uint src = pos * 3 * dim;\n const uint dst = h * seq_len * head_dim + pos * head_dim + hd;\n\n // V and V_t: straight copy (no RoPE)\n half v_val = qkv[src + 2 * dim + d];\n V[dst] = v_val;\n V_t[h * head_dim * seq_len + hd * seq_len + pos] = v_val;\n\n // Q and K: split + RoPE NeoX\n float q_raw = float(qkv[src + d]);\n float k_raw = float(qkv[src + dim + d]);\n\n if (hd < hd2) {\n // First half: v0 position\n float q_pair = float(qkv[src + h * head_dim + hd + hd2]); // wrong: need full dim offset\n float k_pair = float(qkv[src + dim + h * head_dim + hd + hd2]);\n // Actually the pair element is at the same head, same pos, but hd+hd2\n // qkv layout: [pos * 3*dim + d] where d = h*head_dim + hd\n // pair: d_pair = h*head_dim + hd + hd2\n uint d_pair = h * head_dim + hd + hd2;\n float q1 = float(qkv[src + d_pair]);\n float k1 = float(qkv[src + dim + d_pair]);\n float c = cos_t[pos * hd2 + hd];\n float s = sin_t[pos * hd2 + hd];\n Q[dst] = half(q_raw * c - q1 * s);\n K[dst] = half(k_raw * c - k1 * s);\n } else {\n // Second half: v1 position\n uint hd_lo = hd - hd2;\n uint d_pair = h * head_dim + hd_lo;\n float q0 = float(qkv[src + d_pair]);\n float k0 = float(qkv[src + dim + d_pair]);\n float c = cos_t[pos * hd2 + hd_lo];\n float s = sin_t[pos * hd2 + hd_lo];\n Q[dst] = half(q0 * s + q_raw * c);\n K[dst] = half(k0 * s + k_raw * c);\n }\n}\n\n// ============================================================================\n// RoPE NeoX in-place on half Q/K\n// ============================================================================\nkernel void rope_neox_inplace(\n device half* qk [[buffer(0)]],\n device const float* cos_t [[buffer(1)]],\n device const float* sin_t [[buffer(2)]],\n constant uint& seq_len [[buffer(3)]],\n constant uint& n_heads [[buffer(4)]],\n constant uint& head_dim [[buffer(5)]],\n uint tid [[thread_position_in_grid]])\n{\n if (tid >= seq_len * n_heads) return;\n const uint h = tid / seq_len;\n const uint pos = tid % seq_len;\n const uint hd2 = head_dim / 2;\n const uint base = h * (seq_len * head_dim) + pos * head_dim;\n const uint rope_off = pos * hd2;\n for (uint i = 0; i < hd2; i++) {\n float c = cos_t[rope_off + i];\n float s = sin_t[rope_off + i];\n float v0 = float(qk[base + i]);\n float v1 = float(qk[base + i + hd2]);\n qk[base + i] = half(v0 * c - v1 * s);\n qk[base + i + hd2] = half(v0 * s + v1 * c);\n }\n}\n\n// ============================================================================\n// Attention: shared scores + V_t float4, FP16 I/O\n// ============================================================================\nconstant uint N_QR = 8;\n\nkernel void attention_forward(\n device const half* Q [[buffer(0)]],\n device const half* K [[buffer(1)]],\n device const half* V_t [[buffer(2)]], // [n_heads, head_dim, seq] transposed\n device half* output [[buffer(3)]],\n constant uint& seq_len [[buffer(4)]],\n constant uint& n_heads [[buffer(5)]],\n constant uint& head_dim [[buffer(6)]],\n constant float& scale [[buffer(7)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]],\n threadgroup float* shared_base [[threadgroup(0)]])\n{\n const uint h = tgpig.x;\n const uint i = tgpig.y * N_QR + sgitg;\n if (h >= n_heads || i >= seq_len) return;\n\n const uint h_off = h * seq_len * head_dim;\n const uint lane = tiisg;\n const uint hd4 = head_dim / 4;\n threadgroup float* shared = shared_base + sgitg * seq_len;\n\n // Cache Q in registers as float (promote from half)\n device const half* qi = Q + h_off + i * head_dim;\n float qr[64]; // max head_dim\n for (uint d = 0; d < head_dim; d++) qr[d] = float(qi[d]);\n\n // Q·K dot products (half4 K reads)\n float local_max = -1e30f;\n for (uint j = lane; j < seq_len; j += 32) {\n device const half4* kj4 = (device const half4*)(K + h_off + j * head_dim);\n float dot = 0.0f;\n for (uint d4 = 0; d4 < hd4; d4++) {\n half4 k4 = kj4[d4];\n dot += qr[d4*4]*float(k4.x) + qr[d4*4+1]*float(k4.y) + qr[d4*4+2]*float(k4.z) + qr[d4*4+3]*float(k4.w);\n }\n float s = dot * scale;\n shared[j] = s;\n local_max = max(local_max, s);\n }\n\n float global_max = simd_max(local_max);\n float local_sum = 0.0f;\n for (uint j = lane; j < seq_len; j += 32) {\n float e = exp(shared[j] - global_max);\n shared[j] = e;\n local_sum += e;\n }\n float inv_sum = 1.0f / simd_sum(local_sum);\n for (uint j = lane; j < seq_len; j += 32) shared[j] *= inv_sum;\n simdgroup_barrier(mem_flags::mem_threadgroup);\n\n // V accumulation from V_t (half4 reads for bandwidth, F32 accumulate)\n const uint vt_h_off = h * head_dim * seq_len;\n for (uint d = lane; d < head_dim; d += 32) {\n device const half* vt_row = V_t + vt_h_off + d * seq_len;\n float val = 0.0f;\n uint j = 0;\n for (; j + 3 < seq_len; j += 4) {\n half4 v4 = *(device const half4*)(vt_row + j);\n val += shared[j]*float(v4.x) + shared[j+1]*float(v4.y) + shared[j+2]*float(v4.z) + shared[j+3]*float(v4.w);\n }\n for (; j < seq_len; j++) val += shared[j] * float(vt_row[j]);\n output[i * (n_heads * head_dim) + h * head_dim + d] = half(val);\n }\n}\n\n// ============================================================================\n// Fused residual + layernorm (half I/O, F32 compute)\n// ============================================================================\nkernel void residual_layernorm(\n device half* x [[buffer(0)]],\n device const half* y [[buffer(1)]],\n device const float* w [[buffer(2)]],\n device const float* b [[buffer(3)]],\n constant uint& dim [[buffer(4)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]])\n{\n const uint pos = tgpig.x;\n const uint lane = tiisg;\n device half* row = x + pos * dim;\n device const half* y_row = y + pos * dim;\n\n float local_sum = 0.0f;\n for (uint j = lane; j < dim; j += 32) {\n float v = float(row[j]) + float(y_row[j]);\n row[j] = half(v);\n local_sum += v;\n }\n float mean = simd_sum(local_sum) / float(dim);\n\n float local_var = 0.0f;\n for (uint j = lane; j < dim; j += 32) { float d = float(row[j]) - mean; local_var += d * d; }\n float inv_std = rsqrt(simd_sum(local_var) / float(dim) + 1e-5f);\n\n for (uint j = lane; j < dim; j += 32) {\n row[j] = half((float(row[j]) - mean) * inv_std * w[j] + b[j]);\n }\n}\n\n// Variant: reads f32 residual (for post-atomic-scatter MoE norm2, skips f32→f16 dispatch)\nkernel void residual_layernorm_f32(\n device half* x [[buffer(0)]],\n device const float* y_f32 [[buffer(1)]],\n device const float* w [[buffer(2)]],\n device const float* b [[buffer(3)]],\n constant uint& dim [[buffer(4)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]])\n{\n const uint pos = tgpig.x;\n const uint lane = tiisg;\n device half* row = x + pos * dim;\n\n float local_sum = 0.0f;\n for (uint j = lane; j < dim; j += 32) {\n float v = float(row[j]) + y_f32[pos * dim + j];\n row[j] = half(v);\n local_sum += v;\n }\n float mean = simd_sum(local_sum) / float(dim);\n\n float local_var = 0.0f;\n for (uint j = lane; j < dim; j += 32) { float d = float(row[j]) - mean; local_var += d * d; }\n float inv_std = rsqrt(simd_sum(local_var) / float(dim) + 1e-5f);\n\n for (uint j = lane; j < dim; j += 32) {\n row[j] = half((float(row[j]) - mean) * inv_std * w[j] + b[j]);\n }\n}\n\n// ============================================================================\n// Residual + layernorm with COPY: out = layernorm(x + y) — different output buffer\n// ============================================================================\nkernel void residual_layernorm_copy(\n device const half* x [[buffer(0)]], // input 1 (read-only)\n device const half* y [[buffer(1)]], // input 2 (residual)\n device half* out [[buffer(2)]], // output (different from x and y)\n device const float* w [[buffer(3)]],\n device const float* b [[buffer(4)]],\n constant uint& dim [[buffer(5)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]])\n{\n const uint pos = tgpig.x;\n const uint lane = tiisg;\n\n float local_sum = 0.0f;\n for (uint j = lane; j < dim; j += 32) {\n float v = float(x[pos * dim + j]) + float(y[pos * dim + j]);\n out[pos * dim + j] = half(v);\n local_sum += v;\n }\n float mean = simd_sum(local_sum) / float(dim);\n\n float local_var = 0.0f;\n for (uint j = lane; j < dim; j += 32) { float d = float(out[pos * dim + j]) - mean; local_var += d * d; }\n float inv_std = rsqrt(simd_sum(local_var) / float(dim) + 1e-5f);\n\n for (uint j = lane; j < dim; j += 32) {\n out[pos * dim + j] = half((float(out[pos * dim + j]) - mean) * inv_std * w[j] + b[j]);\n }\n}\n\n// ============================================================================\n// GELU in-place (half)\n// ============================================================================\nkernel void gelu_inplace(\n device half* x [[buffer(0)]],\n constant uint& count [[buffer(1)]],\n uint tid [[thread_position_in_grid]])\n{\n if (tid >= count) return;\n float v = float(x[tid]);\n if (v > 10.0f) { x[tid] = half(v); return; }\n if (v < -10.0f) { x[tid] = half(0); return; }\n float t = 0.7978845608f * (v + 0.044715f * v * v * v);\n x[tid] = half(0.5f * v * (1.0f + tanh(t)));\n}\n\n// ============================================================================\n// Fused SIMD gate + softmax + top-k + expert_count\n// 1 simdgroup (32 threads) per token — parallel gate matmul via simd_sum.\n// Dispatch: threadgroups = {seq_len, 1, 1}, threads = {32, 1, 1}\n// ============================================================================\nkernel void gate_softmax_topk_count(\n device const half* hidden [[buffer(0)]],\n device const float* gate_w [[buffer(1)]],\n device int* routing_ids [[buffer(2)]],\n device float* routing_wts [[buffer(3)]],\n device atomic_int* expert_counts [[buffer(4)]],\n constant uint& dim [[buffer(5)]],\n constant uint& n_experts [[buffer(6)]],\n constant uint& k [[buffer(7)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]])\n{\n const uint pos = tgpig.x;\n const uint lane = tiisg;\n\n // Gate matmul: 32 threads cooperatively compute 8 dot products via simd_sum\n float partial[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f};\n device const half* h_row = hidden + pos * dim;\n for (uint j = lane; j < dim; j += 32) {\n float hv = float(h_row[j]);\n for (uint e = 0; e < 8; e++) {\n partial[e] += hv * gate_w[e * dim + j];\n }\n }\n // Reduce across simdgroup → lane 0 gets final logits\n float logits[8];\n for (uint e = 0; e < 8; e++) {\n logits[e] = simd_sum(partial[e]);\n }\n\n // Softmax + top-k (lane 0 only, trivial for 8 elements)\n if (lane == 0) {\n float max_g = logits[0];\n for (uint e = 1; e < n_experts; e++) max_g = max(max_g, logits[e]);\n float sum_exp = 0.0f;\n for (uint e = 0; e < n_experts; e++) { logits[e] = exp(logits[e] - max_g); sum_exp += logits[e]; }\n float inv_sum = 1.0f / sum_exp;\n for (uint e = 0; e < n_experts; e++) logits[e] *= inv_sum;\n\n uint out_base = pos * k;\n for (uint i = 0; i < k; i++) {\n float best_p = -1.0f; int best_e = 0;\n for (uint e = 0; e < n_experts; e++) {\n if (logits[e] > best_p) { best_p = logits[e]; best_e = (int)e; }\n }\n routing_ids[out_base + i] = best_e;\n routing_wts[out_base + i] = best_p;\n logits[best_e] = -1.0f;\n atomic_fetch_add_explicit(&expert_counts[best_e], 1, memory_order_relaxed);\n }\n }\n}\n\n// ============================================================================\n// Gate matmul: hidden(half) @ gate_w(F32) → logits(F32)\n// ============================================================================\nkernel void gate_matmul(\n device const half* hidden [[buffer(0)]],\n device const float* gate_w [[buffer(1)]],\n device float* output [[buffer(2)]],\n constant uint& dim [[buffer(3)]],\n constant uint& n_experts [[buffer(4)]],\n uint2 gid [[thread_position_in_grid]])\n{\n const uint e = gid.x;\n const uint pos = gid.y;\n float sum = 0.0f;\n for (uint j = 0; j < dim; j++) {\n sum += float(hidden[pos * dim + j]) * gate_w[e * dim + j];\n }\n output[pos * n_experts + e] = sum;\n}\n\n// ============================================================================\n// Softmax + Top-K (same as before — operates on F32 logits)\n// ============================================================================\nkernel void softmax_topk(\n device const float* gate_logits [[buffer(0)]],\n device int* routing_ids [[buffer(1)]],\n device float* routing_wts [[buffer(2)]],\n constant uint& n_experts [[buffer(3)]],\n constant uint& k [[buffer(4)]],\n uint tid [[thread_position_in_grid]])\n{\n device const float* row = gate_logits + tid * n_experts;\n float max_g = row[0];\n for (uint e = 1; e < n_experts; e++) max_g = max(max_g, row[e]);\n float sum_exp = 0.0f;\n float probs[8];\n for (uint e = 0; e < n_experts; e++) { probs[e] = exp(row[e] - max_g); sum_exp += probs[e]; }\n float inv_sum = 1.0f / sum_exp;\n for (uint e = 0; e < n_experts; e++) probs[e] *= inv_sum;\n uint out_base = tid * k;\n for (uint i = 0; i < k; i++) {\n float best_p = -1.0f; int best_e = 0;\n for (uint e = 0; e < n_experts; e++) {\n if (probs[e] > best_p) { best_p = probs[e]; best_e = (int)e; }\n }\n routing_ids[out_base + i] = best_e;\n routing_wts[out_base + i] = best_p;\n probs[best_e] = -1.0f;\n }\n}\n\n// ============================================================================\n// Zero int32 buffer\n// ============================================================================\nkernel void zero_int(\n device int* x [[buffer(0)]],\n uint tid [[thread_position_in_grid]])\n{\n x[tid] = 0;\n}\n\n// ============================================================================\n// GPU MoE routing: build gather_map + scatter_wts from routing_ids/wts\n// Uses atomic counters per expert to build contiguous expert groups.\n// Grid: [seq_len] (one thread per token position)\n// ============================================================================\n\nkernel void moe_build_routing(\n device const int* routing_ids [[buffer(0)]], // [seq, k] expert indices\n device const float* routing_wts [[buffer(1)]], // [seq, k] weights\n device int* gather_map [[buffer(2)]], // [total_routing] output: pos indices\n device float* scatter_wts [[buffer(3)]], // [total_routing] output: weights\n device int* expert_offsets [[buffer(4)]], // [n_experts+1] prefix sums (pre-computed)\n device atomic_int* expert_counts [[buffer(5)]], // [n_experts] atomic counters\n constant uint& k [[buffer(6)]], // n_experts_used (=2)\n constant uint& n_experts [[buffer(7)]],\n uint tid [[thread_position_in_grid]])\n{\n const uint pos = tid;\n for (uint ki = 0; ki < k; ki++) {\n int ei = routing_ids[pos * k + ki];\n float w = routing_wts[pos * k + ki];\n // Atomically get slot within this expert's group\n int slot = atomic_fetch_add_explicit(&expert_counts[ei], 1, memory_order_relaxed);\n int dest = expert_offsets[ei] + slot;\n gather_map[dest] = (int)pos;\n scatter_wts[dest] = w;\n }\n}\n\n// ============================================================================\n// Count tokens per expert (for prefix sum)\n// Grid: [seq_len]\n// ============================================================================\n\nkernel void moe_count_experts(\n device const int* routing_ids [[buffer(0)]],\n device atomic_int* expert_counts [[buffer(1)]],\n constant uint& k [[buffer(2)]],\n uint tid [[thread_position_in_grid]])\n{\n for (uint ki = 0; ki < k; ki++) {\n int ei = routing_ids[tid * k + ki];\n atomic_fetch_add_explicit(&expert_counts[ei], 1, memory_order_relaxed);\n }\n}\n\n// ============================================================================\n// Prefix sum for expert offsets (single-thread, tiny: 8 experts)\n// Grid: [1]\n// ============================================================================\n\nkernel void moe_prefix_sum(\n device const int* expert_counts [[buffer(0)]],\n device int* expert_offsets [[buffer(1)]],\n constant uint& n_experts [[buffer(2)]],\n uint tid [[thread_position_in_grid]])\n{\n int sum = 0;\n for (uint e = 0; e < n_experts; e++) {\n expert_offsets[e] = sum;\n sum += expert_counts[e];\n }\n expert_offsets[n_experts] = sum;\n}\n\n// ============================================================================\n// Fused: prefix_sum + zero_counts + build_routing + write_dispatch_args\n// One dispatch replaces 4. Thread 0 computes prefix sum, then all build routing.\n// Dispatch: {seq_len, 1, 1} threadgroups, {32, 1, 1} threads (or 1 TG if seq<32)\n// ============================================================================\nkernel void moe_route_and_dispatch(\n device const int* routing_ids [[buffer(0)]],\n device const float* routing_wts [[buffer(1)]],\n device int* gather_map [[buffer(2)]],\n device float* scatter_wts [[buffer(3)]],\n device int* expert_counts [[buffer(4)]], // input: counts from gate kernel\n device int* expert_offsets [[buffer(5)]], // output: prefix sums\n device uint* dispatch_args [[buffer(6)]], // output: indirect dispatch args\n constant uint& k [[buffer(7)]],\n constant uint& n_experts [[buffer(8)]],\n constant uint& seq_len [[buffer(9)]],\n constant uint& up_out_dim [[buffer(10)]],\n constant uint& down_out_dim [[buffer(11)]],\n constant uint& dim [[buffer(12)]],\n uint tid [[thread_position_in_grid]],\n uint tiisg [[thread_index_in_simdgroup]])\n{\n // Step 1: Thread 0 computes prefix sum + dispatch args + zeros counts\n if (tid == 0) {\n int sum = 0;\n for (uint e = 0; e < n_experts; e++) {\n int ec = atomic_load_explicit((device atomic_int*)&expert_counts[e], memory_order_relaxed);\n expert_offsets[e] = sum;\n atomic_store_explicit((device atomic_int*)&expert_counts[e], 0, memory_order_relaxed);\n sum += ec;\n\n uint eb_u = (uint)max(ec, 0);\n // UP dispatch args\n dispatch_args[e * 3 + 0] = (eb_u + 31) / 32;\n dispatch_args[e * 3 + 1] = (up_out_dim + 63) / 64;\n dispatch_args[e * 3 + 2] = 1;\n // DOWN dispatch args\n dispatch_args[(n_experts + e) * 3 + 0] = (eb_u + 31) / 32;\n dispatch_args[(n_experts + e) * 3 + 1] = (down_out_dim + 63) / 64;\n dispatch_args[(n_experts + e) * 3 + 2] = 1;\n }\n expert_offsets[n_experts] = sum;\n }\n\n // All threads wait for prefix sum + zeros to complete\n threadgroup_barrier(mem_flags::mem_device);\n\n // Step 2: All threads build routing (same as moe_build_routing)\n if (tid < seq_len) {\n for (uint ki = 0; ki < k; ki++) {\n int ei = routing_ids[tid * k + ki];\n float w = routing_wts[tid * k + ki];\n int slot = atomic_fetch_add_explicit(\n (device atomic_int*)&expert_counts[ei], 1, memory_order_relaxed);\n int dest = expert_offsets[ei] + slot;\n gather_map[dest] = (int)tid;\n scatter_wts[dest] = w;\n }\n }\n}\n\n// ============================================================================\n// Write indirect dispatch args from expert_offsets (GPU-side)\n// Grid: [n_experts]\n// Layout: [8 UP args, 8 DOWN args, 8 scatter args] x 3 uint32 each\n// ============================================================================\n// Write indirect dispatch args for MoE mm kernels (simdgroup_matrix)\n// Grid for mm: {ceil(eb/32), ceil(out_dim/64), 1}\n// Grid for scatter: {ceil(eb*dim/256), 1, 1}\nkernel void moe_write_dispatch_args(\n device const int* expert_offsets [[buffer(0)]],\n device uint* dispatch_args [[buffer(1)]],\n constant uint& up_out_dim [[buffer(2)]], // ffn_dim (for UP matmul)\n constant uint& down_out_dim [[buffer(3)]], // dim (for DOWN matmul)\n constant uint& dim [[buffer(4)]], // dim (for scatter)\n constant uint& n_experts [[buffer(5)]],\n uint tid [[thread_position_in_grid]])\n{\n if (tid >= n_experts) return;\n int eb = expert_offsets[tid + 1] - expert_offsets[tid];\n uint eb_u = (uint)max(eb, 0);\n\n // UP mm args: {ceil(eb/32), ceil(ffn_dim/64), 1}\n uint up_idx = tid * 3;\n dispatch_args[up_idx + 0] = (eb_u + 31) / 32;\n dispatch_args[up_idx + 1] = (up_out_dim + 63) / 64;\n dispatch_args[up_idx + 2] = 1;\n\n // DOWN mm args: {ceil(eb/32), ceil(dim/64), 1}\n uint down_idx = (n_experts + tid) * 3;\n dispatch_args[down_idx + 0] = (eb_u + 31) / 32;\n dispatch_args[down_idx + 1] = (down_out_dim + 63) / 64;\n dispatch_args[down_idx + 2] = 1;\n\n // Scatter args: {ceil(eb*dim/256), 1, 1}\n uint sc_idx = (2 * n_experts + tid) * 3;\n uint sc_threads = eb_u * dim;\n dispatch_args[sc_idx + 0] = (sc_threads + 255) / 256;\n dispatch_args[sc_idx + 1] = 1;\n dispatch_args[sc_idx + 2] = 1;\n}\n\n// ============================================================================\n// Batched expert dispatch args: compute expert_tg_offsets (prefix sum of per-expert TG counts)\n// + total grid size for batched_mm kernels. Single-thread kernel.\n// Grid: [1]\n// ============================================================================\nkernel void moe_write_batched_args(\n device const int* expert_offsets [[buffer(0)]], // [n_experts+1] token offsets\n device int* expert_tg_offs [[buffer(1)]], // [n_experts+1] output: TG prefix sums\n device uint* up_grid [[buffer(2)]], // [3] output: indirect dispatch args for batched UP\n device uint* down_grid [[buffer(3)]], // [3] output: indirect dispatch args for batched DOWN\n constant uint& n_experts [[buffer(4)]],\n constant uint& up_out_dim [[buffer(5)]], // ffn_dim\n constant uint& down_out_dim [[buffer(6)]], // dim\n uint tid [[thread_position_in_grid]])\n{\n int sum_tg = 0;\n for (uint e = 0; e < n_experts; e++) {\n expert_tg_offs[e] = sum_tg;\n int eb = expert_offsets[e + 1] - expert_offsets[e];\n sum_tg += (max(eb, 0) + 31) / 32; // ceil(eb / MM_NR1=32)\n }\n expert_tg_offs[n_experts] = sum_tg;\n // UP indirect dispatch: {total_batch_tgs, ceil(ffn_dim/64), 1}\n up_grid[0] = (uint)sum_tg;\n up_grid[1] = (up_out_dim + 63) / 64;\n up_grid[2] = 1;\n // DOWN indirect dispatch: {total_batch_tgs, ceil(dim/64), 1}\n down_grid[0] = (uint)sum_tg;\n down_grid[1] = (down_out_dim + 63) / 64;\n down_grid[2] = 1;\n}\n\n// ============================================================================\n// MoE gather: hidden(half) → moe_input(half)\n// ============================================================================\nkernel void moe_gather(\n device const half* hidden [[buffer(0)]],\n device half* moe_input [[buffer(1)]],\n device const int* gather_map [[buffer(2)]],\n constant uint& dim [[buffer(3)]],\n uint tid [[thread_position_in_grid]])\n{\n const uint ri = tid / dim;\n const uint j = tid % dim;\n const uint pos = (uint)gather_map[ri];\n moe_input[ri * dim + j] = hidden[pos * dim + j];\n}\n\n// ============================================================================\n// Scatter weighted add (half)\n// ============================================================================\nkernel void scatter_weighted_add(\n device half* ffn_out [[buffer(0)]],\n device const half* expert_out [[buffer(1)]],\n device const int* scatter_map [[buffer(2)]],\n device const float* weights [[buffer(3)]],\n constant uint& dim [[buffer(4)]],\n uint tid [[thread_position_in_grid]])\n{\n const uint ri = tid / dim;\n const uint j = tid % dim;\n const uint pos = (uint)scatter_map[ri];\n ffn_out[pos * dim + j] = half(float(ffn_out[pos * dim + j]) + weights[ri] * float(expert_out[ri * dim + j]));\n}\n\n// ============================================================================\n// Atomic MoE scatter — ALL routing slots in ONE dispatch, no sequential barriers\n// Uses atomic_fetch_add on float buffer (Metal 3). After scatter, convert f32→f16.\n// Dispatch: dispatch_1d(total_routing * dim, 256)\n// ============================================================================\nkernel void moe_scatter_atomic(\n device atomic_float* ffn_out_f32 [[buffer(0)]], // float accumulator (zeroed before)\n device const half* expert_out [[buffer(1)]], // packed expert output\n device const int* gather_map [[buffer(2)]], // [total_routing] → token index\n device const float* scatter_wts [[buffer(3)]], // [total_routing] weights\n constant uint& dim [[buffer(4)]],\n constant uint& total_routing [[buffer(5)]],\n uint tid [[thread_position_in_grid]])\n{\n if (tid >= total_routing * dim) return;\n const uint ri = tid / dim;\n const uint d = tid % dim;\n const int pos = gather_map[ri];\n const float val = scatter_wts[ri] * float(expert_out[ri * dim + d]);\n atomic_fetch_add_explicit(&ffn_out_f32[pos * dim + d], val, memory_order_relaxed);\n}\n\n// Convert float buffer → half buffer (after atomic scatter)\n// Dispatch: dispatch_1d(count, 256)\nkernel void f32_to_f16(\n device const float* src [[buffer(0)]],\n device half* dst [[buffer(1)]],\n uint tid [[thread_position_in_grid]])\n{\n dst[tid] = half(src[tid]);\n}\n\n// ============================================================================\n// MoE scatter — GPU-side expert_offsets bounds, zero CPU sync\n// Dispatch: dispatch_1d(seq_len * dim, 256) per expert (excess threads exit)\n// ============================================================================\nkernel void scatter_weighted_add_moe(\n device half* ffn_out [[buffer(0)]],\n device const half* expert_out [[buffer(1)]],\n device const int* scatter_map [[buffer(2)]],\n device const float* weights [[buffer(3)]],\n device const int* expert_offs [[buffer(4)]],\n constant uint& expert_id [[buffer(5)]],\n constant uint& dim [[buffer(6)]],\n uint tid [[thread_position_in_grid]])\n{\n const int base = expert_offs[expert_id];\n const int eb = expert_offs[expert_id + 1] - base;\n const int total = eb * (int)dim;\n if ((int)tid >= total) return;\n\n const uint ri = base + tid / dim;\n const uint j = tid % dim;\n const uint pos = (uint)scatter_map[ri];\n ffn_out[pos * dim + j] = half(float(ffn_out[pos * dim + j]) + weights[ri] * float(expert_out[ri * dim + j]));\n}\n\n// ============================================================================\n// MoE weighted scatter (half) — for sync-free path\n// ============================================================================\nkernel void moe_weighted_scatter(\n device half* ffn_out [[buffer(0)]],\n device const half* expert_out [[buffer(1)]],\n device const int* routing_ids [[buffer(2)]],\n device const float* routing_wts [[buffer(3)]],\n constant uint& dim [[buffer(4)]],\n constant uint& seq_len [[buffer(5)]],\n constant uint& k [[buffer(6)]],\n constant uint& n_experts [[buffer(7)]],\n uint tid [[thread_position_in_grid]])\n{\n const uint pos = tid / dim;\n const uint j = tid % dim;\n if (pos >= seq_len) return;\n float sum = 0.0f;\n for (uint ki = 0; ki < k; ki++) {\n const uint ei = (uint)routing_ids[pos * k + ki];\n sum += routing_wts[pos * k + ki] * float(expert_out[ei * seq_len * dim + pos * dim + j]);\n }\n ffn_out[pos * dim + j] = half(sum);\n}\n\n// ============================================================================\n// Residual add (half)\n// ============================================================================\nkernel void residual_add(\n device half* x [[buffer(0)]],\n device const half* y [[buffer(1)]],\n uint tid [[thread_position_in_grid]])\n{\n x[tid] = half(float(x[tid]) + float(y[tid]));\n}\n\n// ============================================================================\n// Zero region (half)\n// ============================================================================\nkernel void zero_region(\n device half* x [[buffer(0)]],\n constant uint& off [[buffer(1)]],\n uint tid [[thread_position_in_grid]])\n{\n x[off + tid] = half(0);\n}\n\n// ============================================================================\n// Weighted add (half) — for sync-based MoE path\n// ============================================================================\nkernel void weighted_add(\n device half* dst [[buffer(0)]],\n device const half* src [[buffer(1)]],\n constant uint& pos_offset [[buffer(2)]],\n constant float& weight [[buffer(3)]],\n uint tid [[thread_position_in_grid]])\n{\n dst[pos_offset + tid] = half(float(dst[pos_offset + tid]) + weight * float(src[tid]));\n}\n\n// ============================================================================\n// Mean pool + L2 normalize (half input → F32 output)\n// ============================================================================\nkernel void mean_pool_l2(\n device const half* hidden [[buffer(0)]],\n device float* output [[buffer(1)]],\n constant uint& seq_len [[buffer(2)]],\n constant uint& dim [[buffer(3)]],\n uint tid [[thread_position_in_grid]])\n{\n // Mean pool\n for (uint d = 0; d < dim; d++) {\n float sum = 0.0f;\n for (uint p = 0; p < seq_len; p++) sum += float(hidden[p * dim + d]);\n output[d] = sum / float(seq_len);\n }\n // L2 normalize\n float norm = 0.0f;\n for (uint d = 0; d < dim; d++) norm += output[d] * output[d];\n norm = rsqrt(norm + 1e-12f);\n for (uint d = 0; d < dim; d++) output[d] *= norm;\n}\n\n// ============================================================================\n// LayerNorm in-place (half, for embedding norm before layers)\n// ============================================================================\nkernel void layernorm_inplace(\n device half* x [[buffer(0)]],\n device const float* w [[buffer(1)]],\n device const float* b [[buffer(2)]],\n constant uint& dim [[buffer(3)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]])\n{\n const uint pos = tgpig.x;\n const uint lane = tiisg;\n device half* row = x + pos * dim;\n\n float local_sum = 0.0f;\n for (uint j = lane; j < dim; j += 32) local_sum += float(row[j]);\n float mean = simd_sum(local_sum) / float(dim);\n\n float local_var = 0.0f;\n for (uint j = lane; j < dim; j += 32) { float d = float(row[j]) - mean; local_var += d * d; }\n float inv_std = rsqrt(simd_sum(local_var) / float(dim) + 1e-5f);\n\n for (uint j = lane; j < dim; j += 32) {\n row[j] = half((float(row[j]) - mean) * inv_std * w[j] + b[j]);\n }\n}\n"

FLASH_ATTN_SOURCE = "// Flash Attention with online softmax + shared-memory tile scores\n// Processes K/V in tiles of 32. Each tile's scores stored in small shared buffer.\n// No O(n) shared memory — only O(32) per simdgroup.\n//\n// Dispatch: threadgroups = [n_heads, ceil(seq_len / N_FA_ROWS)]\n// threads_per_threadgroup = [32, N_FA_ROWS]\n// shared memory = N_FA_ROWS * 32 * sizeof(float)\n\n#include <metal_stdlib>\nusing namespace metal;\n\nconstant uint N_FA_ROWS = 4; // simdgroups per threadgroup\n\nkernel void attention_flash(\n device const float* Q [[buffer(0)]],\n device const float* K [[buffer(1)]],\n device const float* V_t [[buffer(2)]], // [n_heads, head_dim, seq_len] TRANSPOSED\n device float* output [[buffer(3)]],\n constant uint& seq_len [[buffer(4)]],\n constant uint& n_heads [[buffer(5)]],\n constant uint& head_dim [[buffer(6)]],\n constant float& scale [[buffer(7)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]],\n threadgroup float* shared_base [[threadgroup(0)]])\n{\n const uint h = tgpig.x;\n const uint i = tgpig.y * N_FA_ROWS + sgitg;\n if (h >= n_heads || i >= seq_len) return;\n\n const uint h_off = h * seq_len * head_dim;\n const uint lane = tiisg;\n const uint hd4 = head_dim / 4;\n\n // Per-simdgroup shared tile scores (32 floats)\n threadgroup float* tile_scores = shared_base + sgitg * 32;\n\n // Cache Q in registers\n device const float4* qi4 = (device const float4*)(Q + h_off + i * head_dim);\n float4 qr[16];\n for (uint d = 0; d < hd4; d++) qr[d] = qi4[d];\n\n // Online softmax state\n float m = -1e30f;\n float l = 0.0f;\n float o[2] = {0.0f, 0.0f}; // 2 output dims per lane\n\n const uint vt_h_off = h * head_dim * seq_len;\n\n for (uint tile_start = 0; tile_start < seq_len; tile_start += 32) {\n uint j = tile_start + lane;\n\n // Q·K dot product\n float score = -1e30f;\n if (j < seq_len) {\n device const float4* kj4 = (device const float4*)(K + h_off + j * head_dim);\n float dot = 0.0f;\n for (uint d = 0; d < hd4; d++) {\n dot += metal::dot(qr[d], kj4[d]);\n }\n score = dot * scale;\n }\n\n // Online softmax\n float tile_max = simd_max(score);\n float m_new = max(m, tile_max);\n float correction = exp(m - m_new);\n float p = (j < seq_len) ? exp(score - m_new) : 0.0f;\n l = l * correction + simd_sum(p);\n\n // Store normalized p in shared tile (32 floats — tiny)\n tile_scores[lane] = p;\n simdgroup_barrier(mem_flags::mem_threadgroup);\n\n // V accumulation: lanes split over head_dim, read tile_scores for all 32 keys\n for (uint dl = 0; dl < 2; dl++) {\n uint d = lane + dl * 32;\n if (d >= head_dim) continue;\n device const float* vt_row = V_t + vt_h_off + d * seq_len + tile_start;\n float acc = 0.0f;\n uint tile_end = min(tile_start + 32, seq_len) - tile_start;\n for (uint s = 0; s < tile_end; s++) {\n acc += tile_scores[s] * vt_row[s];\n }\n o[dl] = o[dl] * correction + acc;\n }\n simdgroup_barrier(mem_flags::mem_threadgroup);\n\n m = m_new;\n }\n\n // Finalize: output = o / l\n float inv_l = 1.0f / l;\n uint out_base = i * (n_heads * head_dim) + h * head_dim;\n for (uint dl = 0; dl < 2; dl++) {\n uint d = lane + dl * 32;\n if (d < head_dim) {\n output[out_base + d] = o[dl] * inv_l;\n }\n }\n}\n"

GEMM_MM_F16_SOURCE = "// FP16×FP16 Matrix-Matrix GEMM using simdgroup_matrix_multiply_accumulate\n// Pre-dequantized weights — no Q5K/Q6K dequant in the hot loop.\n//\n// Each threadgroup computes a 64×32 output tile.\n// Weights are FP16 in row-major layout: w[out_dim, in_dim]\n// Input is FP16: x[batch, in_dim]\n// Output is FP16: output[batch, out_dim] (with F32 bias + optional GELU)\n//\n// Dispatch: threadgroups = [ceil(batch/32), ceil(out_dim/64), 1]\n// threads_per_threadgroup = [128, 1, 1]\n// threadgroup_memory = 8192 bytes\n\n#include <metal_stdlib>\n#include <metal_simdgroup_matrix>\nusing namespace metal;\n\n#define FOR_UNROLL _Pragma(\"clang loop unroll(full)\")\n\nconstant int MM_NR0 = 64; // output rows per threadgroup\nconstant int MM_NR1 = 32; // batch elements per threadgroup\nconstant int MM_NK = 32; // K elements per iteration\n\nkernel void simd_mm_f16(\n device const half* w [[buffer(0)]], // pre-dequantized FP16 [out_dim, in_dim]\n device const half* x [[buffer(1)]], // FP16 input [batch, in_dim]\n device const float* bias [[buffer(2)]], // F32 bias [out_dim]\n device half* output [[buffer(3)]], // FP16 output [batch, out_dim]\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n // Shared memory: sa for weights [NR0, NK], sb for input [NR1, NK]\n // Layout: 64 elements per row (8-aligned for simdgroup_load)\n threadgroup half * sa = (threadgroup half *)(shmem); // 64*64 = 4096 half = 8192 bytes? No.\n threadgroup half * sb = (threadgroup half *)(shmem + 4096); // 4096 bytes for B\n\n const int r0 = tgpig.y * MM_NR0; // first output row\n const int r1 = tgpig.x * MM_NR1; // first batch element\n\n const short nr0 = min(MM_NR0, (int)out_dim - r0);\n const short nr1 = min(MM_NR1, (int)batch - r1);\n\n // Thread decomposition for loading\n // 128 threads load NR0×NK = 64×32 = 2048 weight elements (16 per thread)\n // And NR1×NK = 32×32 = 1024 input elements (8 per thread)\n const short lr0 = min((short)(tiitg / 2), (short)(nr0 - 1)); // weight row (0..63)\n const short lr1 = min((short)(tiitg / 4), (short)(nr1 - 1)); // input row (0..31)\n\n // Weight pointer for this thread's row\n device const half * w_row = w + (r0 + lr0) * in_dim;\n // Input pointer for this thread's row\n device const half * x_row = x + (r1 + lr1) * in_dim;\n\n // Simdgroup matrices\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 // Main K-loop\n for (uint k = 0; k < in_dim; k += MM_NK) {\n // Load weight tile [NR0, NK] into sa\n // 128 threads, each loads 16 elements (2 rows × 8 elements)\n threadgroup_barrier(mem_flags::mem_threadgroup);\n {\n const short row_in_tile = tiitg / 2; // 0..63\n const short col_start = (tiitg % 2) * 16; // 0 or 16\n if (row_in_tile < nr0) {\n device const half * src = w + (r0 + row_in_tile) * in_dim + k + col_start;\n threadgroup half * dst = sa + 64 * (row_in_tile / 8) + 8 * (row_in_tile % 8);\n // Store in simdgroup-friendly layout: 64 * (row/8) + 8 * (row%8) + col_block*64\n // Actually, use the same layout as llama.cpp for simdgroup_load compatibility\n for (short j = 0; j < 16; j++) {\n short sx = col_start/8 + j/8; // 0..3 (K block of 8)\n short sy = row_in_tile / 8; // 0..7 (row block of 8)\n short lx = row_in_tile % 8; // 0..7 (row within block)\n short ly = j % 8; // 0..7 (col within K-block of 8)\n short ib = 8 * sx + sy;\n *(sa + 64*ib + 8*ly + lx) = (k + col_start + j < in_dim) ? src[j] : half(0);\n }\n }\n }\n\n // Load input tile [NR1, NK] into sb\n {\n const short row_in_tile = tiitg / 4; // 0..31\n const short col_start = (tiitg % 4) * 8; // 0, 8, 16, 24\n if (row_in_tile < nr1) {\n device const half * src = x + (r1 + row_in_tile) * in_dim + k + col_start;\n short sx = col_start / 8; // 0..3\n short sy = row_in_tile / 8; // 0..3\n short ly = row_in_tile % 8; // 0..7\n short ib = 4*sx + sy;\n *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = (k + col_start < in_dim) ?\n *(device const half2x4 *)src : half2x4(0);\n }\n }\n\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n // Compute: 4 simdgroups, each handles a 32×16 quadrant of the 64×32 tile\n threadgroup const half * lsma = sa + 4*64*(sgitg % 2);\n threadgroup const half * lsmb = sb + 2*64*(sgitg / 2);\n\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL 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_UNROLL 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_UNROLL for (short i = 0; i < 8; i++) {\n simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);\n }\n lsma += 8*64;\n lsmb += 4*64;\n }\n }\n\n // Store results to shared memory, apply bias + GELU, write FP16 output\n threadgroup float * temp = (threadgroup float *)shmem;\n threadgroup float * sg_out = temp + 32*(sgitg & 1) + 16*(sgitg >> 1)*MM_NR0;\n for (short i = 0; i < 8; i++) {\n simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false);\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n {\n const int total_out = nr0 * nr1;\n for (int idx = (int)tiitg; idx < total_out; idx += 128) {\n const int i = idx % nr0;\n const int j = idx / nr0;\n float val = temp[j * MM_NR0 + i] + bias[r0 + i];\n if (apply_gelu) {\n if (val > 10.0f) { }\n else if (val < -10.0f) { val = 0.0f; }\n else {\n float t = 0.7978845608f * (val + 0.044715f * val * val * val);\n val = 0.5f * val * (1.0f + tanh(t));\n }\n }\n output[(r1 + j) * out_dim + r0 + i] = half(val);\n }\n }\n}\n\n// MoE variant — reads expert_offsets from GPU for base pointer computation\nkernel void simd_mm_f16_moe(\n device const half* w [[buffer(0)]], // weights for THIS expert\n device const half* x_packed [[buffer(1)]], // full packed input\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]], // full packed output\n device const int* expert_offs [[buffer(4)]],\n constant uint& expert_id [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half * sa = (threadgroup half *)(shmem);\n threadgroup half * sb = (threadgroup half *)(shmem + 4096);\n\n const int base = expert_offs[expert_id];\n const int eb = expert_offs[expert_id + 1] - base;\n if (eb <= 0) return;\n\n const int r0 = tgpig.y * MM_NR0;\n const int r1 = tgpig.x * MM_NR1;\n\n const short nr0 = min(MM_NR0, (int)out_dim - r0);\n const short nr1 = min(MM_NR1, eb - r1);\n if (nr0 <= 0 || nr1 <= 0) return;\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb[2];\n simdgroup_float8x8 mc[8];\n for (short i = 0; i < 8; i++) mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n\n for (uint k = 0; k < in_dim; k += MM_NK) {\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n // Load weight tile\n {\n const short row_in_tile = tiitg / 2;\n const short col_start = (tiitg % 2) * 16;\n if (row_in_tile < nr0) {\n device const half * src = w + (r0 + row_in_tile) * in_dim + k + col_start;\n for (short j = 0; j < 16; j++) {\n short sx = col_start/8 + j/8;\n short sy = row_in_tile / 8;\n short lx = row_in_tile % 8;\n short ly = j % 8;\n *(sa + 64*(8*sx + sy) + 8*ly + lx) = (k + col_start + j < in_dim) ? src[j] : half(0);\n }\n }\n }\n\n // Load input tile (from packed buffer at expert base offset)\n {\n const short row_in_tile = tiitg / 4;\n const short col_start = (tiitg % 4) * 8;\n if (row_in_tile < nr1) {\n device const half * src = x_packed + (base + r1 + row_in_tile) * in_dim + k + col_start;\n short sx = col_start / 8;\n short sy = row_in_tile / 8;\n short ly = row_in_tile % 8;\n *(threadgroup half2x4 *)(sb + 64*(4*sx + sy) + 8*ly) = (k + col_start < in_dim) ?\n *(device const half2x4 *)src : half2x4(0);\n }\n }\n\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n threadgroup const half * lsma = sa + 4*64*(sgitg % 2);\n threadgroup const half * lsmb = sb + 2*64*(sgitg / 2);\n\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 4; i++) simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 2; i++) simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 8; i++) simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);\n lsma += 8*64; lsmb += 4*64;\n }\n }\n\n threadgroup float * temp = (threadgroup float *)shmem;\n threadgroup float * sg_out = temp + 32*(sgitg & 1) + 16*(sgitg >> 1)*MM_NR0;\n for (short i = 0; i < 8; i++) simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n {\n const int total_out = nr0 * nr1;\n for (int idx = (int)tiitg; idx < total_out; idx += 128) {\n const int i = idx % nr0;\n const int j = idx / nr0;\n float val = temp[j * MM_NR0 + i] + bias[r0 + i];\n if (apply_gelu) {\n if (val > 10.0f) {} else if (val < -10.0f) { val = 0.0f; }\n else { float t = 0.7978845608f * (val + 0.044715f * val * val * val); val = 0.5f * val * (1.0f + tanh(t)); }\n }\n out_packed[(base + r1 + j) * out_dim + r0 + i] = half(val);\n }\n }\n}\n"

GEMM_MM_SOURCE = "// Matrix-matrix GEMM for Q5_K / Q6_K using simdgroup_matrix_multiply_accumulate\n// Adapted from llama.cpp's kernel_mul_mm — hardware-accelerated 8×8 tiles\n//\n// Each threadgroup computes a 64×32 output tile (64 output rows × 32 batch elements)\n// using 4 simdgroups (128 threads), dequantizing weights to FP16 shared memory.\n//\n// For batch > ~8, this is ~5× faster than the scalar SIMD-group GEMM.\n//\n// Dispatch: threadgroups = [ceil(batch/32), ceil(out_dim/64), 1]\n// threads_per_threadgroup = [128, 1, 1]\n// threadgroup_memory = 8192 bytes\n\n#include <metal_stdlib>\n#include <metal_simdgroup_matrix>\nusing namespace metal;\n\n#define FOR_UNROLL _Pragma(\"clang loop unroll(full)\")\n\nconstant uint QK_K = 256;\n\n// ============================================================================\n// Q5_K dequantization: 16 elements from one sub-block → 4×4 half register\n// Block layout: [d:2B][dmin:2B][scales:12B][qh:32B][qs:128B] = 176 bytes\n// il = 0..15 selects which 16-element chunk of the 256-element super-block\n// ============================================================================\nstruct block_q5_K {\n half d;\n half dmin;\n uint8_t scales[12];\n uint8_t qh[32];\n uint8_t qs[128];\n};\n\nstatic inline uchar2 get_scale_min_k4_just2(int j, int k, device const uchar * q) {\n return j < 4 ? uchar2{uchar(q[j+0+k] & 63), uchar(q[j+4+k] & 63)}\n : uchar2{uchar((q[j+4+k] & 0xF) | ((q[j-4+k] & 0xc0) >> 2)),\n uchar((q[j+4+k] >> 4) | ((q[j-0+k] & 0xc0) >> 2))};\n}\n\nvoid dequantize_q5_K_fn(device const block_q5_K *xb, short il, thread half4x4 & reg) {\n device const uint8_t * q = xb->qs;\n device const uint8_t * qh = xb->qh;\n\n short is = (il/4) * 2;\n q = q + 32 * (il/4) + 16 * (il&1);\n qh = qh + 16 * (il&1);\n uint8_t ul = 1 << (il/2);\n il = il & 3;\n const uchar2 sc = get_scale_min_k4_just2(is, il/2, xb->scales);\n const float d = il < 2 ? xb->d : xb->d / 16.f;\n const float min = xb->dmin;\n const float dl = d * sc[0];\n const float ml = min * sc[1];\n\n const ushort mask = il<2 ? 0x0F : 0xF0;\n const float qh_val = il<2 ? 16.f : 256.f;\n for (int i = 0; i < 16; ++i) {\n reg[i/4][i%4] = dl * ((q[i] & mask) + (qh[i] & ul ? qh_val : 0)) - ml;\n }\n}\n\n// ============================================================================\n// Q6_K dequantization: 16 elements from one sub-block → 4×4 half register\n// Block layout: [ql:128B][qh:64B][scales:16B][d:2B] = 210 bytes\n// ============================================================================\nstruct block_q6_K {\n uint8_t ql[128];\n uint8_t qh[64];\n int8_t scales[16];\n half d;\n};\n\nvoid dequantize_q6_K_fn(device const block_q6_K *xb, short il, thread half4x4 & reg) {\n const half d_all = xb->d;\n device const uint16_t * ql = (device const uint16_t *)xb->ql;\n device const uint16_t * qh = (device const uint16_t *)xb->qh;\n device const int8_t * scales = (device const int8_t *)xb->scales;\n\n ql = ql + 32*(il/8) + 16*((il/2)&1) + 8*(il&1);\n qh = qh + 16*(il/8) + 8*(il&1);\n float sc = scales[(il%2) + 2 * ((il/2))];\n il = (il/2) & 3;\n\n const uint32_t kmask1 = il>1 ? (il>2 ? 0xC0C0C0C0 : 0x30303030) : (il>0 ? 0x0C0C0C0C : 0x03030303);\n const uint32_t kmask2 = il>1 ? 0xF0F0F0F0 : 0x0F0F0F0F;\n const float ml = d_all * sc * 32.f;\n const float dl0 = d_all * sc;\n const float dl1 = dl0 / 256.f;\n const float dl2 = dl0 / (256.f * 256.f);\n const float dl3 = dl0 / (256.f * 256.f * 256.f);\n const uint8_t shr_h = il>2 ? 2 : 0;\n const uint8_t shl_h = il>1 ? 0 : (il>0 ? 2 : 4);\n const uint8_t shr_l = il>1 ? 4 : 0;\n for (int i = 0; i < 4; ++i) {\n const uint32_t low = (ql[2*i] | (uint32_t)(ql[2*i+1] << 16)) & kmask2;\n const uint32_t high = (qh[2*i] | (uint32_t)(qh[2*i+1] << 16)) & kmask1;\n const uint32_t q = ((high << shl_h) >> shr_h) | (low >> shr_l);\n reg[i][0] = dl0 * ((half)(q & 0xFF)) - ml;\n reg[i][1] = dl1 * ((float)(q & 0xFF00)) - ml;\n reg[i][2] = dl2 * ((float)(q & 0xFF0000)) - ml;\n reg[i][3] = dl3 * ((float)(q & 0xFF000000)) - ml;\n }\n}\n\n// ============================================================================\n// Shared matmul core: dequant to shared mem → simdgroup_matrix_multiply_accumulate\n// ============================================================================\nconstant int MM_NR0 = 64; // output rows per threadgroup\nconstant int MM_NR1 = 32; // batch elements per threadgroup\nconstant int MM_NK = 32; // K elements per iteration\nconstant int MM_NL0 = 2; // NK/16 — threads sharing weight dequantization\nconstant int MM_NL1 = 4; // NK/8 — threads sharing input loading\nconstant int MM_NL = 16; // QK_NL — sub-blocks per super-block (256/16)\n// Double-buffered shared memory: sa[0]/sb[0] and sa[1]/sb[1]\n// Eliminates 1 of 2 threadgroup_barriers per K-iteration (50% fewer inner barriers)\nconstant int MM_SA_SIZE = 4096; // bytes per weight tile in shmem\nconstant int MM_SB_SIZE = 2048; // bytes per input tile in shmem\nconstant int MM_TILE_SIZE = MM_SA_SIZE + MM_SB_SIZE; // 6144 bytes per tile\n\n// Q5_K matrix-matrix multiply with simdgroup_matrix\nkernel void simd_mm_q5k(\n device const uint8_t* w_raw [[buffer(0)]],\n device const half* x [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* output [[buffer(3)]],\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n // Double-buffered shared memory: tile 0 and tile 1 alternate\n threadgroup half * sa_buf[2] = {\n (threadgroup half *)(shmem),\n (threadgroup half *)(shmem + MM_TILE_SIZE)\n };\n threadgroup half * sb_buf[2] = {\n (threadgroup half *)(shmem + MM_SA_SIZE),\n (threadgroup half *)(shmem + MM_TILE_SIZE + MM_SA_SIZE)\n };\n\n const int r0 = tgpig.y * MM_NR0;\n const int r1 = tgpig.x * MM_NR1;\n\n const short nr0 = min(MM_NR0, (int)out_dim - r0);\n const short nr1 = min(MM_NR1, (int)batch - r1);\n\n const short lr0 = min((short)(tiitg/MM_NL0), (short)(nr0 - 1));\n const short lr1 = min((short)(tiitg/MM_NL1), (short)(nr1 - 1));\n\n const short il0 = tiitg % MM_NL0;\n short il = il0;\n\n const uint row_bytes = (in_dim / QK_K) * 176;\n device const block_q5_K * xw = (device const block_q5_K *)(w_raw + (r0 + lr0) * row_bytes) + il0 / MM_NL;\n\n const short iy = 8 * (tiitg % MM_NL1);\n device const half * y = x + (r1 + lr1) * in_dim + iy;\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 // Load first tile into buffer 0\n {\n threadgroup half * sa = sa_buf[0];\n threadgroup half * sb = sb_buf[0];\n half4x4 temp_a;\n dequantize_q5_K_fn(xw, il, temp_a);\n FOR_UNROLL for (short i = 0; i < 16; i++) {\n const short sx = 2*il0 + i/8;\n const short sy = (tiitg/MM_NL0)/8;\n const short lx = (tiitg/MM_NL0)%8;\n const short ly = i%8;\n *(sa + 64*(8*sx + sy) + 8*ly + lx) = temp_a[i/4][i%4];\n }\n {\n const short sx = (tiitg % MM_NL1);\n const short sy = (tiitg/MM_NL1)/8;\n const short ly = (tiitg/MM_NL1)%8;\n *(threadgroup half2x4 *)(sb + 64*(4*sx + sy) + 8*ly) = *(device const half2x4 *)y;\n }\n il = (il + 2 < MM_NL) ? il + 2 : il % 2;\n xw = (il < 2) ? xw + (2 + MM_NL - 1)/MM_NL : xw;\n y += MM_NK;\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n const uint n_iter = (in_dim + MM_NK - 1) / MM_NK;\n for (uint iter = 0; iter < n_iter; iter++) {\n short cur = iter % 2;\n short nxt = 1 - cur;\n threadgroup half * sa = sa_buf[cur];\n threadgroup half * sb = sb_buf[cur];\n\n // Compute from current buffer\n threadgroup const half * lsma = sa + 4*64*(sgitg % 2);\n threadgroup const half * lsmb = sb + 2*64*(sgitg / 2);\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 4; i++) simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 2; i++) simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 8; i++) simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);\n lsma += 8*64; lsmb += 4*64;\n }\n\n // Load next tile into alternate buffer (overlap with compute above on next iteration)\n if (iter + 1 < n_iter) {\n threadgroup half * sa_n = sa_buf[nxt];\n threadgroup half * sb_n = sb_buf[nxt];\n half4x4 temp_a;\n dequantize_q5_K_fn(xw, il, temp_a);\n FOR_UNROLL for (short i = 0; i < 16; i++) {\n const short sx = 2*il0 + i/8;\n const short sy = (tiitg/MM_NL0)/8;\n const short lx = (tiitg/MM_NL0)%8;\n const short ly = i%8;\n *(sa_n + 64*(8*sx + sy) + 8*ly + lx) = temp_a[i/4][i%4];\n }\n {\n const short sx = (tiitg % MM_NL1);\n const short sy = (tiitg/MM_NL1)/8;\n const short ly = (tiitg/MM_NL1)%8;\n *(threadgroup half2x4 *)(sb_n + 64*(4*sx + sy) + 8*ly) = *(device const half2x4 *)y;\n }\n il = (il + 2 < MM_NL) ? il + 2 : il % 2;\n xw = (il < 2) ? xw + (2 + MM_NL - 1)/MM_NL : xw;\n y += MM_NK;\n }\n threadgroup_barrier(mem_flags::mem_threadgroup); // ONE barrier (was TWO)\n }\n\n // Write output: ALL 128 threads participate (4× faster than sgitg==0 only)\n // Store accumulators to shared, barrier, then cooperative bias+GELU+write\n threadgroup float * temp = (threadgroup float *)shmem;\n {\n threadgroup float * sg_out = temp + 32*(sgitg & 1) + 16*(sgitg >> 1)*MM_NR0;\n for (short i = 0; i < 8; i++) {\n simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false);\n }\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n // ALL 128 threads cooperatively write output with bias + optional GELU\n // nr0/nr1 already computed above\n // Total elements: nr0 * nr1 ≤ 64*32 = 2048. Each of 128 threads handles ~16 elements.\n const int total_out = nr0 * nr1;\n for (int idx = (int)tiitg; idx < total_out; idx += 128) {\n const int i = idx % nr0; // output row within tile\n const int j = idx / nr0; // batch element within tile\n float val = temp[j * MM_NR0 + i] + bias[r0 + i];\n if (apply_gelu) {\n if (val > 10.0f) { }\n else if (val < -10.0f) { val = 0.0f; }\n else {\n float t = 0.7978845608f * (val + 0.044715f * val * val * val);\n val = 0.5f * val * (1.0f + tanh(t));\n }\n }\n output[(r1 + j) * out_dim + r0 + i] = half(val);\n }\n\n}\n\n// Q6_K matrix-matrix multiply with simdgroup_matrix\nkernel void simd_mm_q6k(\n device const uint8_t* w_raw [[buffer(0)]],\n device const half* x [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* output [[buffer(3)]],\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half * sa = (threadgroup half *)(shmem);\n threadgroup half * sb = (threadgroup half *)(shmem + 4096);\n\n const int r0 = tgpig.y * MM_NR0;\n const int r1 = tgpig.x * MM_NR1;\n\n const short nr0 = min(MM_NR0, (int)out_dim - r0);\n const short nr1 = min(MM_NR1, (int)batch - r1);\n\n const short lr0 = min((short)(tiitg/MM_NL0), (short)(nr0 - 1));\n const short lr1 = min((short)(tiitg/MM_NL1), (short)(nr1 - 1));\n\n const short il0 = tiitg % MM_NL0;\n short il = il0;\n\n const uint row_bytes = (in_dim / QK_K) * 210;\n const short offset1 = il0 / MM_NL;\n\n device const block_q6_K * xw = (device const block_q6_K *)(w_raw + (r0 + lr0) * row_bytes) + offset1;\n\n const short iy = 8 * (tiitg % MM_NL1);\n device const half * y = x + (r1 + lr1) * in_dim + iy;\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 < in_dim; loop_k += MM_NK) {\n half4x4 temp_a;\n dequantize_q6_K_fn(xw, il, temp_a);\n\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n FOR_UNROLL for (short i = 0; i < 16; i++) {\n const short sx = 2*il0 + i/8;\n const short sy = (tiitg/MM_NL0)/8;\n const short lx = (tiitg/MM_NL0)%8;\n const short ly = i%8;\n const short ib = 8*sx + sy;\n *(sa + 64*ib + 8*ly + lx) = temp_a[i/4][i%4];\n }\n\n {\n const short sx = (tiitg % MM_NL1);\n const short sy = (tiitg/MM_NL1)/8;\n const short ly = (tiitg/MM_NL1)%8;\n const short ib = 4*sx + sy;\n *(threadgroup half2x4 *)(sb + 64*ib + 8*ly) = *(device const half2x4 *)y;\n }\n\n il = (il + 2 < MM_NL) ? il + 2 : il % 2;\n xw = (il < 2) ? xw + (2 + MM_NL - 1)/MM_NL : xw;\n y += MM_NK;\n\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n threadgroup const half * lsma = sa + 4*64*(sgitg % 2);\n threadgroup const half * lsmb = sb + 2*64*(sgitg / 2);\n\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL 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_UNROLL 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_UNROLL for (short i = 0; i < 8; i++) {\n simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);\n }\n lsma += 8*64;\n lsmb += 4*64;\n }\n }\n\n // Write output with bias + optional GELU\n threadgroup float * temp = (threadgroup float *)shmem;\n threadgroup float * sg_out = temp + 32*(sgitg & 1) + 16*(sgitg >> 1)*MM_NR0;\n for (short i = 0; i < 8; i++) {\n simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false);\n }\n threadgroup_barrier(mem_flags::mem_threadgroup);\n\n // ALL 128 threads cooperatively write output with bias + GELU\n {\n const int total_out = nr0 * nr1;\n for (int idx = (int)tiitg; idx < total_out; idx += 128) {\n const int i = idx % nr0;\n const int j = idx / nr0;\n float val = temp[j * MM_NR0 + i] + bias[r0 + i];\n if (apply_gelu) {\n if (val > 10.0f) { }\n else if (val < -10.0f) { val = 0.0f; }\n else {\n float t = 0.7978845608f * (val + 0.044715f * val * val * val);\n val = 0.5f * val * (1.0f + tanh(t));\n }\n }\n output[(r1 + j) * out_dim + r0 + i] = half(val);\n }\n }\n}\n\n// ============================================================================\n// MoE variants of matrix-matrix GEMM — GPU-side expert_offsets, zero CPU sync\n// Same simdgroup_matrix approach, but input/output are packed by expert.\n// Dispatch: indirect with grid = [ceil(eb/32), ceil(out_dim/64), 1]\n// ============================================================================\n\nkernel void simd_mm_q5k_moe(\n device const uint8_t* w_raw [[buffer(0)]], // weights for THIS expert\n device const half* x_packed [[buffer(1)]], // full packed input\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]], // full packed output\n device const int* expert_offs [[buffer(4)]],\n constant uint& expert_id [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half * sa = (threadgroup half *)(shmem);\n threadgroup half * sb = (threadgroup half *)(shmem + 4096);\n\n const int base = expert_offs[expert_id];\n const int eb = expert_offs[expert_id + 1] - base;\n if (eb <= 0) return;\n\n const int r0 = tgpig.y * MM_NR0; // output row tile\n const int r1 = tgpig.x * MM_NR1; // batch tile within this expert\n\n const short nr0 = min(MM_NR0, (int)out_dim - r0);\n const short nr1 = min(MM_NR1, eb - r1);\n if (nr0 <= 0 || nr1 <= 0) return;\n\n const short lr0 = min((short)(tiitg/MM_NL0), (short)(nr0 - 1));\n const short lr1 = min((short)(tiitg/MM_NL1), (short)(nr1 - 1));\n\n const short il0 = tiitg % MM_NL0;\n short il = il0;\n\n const uint row_bytes = (in_dim / QK_K) * 176;\n device const block_q5_K * xw = (device const block_q5_K *)(w_raw + (r0 + lr0) * row_bytes) + il0/MM_NL;\n\n const short iy = 8 * (tiitg % MM_NL1);\n device const half * y = x_packed + (base + r1 + lr1) * in_dim + iy;\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb[2];\n simdgroup_float8x8 mc[8];\n for (short i = 0; i < 8; i++) mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n\n for (uint loop_k = 0; loop_k < in_dim; loop_k += MM_NK) {\n half4x4 temp_a;\n dequantize_q5_K_fn(xw, il, temp_a);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n FOR_UNROLL for (short i = 0; i < 16; i++) {\n const short sx = 2*il0 + i/8;\n const short sy = (tiitg/MM_NL0)/8;\n const short lx = (tiitg/MM_NL0)%8;\n const short ly = i%8;\n *(sa + 64*(8*sx + sy) + 8*ly + lx) = temp_a[i/4][i%4];\n }\n {\n const short sx = (tiitg % MM_NL1);\n const short sy = (tiitg/MM_NL1)/8;\n const short ly = (tiitg/MM_NL1)%8;\n *(threadgroup half2x4 *)(sb + 64*(4*sx + sy) + 8*ly) = *(device const half2x4 *)y;\n }\n il = (il + 2 < MM_NL) ? il + 2 : il % 2;\n xw = (il < 2) ? xw + (2 + MM_NL - 1)/MM_NL : xw;\n y += MM_NK;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n threadgroup const half * lsma = sa + 4*64*(sgitg % 2);\n threadgroup const half * lsmb = sb + 2*64*(sgitg / 2);\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 4; i++) simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 2; i++) simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 8; i++) simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);\n lsma += 8*64; lsmb += 4*64;\n }\n }\n\n threadgroup float * temp = (threadgroup float *)shmem;\n threadgroup float * sg_out = temp + 32*(sgitg & 1) + 16*(sgitg >> 1)*MM_NR0;\n for (short i = 0; i < 8; i++) simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n {\n const int total_out = nr0 * nr1;\n for (int idx = (int)tiitg; idx < total_out; idx += 128) {\n const int i = idx % nr0;\n const int j = idx / nr0;\n float val = temp[j * MM_NR0 + i] + bias[r0 + i];\n if (apply_gelu) {\n if (val > 10.0f) {} else if (val < -10.0f) { val = 0.0f; }\n else { float t = 0.7978845608f * (val + 0.044715f * val * val * val); val = 0.5f * val * (1.0f + tanh(t)); }\n }\n out_packed[(base + r1 + j) * out_dim + r0 + i] = half(val);\n }\n }\n}\n\nkernel void simd_mm_q6k_moe(\n device const uint8_t* w_raw [[buffer(0)]],\n device const half* x_packed [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]],\n device const int* expert_offs [[buffer(4)]],\n constant uint& expert_id [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n threadgroup half * sa = (threadgroup half *)(shmem);\n threadgroup half * sb = (threadgroup half *)(shmem + 4096);\n\n const int base = expert_offs[expert_id];\n const int eb = expert_offs[expert_id + 1] - base;\n if (eb <= 0) return;\n\n const int r0 = tgpig.y * MM_NR0;\n const int r1 = tgpig.x * MM_NR1;\n\n const short nr0 = min(MM_NR0, (int)out_dim - r0);\n const short nr1 = min(MM_NR1, eb - r1);\n if (nr0 <= 0 || nr1 <= 0) return;\n\n const short lr0 = min((short)(tiitg/MM_NL0), (short)(nr0 - 1));\n const short lr1 = min((short)(tiitg/MM_NL1), (short)(nr1 - 1));\n\n const short il0 = tiitg % MM_NL0;\n short il = il0;\n\n const uint row_bytes = (in_dim / QK_K) * 210;\n device const block_q6_K * xw = (device const block_q6_K *)(w_raw + (r0 + lr0) * row_bytes) + il0/MM_NL;\n\n const short iy = 8 * (tiitg % MM_NL1);\n device const half * y = x_packed + (base + r1 + lr1) * in_dim + iy;\n\n simdgroup_half8x8 ma[4];\n simdgroup_half8x8 mb[2];\n simdgroup_float8x8 mc[8];\n for (short i = 0; i < 8; i++) mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f);\n\n for (uint loop_k = 0; loop_k < in_dim; loop_k += MM_NK) {\n half4x4 temp_a;\n dequantize_q6_K_fn(xw, il, temp_a);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n FOR_UNROLL for (short i = 0; i < 16; i++) {\n const short sx = 2*il0 + i/8;\n const short sy = (tiitg/MM_NL0)/8;\n const short lx = (tiitg/MM_NL0)%8;\n const short ly = i%8;\n *(sa + 64*(8*sx + sy) + 8*ly + lx) = temp_a[i/4][i%4];\n }\n {\n const short sx = (tiitg % MM_NL1);\n const short sy = (tiitg/MM_NL1)/8;\n const short ly = (tiitg/MM_NL1)%8;\n *(threadgroup half2x4 *)(sb + 64*(4*sx + sy) + 8*ly) = *(device const half2x4 *)y;\n }\n il = (il + 2 < MM_NL) ? il + 2 : il % 2;\n xw = (il < 2) ? xw + (2 + MM_NL - 1)/MM_NL : xw;\n y += MM_NK;\n threadgroup_barrier(mem_flags::mem_threadgroup);\n threadgroup const half * lsma = sa + 4*64*(sgitg % 2);\n threadgroup const half * lsmb = sb + 2*64*(sgitg / 2);\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) {\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 4; i++) simdgroup_load(ma[i], lsma + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 2; i++) simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false);\n simdgroup_barrier(mem_flags::mem_none);\n FOR_UNROLL for (short i = 0; i < 8; i++) simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]);\n lsma += 8*64; lsmb += 4*64;\n }\n }\n\n threadgroup float * temp = (threadgroup float *)shmem;\n threadgroup float * sg_out = temp + 32*(sgitg & 1) + 16*(sgitg >> 1)*MM_NR0;\n for (short i = 0; i < 8; i++) simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false);\n threadgroup_barrier(mem_flags::mem_threadgroup);\n {\n const int total_out = nr0 * nr1;\n for (int idx = (int)tiitg; idx < total_out; idx += 128) {\n const int i = idx % nr0;\n const int j = idx / nr0;\n float val = temp[j * MM_NR0 + i] + bias[r0 + i];\n if (apply_gelu) {\n if (val > 10.0f) {} else if (val < -10.0f) { val = 0.0f; }\n else { float t = 0.7978845608f * (val + 0.044715f * val * val * val); val = 0.5f * val * (1.0f + tanh(t)); }\n }\n out_packed[(base + r1 + j) * out_dim + r0 + i] = half(val);\n }\n }\n}\n\n// ============================================================================\n// Batched expert GEMM — ALL experts in ONE dispatch (LTP Diamond surgery)\n// tgpig.x = flattened batch tiles across ALL experts\n// tgpig.y = output row tile\n// Binary search expert_tg_offsets to find expert_id from tgpig.x\n// Grid: {sum_of_all_expert_batch_tgs, ceil(out_dim/64), 1}\n// ============================================================================\n\ninline int find_expert(device const int* tg_offsets, int flat_x, int n) {\n int lo = 0, hi = n - 1;\n while (lo < hi) {\n int mid = (lo + hi + 1) / 2;\n if (tg_offsets[mid] <= flat_x) lo = mid; else hi = mid - 1;\n }\n return lo;\n}\n\n#define BATCHED_MM_BODY(BLOCK_T, BLOCK_BYTES, DEQUANT_FN) \\\n threadgroup half * sa = (threadgroup half *)(shmem); \\\n threadgroup half * sb = (threadgroup half *)(shmem + 4096); \\\n const int eid = find_expert(expert_tg_offs, (int)tgpig.x, (int)n_experts); \\\n const int local_x = (int)tgpig.x - expert_tg_offs[eid]; \\\n const int base = expert_offs[eid]; \\\n const int eb = expert_offs[eid + 1] - base; \\\n if (eb <= 0) return; \\\n const int r0 = tgpig.y * MM_NR0; \\\n const int r1 = local_x * MM_NR1; \\\n const short nr0 = min(MM_NR0, (int)out_dim - r0); \\\n const short nr1 = min(MM_NR1, eb - r1); \\\n if (nr0 <= 0 || nr1 <= 0) return; \\\n const short lr0 = min((short)(tiitg/MM_NL0), (short)(nr0 - 1)); \\\n const short lr1 = min((short)(tiitg/MM_NL1), (short)(nr1 - 1)); \\\n const short il0 = tiitg % MM_NL0; \\\n short il = il0; \\\n const uint row_bytes = (in_dim / QK_K) * BLOCK_BYTES; \\\n device const BLOCK_T * xw = (device const BLOCK_T *)(all_weights + eid * weight_stride + (r0 + lr0) * row_bytes) + il0/MM_NL; \\\n const short iy = 8 * (tiitg % MM_NL1); \\\n device const half * y = x_packed + (base + r1 + lr1) * in_dim + iy; \\\n simdgroup_half8x8 ma[4]; simdgroup_half8x8 mb[2]; simdgroup_float8x8 mc[8]; \\\n for (short i = 0; i < 8; i++) mc[i] = make_filled_simdgroup_matrix<float, 8>(0.f); \\\n for (uint loop_k = 0; loop_k < in_dim; loop_k += MM_NK) { \\\n half4x4 temp_a; DEQUANT_FN(xw, il, temp_a); \\\n threadgroup_barrier(mem_flags::mem_threadgroup); \\\n FOR_UNROLL for (short i = 0; i < 16; i++) { \\\n const short sx = 2*il0 + i/8, sy = (tiitg/MM_NL0)/8, lx = (tiitg/MM_NL0)%8, ly = i%8; \\\n *(sa + 64*(8*sx + sy) + 8*ly + lx) = temp_a[i/4][i%4]; } \\\n { const short sx = tiitg%MM_NL1, sy = (tiitg/MM_NL1)/8, ly = (tiitg/MM_NL1)%8; \\\n *(threadgroup half2x4 *)(sb + 64*(4*sx + sy) + 8*ly) = *(device const half2x4 *)y; } \\\n il = (il + 2 < MM_NL) ? il + 2 : il % 2; \\\n xw = (il < 2) ? xw + (2 + MM_NL - 1)/MM_NL : xw; y += MM_NK; \\\n threadgroup_barrier(mem_flags::mem_threadgroup); \\\n threadgroup const half * lsma = sa + 4*64*(sgitg%2), * lsmb = sb + 2*64*(sgitg/2); \\\n FOR_UNROLL for (short ik = 0; ik < MM_NK/8; ik++) { \\\n simdgroup_barrier(mem_flags::mem_none); \\\n FOR_UNROLL for (short i = 0; i < 4; i++) simdgroup_load(ma[i], lsma + 64*i, 8, 0, false); \\\n simdgroup_barrier(mem_flags::mem_none); \\\n FOR_UNROLL for (short i = 0; i < 2; i++) simdgroup_load(mb[i], lsmb + 64*i, 8, 0, false); \\\n simdgroup_barrier(mem_flags::mem_none); \\\n FOR_UNROLL for (short i = 0; i < 8; i++) simdgroup_multiply_accumulate(mc[i], mb[i/4], ma[i%4], mc[i]); \\\n lsma += 8*64; lsmb += 4*64; } \\\n } \\\n threadgroup float * temp = (threadgroup float *)shmem; \\\n { threadgroup float * sg_out = temp + 32*(sgitg&1) + 16*(sgitg>>1)*MM_NR0; \\\n for (short i = 0; i < 8; i++) simdgroup_store(mc[i], sg_out + 8*(i%4) + 8*MM_NR0*(i/4), MM_NR0, 0, false); } \\\n threadgroup_barrier(mem_flags::mem_threadgroup); \\\n { const int total_out = nr0 * nr1; \\\n for (int idx = (int)tiitg; idx < total_out; idx += 128) { \\\n const int i = idx%nr0, j = idx/nr0; \\\n float val = temp[j*MM_NR0 + i] + bias[r0 + i]; \\\n if (apply_gelu) { if (val > 10.0f) {} else if (val < -10.0f) { val = 0.0f; } \\\n else { float t = 0.7978845608f*(val + 0.044715f*val*val*val); val = 0.5f*val*(1.0f + tanh(t)); } } \\\n out_packed[(base + r1 + j)*out_dim + r0 + i] = half(val); } }\n\nkernel void batched_mm_q5k(\n device const uint8_t* all_weights [[buffer(0)]],\n device const half* x_packed [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]],\n device const int* expert_offs [[buffer(4)]],\n device const int* expert_tg_offs [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n constant uint& n_experts [[buffer(9)]],\n constant uint& weight_stride [[buffer(10)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{ BATCHED_MM_BODY(block_q5_K, 176, dequantize_q5_K_fn) }\n\nkernel void batched_mm_q6k(\n device const uint8_t* all_weights [[buffer(0)]],\n device const half* x_packed [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]],\n device const int* expert_offs [[buffer(4)]],\n device const int* expert_tg_offs [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n constant uint& n_experts [[buffer(9)]],\n constant uint& weight_stride [[buffer(10)]],\n threadgroup char* shmem [[threadgroup(0)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiitg [[thread_index_in_threadgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{ BATCHED_MM_BODY(block_q6_K, 210, dequantize_q6_K_fn) }\n"

GEMM_SOURCE = "// Tiled GEMM for Q5_K quantized weights × F32 input.\n// Uses shared memory for dequantized weight tiles and input tiles.\n// Much better precision than per-element sequential accumulation.\n//\n// C[m, n] = Σ_k dequant(A[m, k]) * B[n, k] + bias[m]\n//\n// A = weight [M, K] quantized Q5_K (M = out_dim, K = in_dim)\n// B = input [N, K] F32 (N = batch/seq_len, K = in_dim)\n// C = output [N, M] F32\n//\n// Tile: TILE_M rows of A × TILE_N rows of B, accumulated over K in chunks of QK_K=256\n\n#include <metal_stdlib>\nusing namespace metal;\n\nconstant uint QK_K = 256;\nconstant uint TILE_M = 4; // Output rows per threadgroup\nconstant uint TILE_N = 4; // Input rows (batch) per threadgroup\nconstant uint THREADS_PER_ROW = 64; // Threads accumulating one output element\n\ninline float2 get_scale_min_k4(int j, const device uint8_t* scales) {\n float sc, m;\n if (j < 4) { sc = float(scales[j] & 63); m = float(scales[j + 4] & 63); }\n else { sc = float((scales[j+4] & 0x0F) | ((scales[j-4] >> 6) << 4)); m = float((scales[j+4] >> 4) | ((scales[j] >> 6) << 4)); }\n return float2(sc, m);\n}\n\n// Tiled Q5_K matmul + optional GELU\n// Grid: [ceil(out_dim/TILE_M), ceil(batch/TILE_N)]\n// Threadgroup: [TILE_M, TILE_N] = 16 threads per group\nkernel void gemm_q5k(\n device const uint8_t* w_raw [[buffer(0)]], // [out_dim rows] Q5_K\n device const float* x [[buffer(1)]], // [batch, in_dim] F32\n device const float* bias [[buffer(2)]], // [out_dim] F32\n device float* output [[buffer(3)]], // [batch, out_dim] F32\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n uint2 tgid [[threadgroup_position_in_grid]], // tile position\n uint2 tid [[thread_position_in_threadgroup]]) // thread in tile\n{\n const uint m = tgid.x * TILE_M + tid.x; // output row\n const uint n = tgid.y * TILE_N + tid.y; // batch index\n if (m >= out_dim || n >= batch) return;\n\n const uint blocks_per_row = in_dim / QK_K;\n const uint row_bytes = blocks_per_row * 176;\n\n device const uint8_t* row_ptr = w_raw + m * row_bytes;\n device const float* x_row = x + n * in_dim;\n\n // Accumulate with float — but over fewer iterations (blocks_per_row = 3 for dim=768)\n // Each block = 256 elements, so only 3 accumulation \"tiles\"\n float sum = bias[m];\n\n for (uint blk = 0; blk < blocks_per_row; blk++) {\n device const uint8_t* bp = row_ptr + blk * 176;\n float d = float(as_type<half>(*(device const ushort*)(bp)));\n float dmin = float(as_type<half>(*(device const ushort*)(bp + 2)));\n // Guard against NaN/Inf in half-precision scale factors\n if (isnan(d) || isinf(d)) d = 0.0f;\n if (isnan(dmin) || isinf(dmin)) dmin = 0.0f;\n device const uint8_t* scales = bp + 4;\n device const uint8_t* qh = bp + 16;\n device const uint8_t* ql = bp + 48;\n\n const uint base_j = blk * QK_K;\n uint8_t u1 = 1, u2 = 2;\n int is = 0;\n uint ql_off = 0;\n\n // Process 4 sub-blocks of 64 elements each\n // Use float4 partial sums for better precision\n float4 partial = float4(0.0f);\n\n for (int iter = 0; iter < 4; iter++) {\n float2 sm0 = get_scale_min_k4(is, scales);\n float d1 = d * sm0.x, m1 = dmin * sm0.y;\n float2 sm1 = get_scale_min_k4(is + 1, scales);\n float d2 = d * sm1.x, m2 = dmin * sm1.y;\n\n // First 32 elements\n float sub_sum0 = 0.0f;\n for (uint l = 0; l < 32; l++) {\n uint j = base_j + (is / 2) * 64 + l;\n sub_sum0 += x_row[j] * (d1 * float((ql[ql_off + l] & 0x0F) + ((qh[l] & u1) ? 16 : 0)) - m1);\n }\n\n // Second 32 elements\n float sub_sum1 = 0.0f;\n for (uint l = 0; l < 32; l++) {\n uint j = base_j + (is / 2) * 64 + 32 + l;\n sub_sum1 += x_row[j] * (d2 * float(((ql[ql_off + l] >> 4) & 0x0F) + ((qh[l] & u2) ? 16 : 0)) - m2);\n }\n\n partial[iter] = sub_sum0 + sub_sum1;\n ql_off += 32;\n is += 2;\n u1 <<= 2;\n u2 <<= 2;\n }\n\n // Reduce float4 — fewer rounding steps than sequential\n sum += (partial[0] + partial[2]) + (partial[1] + partial[3]);\n }\n\n // Double NaN guard: before and after GELU\n if (isnan(sum) || isinf(sum)) sum = 0.0f;\n\n if (apply_gelu) {\n float v = clamp(sum, -20.0f, 20.0f);\n float t = 0.7978845608f * (v + 0.044715f * v * v * v);\n float th = tanh(t);\n sum = 0.5f * v * (1.0f + th);\n if (isnan(sum) || isinf(sum)) sum = 0.0f;\n }\n\n output[n * out_dim + m] = sum;\n}\n\n// Same for Q6_K\nkernel void gemm_q6k(\n device const uint8_t* w_raw [[buffer(0)]],\n device const float* x [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device float* output [[buffer(3)]],\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n uint2 tgid [[threadgroup_position_in_grid]],\n uint2 tid [[thread_position_in_threadgroup]])\n{\n const uint m = tgid.x * TILE_M + tid.x;\n const uint n = tgid.y * TILE_N + tid.y;\n if (m >= out_dim || n >= batch) return;\n\n const uint blocks_per_row = in_dim / QK_K;\n const uint row_bytes = blocks_per_row * 210;\n\n device const uint8_t* row_ptr = w_raw + m * row_bytes;\n device const float* x_row = x + n * in_dim;\n\n float sum = bias[m];\n\n for (uint blk = 0; blk < blocks_per_row; blk++) {\n device const uint8_t* bp = row_ptr + blk * 210;\n device const uint8_t* ql = bp;\n device const uint8_t* qh = bp + 128;\n device const int8_t* sc = (device const int8_t*)(bp + 192);\n float d = float(as_type<half>(*(device const ushort*)(bp + 208)));\n if (isnan(d) || isinf(d)) d = 0.0f;\n\n const uint base_j = blk * QK_K;\n uint ql_off = 0, qh_off = 0, sc_off = 0;\n\n // float4 partial sums per 128-element half-block\n float2 partial = float2(0.0f);\n\n for (int n_iter = 0; n_iter < 2; n_iter++) {\n float sub_sum = 0.0f;\n for (uint l = 0; l < 32; l++) {\n uint is_val = l / 16;\n int q1 = (int(ql[ql_off + l] & 0xF) | ((int(qh[qh_off + l] >> 0) & 3) << 4)) - 32;\n int q2 = (int(ql[ql_off + l + 32] & 0xF) | ((int(qh[qh_off + l] >> 2) & 3) << 4)) - 32;\n int q3 = (int(ql[ql_off + l] >> 4) | ((int(qh[qh_off + l] >> 4) & 3) << 4)) - 32;\n int q4 = (int(ql[ql_off + l + 32] >> 4) | ((int(qh[qh_off + l] >> 6) & 3) << 4)) - 32;\n\n float s0 = float(sc[sc_off + is_val]);\n float s2 = float(sc[sc_off + is_val + 2]);\n float s4 = float(sc[sc_off + is_val + 4]);\n float s6 = float(sc[sc_off + is_val + 6]);\n\n uint j_base = base_j + n_iter * 128;\n sub_sum += x_row[j_base + l] * (d * s0 * float(q1));\n sub_sum += x_row[j_base + l + 32] * (d * s2 * float(q2));\n sub_sum += x_row[j_base + l + 64] * (d * s4 * float(q3));\n sub_sum += x_row[j_base + l + 96] * (d * s6 * float(q4));\n }\n partial[n_iter] = sub_sum;\n ql_off += 64; qh_off += 32; sc_off += 8;\n }\n\n sum += partial[0] + partial[1];\n }\n\n if (isnan(sum) || isinf(sum)) sum = 0.0f;\n\n if (apply_gelu) {\n float v = clamp(sum, -20.0f, 20.0f);\n float t = 0.7978845608f * (v + 0.044715f * v * v * v);\n sum = 0.5f * v * (1.0f + tanh(t));\n if (isnan(sum) || isinf(sum)) sum = 0.0f;\n }\n\n output[n * out_dim + m] = sum;\n}\n"

MAGIC = "GGUF"

SIMD_GEMM_SOURCE = "// SIMD-group GEMM for Q5_K / Q6_K — matches llama.cpp's accumulation pattern.\n// Each output element computed by 32 threads (1 SIMD group).\n// Integer quant values accumulated separately from scale multiplication\n// to avoid large intermediate products that lose F32 precision.\n//\n// Standard: threadgroups = [ceil(out_dim/(N_ROWS*NR0)), batch, 1]\n// MoE: threadgroups = [ceil(out_dim/(N_ROWS*NR0)), seq_len, 1]\n// (kernel reads expert_offsets to compute actual batch + offset)\n\n#include <metal_stdlib>\nusing namespace metal;\n\n#define FOR_UNROLL _Pragma(\"clang loop unroll(full)\")\n\nconstant uint QK_K = 256;\nconstant uint N_ROWS = 2; // simdgroups per threadgroup\nconstant uint NR0 = 2; // output rows per simdgroup (amortize input reads) // output rows per threadgroup (2 simdgroups)\n\n// ============================================================================\n// Q5_K SIMD matmul — llama.cpp style accumulation\n// ============================================================================\nkernel void simd_gemm_q5k(\n device const uint8_t* w_raw [[buffer(0)]],\n device const half* x [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* output [[buffer(3)]],\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n // Each simdgroup processes NR0 consecutive output rows, sharing input reads\n const uint first_row = (tgpig.x * N_ROWS + sgitg) * NR0;\n const uint n = tgpig.y;\n if (first_row >= out_dim || n >= batch) return;\n\n const uint nb = in_dim / QK_K;\n const uint row_bytes = nb * 176;\n\n device const half* y1 = x + n * in_dim;\n\n const short tid = tiisg / 4;\n const short ix = tiisg % 4;\n const short iq = tid / 4;\n const short ir = tid % 4;\n\n const short l0 = 8 * ir;\n const short q_offset = 32 * iq + l0;\n const short y_offset = 64 * iq + l0;\n\n const uint8_t hm1 = 1u << (2*iq);\n const uint8_t hm2 = hm1 << 1;\n const uint8_t hm3 = hm1 << 4;\n const uint8_t hm4 = hm2 << 4;\n\n constexpr uint16_t kmask1 = 0x3f3f;\n constexpr uint16_t kmask2 = 0x0f0f;\n constexpr uint16_t kmask3 = 0xc0c0;\n\n float sumf[NR0] = {0.f}; // accumulator per row\n device const half* yp = y1 + ix * QK_K + y_offset;\n\n for (uint i = ix; i < nb; i += 4) {\n // Load y data ONCE, reuse for all NR0 rows\n device const half* y2 = yp + 128;\n float yl[16], yh[16];\n float4 sumy = {0.f, 0.f, 0.f, 0.f};\n FOR_UNROLL for (short l = 0; l < 8; ++l) {\n yl[l+0] = yp[l+ 0]; sumy[0] += yl[l+0];\n yl[l+8] = yp[l+32]; sumy[1] += yl[l+8];\n yh[l+0] = y2[l+ 0]; sumy[2] += yh[l+0];\n yh[l+8] = y2[l+32]; sumy[3] += yh[l+8];\n }\n\n // Process NR0 rows with shared y data\n for (short row = 0; row < (short)NR0; ++row) {\n device const uint8_t* bp = w_raw + (first_row + row) * row_bytes + i * 176;\n device const uint8_t* q1 = bp + 48 + q_offset;\n device const uint8_t* qh = bp + 16 + l0;\n device const half* dh = (device const half*)bp;\n device const uint16_t* a = (device const uint16_t*)(bp + 4) + iq;\n device const uint8_t* q2 = q1 + 64;\n\n uint16_t sc16[4];\n sc16[0] = a[0] & kmask1;\n sc16[1] = a[2] & kmask1;\n sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2);\n sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2);\n thread const uint8_t* sc8 = (thread const uint8_t*)sc16;\n\n float4 acc1 = {0.f}, acc2 = {0.f};\n FOR_UNROLL for (short l = 0; l < 8; ++l) {\n uint8_t h = qh[l];\n acc1[0] += yl[l+0] * float(q1[l] & 0x0F);\n acc1[1] += yl[l+8] * float(q1[l] & 0xF0);\n acc1[2] += yh[l+0] * float(q2[l] & 0x0F);\n acc1[3] += yh[l+8] * float(q2[l] & 0xF0);\n acc2[0] += (h & hm1) ? yl[l+0] : 0.f;\n acc2[1] += (h & hm2) ? yl[l+8] : 0.f;\n acc2[2] += (h & hm3) ? yh[l+0] : 0.f;\n acc2[3] += (h & hm4) ? yh[l+8] : 0.f;\n }\n\n sumf[row] +=\n dh[0] * ( sc8[0] * (acc1[0] + 16.f*acc2[0])\n + sc8[1] * (acc1[1]/16.f + 16.f*acc2[1])\n + sc8[4] * (acc1[2] + 16.f*acc2[2])\n + sc8[5] * (acc1[3]/16.f + 16.f*acc2[3]) )\n - dh[1] * ( sumy[0]*sc8[2] + sumy[1]*sc8[3]\n + sumy[2]*sc8[6] + sumy[3]*sc8[7] );\n }\n\n yp += 4 * QK_K;\n }\n\n // SIMD reduction + output for each row\n for (short row = 0; row < (short)NR0 && first_row + row < out_dim; ++row) {\n float sum = simd_sum(sumf[row]) + bias[first_row + row];\n if (apply_gelu) {\n if (sum > 10.0f) { }\n else if (sum < -10.0f) { sum = 0.0f; }\n else {\n float t = 0.7978845608f * (sum + 0.044715f * sum * sum * sum);\n sum = 0.5f * sum * (1.0f + tanh(t));\n }\n }\n if (tiisg == 0) {\n output[n * out_dim + first_row + row] = half(sum);\n }\n }\n}\n\n// ============================================================================\n// Q6_K SIMD matmul — exact llama.cpp pattern\n// Block layout: [ql:128B][qh:64B][sc:16B][d:2B] = 210 bytes\n// ============================================================================\nkernel void simd_gemm_q6k(\n device const uint8_t* w_raw [[buffer(0)]],\n device const half* x [[buffer(1)]], // FP16 input\n device const float* bias [[buffer(2)]], // F32 bias\n device half* output [[buffer(3)]], // FP16 output\n constant uint& in_dim [[buffer(4)]],\n constant uint& out_dim [[buffer(5)]],\n constant uint& batch [[buffer(6)]],\n constant uint& apply_gelu [[buffer(7)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n // Each simdgroup processes NR0 consecutive output rows, sharing input reads\n const uint first_row = (tgpig.x * N_ROWS + sgitg) * NR0;\n const uint n = tgpig.y;\n if (first_row >= out_dim || n >= batch) return;\n\n const uint nb = in_dim / QK_K;\n const uint row_bytes = nb * 210;\n\n device const half* yy = x + n * in_dim;\n\n constexpr uint8_t kmask1 = 0x03;\n constexpr uint8_t kmask2 = 0x0C;\n constexpr uint8_t kmask3 = 0x30;\n constexpr uint8_t kmask4 = 0xC0;\n\n // Thread decomposition (matches llama.cpp)\n const short tid = tiisg / 2; // 0..15\n const short ix = tiisg % 2; // 0..1 stride\n const short ip = tid / 8; // 0..1 which 128-elem half\n const short il = tid % 8; // 0..7\n const short l0 = 4 * il;\n const short is = 8 * ip + l0 / 16;\n\n const short y_offset = 128 * ip + l0;\n const short q_offset_l = 64 * ip + l0;\n const short q_offset_h = 32 * ip + l0;\n\n float sumf[NR0] = {0.f}; // accumulator per row\n float yl[16];\n\n for (uint i = ix; i < nb; i += 2) {\n // Load y data ONCE, reuse for all NR0 rows\n device const half* y = yy + i * QK_K + y_offset;\n\n FOR_UNROLL for (short l = 0; l < 4; ++l) {\n yl[4*l + 0] = y[l + 0];\n yl[4*l + 1] = y[l + 32];\n yl[4*l + 2] = y[l + 64];\n yl[4*l + 3] = y[l + 96];\n }\n\n // Process NR0 rows with shared y data\n for (short row = 0; row < (short)NR0; ++row) {\n device const uint8_t* bp = w_raw + (first_row + row) * row_bytes + i * 210;\n device const uint8_t* q1 = bp + q_offset_l;\n device const uint8_t* q2 = q1 + 32;\n device const uint8_t* qh = bp + 128 + q_offset_h;\n device const int8_t* sc = (device const int8_t*)(bp + 192) + is;\n device const half* dh = (device const half*)(bp + 208);\n\n float4 sums = {0.f, 0.f, 0.f, 0.f};\n\n FOR_UNROLL for (short l = 0; l < 4; ++l) {\n sums[0] += yl[4*l + 0] * float((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32);\n sums[1] += yl[4*l + 1] * float((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32);\n sums[2] += yl[4*l + 2] * float((int8_t)((q1[l] >> 4) | ((qh[l] & kmask3) << 0)) - 32);\n sums[3] += yl[4*l + 3] * float((int8_t)((q2[l] >> 4) | ((qh[l] & kmask4) >> 2)) - 32);\n }\n\n sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]);\n }\n }\n\n for (short row = 0; row < (short)NR0 && first_row + row < out_dim; ++row) {\n float sum = simd_sum(sumf[row]) + bias[first_row + row];\n if (apply_gelu) {\n if (sum > 10.0f) { }\n else if (sum < -10.0f) { sum = 0.0f; }\n else {\n float t = 0.7978845608f * (sum + 0.044715f * sum * sum * sum);\n sum = 0.5f * sum * (1.0f + tanh(t));\n }\n }\n if (tiisg == 0) {\n output[n * out_dim + first_row + row] = half(sum);\n }\n }\n}\n\n// ============================================================================\n// MoE variants — GPU-side expert_offsets, zero CPU-GPU sync\n// Input/output are packed contiguously: [expert0_tokens..., expert1_tokens..., ...]\n// Kernel reads expert_offsets to compute base offset + batch bounds\n// Dispatch: threadgroups = [ceil(out_dim/(N_ROWS*NR0)), seq_len, 1]\n// ============================================================================\nkernel void simd_gemm_q5k_moe(\n device const uint8_t* w_raw [[buffer(0)]], // weights for THIS expert\n device const half* x_packed [[buffer(1)]], // full packed input\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]], // full packed output\n device const int* expert_offs [[buffer(4)]], // [n_experts+1] offsets\n constant uint& expert_id [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n const uint first_row = (tgpig.x * N_ROWS + sgitg) * NR0;\n const uint n_local = tgpig.y;\n const uint base = expert_offs[expert_id];\n const uint eb = expert_offs[expert_id + 1] - base;\n if (first_row >= out_dim || n_local >= eb) return;\n\n const uint n = base + n_local;\n const uint nb = in_dim / QK_K;\n const uint row_bytes = nb * 176;\n device const half* y1 = x_packed + n * in_dim;\n\n const short tid = tiisg / 4;\n const short ix = tiisg % 4;\n const short iq = tid / 4;\n const short ir = tid % 4;\n const short l0 = 8 * ir;\n const short q_offset = 32 * iq + l0;\n const short y_offset = 64 * iq + l0;\n const uint8_t hm1 = 1u << (2*iq), hm2 = hm1 << 1, hm3 = hm1 << 4, hm4 = hm2 << 4;\n constexpr uint16_t kmask1 = 0x3f3f, kmask2 = 0x0f0f, kmask3 = 0xc0c0;\n\n float sumf[NR0] = {0.f};\n device const half* yp = y1 + ix * QK_K + y_offset;\n\n for (uint i = ix; i < nb; i += 4) {\n device const half* y2 = yp + 128;\n float yl[16], yh[16];\n float4 sumy = {0.f, 0.f, 0.f, 0.f};\n FOR_UNROLL for (short l = 0; l < 8; ++l) {\n yl[l+0] = yp[l+ 0]; sumy[0] += yl[l+0];\n yl[l+8] = yp[l+32]; sumy[1] += yl[l+8];\n yh[l+0] = y2[l+ 0]; sumy[2] += yh[l+0];\n yh[l+8] = y2[l+32]; sumy[3] += yh[l+8];\n }\n for (short row = 0; row < (short)NR0; ++row) {\n device const uint8_t* bp = w_raw + (first_row + row) * row_bytes + i * 176;\n device const uint8_t* q1 = bp + 48 + q_offset;\n device const uint8_t* qh = bp + 16 + l0;\n device const half* dh = (device const half*)bp;\n device const uint16_t* a = (device const uint16_t*)(bp + 4) + iq;\n device const uint8_t* q2 = q1 + 64;\n uint16_t sc16[4];\n sc16[0] = a[0] & kmask1; sc16[1] = a[2] & kmask1;\n sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2);\n sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2);\n thread const uint8_t* sc8 = (thread const uint8_t*)sc16;\n float4 acc1 = {0.f}, acc2 = {0.f};\n FOR_UNROLL for (short l = 0; l < 8; ++l) {\n uint8_t h = qh[l];\n acc1[0] += yl[l+0] * float(q1[l] & 0x0F);\n acc1[1] += yl[l+8] * float(q1[l] & 0xF0);\n acc1[2] += yh[l+0] * float(q2[l] & 0x0F);\n acc1[3] += yh[l+8] * float(q2[l] & 0xF0);\n acc2[0] += (h & hm1) ? yl[l+0] : 0.f;\n acc2[1] += (h & hm2) ? yl[l+8] : 0.f;\n acc2[2] += (h & hm3) ? yh[l+0] : 0.f;\n acc2[3] += (h & hm4) ? yh[l+8] : 0.f;\n }\n sumf[row] +=\n dh[0] * ( sc8[0] * (acc1[0] + 16.f*acc2[0])\n + sc8[1] * (acc1[1]/16.f + 16.f*acc2[1])\n + sc8[4] * (acc1[2] + 16.f*acc2[2])\n + sc8[5] * (acc1[3]/16.f + 16.f*acc2[3]) )\n - dh[1] * ( sumy[0]*sc8[2] + sumy[1]*sc8[3]\n + sumy[2]*sc8[6] + sumy[3]*sc8[7] );\n }\n yp += 4 * QK_K;\n }\n\n for (short row = 0; row < (short)NR0 && first_row + row < out_dim; ++row) {\n float sum = simd_sum(sumf[row]) + bias[first_row + row];\n if (apply_gelu) {\n if (sum > 10.0f) { }\n else if (sum < -10.0f) { sum = 0.0f; }\n else {\n float t = 0.7978845608f * (sum + 0.044715f * sum * sum * sum);\n sum = 0.5f * sum * (1.0f + tanh(t));\n }\n }\n if (tiisg == 0) {\n out_packed[n * out_dim + first_row + row] = half(sum);\n }\n }\n}\n\nkernel void simd_gemm_q6k_moe(\n device const uint8_t* w_raw [[buffer(0)]],\n device const half* x_packed [[buffer(1)]],\n device const float* bias [[buffer(2)]],\n device half* out_packed [[buffer(3)]],\n device const int* expert_offs [[buffer(4)]],\n constant uint& expert_id [[buffer(5)]],\n constant uint& in_dim [[buffer(6)]],\n constant uint& out_dim [[buffer(7)]],\n constant uint& apply_gelu [[buffer(8)]],\n uint3 tgpig [[threadgroup_position_in_grid]],\n ushort tiisg [[thread_index_in_simdgroup]],\n ushort sgitg [[simdgroup_index_in_threadgroup]])\n{\n const uint first_row = (tgpig.x * N_ROWS + sgitg) * NR0;\n const uint n_local = tgpig.y;\n const uint base = expert_offs[expert_id];\n const uint eb = expert_offs[expert_id + 1] - base;\n if (first_row >= out_dim || n_local >= eb) return;\n\n const uint n = base + n_local;\n const uint nb = in_dim / QK_K;\n const uint row_bytes = nb * 210;\n device const half* yy = x_packed + n * in_dim;\n\n constexpr uint8_t kmask1 = 0x03, kmask2 = 0x0C, kmask3 = 0x30, kmask4 = 0xC0;\n const short tid = tiisg / 2, ix = tiisg % 2;\n const short ip = tid / 8, il = tid % 8;\n const short l0 = 4 * il, is = 8 * ip + l0 / 16;\n const short y_offset = 128 * ip + l0;\n const short q_offset_l = 64 * ip + l0;\n const short q_offset_h = 32 * ip + l0;\n\n float sumf[NR0] = {0.f};\n float yl[16];\n\n for (uint i = ix; i < nb; i += 2) {\n device const half* y = yy + i * QK_K + y_offset;\n FOR_UNROLL for (short l = 0; l < 4; ++l) {\n yl[4*l + 0] = y[l + 0]; yl[4*l + 1] = y[l + 32];\n yl[4*l + 2] = y[l + 64]; yl[4*l + 3] = y[l + 96];\n }\n for (short row = 0; row < (short)NR0; ++row) {\n device const uint8_t* bp = w_raw + (first_row + row) * row_bytes + i * 210;\n device const uint8_t* q1 = bp + q_offset_l;\n device const uint8_t* q2 = q1 + 32;\n device const uint8_t* qh = bp + 128 + q_offset_h;\n device const int8_t* sc = (device const int8_t*)(bp + 192) + is;\n device const half* dh = (device const half*)(bp + 208);\n float4 sums = {0.f, 0.f, 0.f, 0.f};\n FOR_UNROLL for (short l = 0; l < 4; ++l) {\n sums[0] += yl[4*l + 0] * float((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32);\n sums[1] += yl[4*l + 1] * float((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32);\n sums[2] += yl[4*l + 2] * float((int8_t)((q1[l] >> 4) | ((qh[l] & kmask3) << 0)) - 32);\n sums[3] += yl[4*l + 3] * float((int8_t)((q2[l] >> 4) | ((qh[l] & kmask4) >> 2)) - 32);\n }\n sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]);\n }\n }\n\n for (short row = 0; row < (short)NR0 && first_row + row < out_dim; ++row) {\n float sum = simd_sum(sumf[row]) + bias[first_row + row];\n if (apply_gelu) {\n if (sum > 10.0f) { }\n else if (sum < -10.0f) { sum = 0.0f; }\n else {\n float t = 0.7978845608f * (sum + 0.044715f * sum * sum * sum);\n sum = 0.5f * sum * (1.0f + tanh(t));\n }\n }\n if (tiisg == 0) {\n out_packed[n * out_dim + first_row + row] = half(sum);\n }\n }\n}\n"

VERSION = 3

f32_to_f16(f : Float32) : UInt16