research/mlx_speed_deepdive.md

# Can we code MLX ourselves for more decode speed? (June 2026 deep dive)

**Question:** the agentic coder's one weakness is decode speed (~11–14 tok/s). Can custom MLX/Metal code beat it?

## The wall (unchanged, re-confirmed)
Decode is **memory-bandwidth-bound** — generation speed ≈ *bytes moved per token* (every token streams the active
expert weights through the memory bus). TTFT is compute-bound (and the M5 fixes that); decode is not. Our **#33 fused
MoE dequant-matmul Metal kernel already IS the 11–14 tok/s** — it maxed the compute side. You cannot out-*compute*
bandwidth. "More speed" = fewer bytes/token, or faster per-byte processing on M5 hardware.

## Lever 1 — NVFP4 on the M5 hardware path *(easy, no custom code, biggest, do first)*
- **We already have it** (probed: `mx.quantize(..., mode="nvfp4", group_size=16)` works on our MLX 0.31.2).
- **Round-trip fidelity:** affine-3bit mean|err| **0.159** vs **NVFP4-4bit 0.077** — *half the quantization damage.*
- **M5 hardware path:** the Neural Accelerators have a native NVFP4 dequant+matmul; a 35B-A3B MoE went **58→112 tok/s
  (2×)** on M5 Max with NVFP4. It's 4-bit (more bytes than our 3-bit) but the hardware path + better fidelity net a win.
- **Action:** re-quantize the 4-bit layers in **#59** as NVFP4 (group 16), keep the middle 3-bit affine; measure decode.
  Zero custom code — just the quant mode + a `04b/24b` NVFP4 path.

## Lever 2 — a Metal-4 **TensorOps** fused-MoE kernel *(the "code it ourselves" win, ~30–60%)*
- **#33's kernel predates the M5** — it uses plain Metal, not the **Neural Accelerators' hardware matmul** (Metal 4
  TensorOps). MLX gets 30–60% on M5 by using TensorOps; a *custom* MoE kernel that does too captures that for our model.
- **How:** `mx.fast.metal_kernel` + Metal 4 **TensorOps** — slice the gather'd expert weights into tiles, tile-wise
  matmul on the Neural Accelerators, keep data in cache (WWDC26 session 330 "Optimize custom ML operations with Metal
  tensors" is the recipe). Fuse: expert-gather → NVFP4-dequant → matmul → accumulate, in one dispatch.
- **Effort:** real Metal work (~weeks). Reference impls exist (TurboQuant-MLX fused kernels; simdgroup-MMA int4-MoE kernels).

## Honest ceiling
Neither lever breaks the bandwidth wall — they reduce bytes (NVFP4) and speed per-byte (TensorOps on the Accelerators).
Realistic: **11–14 → ~20–28 tok/s** from NVFP4+M5 alone; **~30+** with the custom TensorOps kernel. A genuine **2–3×**,
not 100+. Speculative decode stays dead on this MoE (any multi-token verify reloads ~all experts — the bandwidth wall again).

## Plan (CPU now → GPU later)
1. **NVFP4 re-quant** — wire `nvfp4` (group 16) into `04b`/`24b` as the 4-bit-layer format for #59; re-quant; **measure decode** (the proof).
2. If the gain lands and we want more → **write the Metal-4 TensorOps fused-MoE kernel** (the 30–60% lever); benchmark vs #33.
3. Keep the verify-everything stack (constrained decode, compiler-steer) intact — none of this touches it.

*Sources: MLX custom Metal kernels (ml-explore docs 0.31.2) · `mx.fast.metal_kernel` · WWDC26 #330 Metal tensors ·
Apple "LLMs with MLX + M5 Neural Accelerators" · TurboQuant-MLX fused kernels · MLX-vs-llama.cpp M5 benchmarks (yage.ai 2026) · our SPEED.md root-cause.*