perf: sparse DeepSeek V4 prefill attention (25K-token TTFT 112 s → 31 s) - #2402

Motivation

Long prompts on mlx-community/DeepSeek-V4-Flash prefill slowly on a 4-node tensor-parallel (MlxJaccl) instance (4× Mac mini M5 Pro, Thunderbolt 5 RDMA). A 25.5K-token prompt took 112 s to the first token. The cost per 4096-token chunk grew from 11 s at the start of the prompt to 25 s at 20K tokens of context.

Profiling one chunk found three causes:

• Attention computes over every key in the chunk.
  • mlx_lm's prefill attends each query to all S + window raw keys and the whole compressed pool, then masks almost all of them away. Each query sees only 128 window keys plus at most 512 Indexer-selected rows.
  • Head_dim 512 has no fused SDPA in MLX, so this runs as a dense [S, S + window + pool] matmul.
  • On ratio-4 layers, mlx_lm also builds the Indexer's selection mask from a one-hot [S, top_k, pool] comparison. That is 13 G elements per layer at 25K context.
• Indexer scoring materializes [S, n_heads, pool] scores twice. The ReLU rewrites the whole tensor (2.7 GB at 20K) and the weighted head sum then reads it again. This runs replicated on every rank.
• The hyper-connection ops run mlx_lm's float32 paths at full length. hc_pre casts to fp32, normalizes and matmuls. hc_post uses an fp32 matmul. Together they take about 1.2 s per chunk.

Stacked on #2400 (which includes #2398).

Changes

1. patches/deepseek_v4_prefill_attention.py (new): sparse prefill attention.
   • Takes uniform prefill chunks (S > 4, RotatingKVCache, uniform pool) and leaves other cases to the existing path.
   • A Metal kernel visits only the keys each query can see: its sliding window, plus the visible compressed rows (on ratio-4 layers, only the Indexer's selected rows).
   • The kernel uses the same visibility rules as mlx_lm's masks. Rows the Indexer selects that are not visible yet stay excluded.
   • Scores, the online softmax with sinks and the weighted sum are float32. Each staged key is used for two heads per load (MQA), which is 1.5× faster than one head per simdgroup.
   • The rest of V4Attention.__call__ is mlx_lm's code: projections, compressor, Indexer and output projection.
2. deepseek_v4_indexer: fused ReLU + weighted head sum for prefill.
   • Scoring keeps mlx_lm's formula and selection semantics.
   • The ReLU and the head reduction become one kernel that reads the score tensor once and accumulates in float32, like the matmul it replaces.
3. deepseek_v4_decode_kernels: hc kernels at every length.
   • hc_post uses the existing fused kernel at prefill sizes too (3.5× faster).
   • hc_pre gets a many-row mix kernel: two rows per threadgroup, so x is read once and fn is reused across rows (5.3 → 3.0 ms per call).
   • Both match a float64 reference more closely than mlx_lm's fp32 matmuls do on this GPU.

Results

End to end, measured through the API with a single request (MTP on):

Decode speed is unchanged (59 tok/s, short prompt). Needle-in-a-haystack recall at 10%, 50% and 90% depth is 3/3 at every length.

Per 4096-token chunk at 20K tokens of context, standalone TP: 25.4 s before, 5.5 s after.

After, the time splits as follows: attention 3.35 s (including the Indexer's 1.25 s), MoE 1.6 s, hc_pre + hc_post 0.5 s.

Correctness

• Kernel vs float32 masked SDPA:
  • Synthetic tests use mlx_lm's own mask builders as the reference.
  • Mean absolute error vs the float32 reference is the same as for the existing bf16 path (6.0e-4 vs 6.0e-4, at S=4096 with 6144 compressed rows).
• Full model, 4-rank TP: I compared per-token output distributions between the old and new prefill over two 20K-token code corpora, on the real Flash weights.
  • Mean KL is 0.0026 / 0.0033. Changing only the chunk size (4096 → 2048) gives 0.0025, so the difference is at the level of summation-order noise. That baseline was measured with a chunk-size-independent Indexer selection; see the note below.
  • Argmax agreement is 99.2% / 99.3%.
  • Mean NLL changes by ±0.0001 nats.
• hc_post, hc_pre and the ReLU-weighted sum: the new kernels match float64 references more closely than the paths they replace.

Note: Indexer RoPE layout in mlx_lm

While testing I found that mlx_lm's Indexer.__call__ rotates its queries in [B, S, H, D] layout. mx.fast.rope takes positions from the second-to-last axis, so head h is rotated to position offset + h. DeepSeek's reference rotates by token position.

Its prefill top-k also ranks rows that complete later in the same chunk, which makes results depend on the chunk size.

I tried fixing both. KL vs the current behaviour was 0.005, but NLL did not improve measurably on three corpora. This PR keeps mlx_lm's Indexer semantics unchanged, and the issue is left for a separate discussion.

Tests

src/exo/worker/engines/mlx/tests/test_deepseek_v4_prefill.py uses the tiny DeepSeek V4 from #2400. It covers:

• The prefill kernel vs float32 masked SDPA for ratio 0, 4 and 128, with offsets inside and past the window, Indexer selections that include not-yet-visible rows, batch 2, and 8, 16 and 32 heads.
• The ReLU-weighted sum and the many-row hc_mix vs float64.
• Patched vs mlx_lm chunked prefill per attention layer, with chunks of 64 and 320 over 320 tokens, past the window and with the Indexer active.
• The full model's chunked prefill plus decode.
• Batched caches falling back.

uv run basedpyright (0 errors), uv run ruff check, nix fmt, and uv run pytest (546 passed).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XSvpsPpNrx5ZY24MwZ9xtd