Writing Speed-of-Light Flash Attention for 5090 in Cuda C++

Hmm, but supposing the accelerated NVIDIA specific inference data types were available for Triton, then you would just use that? Why not contribute to Triton, they accept PRs? Like so what if you do free product ecosystem development for NVIDIA and giant corporations by contributing to Triton?

I had a 5090 some months ago but couldnt get flash attention to work. Does it now work natively? What about 5080?

I was surprised to see 5090's theoretical BF16 TFLOPs at just 209.5. That's not even 10% of the server Blackwell (B200 is 2250, and GB200 is 2500). B200 costs around $30-40k per GPU, so they are pretty close in performance per dollar.

Starting with 4090, NVIDIA limits the performance of tensor cores on gaming cards, specifically for ops that might be used in ML training. FP8 and FP16 matmuls run at full speed if accumulating in FP16 (I've never seen anyone use this), but only half speed when accumulating in FP32. This restriction is not present for lower precision matmuls like FP4, and is removed entirely on the workstation-class cards like RTX Pro 6000.

It doesn't seem worth it to use NVIDIA gaming cards as a "cheaper FLOPs" alternative anymore (e.g. diffusion models could have been cheaper to run on 3090 than A100). They are generous with memory bandwidth though, nearly 2TB/s on 5090 is amazing!

Damn awesome. This going to take me 3 reads and a week to digest

> Due to improvements in newer hardware, you might need to use more tricks to reach Speed-of-Light on older GPUs e.g. pipeline shared memory to register memory data movements.

On the contrary, older GPUs are a lot easier to hit rooflines on. Newer GPUs run so fast that they keep adding new tricks to remove bottlenecks. Not to discount the author's work here but a 5090 is pretty bad on the FLOPs/memory bandwidth ratio so it's comparatively easier to get throttled by tensor cores there; on datacenter hardware your tensor cores are so fast that you'll hit limits that were glossed over here.

For example, using Ampere "mma" instructions won't cut it, because they compute a really small MMA and force your input to live in registers. You'll need TMA to get data into shared memory and wmma to do a matrix multiply out of them. At those speeds you will run into issues with dispatching instructions and computing addresses (and doing out-of-bounds calculation) fast enough that you will need to offload it to specialized hardware to keep up with the tensor cores.

My issue with upgrading to the 5090 for workstation ML use is that it both has higher TDP than the 4090 and it can only be limited to 70% power (not 50% like the 4090).

Definitely going to save this for later and come back to it after I get some more CUDA experience under my belt. It feels so nice right now making nice beautiful to use pipeline code w/ npp and some CUDA kernels here and there, the code is much faster than what it's replacing, but then I look at this guy getting down into the weeds of memory bank contention, prefetching, loop invariance, etc. Makes me feel like I'm playing with LEGO, I'm a little jealous.

The tip that Nsight can run on Mac over SSH is great, too. I've been capturing and viewing data over RDP, basically, will have to give it a shot next week.