---
title: "5.5 real sparse MoE kernels (Metal)"
sidebar:
  order: 505
---

<!-- PRD metadata (moved from frontmatter for Blume schema compatibility) -->
**status:** blocked-upstream  
**owner:** unassigned  
**created:** 2026-06-13  
**parent_plan:** docs/PLAN.md §3 Tier 5 (5.5)  
**related_prds:** docs/learn/moe.md (the design rationale already documented), docs/research/wave_2_5_kernel_audit.md (kernel landscape audit)

# 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 posttrainllm 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 posttrainllm routing; this PRD is kernels, not policy.

## Files to touch (when unblocked)

| File | Change |
|---|---|
| `Sources/TinyGPTModel/MoESparse.swift` | new — Swift wrapper around the Metal kernels |
| `Sources/TinyGPTModel/MoE.swift` | route between dense (existing) and sparse (new) based on a `--moe-routing {dense,sparse}` flag |
| `kernels/moe_topk.metal` | new — Metal shader for top-K + dispatch |
| `kernels/moe_combine.metal` | new — scatter_add equivalent |
| `evals/moe-flops-sparse-vs-dense.sh` | new — measurement |
| `docs/research/sparse-moe-metal.md` | new — publishable doc |
| `docs/PLAN.md` | 5.5 from blocked → ✅ when shipped |

## 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](https://github.com/ml-explore/mlx-swift) — the
  extension surface for adding new MLX ops backed by Metal.
- [DeepSeek-V3 paper](https://arxiv.org/abs/2412.19437) — 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 posttrainllm). **Recommendation:** yes — apple-silicon-
  sparse-MoE kernels are useful to projects beyond posttrainllm.