Skip to content

perf(core): fused sparse-attention decode via page indirection - #1005

Merged
inureyes merged 4 commits into
mainfrom
feature/issue-904-sparse-paged-decode
Aug 3, 2026
Merged

perf(core): fused sparse-attention decode via page indirection#1005
inureyes merged 4 commits into
mainfrom
feature/issue-904-sparse-paged-decode

Conversation

@inureyes

@inureyes inureyes commented Aug 2, 2026

Copy link
Copy Markdown
Member

Summary

Sparse attention is reduced to page indirection over the #898 fused v2 decode kernel: a selection becomes a page_size = 1 page table, the KV allocation becomes the pool through a pure reshape, and the kernel visits the selected rows and nothing else. MiniMax-M3 block-sparse decode is routed onto it. DeepSeek Sparse Attention is not, for a reason that is a property of the kernel rather than of the selection, and that reason is costed rather than asserted.

Two premises in the issue turned out to be wrong and are corrected here rather than worked around. Both are documented in docs/sparse-paged-decode.md.

The reduction, and why it needed no kernel change

The v2 kernel resolves a visible token through indices[indptr[r] + (fpo + i) / page_size] and never assumes those rows are adjacent, ordered, or complete. A sparse selection is therefore a shorter page list, not a new kernel. At page_size = 1 each page holds one token and the page indirection becomes a per-token gather executed inside the attention loop, so no gathered copy is materialized at any point.

ContiguousCacheLayout addresses a dense [B, H, Cap, D] KVCache allocation as a [B * H * Cap, 1, 1, D] pool, which is a pure reshape, and folds the KV-head axis into the request axis so one CTA still owns one KV head and a group of its query heads. The reshape is of the allocation, not of the window update_and_fetch returns: that window is a token-axis slice of a step-padded buffer, and reshaping it would copy the whole cache and defeat the point. KVCache::raw_kv_allocations is the accessor, and it returns None for every mode whose rows are not raw fp16 K/V, for pool-backed caches, and for a head-trimmed window.

The selected content is a device value but the cardinality is not, because a top-k selection picks k of something on every step by construction. SparseCsrStructure is therefore built on the host with no synchronization while indices stays a lazy MLX array that is never read back.

Corrected premise 1: DeepSeek Sparse Attention does not land

