5 ms·

FlashAttention-T: Towards Tensorized Attention

116 points by matt_d 9mo ago

measurablefunc 9mo ago
[flagged]
- dheera 9mo ago
  "Most people" didn't figure this out either, the top 0.01% did.
  - E-Reverance 9mo ago
    I also wouldn't be surprised if they used AI to assist themselves in small ways
  - measurablefunc 9mo ago
    You're just moving the goal post & not addressing the question I asked. Why isn't AI optimizing the kernels in its own code the way people have been optimizing it like in the posted paper?
    - phkahler 9mo ago
      It will, right after it reads the paper.
      - measurablefunc 9mo ago
        I read the paper. All the prerequisites are already available in existing literature & they basically profiled & optimized around the bottlenecks to avoid pipeline stalls w/ instructions that utilize the available tensor & CUDA cores. Seems like something these super duper AIs that don't get tired should be able to do pretty easily.
    - aspenmartin 9mo ago
      They do? https://www.deeplearning.ai/the-batch/alphatensor-for-faster-matrix-multiplication-explained/ https://www.deeplearning.ai/the-batch/alphatensor-for-faster... https://deepmind.google/blog/alphaevolve-a-gemini-powered-coding-agent-for-designing-advanced-algorithms/ https://deepmind.google/blog/alphaevolve-a-gemini-powered-co... https://www.rubrik.com/blog/ai/25/teaching-ai-to-write-gpu-code-a-deep-dive-into-reinforcement-fine-tuning?utm_source=chatgpt.com https://www.rubrik.com/blog/ai/25/teaching-ai-to-write-gpu-c...
semiinfinitely 9mo ago
tri dao isn't on the paper is it even allowed to call it "FlashAttention"???
saagarjha 9mo ago
Less annoying link directly to the paper: https://dl.acm.org/doi/pdf/10.1145/3774934.3786425?download=true https://dl.acm.org/doi/pdf/10.1145/3774934.3786425?download=...
- SpaceManNabs 9mo ago
  link if you don't want to automatically download files https://dl.acm.org/doi/pdf/10.1145/3774934.3786425 https://dl.acm.org/doi/pdf/10.1145/3774934.3786425
