RDNA3 support

[WilliamGazeley]

Great work so far. I'm trying to run vLLM on my 7900XTX cards and was wondering if there were any plans to support RDNA3?

[sdli1995]

A CK disscussion has show a branch which has flash-attention kernel impl and already work in ait ROCm/composable_kernel#1032
are there any barrier on RNDA3 support ?

[dejay-vu]

Hi @WilliamGazeley and @sdli1995, I wanted to update you on the attention kernels for the NAVI platform. My colleague @aska-0096 did activate them. However, these Flash-Attention kernels were initially developed for MI devices, which operate on a distinct set of CK kernels.

In essence, the issue is that we haven't yet integrated these kernels into our current API. I plan to work on this integration in my spare time.

[AlpinDale]

Great work on the fork, @howiejayz

Will it take too long to port the kernels for gfx1100 and other non-MI architectures? I need the kernels for my project and I'd be willing to help out if I can.

[dejay-vu]

Thanks for reaching out and offering to help @AlpinDale! I'm currently tied up with a few other projects, so I can't give an exact timeframe for porting the kernels for gfx1100 and other architectures. But I'm definitely planning to tackle this soon. The first step will be creating a new code path for gfx110x, considering their CK kernels are only for forward ops.

I'm totally open to suggestions or any help you can provide. It'd be great to have some extra hands on this. Let me know if you're interested!

[evshiron]

I am a complete novice in this field, but a few months ago I managed to make Composable Kernel, Flash Attention and PyTorch work together for my RX 7900 XTX (see here, sort by performance, and look for the first one). Although I was able get that Flash Attention implementation "working" in the end, the generated images were meaningless, and I gave up because I didn't know how to fix it. Here are the relevant branch links, and I hope they can be of some help to you:

• https://github.com/are-we-gfx1100-yet/composable_kernel/tree/feat/flash-attn
• https://github.com/are-we-gfx1100-yet/flash-attention/tree/feat/adapter
• https://github.com/are-we-gfx1100-yet/pytorch/tree/feat/flash-attn

I made a grouped fused kernel by Frankensteining the batched fused kernel, which matched the call signatures in this repo at that time. However, that self-made kernel might just be broken.

[dejay-vu]

Hi @evshiron! First off, I must say I'm seriously impressed by your work! It's quite an achievement, and the resources you've provided are invaluable.

I've had the opportunity to build your implementation on gfx1100, and I'm pleased to report that the build was successful. However, I encountered an issue with the unit tests not passing due to incorrect results in the forward pass:

assert (output - output_ref).abs().max().item() <= 2 * (output_pt - output_ref).abs().max().item()
AssertionError: assert 0.109619140625 <= (2 * 0.000244140625)

which is likely stemming from incorrect parameter settings in the CK kernels. I guess this should be the reason why the output image become meaningless.

Despite this, your work has been immensely helpful! This will massively speed up the navi porting process for the v2 implementation.

[evshiron]

@howiejayz

