inkling: shared experts to the GPU, 1.75 to 2.46 tok/s on Apple Silicon

[rgbkrk]

Follow-up to #750, same validation discipline: two changes that take
Inkling-Small on the same Apple Silicon box from 1.75 to 2.46 tok/s decode
and 10.3 to 6.4 s prefill, with greedy output token-identical through eos
and the tiny multimodal oracle exact at every step.

The profile after #750 put the shared experts at 8.0 s of a 24-token probe,
the single biggest line. Both commits attack that.

Commit 1: overlap shared experts with the last GPU round

Routed and shared experts are independent given the layer input: both read
x, both accumulate into out, and the GPU only touches out in
_end's scatter-add after the wait. So the last routed round submits with
the existing moe_block_begin/_end pair and the CPU computes the shared
experts while the GPU is in flight. On a GPU fault _end returns 0, the
CPU redoes the round, and a shared_done flag guards double-compute.
No new math, no new kernel. 1.75 -> 1.96 tok/s.

Commit 2: shared experts on the GPU as a bf16 block

moe_gemv gains fmt=5: raw bf16 rows, no scales, dequant is a 16-bit
shift into the upper half of an f32. The branch is additive and every
existing format's code path is byte-untouched; glm/kimi call sites never
pass fmt=5, so this is dead code for them.

On the inkling side the shared-expert weights move into one page-aligned
slab per sparse layer at load, and the Wt handles repoint INTO the slab,
so the CPU fallback path reads the same bytes it always did. Each MoE layer
then submits one fmt=5 block before the routed rounds: every token uses
every shared expert, so group j is simply all S rows with weight w[K+j],
and the GPU chews it while the CPU routes, fills, and packs the routed
rounds. INK_METAL_SHARED=0 opts out.

Shared experts drop from 8.0 s to 0.1 s on the probe. 1.96 -> 2.46 tok/s
decode, prefill 10.3 -> 6.4 s. Audio prefill (113-token spoken prompt):
22.4 s, from 148.8 s on the day-one CPU path.

Validation

• Tiny multimodal oracle exact on CPU-f32 and container-int4 CPU/Metal
  (the fixture's f32 shared weights correctly skip the bf16 path and take
  the CPU loop).
• Real-model greedy generation token-identical through eos across both
  commits, including the audio transcript.
• make check green.
• One caveat to be transparent about: make -C c metal-test does not
  currently compile on dev (a pre-existing fp8_nblk redefinition in
  tests/test_backend_metal.mm, unrelated to this change), so the kernel
  parity suite could not be run. The fmt=5 branch is exercised by the live
  runs above; happy to add a fmt=5 case to the parity suite once that
  harness builds again.

Benchmark report

• Commit: rebased onto dev at ee0c668
• Hardware: Apple M-series Mac, 128 GB unified memory
• Model/storage: Inkling-Small, 131 GiB int4 container, internal NVMe,
  warm page cache
• Workload: frozen probe (fixed prompt, --cap 131, pins loaded,
  USAGE_SAVE=0), greedy, THINK=0
• Runs: each number reproduced twice, reps within 0.01 tok/s

	Decode	Prefill	Shared experts (phase)
#750 as merged	1.75 tok/s	10.3 s	8.0 s
+ overlap	1.96 tok/s	10.3 s	8.0 s (hidden)
+ bf16 GPU block	2.46 tok/s	6.4 s	0.1 s

Remaining profile for anyone hunting further: attention (7.0 s, CPU) is now
the wall, and the MTP head is still unloaded.

[JustVugg]

The measurements are excellent and the correctness discipline is exactly right — greedy output token-identical through eos across both commits is the claim that matters, and you made it the headline rather than the throughput. Shared experts from 8.0 s to 0.1 s on the probe, audio prefill 148.8 s → 22.4 s. That is a real change in what the model is usable for on that machine.

Commit 1 is a clean piece of reasoning too: routed and shared experts are independent given the layer input, so the CPU can compute one while the GPU is in flight, with shared_done guarding the double-compute on a GPU fault. No new math, and it buys 12%.

One thing I want changed before merging, and it is a number rather than any of the engineering.

fmt=5 is already taken

moe_gemv now handles fmt 2, 5 and 6, with 5 meaning "raw bf16 rows, no scales". But the Metal kernels share the global format numbering — 1 int8, 2 int4, 3 int2, 4 int4-grouped, 6 E8/IQ3, 8 fp8-e4m3, all matching FORMATS.md — and in that registry:

| 5 | int3-g64 | O*ceil(I/64)*24 bytes … |

merged in #465.

Your defence in the description is correct today: "glm/kimi call sites never pass fmt=5, so this is dead code for them". The problem is the day someone wires int3-g64 experts to the Metal MoE path — a perfectly reasonable thing to do, since fmt=5 exists precisely as an expert format — and lands in the bf16 branch. They will not get an error. They will get numbers, and the numbers will be wrong.

That is the failure mode FORMATS.md and the #529 container stamp were built to prevent: a format ordinal meaning two different things depending on which function you are inside.

The fix, and the repo already has the convention

From FORMATS.md:

PRIVATE ORDINAL BLOCK convention (this repo, pending upstream review)

… format ordinals from a private block starting at 100, never from the [shared range]

fmt=8 itself was developed as fmt=100 under exactly this rule before being assigned a real number — and it was then renumbered from 7 to 8 when #705 claimed 7 for MXFP4 while the PR was open. That convention exists because this has already happened once.

Two options, either is fine by me:

1. fmt=100 for the bf16 path — private, unambiguous, no registry change needed.
2. Register bf16 properly as the next free shared ordinal (9) with a FORMATS.md row. Worth it if you think bf16 rows will be wanted beyond Inkling's shared experts, which is plausible — every engine has some tensor it keeps unquantised.

I would take (1) now and (2) later if a second consumer appears, but you have thought about the bf16 path more than I have.

Everything else is a merge as soon as that lands. To be explicit that this is not a review of the engineering: the slab-per-sparse-layer with Wt handles repointing into it, so the CPU fallback reads the same bytes it always did, is the detail that makes me comfortable with the whole thing — the fallback is not a second implementation that can drift.

And the note that the fixture's f32 shared weights correctly skip the bf16 path and take the CPU loop is the kind of thing most people would not have checked, let alone written down.