sigbottle 9mo ago
Oh wow there's still work being done on ampere? I was wondering - I've been thinking about switching to AI systems programming (I know, easy task), but from what I understand, industry cloud GPUs are the main winners, right? Nobody's going to pay me (assuming I even had the skills) to optimize for consumer GPUs? From what I understand, it's not just number + capacity + performance, it's literal core primitives. I don't think any of the "Blackwell" chips like the grace one or rtx 5090 have for example SM pairs in their ISA? And likewise similar fundamental differences between consumer and cloud hopper (where the majority of the perf is the cloud one's ISA?) So I guess I'm wondering if I should buy a GPU myself or should I just rent on the cloud if I wanted to start getting some experience in this field. How do you even get experience in this normally anyways, do you get into really good schools and into their AI labs which have a lot of funding?
- Maxious 9mo ago
  yep, https://github.com/poad42/cuda-fp8-ampere https://github.com/poad42/cuda-fp8-ampere recently another attempt at squeezing whatever's left from ampere
- storus 9mo ago
  I still have 2x NVLinked A6000 and they aren't that bad compared to a single RTX 6000 Pro.
- vlovich123 9mo ago
  Look at am the email addresses. If you’ll recall there’s an embargo on China.
- coolsunglasses 9mo ago
  I do CUDA for a living (not inference) and for the life of me (and a couple of LLMs for that matter) I cannot figure out what you mean by "SM pairs". Do you mean the coupled dies on stuff like the B200? An NVidia chip die has many SMs if so. Do you mean TMEM MMA cooperative execution? I'm guessing that must be it given what the paper is about.
  - 9mo ago
simianwords 9mo ago
OT but instead of quadratic attention can we not have n^10 or something crazier? I feel like we are limiting the intelligence just to save cost. But I can imagine that there might be some questions that may be worth paying higher cost for. I feel like n^10 attention can capture patterns that lower complexity attention may not. So it seems arbitrary that we have n^2 attention.
- eldenring 9mo ago
  This is a common way of thinking. In practice this type of thing is more like optimizing flop allocation. Surely with an infinite compute and parameter budget you could have a better model with more intensive operations. Another thing to consider is that transformers are very general computers. You can encode many many more complex architectures in simpler, multi layer transformers.
- refulgentis 9mo ago
  n^2 isn't a setting someone chose, it's a mathematical consequence of what attention is. Here's what attention does: every token looks at every other token to decide what's relevant. If you have n tokens, and each one looks at n others, you get n * n = n^2 operations. Put another way: n^2 is when every token gets to look at every other token. What would n^3 be? n^10? (sibling comment has same interpretation as you, then handwaves transformers can emulate more complex systems)
  - measurablefunc 9mo ago
    There are lots more complicated operations than comparing every token to every other token & the complexity increases when you start comparing not just token pairs but token bigrams, trigrams, & so on. There is no obvious proof that all those comparisons would be equivalent to the standard attention mechanism of comparing every token to every other one.
    - vlovich123 9mo ago
      While you are correct at a higher level, comparing bigrams/trigrams would be less compute not more because there’s fewer of them in a given text
jmward01 9mo ago
I built guided window attn (literally predict the position of the window) a while ago and that works great. Why are we still stuck on any form of attn that looks at the entire context in any meaningful way? Do humans work this way? Do I need a whole book to predict the next word? Who out there is working on really new unique ways to deal with infinite history, other than me of course :)
- cs702 9mo ago
  > Who out there is working on ... infinite history? Many people are still working on improving RNNs, mostly in academia. Examples off the top of my head: * RWKV: https://arxiv.org/abs/2006.16236 https://arxiv.org/abs/2006.16236 / https://arxiv.org/abs/2404.05892 https://arxiv.org/abs/2404.05892 https://arxiv.org/abs/2305.13048 https://arxiv.org/abs/2305.13048 * Linear attention: https://arxiv.org/abs/2503.14456 https://arxiv.org/abs/2503.14456 * State space models: https://arxiv.org/abs/2312.00752 https://arxiv.org/abs/2312.00752 / https://arxiv.org/abs/2405.21060 https://arxiv.org/abs/2405.21060 * Linear RNNs: https://arxiv.org/abs/2410.01201 https://arxiv.org/abs/2410.01201 Industry OTOH has gone all-in on Transformers.
  - viraptor 9mo ago
    > Industry OTOH has gone all-in on Transformers. It's so annoying. Transformers keep improving and recurrent networks are harder to train so until we hit some real wall, companies don't seem eager to diverge. It's like lithium batteries improving easy faster than it was profitable to work on sodium ones, even though we unfortunately want the sodium ones to be better.
    - cs702 9mo ago
      I'd add this to the list of linear-attention RNNs: https://arxiv.org/abs/2602.00294 https://arxiv.org/abs/2602.00294 Recently saw it on HN.
  - jmward01 9mo ago
    RNNs have two huge issues: - long context. Recurrence degrades the signal for the same reason that 'deep' nn architectures don't go much past 3-4 layers before you need residual connections and the like - (this is the big one) training performance is terrible since you can't parallelize them across a sequence like you can with causal masked attn in transformers On the huge benefit side though you get: - guaranteed state size so perfect batch packing, perfect memory use, easy load/unload from a batch, O(1) of token gen so generally massive performance gains in inference. - unlimited context (well, no need for a concept of a position embedding or similar system) Taking the best of both worlds is definitely where it is at for the future. An architecture that can train parallelized, has a fixed state size so you can load/unload and patch batches perfectly, unlimited context (with perfect recall), etc etc. That is the real architecture to go for.
verytrivial 9mo ago
Tldr: 5% - 17% speedup due to removing a bottleneck by juggling where on a GPU/compute core a computation is done during Flash attention.
findalex 9mo ago
QM would tell us the order of your Hamiltonian (attention operator) doesn’t limit the complexity of the wave function (hidden state). It might be more efficient to explicitly correlate certain many-body interactions, but pair-wise interactions, depth and a basis (hidden state dimension) approaching completeness "are all you need”.
- twothreeone 9mo ago
  The terminology is overloaded.. Tensors in QM are objects obeying transformation laws, in ML Tensors are just data arranged in multidimensional arrays. There are no constraints on how the data transforms.
  - findalex 9mo ago
    Intended as analogy - but it is essentially a description of the DMRG algorithm (quantum chem). Only pair-wise operators there but the theory approaches exact when there are enough terms in your tensor product (iterations ~ depth) and a large enough embedding dimension. > There are no constraints on how the data transforms. Except those implicit in your learned representation. And that representation could be the MB WF.
kioku 9mo ago
> Our key insight is to offload critical softmax primitives to idle tensor units, maximizing hardware utilization and throughput. > … speedups of 1.05–1.17×across diverse attention configurations on Ampere and Hopper GPUs …

