PRD — Custom Metal kernels for sparse MoE routing

Goal

Build the missing-upstream Metal kernels (sparse top-K routing, scatter_add equivalent, expert-parallel matmul gather) so MoE in TinyGPT can run hard routing — actually skipping inactive experts’ compute — instead of the dense soft-routing fallback that currently ships. Measure FLOP reduction vs the dense baseline; publish as “from-scratch sparse-MoE kernels on Apple silicon.”

Tier-5 research arc. Blocked upstream: MLX-Swift has no scatter_add primitive, which is the root reason hard routing isn’t shipped. Unblocks once either (a) MLX adds the op, or (b) we write the Metal shader ourselves through the MLXFast.metallib extension surface.

This PRD documents the design + the gate work that makes the implementation viable when unblocked. Implementing the kernel itself is the actual ship.

Why now

• Filed as ⬜ in PLAN.md Tier 5 with a known upstream-blocked status. Without a PRD, the design + measurement plan exists only in docs/learn/moe.md (rationale) and the wave-2.5 kernel audit (landscape). This PRD anchors the executable plan when MLX-Swift catches up.
• DeepSeek-V3 family is the production target. The kernel audit’s verdict was: MLX quantized_matmul is hand-tuned for dense; sparse-MoE remains “different family” until proven otherwise.
• If we write the Metal kernel ourselves, we end up with a piece of Apple-silicon-sparse-MoE infra not available elsewhere — publishable on its own.

Scope — in

• Design doc + kernel shapes in this PRD (the implementation ships when unblocked):
  • Top-K routing kernel: routing_logits → (top_k_indices, top_k_weights) in one shader, no host roundtrip.
  • Expert-dispatch: (tokens, expert_indices) → grouped_tokens — Metal threadgroup-shared rearrangement. The hard part.
  • scatter_add equivalent: post-expert combine; threadgroup- atomic add into the output buffer.
• Measurement framework: evals/moe-flops-sparse-vs-dense.sh — runs the existing dense MoE forward + the (when-shipped) sparse forward; reports FLOP and wall-time ratio.
• Acceptance gate (when implemented): sparse-MoE forward must be within 5% of dense’s output (KL ε) and ≥ 1.5× faster on a realistic config (8 experts, top-2, d_model=2048).
• Documentation: docs/research/sparse-moe-metal.md — the publishable artifact (shipping path + measurement results) on land.

Scope — out

• Sparse-MoE training (backward pass for sparse routing). V1 is forward-only; training stays on dense routing. The forward win alone justifies the work for inference.
• Multi-device MoE (expert parallelism across multiple Macs). Out of scope.
• Routing-policy research (top-K vs top-P vs learned routers). Use the existing TinyGPT routing; this PRD is kernels, not policy.

Files to touch (when unblocked)

Don’t touch (until unblocked)

• The current dense-MoE path. Sparse is sibling, default off.
• Anything else in TinyGPTModel/ — keep the diff localized.

Acceptance criteria (when implementation ships)

• Sparse-MoE forward matches dense-MoE forward to KL ε on the shipped MoE config (8 experts, top-2, d_model=2048) over 100 random batches.
• Wall-clock for the sparse path is ≥ 1.5× faster than dense on M5 Pro for the same config.
• FLOP-count instrumentation reports the expected (top_k / total_experts) activation ratio.
• moe-flops-sparse-vs-dense.sh passes; rows append to the shared E0 schema.
• docs/research/sparse-moe-metal.md publishable.

Reference patterns

• kernels/ (if it doesn’t exist yet, this PRD creates the directory) — Metal shader source lives here. Build glue lives in Package.swift; coordinate with the maintainer for the addition.
• MLXFast — the extension surface for adding new MLX ops backed by Metal.
• DeepSeek-V3 paper — the reference MoE arch; cite the routing details.

Open questions

• Whether to wait for MLX-Swift to ship scatter_add or write the Metal kernel ourselves. Recommendation: wait until the rest of Tier 5 (5.1, 5.2, 5.3, 5.6, 5.7) ships first. Sparse-MoE kernels are the highest-effort lowest-immediate-product-value Tier-5 item; ordering it last is correct.
• Whether to publish the Metal kernel as a standalone package (separable from TinyGPT). Recommendation: yes — apple-silicon- sparse-MoE kernels are useful to projects beyond TinyGPT.