metal: deepseek v4 sparse prefill - #26512
Closed
tarruda wants to merge 2 commits into
Closed
Conversation
Replace dense compressed-attention scans during DeepSeek V4 prefill with a backend-neutral sparse path that packs the sliding window and Lightning Indexer selections into a compact working set. - Add DSV4_SPARSE_PACK with a CPU reference and backend capability probes. - Extend Flash Attention with broadcast masks and row-indexed sinks. - Select sparse prefill graphs when supported while preserving dense fallback. - Test sparse packing and tiled and padded Flash Attention paths. Assisted-by: Codex
Implement the sparse DeepSeek V4 prefill path on Metal so Flash Attention consumes compact per-token key sets instead of scanning the full compressed cache. - Add a DSV4_SPARSE_PACK kernel for raw-window and selected keys. - Add an exact radix TOP_K specialization for 512 Lightning Indexer results. - Extend tiled Flash Attention with broadcast masks and row-indexed sinks. - Register backend support for the fused packing operation. Assisted-by: Codex
Contributor
Author
|
The commits are split into common/metal, but can also create separate PRs if desired. |
github-actions
Bot
added
model
This was referenced Aug 3, 2026
ggerganov
reviewed
Aug 3, 2026
Comment on lines
+2611
to
+2621
| // Packs per-token raw-window and Lightning-Indexer selections into the | ||
| // strided K + mask storage consumed by DeepSeek V4 sparse flash attention. | ||
| GGML_API struct ggml_tensor * ggml_dsv4_sparse_pack( | ||
| struct ggml_context * ctx, | ||
| struct ggml_tensor * raw_k, | ||
| struct ggml_tensor * comp_k, | ||
| struct ggml_tensor * raw_mask, | ||
| struct ggml_tensor * comp_mask, | ||
| struct ggml_tensor * comp_idx, | ||
| int64_t n_raw); | ||
|
|
Member
There was a problem hiding this comment.
Earlier I didn't notice that this involves a new ggml OP. So, likely will not get merged as I think we are already working on similar functionality in another PR.
Contributor
Author
|
@ggerganov closing it then. My branch still has optimizations on top of this PR. For reference here are the current numbers: After the new op is implemented on master, I can re-evaluate there are still things that can be cherry-picked to contribute |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Picking up initial sparse prefill work from my branch
Additional information
This will greatly improve prefill performance as context grows. Here's
-d 0,10000,20000,30000 -n 128 -p 2048 -b 2048 -ub 2048on a M1 ultra before these changes:After:
The official inference code was used as reference for this.
Requirements