perf(deepseek_v4): faster exact DeepSeek-V4 prefill (wsdpa fusion + scheduler/cache amortization)

[jonathan308]

Summary

Exact-preserving DeepSeek-V4 prefill speedup, rebased onto current main (post-#2559/#2562/#2563/cfec5021). Six commits:

1. c2f55836 fuse windowed+pooled prefill attention — JIT Metal kernel for the exact DSV4 shape (B=1, H=64, D=512, bf16), routed first for LocalAttention, CompressedAttention, and SparseCompressedAttention (small-pool and indexer top-k). Dispatch is gated on a _standard_mask flag threaded from the model loop (same pattern as perf(deepseek-v4): route ratio-128 prefill through native attention #2559): the kernels rebuild causality/sliding-window/pooled visibility from offsets, so direct callers with custom masks stay on the reference path.
2. d372749d setup.py CMake pinning so the native prefill kernels build reliably (the block-alignment half of the original commit is covered by cfec502 and was dropped).
3. 19827b93 skip discarded lm_head projection in intermediate chunked-prefill chunks (logits of non-final chunks are thrown away).
4. a77ab09e append pooled cache in place during prefill (removes per-chunk pooled-prefix re-concatenation).
5. 760fe945 fold the pooled-ratio causal mask into the indexer score kernel epilogue, with a signature probe (_probe_mask_fold): extension builds predating the fold kwargs keep the historical call signature (GLM-5.2 unaffected) and get the same mask in a second pass with identical sentinel semantics. The indexer tile config is fixed at bn=64; the env-gated bn=128 variant measured no win and was removed.
6. 9423196a test fixture compatibility for the skip_lm_head capability check (thanks @PhilipJohnBasile).

The MXFP4 tile-tuning commit from the previous revision was dropped: #2562 tunes the same constant and both pick BM16 at 2048-token chunks.

No precision, selection-quality, sampling, or cache-semantics changes. Decode, DSpark verify, batching, cache update, prompt-boundary snapshots, and fallback behavior are unchanged; every kernel path falls back to stock SDPA/native on shape-gate miss, setup failure, or runtime error. Escape hatches: OMLX_DSV4_WSDPA=0, OMLX_DSV4_WSDPA_TOPK=0.

Performance

@jundot's A/B on current main (M3 Ultra, V4-Flash oQ4e, cache on, cold runs, tg128), cherry-picking the three core commits:

context	main	+3 commits, wsdpa first
4,096	549.4	630.8
32,768	549.7	588.5

tg unchanged. The JIT kernel also beats the precompiled deepseek_v4_sparse_attention head to head and helps installs without the compiled extension.

The remaining high-context taper is architectural: the DSA indexer scores all context/4 pooled rows per query, so prefill FLOPs grow linearly with absolute context. Kernel-traffic variants (K-tile hoist, bn=128) were implemented and measured at 0% gain — excluded.

Output-parity note (release-note worthy)

Greedy/temp-0 outputs after this PR are tolerance-level parity, not bit-identical, with current main's native ratio-128 path: floating-point reduction order differs, so greedy flips can occur where the model is effectively indifferent between top tokens — far below the checkpoint's own quantization noise (the same standard the merged native path landed under). Each change was validated temp-0 exact-match against its own pre-change build at 10k/40k/75k-token prompt sets, and the fused kernels are bit-exact against explicit fp32 references (including trimmed-RotatingKVCache coordinate regression tests — both fail pre-fix, pass post-fix).

Upgrade note for the release log: identical prompts can produce different (equivalent-quality) outputs after this change.

Validation

• Focused suite on this branch: 618 passed (test_deepseek_v4_wsdpa, test_dsa_indexer_fused_mask (bit-identity incl. fresh bn=64 build), test_custom_kernel_abi_probe (old/new/missing mask-fold signatures + exact second-pass masking), test_pooling_cache_*, test_scheduler*, test_deepseek_v4_patch, test_deepseek_v4_dspark, test_glm_*, test_prefill_oom_graceful)
• Prefix cache intact: 96.7–96.9% hit at ~38k turns (2048-block-aligned flooring); warm repeated turn 25.6 s → 3.6 s
• KV footprint flat across a full 100k cold prefill (0.0 GB settled delta); ~13.1 GB accounted KV+gate at 1,048,576 tokens

Evaluated and deliberately excluded (with data): indexer 64-alignment padding (~0.5% gain, shifts temp-0 numerics), mx.compile MoE wrapper (0%), mx.compile decode step (unvalidated state risk), K-tile/bn=128 indexer variants (0%, FLOP-bound), MXFP4 tile re-tune (superseded by #2562).

[jonathan308]

Follow-up validation / fix in adbf4f54:

• Root cause for the missing long-context acceleration was an install-path defect: the documented native-kernel build could let CMake choose the system Python instead of the active oMLX environment, where nanobind was unavailable.
• setup.py now pins CMake to the interpreter running pip. OMLX_WITH_CUSTOM_KERNEL=1 python -m pip install -e . --no-build-isolation builds successfully without manual CMake overrides; all native kernel modules load.
• DeepSeek rotating-cache alignment now defaults to the validated 2,048-token block size, matching the WSDPA benchmark configuration.
• Focused validation: 290 passed, 1 skipped (test_deepseek_v4_wsdpa, DeepSeek patch/DSpark/MTP regressions, and scheduler alignment).

Current cold-server measurements on this machine with the native kernel path, WSDPA top-k enabled, 2,048-token blocks, and cached_tokens=0 were 406 tok/s at 9,825 prompt tokens and 423 tok/s at 48,899 prompt tokens. I also tested LM-head skipping, MLX graph compilation, alternate native sparse attention dispatch, and 4k chunk boundaries; none produced a repeatable gain, so none are included. These figures do not substantiate a 600 tok/s claim.

[jonathan308]

Follow-up validation after the native-kernel build fix:

• M3 Ultra, DeepSeek-V4-Flash-0731-MXFP4-MLX, cold 48,671-token prompt (cached_tokens=0).
• Controlled no-cache A/B: existing 32x32 MXFP4 tile = 425.0 / 425.1 tok/s; new 16x32 tile = 427.5 / 427.7 tok/s (+0.6%).
• Cache-enabled end-to-end confirmation: 48,671 tokens / 113.97 s = 427.1 tok/s, still cached_tokens=0.
• This only changes tile selection among already-built native kernels; it does not change routing, weights, cache behavior, or numerical operations.

Validation: native MoE parity tests passed; all five WSDPA/trimmed-rotating-cache regression tests passed; a live temperature-0 arithmetic response remained correct.

The older ~470 tok/s values in the original draft could not be reproduced under this clean, cache-safe contract, so please use the measured 427.1 tok/s high-context figure for release/PR notes. The PR remains open with the build fix, the long-context gibberish fix, and this verified incremental MoE improvement.

[jonathan308]

Validation complete — reopening. Summary of what changed since this was closed:

What was validated

• All 6 commits re-validated end-to-end on a clean from-scratch build (M3 Ultra 256 GB, DeepSeek-V4-Flash-0731-MXFP4-MLX)
• Temp-0 deterministic output equivalence old vs new path: exact match at 10k / 40k / 75k token prompt sets, per change
• Kernel bit-exactness vs fp32 references incl. trimmed-RotatingKVCache coordinate regression tests (fail pre-fix / pass post-fix)
• Focused pytest suite: 397 passed
• Prefix-cache behavior intact (96.7–96.9% hits at ~38k, block-aligned flooring); KV footprint flat across a 100k cold prefill (0.0 GB settled delta)

Honest benchmark correction
The ~470 tok/s figures originally posted here were not reproducible — they came from prefix-cache contamination in the benchmark harness (cached_tokens > 0 on supposedly cold runs) plus a local dev state that no longer exists. Reproducible cold-prefill numbers from a clean build of this branch:

Prompt tokens	upstream main	this PR
~10,350	—	430–448 tok/s
~41,100	328–345 tok/s	427–433 tok/s
~77,000	~318 tok/s cap	404–405 tok/s
~102,700	~318 tok/s cap	380–391 tok/s
~153,900	~318 tok/s cap	335 tok/s

The upstream cap comes from a measured ~1.6 s fixed cost per 512-token rotating-block boundary (snapshot + stream sync + cache clear + scheduler step); the 2048 alignment in adbf4f54 amortizes it losslessly. The residual taper at high context is the DSA indexer's all-context scoring — FLOP-inherent, not a regression (kernel-traffic variants measured 0% and were excluded).

Everything in the PR body is updated to match. Happy to run any additional validation you'd like to see.

[PhilipJohnBasile]

I traced the three failing CI jobs to the same test-fixture omission and opened a one-line patch against this PR's head branch: jonathan308#1

The patch adds _supports_skip_lm_head=lambda: False to the lightweight SimpleNamespace used by test_step_prefill_reclaims_before_first_guard. It changes no production code and keeps the test on its historical stock-model path.

Verification on the exact patch:

• formerly failing regression: 1 passed
• full tests/test_prefill_oom_graceful.py: 45 passed
• git diff --check: passed
• frozen one-line diff independently reviewed: approved with no blockers

The original CI run is the negative control: the same test fails identically on Python 3.11, 3.12, and 3.13 after 7,982 other tests pass.

[jonathan308]

CI fix merged: @PhilipJohnBasile correctly traced the 3 failing jobs to a test-fixture omission — test_step_prefill_reclaims_before_first_guard invokes the unbound scheduler method on a SimpleNamespace that was missing the new _supports_skip_lm_head capability stub from e3c6912b. Merged their one-line fixture patch (jonathan308#1) after local verification: the previously failing test passes, the full test_prefill_oom_graceful.py file passes 45/45, and the focused suite passes 442/442. No production code changed.

[jundot]

@jonathan308, Thank you for all the work on this one. The wsdpa kernel in particular turned out to be the fastest DeepSeek V4 attention path I have measured so far, so I would really like to get the good parts of this PR in. Since it was opened, #2559, #2562, #2563 and a follow-up (cfec502, DeepSeek V4 cache blocks aligned to 2048) landed on main, which changes the picture a bit, so here is where each commit stands.

adbf4f5: main now covers this through cfec502, so only the setup.py CMake pinning part is still needed, and that part is a nice fix. One note from my machine: the ~1.6 s per 512-boundary cost did not reproduce here, the win seems to come from chunk size itself.

981f1d6: I think this one can go. #2562 tunes the same constant, and at 2048-token chunks both versions pick BM16 anyway, so it would not add anything on top.

95b27cb + e3c6912 + f601239: these three are the keepers. I cherry-picked them onto current main and measured (M3 Ultra, V4-Flash oQ4e, cache on, cold runs, tg128):

context	main	+3 commits, wsdpa as fallback	+3 commits, wsdpa first
4,096	549.4	599.6	630.8
32,768	549.7	564.1	588.5

tg unchanged. Your JIT kernel beats the precompiled deepseek_v4_sparse_attention head to head, so the dispatch should try wsdpa first, and it also helps installs without the compiled extension. Two small asks for the rebase: a standard-mask guard like the one #2559 threads through (the kernel rebuilds causality from offsets, so custom-mask callers need to stay on the reference path), and dropping the env-gated bn=128 variant.

a50836a: this one is close, it just needs a compatibility gate. fast.py currently passes the new kwargs unconditionally, so an older built _ext raises TypeError: GLM-5.2's native path stops working and the V4 indexer quietly falls back while the startup probe still reports the kernels as available. A signature probe in fast.py plus an abi-probe test update would cover it. The mask math itself checked out nicely.

One correction on the claims, which I want to frame carefully: temp-0 outputs do change on top of main's native ratio-128 path, so tolerance-level parity is the accurate description rather than exact. To be clear, that is not a quality loss. Only the floating-point reduction order differs, greedy flips only happen where the model is effectively indifferent between the top tokens, and the error is far below the checkpoint's own quantization noise. The already-merged native path landed under exactly the same standard. It is just worth a release-note line, since identical prompts will produce different outputs after the upgrade.

I know this PR has been through several rounds already, and the remaining pieces are worth the effort. Happy to rerun the same A/B on your rebased branch whenever it is ready.

[jonathan308]

@jundot Rebased onto current main with every ask addressed — ready for your A/B rerun whenever you'd like.

Per-commit:

• adbf4f54 → reduced to the setup.py CMake pinning only (d372749d); the block-alignment half is dropped since cfec502 covers it. Noted on the 1.6 s/512-boundary figure — that was measured on our tree at the time; if it doesn't reproduce on your side, the chunk-size explanation works for us either way, and it's moot now.
• 981f1d69 → dropped entirely; agreed perf(deepseek-v4): raise MXFP4 large-block threshold to 16K #2562 covers it at 2048-token chunks.
• The three keepers are rebased (c2f55836, 19827b93, a77ab09e). Dispatch now tries wsdpa first in all three attention classes. A _standard_mask guard is threaded from the model loop through the block into LocalAttention, CompressedAttention, SparseCompressedAttention, and _sparse_pooled_attention (same pattern as perf(deepseek-v4): route ratio-128 prefill through native attention #2559) — any direct caller with a custom mask stays on the reference path. The env-gated bn=128 variant is removed; bn=64/wm2/wn2 is fixed in the kernel with the measurement note kept in a comment.
• a50836ac → now has the compatibility gate (760fe945): fast.py probes the built extension's dsa_indexer_scores signature once at import (_probe_mask_fold, via the nanobind-rendered __doc__). Pre-fold builds keep the historical call signature, so GLM-5.2's native path is untouched, and when the fold is requested the same mask is applied in a second pass with identical sentinel semantics. test_custom_kernel_abi_probe.py gained cases for fold-aware/pre-fold/missing signatures plus exact second-pass masking on a simulated pre-fold build; the existing bit-identity suite passes on a fresh bn=64-only build.
• Claims framing: accepted and corrected — the body now says tolerance-level parity vs main's native ratio-128 path (reduction-order differences only, greedy flips at token-indifference points, below quantization noise), keeps "exact-match vs own pre-change build" scoped to our validation, and carries an explicit release-note line about differing outputs for identical prompts.

Focused suite on the rebased branch: 618 passed, including the fused-mask bit-identity tests on a freshly built extension and the trimmed-cache coordinate regression tests.

[jundot]

Reran the A/B on the rebased branch (M3 Ultra, V4-Flash oQ4e, cache on, cold runs, tg128):

context	main	this PR
4,096	548.8	630.6
32,768	549.5	588.3

tg unchanged, and the numbers land right where the wsdpa-first slice predicted. I also verified the compatibility gate against a real pre-fold extension build: the probe detects it, the historical call signature is kept so GLM-5.2 stays on its native path, and the second-pass mask is bit-identical to the fused semantics. Greedy outputs shift as expected from the kernel swap and are deterministic per build. Everything I asked for is in, so this is good to merge on my side. Merging now, and I will pick up a couple of cosmetic leftovers in a follow-up (the unused bn=128 metal instantiations and an inline import).