Perf: accelerate DeepSeek V4 Flash prefill sparse attention

[hashiqiqixian]

Summary

• Pass per-token compressed-row counts through CSA, HCA, and SWA prefill so sparse gather, QK/PV, and online-softmax merge skip inactive blocks.
• Fan the DeepSeek V4 Flash prefill indexer score over 8-token groups, cap the useful top-k prefix at 256, and remove the redundant final merge stage.
• Preserve separate QK/PV and online-merge SPMD tasks with an explicit dependency, so AIC-produced statistics are published before the vector merge consumes them.

Performance

• Standalone sparse-attention median run mean: 635.2 -> 609.6 us (ratio 4, 128 tokens, a2a3 auto device, 3 x 20 rounds).
• Standalone indexer latency: 952.9 -> 488.7 us (128 tokens, start-pos 896, a2a3, 20 rounds).

Validation

• On 686/a2a3, serial standalone sparse attention, CSA, HCA, and SWA tests pass.
• The two-device EP2 prefill_fwd.py -p a2a3 test passes without scheduler stalls.
• The exact eight-device DeepSeek serving accuracy test passes: 4 passed.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR adds per-token compressed-slot counts to packed prefill sparse attention, updates CSA, HCA, and SWA call sites, changes indexer score tiling and top-k sorting, and fuses sparse attention scoring with online-softmax merging.

Changes

Sparse prefill count propagation

Layer / File(s)	Summary
Indexer tiling and top-k pipeline models/deepseek/v4-flash/prefill_indexer.py	Score computation uses eight-token groups, requires divisible token counts, adopts INDEXER_TOPK_CAP, and removes the extended merge-sort stage.
Compressed count production models/deepseek/v4-flash/prefill_attention_csa.py, models/deepseek/v4-flash/prefill_attention_hca.py, models/deepseek/v4-flash/prefill_attention_swa.py	CSA and HCA compute clamped cmp_counts; SWA passes a zero-initialized dummy count tensor.
Sparse gather and fused attention models/deepseek/v4-flash/prefill_sparse_attn.py	Sparse KV gathering and online-softmax merging use per-token compressed counts, with the qk_pv and merge stages fused.
Sparse attention test plumbing models/deepseek/v4-flash/prefill_sparse_attn.py	Test signatures, calls, tensor specifications, and count initialization are updated for cmp_counts.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PrefillIndexer
  participant CSA_HCA_SWA
  participant PrefillSparseAttn
  PrefillIndexer->>CSA_HCA_SWA: provide sparse indices and visible compressed range
  CSA_HCA_SWA->>CSA_HCA_SWA: compute per-token cmp_counts
  CSA_HCA_SWA->>PrefillSparseAttn: pass cmp_counts with sparse metadata
  PrefillSparseAttn->>PrefillSparseAttn: gather valid compressed KV slots
  PrefillSparseAttn->>PrefillSparseAttn: merge attention outputs over active blocks

Loading

Possibly related PRs

• hw-native-sys/pypto-lib#823: Refactors SWA dummy compressed metadata tensors used by the same sparse attention call.

Poem

I’m a rabbit with counts in my pack,
Eight bright tokens marching in track.
Sparse keys now know where to play,
Fused merges hop through the day.
Top-k sorts with a nimble twist—
Count every carrot on the list!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.09% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: faster DeepSeek V4 Flash prefill sparse attention.
Description check	✅ Passed	The description is clearly related to the changeset and matches the performance and validation-focused updates.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}