Skip to content

DSA: Better way to build the attention mask#2119

Merged
ikawrakow merged 11 commits into
mainfrom
ik/better_indexer_mask
Jul 13, 2026
Merged

DSA: Better way to build the attention mask#2119
ikawrakow merged 11 commits into
mainfrom
ik/better_indexer_mask

Conversation

@ikawrakow

Copy link
Copy Markdown
Owner

This PR improves the way the attention mask is built when using DSA. The result is a reduction of compute buffer sizes and a very minor improvement of PP performance (1-3%, increasing with context length).

On the main branch (and the pending PR #2109), an f32 mask filled with $-\infty$ is used as a base for preparing the attention mask. When the top_k KV cache entries are selected by the indexer, the corresponding values in this mask are overwritten with zeros, and then the actual attention mask is added. The issue with this approach is that, as far as the scheduler is concerned, the f32 mask filled with $-\infty$ is a graph leaf, so always allocated first and independently of other graph nodes, so there is no compute buffer reuse and the memory price for this mask is always paid. In addition, when running on a multi-GPU system, this mask becomes an input for each graph split, so needs to be copied to each GPU, and each GPU pays the corresponding VRAM price. The $-\infty$ mask is n_ctx x n_ubatch, so for a context of 128k tokens and u-batch size of 2048, this woks out to 1 GiB!

In this PR I have added a new op (GGML_OP_MASK_TOPK) that takes the attention mask and the indexer top_k as input, and builds the resulting mask as needed by the DSA attention. This is a now a "normal" compute graph node, so the buffer allocated for this mask can be used for other graph nodes.

For a context of 100k tokens on the main branch I can only offload 2 MoE layers per GPU to not run into OOM. After this change, I'm able to have 3 MoE layers per GPU offloaded.

Running CPU only with a context of 100k tokens I see a compute buffer of 4178 MiB versus 4914 MiB on the main branch.

Removing the need to copy the $-\infty$ mask to each GPU leads to the observed PP performance improvement (that will obviously depend on the speed of the PCI-E bus, so can be more or less if slower/faster than mine). TG is not affected as in that case the mask is already built differently before this PR.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant