Metal (Apple GPU) backend for Kimi K3

[RDouglasSharp]

Metal (Apple GPU) backend for Kimi K3

Summary

Adds an opt-in Metal GPU path for Kimi K3, covering the KDA (Kimi Delta
Attention) token step and the KDA/MLA/lm_head projection matmuls, with a full
CPU fallback at every dispatch site. The MoE expert matmuls stay on CPU
(matmul_mxfp4), which is why decode — dominated by MoE — remains CPU/I/O-bound;
the GPU win is on the attention/projection compute. Enabled at runtime with K3_METAL=1 (build with
make -C c kimi_k3 METAL=1); default is off and the CPU path is unchanged.

On an M5 Max / 128 GB this gives ~1.7× prefill and ~2.4× attention on the
compute-bound phases. Decode remains I/O-bound (routed-expert streaming), so
wall-clock decode is governed by the expert cache, not the backend — see below.

Safety / semantics contract

• Default off. No behavior changes unless K3_METAL=1 is set at runtime and
  the binary was built with METAL=1. CPU-only builds are untouched.
• CPU path unchanged, and preserved as the correctness oracle. Every Metal
  dispatch is guarded by g_k3_metal && coli_metal_available() and falls through
  to the existing CPU code on any failure (self-healing: KDA fused dispatch flips
  g_k3_metal=0 and continues on CPU).
• Not token-exact vs CPU under greedy. Metal kernels accumulate in a
  different order than the CPU/AVX2 reference, so logits differ at ~FP32 epsilon
  and greedy (COLI_TEMP=0) can flip a near-tie token. Output stays coherent and
  logits track within tolerance; it is not bit-identical. Because the path is
  opt-in and default-off, this changes no default behavior. (Same class as the
  known Metal-prefill near-tie divergence.)

CPU fallback — every dispatch site

Every coli_metal_* call in kimi_k3.c is compiled under #ifdef COLI_METAL
and gated at runtime on g_k3_metal, so a Metal-enabled binary with K3_METAL
unset — or on a machine where the device is unavailable — never dispatches.
Site by site (line numbers as of this branch):

Call	Line(s)	Guard
coli_metal_init	768	inside the K3_METAL env check; result stored in g_k3_metal
coli_metal_matmul (KDA/MLA/lm_head proj, fmt 0/1/4)	335, 339, 343	if (g_k3_metal && coli_metal_available()) (332)
coli_metal_matmul (f32 low-rank)	362	if (g_k3_metal && coli_metal_available()) (361)
coli_metal_kda_fused_token	904	inside if (g_k3_metal) (873); return value checked — on failure sets g_k3_metal = 0 and falls through to the CPU path (self-healing)
coli_metal_shutdown	2029, 2082, 2141	if (g_k3_metal)

MLA KV write/clear are pure CPU — the former Metal calls are gone; only comments
reference them. No coli_metal_* call is reachable with g_k3_metal == 0, and
none is reachable in a non-COLI_METAL build.

What runs where (per layer)

Component	Backend	Notes
KDA projections (q,k,v,g,o)	GPU	w_matmul → coli_metal_matmul
KDA low-rank (fa,fb,bp)	GPU	k3_matmul_f32
KDA conv1d+SiLU, L2-norm, state recurrence	GPU	fused, one command buffer/token
MLA projections (qa,qb,kva,o)	GPU	w_matmul → coli_metal_matmul
lm_head	GPU	coli_metal_matmul
MoE expert matmul	CPU	matmul_mxfp4 / matmul_mxfp4_i8 — not GPU-dispatched
MLA attention loop	CPU	unchanged
MLA KV cache write/clear	CPU	see design note
MoE routing / RMSNorm / residual	CPU	unchanged
Expert loading (disk)	CPU	unchanged, backend-agnostic

Correctness

• Per-layer intermediate tensors (KDA and MLA layers) validated bit-for-bit
  against the CPU reference at the time of authoring (K3_VALIDATE_LAYER).
• KDA token-0 zero-state check: oh[i]/(vh[i]·beta) (= qn·kn) matches CPU to 7
  significant figures.
• End-to-end: full 93-layer inference runs to completion and stays coherent;
  logits track CPU within FP tolerance (greedy may diverge on near-ties, above).
• c/tests/test_backend_metal.mm exercises the primitives.

Performance (M5 Max, 128 GB, unified memory)

Compute-bound phases (prompt prefill, attention), Metal vs CPU at matched expert
hit rate:

Phase	CPU	Metal	Speedup
Prefill (13 tok)	45.4 s	27.2 s	1.7×
Attention (time: attn)	34.5 s	14.4 s	2.4×

Decode is I/O-bound (routed experts stream from disk), so wall-clock decode is
set by the expert LRU cache, not the backend. time: attn/moe/eload are
printed per run so the split is checkable rather than asserted.

Design notes worth a reviewer's eye

• Zero-copy state persistence (afcalloc). KDA state and conv windows are
  16 KB-aligned so Metal's wrap() takes the newBufferWithBytesNoCopy path —
  GPU writes land directly in host memory and persist across tokens. Without
  alignment the write goes to a throwaway copy and the recurrence never
  accumulates (decode degenerates). afcalloc falls back to plain calloc when
  COLI_METAL is undefined, so non-Metal/Windows builds are unaffected.
