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.