2-Pass Sdpa Inference Kernel #1597
Merged
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
awni
approved these changes
Nov 18, 2024
| keys += blocks * stride; | ||
| values += blocks * stride; | ||
| } | ||
| threadgroup_barrier(mem_flags::mem_threadgroup); |
There was a problem hiding this comment.
This is actually a great catch. It should be inside the loop. Same for the 1 pass. The reasoning is it makes sure the thread group all reads the block at the same time so one simdgroup cannot just run ahead. I had seen it provides a small improvement but then in one of the edits it probably got restored back outside the loop.
There was a problem hiding this comment.
I take it back... I tested again and I get mixed results which is probably why I reverted it in some previous commit. I will remove it since indeed conceptually it needn't be there.
This was referenced Nov 26, 2024
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 aims to improve long context generation performance by increasing parallelization for large numbers of keys/values. There are mild benefits for smaller machines and very significant benefits for Ultra machines.
The main benefit for small machines stems from accessing the keys and values in a more cache friendly way when there is GQA and for the Ultra machines it stems from launching more thread groups which allows using more of the chip.
Speedup for M2 Max
The following speedup is in total tokens per second and not attention speedup. Note the phi model which does not improve does not have GQA. The 1 pass SDPA on the M2 Max achieves ~350 to 380 GB/s read for sequence length ~2048 so there isn't really much room for speedup.
Speedup for M2 Ultra
Again the speedup is in total tokens per second and not attention specific. The M2 Ultra is sped up for all cases, no GQA required. The 2048 sequence length without GQA peaks at >800GB/s which also means there is probably little room for improvement (there could be for longer sequences).