Do you support RDNA4

[ouco1986]

May I ask if this graphics card supports RDNA4? gfx1201

[evshiron]

• https://github.com/ROCm/flash-attention?tab=readme-ov-file#triton-backend

[Umio-Yasuno]

Hi, I have confirmed that the CK backend can be compiled for RDNA4/GFX12 by applying the following patch and running the build command.

build command:

GPU_ARCHS="gfx1201" FLASH_ATTENTION_TRITON_AMD_ENABLE="FALSE" python setup.py install

patch:

---
 .gitmodules                           |  2 +-
 csrc/composable_kernel                |  2 +-
 csrc/flash_attn_ck/mha_fwd.cpp        |  8 ++++++--
 csrc/flash_attn_ck/mha_varlen_fwd.cpp |  4 ++++
 setup.py                              | 17 ++++++++++++-----
 5 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/.gitmodules b/.gitmodules
index a6446cc5..7df5611b 100644
--- a/.gitmodules
+++ b/.gitmodules
@@ -4,4 +4,4 @@
 [submodule "csrc/composable_kernel"]
 	path = csrc/composable_kernel
 	url = https://github.com/ROCm/composable_kernel.git
-	branch = amd-master
+	branch = streamhpc/fmha-wmma
diff --git a/csrc/composable_kernel b/csrc/composable_kernel
index 663992e9..6ac2092d 160000
--- a/csrc/composable_kernel
+++ b/csrc/composable_kernel
@@ -1 +1 @@
-Subproject commit 663992e99b412991eab554b0deb89bb916d40161
+Subproject commit 6ac2092d7577e43ac421ee9d915d34d1e312082d
diff --git a/csrc/flash_attn_ck/mha_fwd.cpp b/csrc/flash_attn_ck/mha_fwd.cpp
index 68e28355..ab04e202 100644
--- a/csrc/flash_attn_ck/mha_fwd.cpp
+++ b/csrc/flash_attn_ck/mha_fwd.cpp
@@ -98,12 +98,16 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
                          has_dropout_randval ? dropout_randval.data_ptr() : nullptr,
                          has_lse ? softmax_lse.data_ptr() : nullptr,
                          out.data_ptr(),
+                         nullptr, // cu_seqlen_q_ptr
+                         nullptr, // cu_seqlen_kv_ptr
                          nullptr, // seqstart_q
                          nullptr, // seqstart_k
-                         nullptr,
+                         nullptr, // seqlen_k_ptr
+                         nullptr, // seqstart_padded_q_ptr
+                         nullptr, // seqstart_padded_k_ptr
                          seqlen_q,
                          seqlen_k,
-                         b,
+                         b, // batch
                          seqlen_q,      // max_seqlen_q
                          d,             // hdim_q
                          d,             // hdim_v
diff --git a/csrc/flash_attn_ck/mha_varlen_fwd.cpp b/csrc/flash_attn_ck/mha_varlen_fwd.cpp
index 3e4422ef..dc1525c2 100644
--- a/csrc/flash_attn_ck/mha_varlen_fwd.cpp
+++ b/csrc/flash_attn_ck/mha_varlen_fwd.cpp
@@ -119,9 +119,13 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
                          has_dropout_randval ? dropout_randval.data_ptr() : nullptr,
                          has_lse ? softmax_lse.data_ptr() : nullptr,
                          out.data_ptr(),
+                         nullptr, // cu_seqlen_q_ptr
+                         nullptr, // cu_seqlen_kv_ptr
                          seqlens_q.data_ptr(), // seqstart_q
                          seqlens_k.data_ptr(), // seqstart_k
                          nullptr,              // seqlen_kpads
+                         nullptr, // seqstart_padded_q_ptr
+                         nullptr, // seqstart_padded_k_ptr
                          total_q,
                          total_k,
                          b,
diff --git a/setup.py b/setup.py
index d54e93f6..e2cbee7e 100644
--- a/setup.py
+++ b/setup.py
@@ -132,7 +132,7 @@ def rename_cpp_to_cu(cpp_files):
 
 def validate_and_update_archs(archs):
     # List of allowed architectures
