CUDA: use mma FA kernel for gqa > 4 on RTX 4000 #15035
Merged
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.
For models such as Qwen 2.5 3b with 8 Q heads per K/V head it seems to be better to use the mma FlashAttention kernel than the vector kernel:
Notably if the GPU is frequency limited this difference is even larger and up to +25%. With a frequency limit it would also be better to use the mma kernel in other cases such as with LLaMA which has 4 Q heads per K/V head. And since datacenter GPUs have lower frequencies than consumer GPUs this implies that choosing kernels solely based on compute capability is suboptimal. It's currently unclear to me how to best retrieve the GPU clocks in a way that considers user-defined limits, I opened a thread in the NVIDIA developer forums and will make a PR once I get a reply.