116 points by matt_d 9mo ago

measurablefunc 9mo ago

[flagged]

dheera 9mo ago
"Most people" didn't figure this out either, the top 0.01% did.
- E-Reverance 9mo ago
  I also wouldn't be surprised if they used AI to assist themselves in small ways
- measurablefunc 9mo ago
  You're just moving the goal post & not addressing the question I asked. Why isn't AI optimizing the kernels in its own code the way people have been optimizing it like in the posted paper?
  - phkahler 9mo ago
    It will, right after it reads the paper.
    - measurablefunc 9mo ago
      I read the paper. All the prerequisites are already available in existing literature & they basically profiled & optimized around the bottlenecks to avoid pipeline stalls w/ instructions that utilize the available tensor & CUDA cores. Seems like something these super duper AIs that don't get tired should be able to do pretty easily.
  - aspenmartin 9mo ago
    They do? https://www.deeplearning.ai/the-batch/alphatensor-for-faster-matrix-multiplication-explained/ https://www.deeplearning.ai/the-batch/alphatensor-for-faster... https://deepmind.google/blog/alphaevolve-a-gemini-powered-coding-agent-for-designing-advanced-algorithms/ https://deepmind.google/blog/alphaevolve-a-gemini-powered-co... https://www.rubrik.com/blog/ai/25/teaching-ai-to-write-gpu-code-a-deep-dive-into-reinforcement-fine-tuning?utm_source=chatgpt.com https://www.rubrik.com/blog/ai/25/teaching-ai-to-write-gpu-c...

semiinfinitely 9mo ago

tri dao isn't on the paper is it even allowed to call it "FlashAttention"???

saagarjha 9mo ago

Less annoying link directly to the paper: https://dl.acm.org/doi/pdf/10.1145/3774934.3786425?download=true https://dl.acm.org/doi/pdf/10.1145/3774934.3786425?download=...

SpaceManNabs 9mo ago
link if you don't want to automatically download files https://dl.acm.org/doi/pdf/10.1145/3774934.3786425 https://dl.acm.org/doi/pdf/10.1145/3774934.3786425

sigbottle 9mo ago

Oh wow there's still work being done on ampere? I was wondering - I've been thinking about switching to AI systems programming (I know, easy task), but from what I understand, industry cloud GPUs are the main winners, right? Nobody's going to pay me (assuming I even had the skills) to optimize for consumer GPUs? From what I understand, it's not just number + capacity + performance, it's literal core primitives. I don't think any of the "Blackwell" chips like the grace one or rtx 5090 have for example SM pairs in their ISA? And likewise similar fundamental differences between consumer and cloud hopper (where the majority of the perf is the cloud one's ISA?) So I guess I'm wondering if I should buy a GPU myself or should I just rent on the cloud if I wanted to start getting some experience in this field. How do you even get experience in this normally anyways, do you get into really good schools and into their AI labs which have a lot of funding?

Maxious 9mo ago
yep, https://github.com/poad42/cuda-fp8-ampere https://github.com/poad42/cuda-fp8-ampere recently another attempt at squeezing whatever's left from ampere
storus 9mo ago
I still have 2x NVLinked A6000 and they aren't that bad compared to a single RTX 6000 Pro.
vlovich123 9mo ago
Look at am the email addresses. If you’ll recall there’s an embargo on China.
coolsunglasses 9mo ago
I do CUDA for a living (not inference) and for the life of me (and a couple of LLMs for that matter) I cannot figure out what you mean by "SM pairs". Do you mean the coupled dies on stuff like the B200? An NVidia chip die has many SMs if so. Do you mean TMEM MMA cooperative execution? I'm guessing that must be it given what the paper is about.
- 9mo ago

simianwords 9mo ago

OT but instead of quadratic attention can we not have n^10 or something crazier? I feel like we are limiting the intelligence just to save cost. But I can imagine that there might be some questions that may be worth paying higher cost for. I feel like n^10 attention can capture patterns that lower complexity attention may not. So it seems arbitrary that we have n^2 attention.

eldenring 9mo ago
This is a common way of thinking. In practice this type of thing is more like optimizing flop allocation. Surely with an infinite compute and parameter budget you could have a better model with more intensive operations. Another thing to consider is that transformers are very general computers. You can encode many many more complex architectures in simpler, multi layer transformers.
refulgentis 9mo ago
n^2 isn't a setting someone chose, it's a mathematical consequence of what attention is. Here's what attention does: every token looks at every other token to decide what's relevant. If you have n tokens, and each one looks at n others, you get n * n = n^2 operations. Put another way: n^2 is when every token gets to look at every other token. What would n^3 be? n^10? (sibling comment has same interpretation as you, then handwaves transformers can emulate more complex systems)
- measurablefunc 9mo ago
  There are lots more complicated operations than comparing every token to every other token & the complexity increases when you start comparing not just token pairs but token bigrams, trigrams, & so on. There is no obvious proof that all those comparisons would be equivalent to the standard attention mechanism of comparing every token to every other one.
  - vlovich123 9mo ago
    While you are correct at a higher level, comparing bigrams/trigrams would be less compute not more because there’s fewer of them in a given text

jmward01 9mo ago

I built guided window attn (literally predict the position of the window) a while ago and that works great. Why are we still stuck on any form of attn that looks at the entire context in any meaningful way? Do humans work this way? Do I need a whole book to predict the next word? Who out there is working on really new unique ways to deal with infinite history, other than me of course :)