-    allowed_archs = ["native", "gfx90a", "gfx950", "gfx942"]
+    allowed_archs = ["native", "gfx90a", "gfx950", "gfx942", "gfx1201"]
 
     # Validate if each element in archs is in allowed_archs
     assert all(
@@ -328,10 +328,10 @@ elif not SKIP_CUDA_BUILD and IS_ROCM:
         if not os.path.exists("./build"):
             os.makedirs("build")
 
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd", "--output_dir", "build", "--receipt", "2"], check=True)
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_appendkv", "--output_dir", "build", "--receipt", "2"], check=True)
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_splitkv", "--output_dir", "build", "--receipt", "2"], check=True)
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "bwd", "--output_dir", "build", "--receipt", "2"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd", "--output_dir", "build", "--receipt", "2", "--targets", "gfx12"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_appendkv", "--output_dir", "build", "--receipt", "2", "--targets", "gfx12"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_splitkv", "--output_dir", "build", "--receipt", "2", "--targets", "gfx12"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "bwd", "--output_dir", "build", "--receipt", "2", "--targets", "gfx12"], check=True)
 
         # Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h
         # See https://github.com/pytorch/pytorch/pull/70650
@@ -387,7 +387,14 @@ elif not SKIP_CUDA_BUILD and IS_ROCM:
                     "-DCK_ENABLE_FP8",
                     "-DCK_ENABLE_INT8",
                     "-DCK_USE_XDL",
+                    "-DCK_USE_WMMA",
+                    "-DCK_TILE_USE_WMMA",
+                    "-DCK_USE_WMMA_FP8",
+                    "-DCK_USE_OCP_FP8",
                     "-DUSE_PROF_API=1",
+                    "-DCMAKE_HIP_COMPILER_LAUNCHER=sccache",
+                    "-DCMAKE_CXX_COMPILER_LAUNCHER=sccache",
+                    "-DCMAKE_C_COMPILER_LAUNCHER=sccache",
                     # "-DFLASHATTENTION_DISABLE_BACKWARD",
                     "-D__HIP_PLATFORM_HCC__=1"]
 
-- 
2.51.0

[ouco1986]

@Umio-Yasuno Sorry, I just saw the message that I cloned the latest flash attention source file and applied a patch. The following error is prompted:,
“patching file .gitmodules
File csrc/composable_kernel is not a regular file -- refusing to patch
1 out of 1 hunk ignored -- saving rejects to file csrc/composable_kernel.rej
patching file csrc/flash_attn_ck/mha_fwd.cpp
patching file csrc/flash_attn_ck/mha_varlen_fwd.cpp
patching file setup.py
Hunk #1 succeeded at 185 (offset 53 lines).
Hunk #2 FAILED at 328.
Hunk #3 succeeded at 429 (offset 42 lines).
1 out of 3 hunks FAILED -- saving rejects to file setup.py.rej
”
Continue using commands GPU_ARCHS="gfx1201" FLASH_ATTENTION_TRITON_AMD_ENABLE="FALSE" python setup.py install Compile, after compiling for a while, the build failed with the following error message. I am using py3.13 rocm7.02. I don't know if it's not supported? I hope to receive an answer. Thank you.
“fatal error: too many errors emitted, stopping now [-ferror-limit=]
20 errors generated when compiling for gfx1201.
failed to execute:/opt/rocm-7.0.2/lib/llvm/bin/clang++ --offload-arch=gfx1201 -I/media/wr/LTW/linux/ROCm7/flash-attention/csrc/composable_kernel/include -I/media/wr/LTW/linux/ROCm7/flash-attention/csrc/composable_kernel/library/include -I/media/wr/LTW/linux/ROCm7/flash-attention/csrc/composable_kernel/example/ck_tile/01_fmha -I/media/wr/LTW/linux/ROCm7/venvl/lib/python3.13/site-packages/torch/include -I/media/wr/LTW/linux/ROCm7/venvl/lib/python3.13/site-packages/torch/include/torch/csrc/api/include -I/media/wr/LTW/linux/ROCm7/venvl/lib/python3.13/site-packages/torch/include/THH -I/opt/rocm-7.0.2/include -I/media/wr/LTW/linux/ROCm7/venvl/include -I/usr/include/python3.13 -c -c -x hip /media/wr/LTW/linux/ROCm7/flash-attention/build/fmha_bwd_convert_dq_d128_bf16_b64x128_batch_o2_pd_deterministic.hip -o "/media/wr/LTW/linux/ROCm7/flash-attention/build/temp.linux-x86_64-cpython-313/build/fmha_bwd_convert_dq_d128_bf16_b64x128_batch_o2_pd_deterministic.o" -D__HIP_PLATFORM_AMD__=1 -DUSE_ROCM=1 -DHIPBLAS_V2 -fPIC -DCUDA_HAS_FP16=1 -D__HIP_NO_HALF_OPERATORS__=1 -D__HIP_NO_HALF_CONVERSIONS__=1 -DHIP_ENABLE_WARP_SYNC_BUILTINS=1 -O3 -std=c++17 -DCK_TILE_FMHA_FWD_FAST_EXP2=1 -fgpu-flush-denormals-to-zero -DCK_ENABLE_BF16 -DCK_ENABLE_BF8 -DCK_ENABLE_FP16 -DCK_ENABLE_FP32 -DCK_ENABLE_FP64 -DCK_ENABLE_FP8 -DCK_ENABLE_INT8 -DCK_USE_XDL -DCK_USE_WMMA -DCK_TILE_USE_WMMA -DCK_USE_WMMA_FP8 -DCK_USE_OCP_FP8 -DUSE_PROF_API=1 -DCMAKE_HIP_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache -DCMAKE_C_COMPILER_LAUNCHER=sccache -D__HIP_PLATFORM_HCC__=1 -DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT=3 -mllvm --lsr-drop-solution=1 -fno-offload-uniform-block -mllvm -enable-post-misched=0 -mllvm -amdgpu-early-inline-all=true -mllvm -amdgpu-function-calls=false -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=flash_attn_2_cuda -fno-gpu-rdc”