The issue lists DSA first. Its decode score is scale * (q_latent . kv_latent[t]) + scale * (q_pe . k_pe[t]), and the second term reaches SDPA as an additive per-(head, position) mask. The v2 kernel computes scale * (q[h] . k_pool[row]) and has no input through which such a term can arrive; the selection plumbing applies to DSA unchanged, the kernel does not. Three ways out are costed in the doc: materializing the bias needs the very gather being removed (or 8 MB and 8x the MACs over the full context), folding the positional term into the pool row costs a ~19 MB per-layer-per-step copy at 16K unless the cache field order changes, and the real answer is a row-stride plus second-K-stream kernel generalization (the shape FlashInfer's paged-MLA kernel uses) that composes with everything here and needs no cache change. It was not attempted because it modifies the kernel #899 just shipped into production, for a family whose only checkpoint (~685B) cannot be run on any available machine, and a fused decode path that silently never activates is this epic's most expensive failure mode. One further fact for that follow-up: with Hkv == 1, head folding yields one CSR request per sequence, so a batch-1 MLA decode faces the 4096-row single-request floor while index_topk defaults to 2048, and the floor has to be re-derived for Hkv == 1 before MLA can dispatch at all.

phi3-small and phimoe are also not routed: their block pattern varies per query head ((kb + h + 1) % vert_stride == 0), and a CSR request necessarily shares one page list across the n_rep query heads grouped under its KV head.

Corrected premise 2: a token floor is not sufficient, sparsity has to be gated too

The issue assumes replacing the mask with a fused pass is a win at any sparse configuration. It is not. mlx::fast::scaled_dot_product_attention is a heavily tuned dense kernel and the v2 partial kernel is a scalar per-lane sweep, so reading half the data through a kernel with twice the constant factor loses. examples/sparse_paged_decode_bench at MiniMax-M3 geometry, one repetition of 40 steps on an idle Apple Silicon host (indicative, not a recorded result): 0.67x at 2x sparsity (ctx 4096), 0.77x at 4x (8192), 1.22x at 8x (16384), 1.17x at 16x (32768), 2.06x at 32x (65536), with transient memory dropping at every point.

So a launch must clear MIN_SPARSITY_RATIO (default 8, MLXCEL_SPARSE_PAGED_MIN_SPARSITY, 0 disables) in addition to the #899 token floor. The default sits on top of the measured win rather than interpolated into the unmeasured 4x-to-8x band: #899 argued the opposite way for its token floor because what lay under its weakest point was a missed opportunity, whereas here what lies under the point is a measured regression. The harness itself flatters the mask arm twice over (it builds the mask on the host outside the timed region, and hands the sparse arm a prebuilt block list), so the real crossover should sit lower, which makes 8 conservative in the direction that cannot regress anything. At MiniMax-M3's shipped configuration the gate opens at 16K, which is exactly where the issue requires improvement.

The dispatch floor also had to change shape: the count that must clear it is the selected count rather than the context, and folding the head axis into the request axis puts a batch-1 launch with Hkv > 1 into the batched regime. Both are derived from #899's table rather than measured on a sparse page list, and both say so in the code.

Making the selection inspectable and proving it against dense

Skipping blocks changes results, so the deliverable is not "is it faster" but "which tokens did we stop attending to". The selection is checked from both ends. minimax_m3_tests::the_decode_selection_names_exactly_the_blocks_the_mask_path_keeps asserts set equality between the page table's attended positions and the additive mask's kept columns across four context lengths including a partial final block, which is strictly stronger than comparing attention outputs, because a wrong block choice hides easily inside a plausible output tensor and not at all inside a set comparison. paged_v2::sparse_tests compares the fused output against a host reference restricted to the selected positions (which is the masked-dense implementation being replaced), and then requires it to differ from full dense attention, so a kernel that ignored the page list could not pass. SparseSelection::materialize and MLXCEL_SPARSE_PAGED_DUMP=1 expose the selected rows directly.

Proving which path a benchmark arm took

Every call returns a SparseDecodeOutcome, and report_sparse_outcome_once announces the first occurrence of each kind on stderr as well as through tracing, because the mlxcel CLI installs no tracing subscriber and a tracing::info! there prints nothing at any RUST_LOG. One flag per kind, not one global flag, since a single one-shot reports only the warmup request. The harness prints, by name, whether the sparse arm fused or fell back plus the sparse_decode_stats() delta, checks the two arms agree numerically before timing anything, and refuses to report a timing comparison at all when the sparse arm did not fuse.

What changed

  • src/lib/mlxcel-core/src/cache/sparse_csr.rs (new): SparseSelection (host and device forms), SparseCsrStructure, ContiguousCacheLayout, shared_row_mapping, selection_from_positions, selection_from_blocks. Plus sparse_csr_tests.rs, 16 host-only tests.
  • src/lib/mlxcel-core/src/paged_v2/sparse.rs (new): run_sparse_decode, SparseDecodeOutcome and its nine kinds, the kill switch, the sparsity gate, the counters and the one-shot reporter. Plus sparse_tests.rs, 17 GPU tests.
  • src/lib/mlxcel-core/src/cache.rs: KVCache::raw_kv_allocations.
  • src/models/minimax_m3_indexer.rs: decode_selection_shape and decode_selected_blocks, the host-known-cardinality decode selection. The final block is excluded from the scored set and appended unconditionally, which changes no selection (local_blocks >= 1 already pins the query's own block) but keeps the cardinality a function of kv_len and removes the only way a selected block could name a row past the live window.
  • src/models/minimax_m3_layers.rs: try_fused_sparse_decode at l == 1, falling through to the existing mask path on any decline.
  • examples/sparse_paged_decode_bench.rs (new): the mask-versus-fused harness.
  • docs/sparse-paged-decode.md (new), plus docs/README.md and docs/environment-variables.md.

Scope and limits

Neither target checkpoint fits on the development machine (DeepSeek-V3.2 is ~685B, MiniMax-M3 is 427B, and minimax_m3_indexer.rs has said so since the family landed), so everything here is validated synthetically at the real head geometry and the real sparse configuration. Prefill is untouched on both families. CUDA is unvalidated: the v2 partial kernel's CUDA body was already marked unvalidated by #898 and nothing here was compiled or run against nvcc; the host-side selection code is backend-independent, the kernel it feeds is not.

I did not run the benchmark for the record. The numbers above come from short exploratory runs used to decide the sparsity gate and are labelled indicative throughout; the harness exists so they can be replaced by a serialized, repeated measurement on a quiet machine.

Test plan

Closes #904

@inureyes inureyes added type:performance Performance improvements priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:core mlxcel-core: MLX FFI, primitives, KV cache, layers status:review Under review labels Aug 2, 2026
inureyes added 3 commits August 3, 2026 01:11
@inureyes
feat(core): sparse attention as page indirection over a contiguous cache
d6ae984 
Sparse-attention families select which KV to attend to and then pay dense costs anyway: DSA gathers the selected rows into a fresh contiguous copy, block-sparse models build an additive -inf mask and let the dense kernel compute the dropped blocks before discarding them. The #898 CSR page table already resolves a visible token through an arbitrary physical row and never assumes those rows are adjacent, ordered, or complete, so a sparse selection is not a new kernel, it is a shorter page list.

`cache/sparse_csr.rs` encodes a uniform-cardinality selection as a `page_size = 1` page table: each page is one token, `indices` is the selected-position list itself, and the kernel's page indirection becomes a per-token gather inside the attention loop with no materialized gathered copy. `ContiguousCacheLayout` addresses a dense `[B, H, Cap, D]` allocation as a `[B * H * Cap, 1, 1, D]` pool, which is a pure reshape, and folds the KV-head axis into the request axis so one CTA still owns one KV head and a group of its query heads. `shared_row_mapping` declines the combinations where one row list cannot address both allocations rather than mis-addressing them.

The row ids are a device value but the cardinality is not, because a top-k selection picks k of something on every step by construction. `SparseCsrStructure` is therefore built on the host with no synchronization, and `SparseSelection::materialize` (which does sync) is reserved for tests and the opt-in `MLXCEL_SPARSE_PAGED_DUMP`.

`paged_v2/sparse.rs` runs the v2 partial and merge kernels over the selection unchanged, returns a `SparseDecodeOutcome` for every call, and announces the first occurrence of each outcome kind on stderr as well as through tracing: the CLI installs no tracing subscriber, so a benchmark that cannot see which kernel it measured is how #899 first compared the fallback against itself. The dispatch floor is applied to the selected-row count rather than the context length, and the head-folded request count puts a single-sequence sparse decode in the batched regime; both deviations are derived from the #899 table, not measured on a sparse page list, and are documented as such.

Token-exact selection is the landing granularity. Block-granular pages need the per-request row stride to be a multiple of the block size, which holds for a `PagedBlockPool` but not for the contiguous allocations DSA and MiniMax-M3 actually use.

GPU tests check the selection from both ends: the output matches a host reference restricted to the selected positions (which is the masked-dense implementation this replaces), and it is required to differ from full dense attention, so a kernel that ignored the page list cannot pass.

Refs #904
@inureyes
feat(minimax-m3): route block-sparse decode through page indirection
78e97c2 
MiniMax-M3's block-sparse decode builds an additive mask that is 0 on the selected blocks and -inf everywhere else, then hands it to dense SDPA. The dropped blocks are computed in full and discarded, so the selection saves nothing: a 32K context with a 2048-token selection still pays for 32K columns.

Decode now takes the selected block ids straight to a `page_size = 1` page table and runs the #898 fused kernel over exactly those rows. `decode_selection_shape` states the part of the selection that is knowable without touching device memory, and `decode_selected_blocks` scores the whole blocks the same way `build_block_drop_mask` does. The final block is excluded from the scored set and appended unconditionally: it is the query's own block, so `local_blocks >= 1` already pins it, and excluding it keeps the selected cardinality a function of `kv_len` rather than of the scores while removing the only way a selected block could name a row past the live window.

`KVCache::raw_kv_allocations` exposes the reserved buffers rather than the fetched window, because the window is a token-axis slice of a step-padded buffer and reshaping it into a pool view would copy the whole cache. It returns `None` for every mode whose rows are not raw fp16 K/V, for pool-backed caches, and for a head-trimmed window.

`selection_from_blocks` expands block ids into device-resident row ids in one `O(selected)` graph pass, so nothing synchronizes on the decode path.

Any decline falls through to the existing mask path with the reason announced once per kind, and `MLXCEL_SPARSE_PAGED_ATTENTION=0` restores the mask path wholesale.

The 427B checkpoint does not fit on the development machine, so this is validated synthetically. `the_decode_selection_names_exactly_the_blocks_the_mask_path_keeps` asserts set equality between the page table's attended positions and the mask path's kept columns across four context lengths including a partial final block, which is a stronger statement than comparing two attention outputs: a wrong block choice hides easily inside a plausible output tensor and not at all inside a set comparison.

Refs #904
@inureyes
perf(core): gate fused sparse decode on sparsity, not just launch size
6b43d22 
The #899 token floor asks "is this launch big enough". It does not ask the question that decides a sparse launch, which is "is skipping worth the kernel you have to skip with". `mlx::fast::scaled_dot_product_attention` is a heavily tuned dense kernel and the v2 partial kernel is a scalar per-lane sweep, so reading half the data through a kernel with twice the constant factor is a loss.

`examples/sparse_paged_decode_bench` measures it at MiniMax-M3 geometry, one repetition of 40 steps on an idle Apple Silicon host: 0.67x at 2x sparsity (ctx 4096), 0.77x at 4x (8192), 1.22x at 8x (16384), 1.17x at 16x (32768), 2.06x at 32x (65536). Transient memory drops at every point. Indicative, not a recorded result.

So a launch must also clear `MIN_SPARSITY_RATIO`, default 8, overridable through `MLXCEL_SPARSE_PAGED_MIN_SPARSITY` and disabled at 0. The default sits on top of the measured win rather than interpolated into the unmeasured band below it. #899 argued the opposite way for its token floor and deliberately sat under its weakest measured point, because what lay under that point was a missed opportunity; here what lies under it is a measured regression, and shipping a regression in a narrow band is worse than declining a modest win in one. For MiniMax-M3's shipped configuration the gate opens at a 16K context, which is where the issue requires decode to improve.

`clears_sparsity_gate` uses `checked_mul` rather than `saturating_mul`: saturating the requirement to `usize::MAX` makes an absurd selection pass the gate, since the comparison then reads `usize::MAX >= usize::MAX`. A unit test pins that.

The harness itself is the other half of this commit. It drives both arms off the same selection, checks they agree numerically before timing anything, prints by name whether the sparse arm fused or fell back together with the counter delta, and refuses to report a timing comparison at all when the sparse arm did not fuse. #899 shipped a fused path that silently never activated and then benchmarked the fallback against itself, and the resulting clean-looking null was nearly accepted; an arm that cannot say what it ran is not evidence.

`docs/sparse-paged-decode.md` records the whole design, including two families the issue names that did not land and why. DeepSeek Sparse Attention decode adds a per-(head, position) positional term through the mask argument, and the v2 kernel has no input through which such a term can arrive; the three ways out are costed there, and the recommended one is a row-stride plus second-K-stream kernel generalization that needs no cache change. phi3-small's block pattern varies per *query* head, and a CSR request necessarily shares one page list across the `n_rep` query heads grouped under its KV head.

Refs #904
@inureyes
inureyes force-pushed the feature/issue-904-sparse-paged-decode branch from 168f361 to 6b43d22 Compare August 3, 2026 00:00
@inureyes
docs(benchmarks): record sparse paged decode validation on M1 Ultra
6b906a3 
Adds the performance-validation report issue #904 requires, across three independent invocations at context 4096 through 65536.

The fused sparse path wins from 16x sparsity upward, at 1.39x median, reaching 2.99x at 32x sparsity, with transient peak memory lower at every fused point. Correctness agrees with the dense mask path to 4e-4 in f16 storage.

Corrects the routing threshold's justification. `MIN_SPARSITY_RATIO` defaults to 8, and 8x measures at parity rather than at a win: median 1.02x over a 0.94x to 1.04x range. The first invocation alone read 0.94x, which looks like a regression, and two further invocations put it at break-even, so that single reading was noise. The threshold is left at 8 because parity costs nothing measurable and because the harness flatters both arms in ways whose residual sign is not established, with the note that raising it to 16 is justified if a production workload regresses there.

Records that the harness declines to print timings when the sparse arm did not fuse, rather than printing a comparison of the mask path against itself. That failure mode occurred three times elsewhere in this epic, and here it is structurally impossible to misread because no number appears.

Also records the two premise corrections the implementation surfaced: DeepSeek Sparse Attention did not land because its decode score carries an additive per-head per-position term the v2 kernel has no input for, and phi3-small and phimoe cannot use this at all because their pattern varies per query head while a CSR request shares one page list across the query heads under its KV head.
@inureyes inureyes added status:done Completed and removed status:review Under review labels Aug 3, 2026
@inureyes
inureyes merged commit 230163d into main Aug 3, 2026
5 checks passed
@inureyes inureyes mentioned this pull request Aug 3, 2026
Closed
@inureyes
inureyes deleted the feature/issue-904-sparse-paged-decode branch August 4, 2026 12:08
@inureyes inureyes self-assigned this Aug 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

@inureyes inureyes

Labels

area:core mlxcel-core: MLX FFI, primitives, KV cache, layers area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:done Completed type:performance Performance improvements

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Fused sparse-attention decode via page indirection for DSA and block-sparse models

1 participant

@inureyes