DSA: Better way to build the attention mask#2119
Merged
Merged
Conversation
4 tasks
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
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.
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$-\infty$ is used as a base for preparing the attention mask. When the $-\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
f32mask filled withtop_kKV 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, thef32mask filled withn_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 indexertop_kas 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.