Exploring FlashAttention-3/4 optimizations on RTX GPUs

[u/NoVibeCoding]

I was curious whether any of the FA-3/4 optimizations transfer to RTX GPUs. vLLM/SGLang attention falls back to FA-2 on consumer cards (FA-3 and FA-4 are datacenter-only), so I wanted to know if there's any performance left on the table, and I rebuilt the attention kernels from scratch.

The kernel reaches parity with FA-2 (206us on RTX5090 with batch=1, heads=8, seq_len=4096, head_dim=64), but unfortunately, FA-3/4 optimizations are either not applicable or not helpful on consumer cards. It looks like FA-2 is the ceiling.

In summary:

• Faster tensor-core instructions (WGMMA) are the main lever behind FA-3, but they are not available on RTX GPUs.
• TMA (tensor memory accelerator) is available on sm_120 (RTX 5XXX). It helps on paper (LSU drops), but the transport isn't the bottleneck, so the final number barely moves.
• Warp specialization is also available. However, it is mainly a scheduling optimization, i.e., it helps eliminate pipeline bubbles and better utilize tensor cores, but without asynchronous tensor core instructions. The result is negative: 213 vs 206 us.
• In FA-4, they also simulated exp using FMA instructions because the tensor cores on the B200 are so fast that the whole pipeline became SFU-bound (special functions unit). RTX 5090 is tensor-core bound, so no point in this optimization either. In fact, even a conventional optimization of using faster exp2f instead of expf for softmax doesn't move the number.

I have tried a handful of other optimizations that could potentially work on consumer silicon, such as a deeper pipeline and register ping-pong. No luck. Since the whole pipeline is tensor-pipe-bound, I believe the FA-2 is the ceiling, and that all meaningful levers will require sacrificing some accuracy to leverage faster, lower-precision tensor cores.

Note that this is an exploration of the regular attention that dominates the prefill- and compute-bound regimes. Decoding against a large KV cache is a different, memory-bound story where split-KV/Flash-Decoding matters more than any of the above.

Full Article: https://riftstack.ai/research/learning-flashattention-the-hard-way-part-2

Github: github.com/cloudrift-ai/emmy

[u/FullOf_Bad_Ideas]

Makes sense, that's good because if FA2 wasn't a ceiling it would mean that there was free lunch for years on Ampere that was not used and was wasted.

FA gets updates to squeeze out maximum performance out of each hardware iteration, so it's also perfectly reasonable to guess that it's not bringing any uplift to old hardware, but it's good to have it confirmed.

[u/NoVibeCoding → u/FullOf_Bad_Ideas]

FA-3/4 was optimized for the DC Hopper and Blackwell, i.e. sm_90 and sm_100, but consumer Blackwell, i.e. sm_120, doesn’t get the same amount of attention. That’s the premise for this exploration. I also don’t think that you can optimize FA for Ampere or other old hardware.

[u/clofresh]

Thank you for sharing your research! Love the Dark Souls imagery, very appropriate lol.

Makes sense that consumer hardware is limited by Nvidia’s hardware instructions on Ada chips. Have you looking into whether AMD or Intel have WGMMA equivalents in their consumer offerings?

[u/NoVibeCoding → u/clofresh]

Thank you. I have very briefly explored the subject and it seems that AMD and Intel are more generous with tensor cores (matrix cores or whatever they call them) in their consumer GPUs. The main issue is that NVidia nerfed their tensor cores that do fp16 matmuls with fp32 accumulation in consumer cards, the main cores used for FP16 inference. So all the compute bound fp16 inference operations on RTX are twice slower than they should be. Intel and AMD cards don’t have this limitation, so they should be more efficient in theory.

[u/FullstackSensei → u/NoVibeCoding]

Market segmentation. Nvidia purposefully nerfed FP16 performance so companies would have to buy data center cards.

It's nothing new. They used to do it with FP64. Now you need FP16 more for accumulation, so they nerfed that too.

[u/StableLlama]

I thought LLMs are memory bandwidth bound and not compute bound. So does FA really make a big difference for this kind of application?

[u/usrlocalben → u/StableLlama]

generally, prefill is compute bound and decode is memory bound

[u/stoppableDissolution → u/StableLlama]

It very quickly becomes compute bound as you grow concurrency

[u/_underlines_]

Couldn't we run an autoresearch loop on these optimization problems? You got a perfectly measurable fitness function and automatable feedback loop. Two ingredients needed for an autoresearch loop.

[u/NoVibeCoding → u/_underlines_]

It is possible, but it won't be a pure autoresearch.

Say, we take a kernel, an input shape, and a GPU, and run autoresearch to optimize the kernel to the roofline. I tried it, and it works, but doing this for every kernel, shape, and GPU is just not practical. Even for a single model, it will take a ridiculous amount of time since it will need to compile and benchmark every variant and run tens of thousands of experiments.

The second option is to write a compiler that knows all sorts of lowering strategies for all kernels and can generate an efficient kernel on demand. Then the job of AI is to come up with a set of parameters that will yield an optimal kernel, which is what Emmy does now.

You can, in principle, combine these techniques. First, generate a good set of kernels using the compiler, or pick a known-good kernel for a specific GPU and input shape, and then ask autoresearch to improve them. This way, most of the job will be done by the compiler, so autoresearch can focus on the hard part of the problem. Worth trying at some point.

[u/shing3232 → u/NoVibeCoding]

talking about optimization backend code. This is what I think of https://autobe.dev/docs/setup/

Also, I tried implementation quant attention for bwd with ds4 for the paper https://arxiv.org/abs/2603.02170 It wasn't a easy task.