I'm glad that my humble work could be of some help. I am indeed unfamiliar with this field, so I can only leave it to professionals.
Furthermore, as you can see, even though I managed to compile it, the improvement in the benchmark is quite limited (I didn't use the specific commit showed here). I hope it's just an issue with my implementation and I look forward to better performance in future implementations.

[dejay-vu]

Guys I have added the batched forward(consistent sequence lengths) support for gfx1100, gfx1101, gfx1102 under this branch. Thanks to @aska-0096's CK kernels. The implementation is still under development and there are a lot of things to fine-tune. For now I see the performance is generally better when head dim = 64

To install just use pip install .

I only had the chance to test it on gfx1100 but I expect it works as well for the other two. Let me know if there is any issue! The docker I used to test is rocm/pytorch:latest where torch==2.1.0

[xzuyn]

under this branch.

benchmark_flash_attention_forward.py works, but benchmark_flash_attention.py doesn't. Forward speeds look pretty nice.

Using a 7900XTX with torch 2.2.0.dev20231209+rocm5.7.

Results for `benchmark_flash_attention_forward.py`

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 38.52 TFLOPs/s
Pytorch fwd: 13.37 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 39.33 TFLOPs/s
Pytorch fwd: 14.22 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 41.23 TFLOPs/s
Pytorch fwd: 16.04 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 42.06 TFLOPs/s
Pytorch fwd: 18.02 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 42.02 TFLOPs/s
Pytorch fwd: 19.16 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 37.27 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=128, batch_size=32, seqlen=512 ###
Flash2 fwd: 28.27 TFLOPs/s
Pytorch fwd: 20.43 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=128, batch_size=16, seqlen=1024 ###
Flash2 fwd: 29.38 TFLOPs/s
Pytorch fwd: 21.05 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=128, batch_size=8, seqlen=2048 ###
Flash2 fwd: 30.49 TFLOPs/s
Pytorch fwd: 25.23 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=128, batch_size=4, seqlen=4096 ###
Flash2 fwd: 31.00 TFLOPs/s
Pytorch fwd: 26.99 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=128, batch_size=2, seqlen=8192 ###
Flash2 fwd: 27.50 TFLOPs/s
Pytorch fwd: 28.47 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=False, headdim=128, batch_size=1, seqlen=16384 ###
Flash2 fwd: 20.67 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 24.02 TFLOPs/s
Pytorch fwd: 5.07 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 29.08 TFLOPs/s
Pytorch fwd: 5.48 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 33.49 TFLOPs/s
Pytorch fwd: 5.84 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 36.44 TFLOPs/s
Pytorch fwd: 6.21 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 38.54 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 39.70 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=128, batch_size=32, seqlen=512 ###
Flash2 fwd: 17.89 TFLOPs/s
Pytorch fwd: 8.42 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=128, batch_size=16, seqlen=1024 ###
Flash2 fwd: 21.69 TFLOPs/s
Pytorch fwd: 8.68 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=128, batch_size=8, seqlen=2048 ###
Flash2 fwd: 25.64 TFLOPs/s
Pytorch fwd: 9.78 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=128, batch_size=4, seqlen=4096 ###
Flash2 fwd: 27.06 TFLOPs/s
Pytorch fwd: 10.01 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=128, batch_size=2, seqlen=8192 ###
Flash2 fwd: 27.43 TFLOPs/s
Pytorch fwd: 9.87 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s
### causal=True, headdim=128, batch_size=1, seqlen=16384 ###
Flash2 fwd: 24.68 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s
Triton fwd: 0.00 TFLOPs/s

Results for `test_flash_attn_wmma_rocm.py`

=============== 125 failed, 2148 passed, 4606 skipped in 46.01s ================

Full Log: test_flash_attn_wmma_rocm.log

Error for `benchmark_flash_attention.py`

> python benchmarks/benchmark_flash_attention.py

Traceback (most recent call last):
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/flash-attention/benchmarks/benchmark_flash_attention.py", line 97, in <module>
    f, b = time_fwd_bwd(
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/flash-attention/benchmarks/benchmark_flash_attention.py", line 66, in time_fwd_bwd
    time_f, time_b = benchmark_fwd_bwd(func, *args, **kwargs)
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/flash_attn/utils/benchmark.py", line 99, in benchmark_fwd_bwd
    benchmark_backward(fn, *inputs, grad=grad, repeats=repeats, desc=desc, verbose=verbose,
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/flash_attn/utils/benchmark.py", line 53, in benchmark_backward
    m = t.timeit(repeats)
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/torch/utils/benchmark/utils/timer.py", line 274, in timeit
    self._timeit(number=max(int(number // 100), 2))
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/torch/utils/benchmark/utils/timer.py", line 264, in _timeit
    return max(self._timer.timeit(number), 1e-9)
  File "/usr/lib/python3.10/timeit.py", line 178, in timeit
    timing = self.inner(it, self.timer)
  File "<timeit-src>", line 6, in inner
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/flash_attn/utils/benchmark.py", line 46, in f
    y.backward(grad, retain_graph=True)
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/torch/_tensor.py", line 503, in backward
    torch.autograd.backward(
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 266, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/torch/autograd/function.py", line 289, in apply
    return user_fn(self, *args)
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 109, in backward
    _flash_attn_backward(
  File "/home/USER/clones/LLaMA-Efficient-Tuning/venv/lib/python3.10/site-packages/flash_attn/flash_attn_interface.py", line 66, in _flash_attn_backward
    dq, dk, dv, softmax_d, = flash_attn_cuda.bwd(
TypeError: bwd(): incompatible function arguments. The following argument types are supported:
    1. () -> None

[AlpinDale]

Some benchmark results.

RTX 4090

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 150.90 TFLOPs/s,
Pytorch fwd: 20.23 TFLOPs/s,
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 154.49 TFLOPs/s,
Pytorch fwd: 23.69 TFLOPs/s,
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 171.80 TFLOPs/s,
Pytorch fwd: 26.21 TFLOPs/s,
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 172.81 TFLOPs/s,
Pytorch fwd: 27.89 TFLOPs/s,
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 172.96 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 173.04 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,

7900 XTX

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 42.98 TFLOPs/s
Pytorch fwd: 14.38 TFLOPs/s
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 44.31 TFLOPs/s
Pytorch fwd: 14.83 TFLOPs/s
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 48.25 TFLOPs/s
Pytorch fwd: 17.32 TFLOPs/s
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 47.55 TFLOPs/s
Pytorch fwd: 19.40 TFLOPs/s
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 38.40 TFLOPs/s
Pytorch fwd: 20.19 TFLOPs/s
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 41.01 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s

[sdli1995]

Some benchmark results.

RTX 4090

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 150.90 TFLOPs/s,
Pytorch fwd: 20.23 TFLOPs/s,
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 154.49 TFLOPs/s,
Pytorch fwd: 23.69 TFLOPs/s,
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 171.80 TFLOPs/s,
Pytorch fwd: 26.21 TFLOPs/s,
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 172.81 TFLOPs/s,
Pytorch fwd: 27.89 TFLOPs/s,
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 172.96 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 173.04 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,

7900 XTX

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 42.98 TFLOPs/s
Pytorch fwd: 14.38 TFLOPs/s
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 44.31 TFLOPs/s
Pytorch fwd: 14.83 TFLOPs/s
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 48.25 TFLOPs/s
Pytorch fwd: 17.32 TFLOPs/s
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 47.55 TFLOPs/s
Pytorch fwd: 19.40 TFLOPs/s
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 38.40 TFLOPs/s
Pytorch fwd: 20.19 TFLOPs/s
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 41.01 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s

4090 fp16 accumulate fp16 tensorcore performance is 330T ,while 7900xtx is 120T the better reference nvidia card is rtx3090

[aska-0096]

Some benchmark results.
RTX 4090

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 150.90 TFLOPs/s,
Pytorch fwd: 20.23 TFLOPs/s,
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 154.49 TFLOPs/s,
Pytorch fwd: 23.69 TFLOPs/s,
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 171.80 TFLOPs/s,
Pytorch fwd: 26.21 TFLOPs/s,
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 172.81 TFLOPs/s,
Pytorch fwd: 27.89 TFLOPs/s,
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 172.96 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 173.04 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,

7900 XTX

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 42.98 TFLOPs/s
Pytorch fwd: 14.38 TFLOPs/s
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 44.31 TFLOPs/s
Pytorch fwd: 14.83 TFLOPs/s
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 48.25 TFLOPs/s
Pytorch fwd: 17.32 TFLOPs/s
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 47.55 TFLOPs/s
Pytorch fwd: 19.40 TFLOPs/s
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 38.40 TFLOPs/s
Pytorch fwd: 20.19 TFLOPs/s
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 41.01 TFLOPs/s
Pytorch fwd: 0.00 TFLOPs/s

4090 fp16 accumulate fp16 tensorcore performance is 330T ,while 7900xtx is 120T the better reference nvidia card is rtx3090

Thanks for the benchmark data. We are going to launch a new version of Composable Kernel with better flash-attention performance. Adapt the optimization on RDNA3 is in my plan.

[AlpinDale]

@sdli1995 here's the benchmarks with a 3090:

### causal=False, headdim=64, batch_size=32, seqlen=512 ###
Flash2 fwd: 65.38 TFLOPs/s,
Pytorch fwd: 18.38 TFLOPs/s,
### causal=False, headdim=64, batch_size=16, seqlen=1024 ###
Flash2 fwd: 72.94 TFLOPs/s,
Pytorch fwd: 21.69 TFLOPs/s,
### causal=False, headdim=64, batch_size=8, seqlen=2048 ###
Flash2 fwd: 74.11 TFLOPs/s,
Pytorch fwd: 18.92 TFLOPs/s,
### causal=False, headdim=64, batch_size=4, seqlen=4096 ###
Flash2 fwd: 74.98 TFLOPs/s,
Pytorch fwd: 22.27 TFLOPs/s,
### causal=False, headdim=64, batch_size=2, seqlen=8192 ###
Flash2 fwd: 75.06 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,
### causal=False, headdim=64, batch_size=1, seqlen=16384 ###
Flash2 fwd: 75.12 TFLOPs/s,
Pytorch fwd: 0.00 TFLOPs/s,

[AlpinDale]

Any updates on this?

[Wintoplay]

We need official support for flash attention