• Wrap-once buffer cache (wrap_persistent). Model-lifetime buffers (state,
  conv windows, taps) are wrapped once and reused instead of re-wrapping ~13
  buffers per token across 69 KDA layers. Keyed by host pointer, valid because
  those allocations are never freed/resized during inference.
• MLA KV cache kept on CPU. The MLA attention loop reads the cache from host
  memory, and the KV write is a trivial per-row rmsnorm+copy (negligible next to
  MoE). Keeping write+clear on CPU avoids an unaligned GPU round-trip and keeps
  the cache path consistent with the CPU attention that consumes it.
• Fused KDA token. conv_silu×3 + l2_norm + state recurrence encode into a
  single MTLCommandBuffer per token (~4× fewer submits than one-CB-per-kernel).

Config

• K3_METAL=0|1 — enable the Metal path (default 0).
• K3_EXPERT_GB=N — routed-expert LRU budget; the decode lever (I/O-bound).
• K3_LOAD_THREADS=N, K3_PIPE=1, K3_DIRECT=1 — expert I/O (backend-agnostic).
• OMP_NUM_THREADS — CPU compute team for the CPU-resident work.

Known limitations / follow-ups

• Not token-exact under greedy (FP reordering); opt-in mitigates.
• DSA indexer Metal kernels not implemented (CPU only).
• MLA attention loop and MoE routing remain CPU.
• Decode is expert-cache-bound; the cache/prefetch is the next lever, not the
  kernels.

Files

c/backend_metal.mm, c/backend_metal.h, c/kimi_k3.c, c/Makefile,
c/tests/test_backend_metal.mm, and docs (docs/metal_implementation.md,
docs/kimi_metal_gap_analysis.md, docs/METAL.txt).

[JustVugg]

Two housekeeping notes, then the review — which I want to do properly rather than quickly, because this is 3,800 lines across every engine.

Retargeted to dev. It was opened against main, which is protected here; everything lands on dev first. No action needed from you, and CI is green on the new base.

#763 is merged, with the default off. Thank you for turning that around without argument — the measurement was never in question, only which way the switch points on a project whose one unconditional promise is semantics.

Merge order, so you rebase once instead of twice

This PR overlaps two others:

	overlap with #790
#763 (yours)	backend_metal.mm — merged, done
#787 (KV prefix reuse)	kimi_k3.c — CI running now, small (+71/-4)

So the order is #763 → #787 → #790, which costs you one rebase at the end rather than a rebase after each. The alternative — merging 3,800 lines first and making a 71-line PR rebase around it — would be the wrong way round, and #787 also fixes a leak in kv_alloc that your kimi_k3.c changes sit next to.

I will ping you the moment #787 lands.

What I will be reading for

Not a rubber stamp — the parts I want to understand before merging:

• K3_METAL=1 default-off and the CPU fallback at every dispatch site. You state both; I want to check the second one is true at every site, because a single unguarded dispatch turns "opt-in" into "crashes on a Mac without the feature".
• compat.h, Makefile, openai_server.py — shared files touched by an engine-specific PR. Usually correct and worth a second look, since a change there reaches GLM, Inkling and OLMoE too.
• test_serve_sentinel.c — that test exists because of [Bug]: Loading Kimi K3 layers is slow, and then Colibri stops doing anything. #748, where the READY bytes got rewritten and Kimi hung forever with no error. I want to see what changed and why.
• The decode claim. You say decode stays I/O-bound and governed by the expert cache, not the backend. That matches everything we have measured, and stating it rather than implying a speedup you did not get is the reason I trust the ~1.7×/~2.4× on the compute-bound phases.

One question while it waits: on an M5 Max with 128 GB, roughly what expert residency did you have during those runs? K3 is ~1.6 TB, so I want to know whether the prefill numbers came from a warm cache or a cold one — it changes how they generalise to someone with 36 GB.

[RDouglasSharp]

Residency was low — these were effectively cold runs. ~65 GB resident (≈35 GB weights + ~30 GB expert LRU at K3_EXPERT_GB=44), so ~4% of the ~1.6 TB model, ~40–50% routed-expert hit, ~300 GB streamed per 16-token decode. So the prefill/attn numbers aren't from a warm cache.

They should still generalize down to 36 GB, because the 1.7×/2.4× are on the compute-bound phases (KDA attention + projections dispatched to the GPU), which don't depend on expert residency. MoE experts stay on CPU (matmul_mxfp4) on both paths, so residency doesn't enter the Metal-vs-CPU delta at all — it only affects decode wall-clock, which I deliberately didn't claim a speedup on.

Two clarifications while you review: (1) against dev the PR touches only Makefile, backend_metal.{mm,h}, kimi_k3.c, test_backend_metal.mm and docs — the compat.h/openai_server.py/test_serve_sentinel.c deltas were from the original against-main diff, not my changes. (2) MoE expert matmuls are CPU, not GPU — the GPU coverage is the KDA token step plus the KDA/MLA/lm_head projections.

Standing by for the #787 rebase.