FA3 kvcache + split kv + gqa parallelization #1236
Draft
+2,259
−311
Conversation
This file contains 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
KV cache functionality not added yet.
ipiszy
reviewed
Sep 19, 2024
| @@ -174,7 +175,8 @@ def forward( | |||
| causal, | |||
| descale_q=descale_q, | |||
| descale_k=descale_k, | |||
| descale_v=descale_v, | |||
| descale_v=descale_v, | |||
| gqa_decoding=False, | |||
There was a problem hiding this comment.
I wonder does it make sense to give user an option to enable GQA optimization for general use cases outside of decoding?
e.g. It's generally useful for small seq_len prefill. In this case we don't really need split-kv, but we want to have each threadblock handle multiple Q heads with the same KV head.
There was a problem hiding this comment.
Furthermore, does it make sense to just enable GQA optimization by default when input is GQA? I feel it won't cause perf regressions even for long sequence length.
There was a problem hiding this comment.
I feel it might slow things down a bit, but I haven't tried
…ded template params
…ly matters for fp8 support
…n into fa3-kvcache-gqa
…n into fa3-kvcache-gqa
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 adds split KV ("Flash decoding") and GQA parallelization improvements for FA3. Some essential parts of the KV cache API are added as well, including the
cache_seqlensandcache_batch_idxarguments.Up to 15x improvement over FA2 measured on my H100 PCIe in exceptional cases, e.g.
Times given in microseconds. GB/s is measured in terms of loading the KV cache. Note that theoretical max bandwidth is 2 TB/s for H100 PCIe.
TODO on this PR before merge: add split kv heuristic, implement for FP8.
fa3-decoding-times-091724.log