perf(deepseek-v4): skip redundant all-pooled indexer scoring

[DiscoStew6082]

Summary

• Skip indexer query projection, scoring, and top-k selection when every pooled row already fits within index_topk.
• Still run the indexer compressor so its PoolingCache advances exactly as before.
• Guard against divergence between the attention and indexer compressor row counts.
• Retain the existing full indexer path for empty pools and pools larger than index_topk.
• Leave downstream dense/native attention dispatch unchanged.

Correctness

The regression test uses the existing full indexer as the reference for a non-aligned ratio-4 prefill: 10 input tokens produce two pooled rows and a two-token remainder.

It verifies that:

• the full indexer scoring path is not invoked;
• the indexer compressor is invoked exactly once;
• the attention and indexer caches advance to the expected pooled offset;
• the remainder and every public element of PoolingCache.state exactly match the reference path;
• the reference selection contains the expected sequential pooled indices.

Performance evidence

Issue #2558 attributes approximately 0.4% marginal throughput improvement to this slice in its controlled leave-one-out testing.

That result is scoped to:

• Mac Studio with M3 Ultra and 512 GB unified memory;
• DeepSeek-V4-Flash-0731 at revision 7872f01b1d1fe23eabc4c98b48bffcef5a386062;
• fixed 17,219-token prompt, SHA-256 a1465f4b5ee68dbd173c138dd65718bd06957439b66e99069e91f50364bf81f1;
• cold cache with one request at a time;
• balanced memory guard and 2,048-token prefill steps;
• MTP enabled with three draft tokens.

This is an M3 Ultra result, not an unconditional cross-chip performance claim.

Evidence package:
https://github.com/DiscoStew6082/omlx/tree/06b515b0fd3543341b346e74cf0d3a13da2d6c06/benchmarks/evidence/deepseek_v4_flash_prefill_m3_ultra

Validation

• tests/test_deepseek_v4_patch.py: 92 passed
• tests/test_deepseek_v4_dspark.py: 64 passed, 11 skipped
• Black formatting check passed
• Python compilation check passed
• git diff --check passed
• Independent code review found no issues

Refs #2558

[jundot]

I reviewed and benchmarked #2559, #2562, and #2563 together since they merge cleanly as a set, so I am posting the same summary on all three. Setup: M3 Ultra 512 GB, DeepSeek-V4-Flash-0731 oQ4e (4-bit top level with mxfp4 experts, so all three changes are active), MTP off, single request, cold prefill with cached_tokens=0 verified on every run, one run per point, current main vs main plus all three PRs, using the built-in throughput benchmark (tg128).

With the prompt cache disabled, which gives 2048-token prefill steps matching the #2558 methodology:

context	main pp tok/s	+3 PRs pp tok/s	delta
4,096	464.1	546.3	+17.7%
16,384	487.8	559.6	+14.7%
65,536	463.4	525.1	+13.3%
131,072	416.1	460.8	+10.7%

tg was unchanged everywhere, which is expected for prefill-only changes. This independently reproduces the #2558 ablation on a different checkpoint, so the numbers hold up well.

One finding on my side rather than yours: with the default server config the paged SSD cache aligns the rotating block size to 512, which clamps prefill to 512-token chunks, and in that regime these gains mostly disappear (the MoE route count stays below the MXFP4 crossover and the ratio-128 win shrinks with the chunk). Raising the DeepSeek V4 block alignment to 2048 with the cache still enabled recovers the full numbers above (549.5 / 558.7 / 524.2 / 461.6 pp tok/s). I will handle that follow-up separately so default-config users actually see the improvement. Merging all three.