cs702 9mo ago
> Who out there is working on ... infinite history? Many people are still working on improving RNNs, mostly in academia. Examples off the top of my head: * RWKV: https://arxiv.org/abs/2006.16236 https://arxiv.org/abs/2006.16236 / https://arxiv.org/abs/2404.05892 https://arxiv.org/abs/2404.05892 https://arxiv.org/abs/2305.13048 https://arxiv.org/abs/2305.13048 * Linear attention: https://arxiv.org/abs/2503.14456 https://arxiv.org/abs/2503.14456 * State space models: https://arxiv.org/abs/2312.00752 https://arxiv.org/abs/2312.00752 / https://arxiv.org/abs/2405.21060 https://arxiv.org/abs/2405.21060 * Linear RNNs: https://arxiv.org/abs/2410.01201 https://arxiv.org/abs/2410.01201 Industry OTOH has gone all-in on Transformers.
- viraptor 9mo ago
  > Industry OTOH has gone all-in on Transformers. It's so annoying. Transformers keep improving and recurrent networks are harder to train so until we hit some real wall, companies don't seem eager to diverge. It's like lithium batteries improving easy faster than it was profitable to work on sodium ones, even though we unfortunately want the sodium ones to be better.
  - cs702 9mo ago
    I'd add this to the list of linear-attention RNNs: https://arxiv.org/abs/2602.00294 https://arxiv.org/abs/2602.00294 Recently saw it on HN.
- jmward01 9mo ago
  RNNs have two huge issues: - long context. Recurrence degrades the signal for the same reason that 'deep' nn architectures don't go much past 3-4 layers before you need residual connections and the like - (this is the big one) training performance is terrible since you can't parallelize them across a sequence like you can with causal masked attn in transformers On the huge benefit side though you get: - guaranteed state size so perfect batch packing, perfect memory use, easy load/unload from a batch, O(1) of token gen so generally massive performance gains in inference. - unlimited context (well, no need for a concept of a position embedding or similar system) Taking the best of both worlds is definitely where it is at for the future. An architecture that can train parallelized, has a fixed state size so you can load/unload and patch batches perfectly, unlimited context (with perfect recall), etc etc. That is the real architecture to go for.

verytrivial 9mo ago

Tldr: 5% - 17% speedup due to removing a bottleneck by juggling where on a GPU/compute core a computation is done during Flash attention.

findalex 9mo ago

QM would tell us the order of your Hamiltonian (attention operator) doesn’t limit the complexity of the wave function (hidden state). It might be more efficient to explicitly correlate certain many-body interactions, but pair-wise interactions, depth and a basis (hidden state dimension) approaching completeness "are all you need”.

twothreeone 9mo ago
The terminology is overloaded.. Tensors in QM are objects obeying transformation laws, in ML Tensors are just data arranged in multidimensional arrays. There are no constraints on how the data transforms.
- findalex 9mo ago
  Intended as analogy - but it is essentially a description of the DMRG algorithm (quantum chem). Only pair-wise operators there but the theory approaches exact when there are enough terms in your tensor product (iterations ~ depth) and a large enough embedding dimension. > There are no constraints on how the data transforms. Except those implicit in your learned representation. And that representation could be the MB WF.

kioku 9mo ago

> Our key insight is to offload critical softmax primitives to idle tensor units, maximizing hardware utilization and throughput. > … speedups of 1.05–1.17×across diverse attention configurations on Ampere and Hopper GPUs …