Faster DeepSeek FA on CUDA#408
Merged
Merged
Conversation
added 3 commits
May 11, 2025 09:58
Does not work because the RoPE portion is stored at the end in our case, while in mainline it is stored at the beginning, and the FA kernel assumes that.
Contributor
|
An RTX 4080 has 76 streaming multiprocessor, the CUDA code assigns KV slices to SMs in chunks of size 256. So every 76*256=19456 tokens the size of biggest workload across all of the SMs increases and there is a dip in performance. These so-called quantization effects are much more noticeable with compute than with I/O so they become more pronounced if the I/O of a kernel is optimized. |
|
Just tested on DeepSeek V3 0324 Q2_K_XL and it seems to have improved my t/s TG by about 1-2% (I guess with offloading there isn't much difference), but tested a smaller models (DeepSeek2 16B) on a single GPU (5090) and got about 8-12% speedup, so pretty nice! This is on top of #405 PR. Now I'm gonna try #409 on top of that PR and this PR. |
5 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 is a port of this PR in mainline
llama.cpp.The main difference to PR #386 is that now the FA kernel takes advantage of the fact that the V tensor contains the same data as the K tensor (it is a view on the K cache with an offset given by the RoPE embedding size). Hence, one can reduce the number of loads by reusing K tiles when processing
V*softmax(K*Q).To take advantage of this new kernel I had to change the way the K cache is organized. In mainline
llama.cppthe K cache stores(RoPE, NoPE)parts in that order, and the FA kernel assumes this arrangement. But inik_llama.cppprior to this PR the K cache was stored as(NoPE, RoPE). As there are several places where the views into the K cache can go wrong when building the graph, the PR should be tested more thoroughly before merging. I have tested all possible combinations ofmlaandfausing DeepSeek-Lite and it appears to work correctly, but still.The next graph shows a TG performance comparison between the main branch (black) and this PR (red). Model is DeepSeek-Lite quantized with
Q4_0, GPU is RTX-4080. We see nice performance improvements, but also a more peculiar behavior as a function ofN_KV, the number of tokens in the KV cache.When
mla = 2ormla = 3this PR has no effect on PP, so the next graph compares PP speed between the main branch (black) and the PR (red) formla = 1. For reference I have also included PP performance formla = 3with blue symbols. In case I ave not shown a graph such as this one, it illustrates what one gives up in terms of PP performance by using a mainlinellama.cppMLA-enabled GGUF for DeepSeek models. The difference is ~25% forN_KV = 0and nearly a factor of 2 at 60k tokens. The PR improvesmla = 1performance by a few percent.Finally, being curious about the peculiar TG behavior as a function of
N_KV, I ransweep-benchwith the llama.cpp PR, and the next graph shows a TG performance comparison between this PR and the mainline PR. We see that the two curves align very closely, so the strange behavior is not due to me screwing up with the port. I wonder if @JohannesGaessler is aware.