[ouco1986]

@Umio-Yasuno Thank you for your tutorial. I have successfully built from the source code of your rdna4 branch. And FP8's flux speed has improved, but WAN video speed is not as fast as PyTorch's attention speed. I don't know what the function of this parameter 'FLASH_TTTENTION-TRITON-AMD_INABLE=' False 'is? Is it due to support issues?

[androiddrew]

@Umio-Yasuno I tried using your patch but there seems to be significant drift since October. Is this still needed for Flash attention support with the gfx1201?

[Umio-Yasuno]

No, just apply the following patch or PR:

diff --git a/setup.py b/setup.py
index 730a190a..bddaf975 100644
--- a/setup.py
+++ b/setup.py
@@ -197,7 +197,7 @@ def rename_cpp_to_cu(cpp_files):
 
 def validate_and_update_archs(archs):
     # List of allowed architectures
-    allowed_archs = ["native", "gfx90a", "gfx950", "gfx942"]
+    allowed_archs = ["native", "gfx90a", "gfx950", "gfx942", "gfx1200", "gfx1201"]
 
     # Validate if each element in archs is in allowed_archs
     assert all(
@@ -382,10 +382,10 @@ elif not SKIP_CUDA_BUILD and IS_ROCM:
             os.makedirs("build")
 
         optdim = os.getenv("OPT_DIM", "32,64,128,256")
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd", "--output_dir", "build", "--receipt", "2", "--optdim", optdim], check=True)
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_appendkv", "--output_dir", "build", "--receipt", "2", "--optdim", optdim], check=True)
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_splitkv", "--output_dir", "build", "--receipt", "2", "--optdim", optdim], check=True)
-        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "bwd", "--output_dir", "build", "--receipt", "2", "--optdim", optdim], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd", "--output_dir", "build", "--receipt", "2", "--optdim", optdim, "--targets", "gfx12"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_appendkv", "--output_dir", "build", "--receipt", "2", "--optdim", optdim, "--targets", "gfx12"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "fwd_splitkv", "--output_dir", "build", "--receipt", "2", "--optdim", optdim, "--targets", "gfx12"], check=True)
+        subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", "bwd", "--output_dir", "build", "--receipt", "2", "--optdim", optdim, "--targets", "gfx12"], check=True)
 
         # Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h
         # See https://github.com/pytorch/pytorch/pull/70650

Dao-AILab#2054