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.

[RDouglasSharp]

Rebased onto current dev (post #787) — the single rebase we'd planned.

The kimi_k3.c / Makefile / backend_metal.h conflicts were mechanical: the Model struct gains both kvp and the DSA fields; I took your #787 grow-don't-restart kv_alloc over the branch's free-and-realloc, since prefix reuse needs those buffers kept; the main cleanup keeps USAGE_SAVE alongside the Metal shutdown; the header takes dev's qgs signature.

backend_metal.mm needed real care rather than a merge. The branch forked ~228 commits back and that file's been rewritten heavily since (fmt=4 grouped, fmt=6, fmt=8, top-8), so I didn't trust git's silent auto-merge of it. I rebuilt it deliberately: current dev's file plus only the K3-additive code (the KDA kernels, the coli_metal_kda_* dispatch, wrap_persistent), with the three shared functions the branch also touches hand-merged onto dev's versions — fmt_scale_bytes, coli_metal_matmul, coli_metal_init.

That was warranted: the auto-merge had silently kept the branch's older fmt_scale_bytes, which drops dev's per-row scale catch-all for fmt 0/2/3 — so int4/int2/f32 GEMV returned garbage (nerr=1.0) while int8 and the grouped/fp8/e8 paths passed. I restored dev's catch-all and kept the fmt=6 line, then audited the whole delta: every dev line the branch removes is a deliberate replacement, nothing else dropped.

Tests: make metal-test green (including the fmt GEMV arms that caught the above), make check green, make kimi_k3 METAL=1 builds.

[JustVugg]

@RDouglasSharp — you rebased on August 8, explained the conflict resolutions in detail (taking #787's grow-don't-restart kv_alloc because prefix reuse needs those buffers kept, USAGE_SAVE alongside the Metal shutdown, dev's qgs signature), and then nobody answered you for eleven days. That is a maintainer failure, not a problem with your PR, and I'm sorry. In that window dev took 71 merges, so the branch is conflicting again — which means the work you did on request has been partly invalidated by our silence.

Where this actually stands, stated plainly so you can decide whether it is worth your time:

• We want it. Kimi K3 is the family with the least GPU coverage, and Apple silicon is where a 2.8 T model streaming from disk is most compelling.
• We cannot test it. No maintainer here has a Mac. That means the merge bar is your measurements plus CI, and I will not pretend otherwise or make you chase a verification we can't perform.
• The honest cost: dev has moved a lot (a sixth engine and its CUDA tier, the expert-matmul path rebuilt, a new ARM CI job). backend_metal.mm is the file most likely to have drifted under you.

Two ways forward, your call:

1. You rebase once more and we merge on green CI without another wait — I'll watch it and merge the same day. If it conflicts again after that, that's our fault for not moving fast enough, and we'll carry the rebase ourselves.
2. We carry the rebase, opening a PR that keeps your commits and authorship (we did exactly this for feat(moe): FUSED3=1 opt-in AVX2 expert matmul — 40% less matmul time, bit-identical output, off by default #1024 → feat(moe): FUSED3=1 opt-in AVX2 expert matmul (olmoe) — rebase of #1024 by @outtodata #1082 last week), and you review the conflict resolution rather than performing it.

Say which and it happens. v1.7.0 is being cut from dev right now and this is not in it — that is a consequence of our delay, not a judgement on the work.

[JustVugg]

Done — we carried the rebase: #1113, your two commits, your authorship. You review the conflict resolution instead of performing it again.

Four conflicts, each decision stated in the PR body: Makefile deps unioned with your METAL_OBJ; struct W unioned with dev's mmap fields; your duplicate g_k3_vk declaration dropped because dev already has one; and dev's Apple-Paravirtual skip guard kept in the test (it's what stops CI hanging 47 minutes) with your printf text. backend_metal.mm did not conflict — your August 8 work still applies cleanly, which says something about how carefully you did it.

All four engines build clean on Linux after the rebase. What we can't do is run Metal — so if you can confirm the branch still works on your M5 Max, that closes it. If you're unavailable, we'll merge anyway: our delay caused this, and it shouldn't cost you a second round.

Targeting v1.7.0, which is being cut today.

[JustVugg]

Merged via #1113 with your commits and authorship intact — GitHub now credits you for the Metal Kimi K3 backend. It ships in v1.7.0, being cut today. Thank you for the patience our silence did not deserve.