Optimize TPU Flash Attention (400x speed-up on 32k long context) #845
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
Use splash attention lazy mask instead of jnp mask, which is O(T^2). The memory for jnp mask is O(T^2), which almost negates the benefits of reducing HBM communication with flash attention. Let’s use splash attention lazy mask, which lazily generates causal masks. In addition, pallas supports CPU simulation (interpret=True), so use same pallas kernel on CPU, which makes it easier to debug the code. * Benchmark: on TPUv5p, (model_dim/heads/kv_heads/seq_len). NumpyMask (ASIS) ---------------------------------------------------------------------------------------- Benchmark Time CPU Iterations ---------------------------------------------------------------------------------------- FlashAttentionBenchmark/256/2/2/512 1.71 ms 1.09 ms 592 (4.43M) FlashAttentionBenchmark/2048/16/2/1024 4.44 ms 1.21 ms 483 (28.62M) FlashAttentionBenchmark/4096/16/2/1024 8.61 ms 1.36 ms 302 (53.27M) FlashAttentionBenchmark/4096/16/2/4096 3264 ms 1537 ms 1 (197.38M) FlashAttentionBenchmark/4096/16/2/8192 7426 ms 5603 ms 1 (389.54M) FlashAttentionBenchmark/4096/16/2/32768 94427 ms 92256 ms 1 (1.50G) CausalMask (Proposed PR): This PR saves both memory and computation. In long context, speed-up (400x) and HBM saving (3x). ---------------------------------------------------------------------------------------- Benchmark Time CPU Iterations ---------------------------------------------------------------------------------------- FlashAttentionBenchmark/256/2/2/512 1.55 ms 1.01 ms 578 (3.43M) FlashAttentionBenchmark/2048/16/2/1024 4.21 ms 1.11 ms 490 (13.57M) FlashAttentionBenchmark/4096/16/2/1024 6.50 ms 1.17 ms 493 (24.22M) FlashAttentionBenchmark/4096/16/2/4096 16.8 ms 1.38 ms 228 (84.33M) FlashAttentionBenchmark/4096/16/2/8192 28.8 ms 1.58 ms 217 (164.50M) FlashAttentionBenchmark/4096/16/2/32768 230 ms 6.36 ms 16 (644.60M)
ds-hwang
changed the title
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.
Optimize TPU Flash Attention (400x speed-up on 32k long context)
Use splash attention lazy mask instead of jnp mask, which is O(T^2).
The memory for jnp mask is O(T^2), which almost negates the benefits of
reducing HBM communication with flash attention. Let’s use splash attention
lazy mask, which lazily generates causal masks.
In addition, pallas supports CPU simulation (interpret=True), so use same
pallas kernel on CPU, which makes it easier to debug the code.
NumpyMask (ASIS)
CausalMask (Proposed PR): This PR saves both memory and computation. In long
context, speed-up (400x) and HBM saving (3x).