[CK_TILE] Update CK and add RDNA build support

README.md

@@ -140,9 +140,10 @@ container from ROCm, which has all the required tools to install FlashAttention.
 
 #### Composable Kernel Backend
 FlashAttention-2 ROCm CK backend currently supports:
-1. MI200x, MI250x, MI300x, and MI355x GPUs.
+1. MI200x, MI250x, MI300x, MI355x, and RDNA 3/4 GPUs.
 2. Datatype fp16 and bf16
 3. Both forward's and backward's head dimensions up to 256.
+4. RDNA 3 GPUs do not currently support backward, and RDNA 4 GPUs support backward only with deterministic=False
 
 #### Triton Backend
 The Triton implementation of [Flash Attention](https://tridao.me/publications/flash2/flash2.pdf) supports AMD's CDNA (MI200, MI300) and RDNA GPUs using fp16, bf16, and fp32 datatypes. It provides forward and backward passes with causal masking, variable sequence lengths, arbitrary Q/KV sequence lengths and head sizes, MQA/GQA, dropout, rotary embeddings, ALiBi, paged attention, and FP8 (via the Flash Attention v3 interface). Sliding window attention is currently a work in progress.

csrc/flash_attn_ck/flash_common.hpp

@@ -9,6 +9,11 @@
 #include <torch/nn/functional.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
+#include <string>
+
+#ifdef USE_ROCM
+#include <hip/hip_runtime.h>
+#endif
 
 #ifdef OLD_GENERATOR_PATH
 #include <ATen/CUDAGeneratorImpl.h>
@@ -73,4 +78,56 @@ inline int num_splits_heuristic_ck(int batch_nheads_mblocks, int num_SMs, int nu
 
 int override_num_splits_if_necessary(int batch, int nhead, int max_seqlen_q, int hdim_v, float p_drop, int num_splits);
 
+inline std::string get_gcn_arch_name() {
+#ifdef USE_ROCM
+    static const std::string cached_arch = []() {
+        int dev = 0;
+        if (hipGetDevice(&dev) != hipSuccess) {
+            return std::string{};
+        }
+        hipDeviceProp_t prop{};
+        if (hipGetDeviceProperties(&prop, dev) != hipSuccess) {
+            return std::string{};
+        }
+        return std::string{prop.gcnArchName};
+    }();
+    return cached_arch;
+#else
+    return "";
+#endif
+}
+
+inline bool is_gfx11_arch() {
+    static const bool cached = []() {
+        const std::string arch = get_gcn_arch_name();
+        return !arch.empty() && arch.rfind("gfx11", 0) == 0;
+    }();
+    return cached;
+}
+
+inline bool is_gfx12_arch() {
+    static const bool cached = []() {
+        const std::string arch = get_gcn_arch_name();
+        return !arch.empty() && arch.rfind("gfx12", 0) == 0;
+    }();
+    return cached;
+}
+
+inline bool is_gfx1x_arch() {
+    static const bool cached = is_gfx11_arch() || is_gfx12_arch();
+    return cached;
+}
+
+inline void check_gfx1x_bwd_supported(bool deterministic) {
+    if (is_gfx11_arch()) {
+        TORCH_CHECK(false, "CK backward is not supported on gfx11.");
+    }
+
+    if (is_gfx12_arch() && deterministic) {
+        TORCH_CHECK(false,
+                    "Deterministic CK backward is not supported on gfx12. "
+                    "Please rerun with deterministic=False.");
+    }
+}
+
 } // namespace flash

csrc/flash_attn_ck/mha_bwd.cpp

@@ -346,6 +346,9 @@ mha_bwd(const at::Tensor &dout,                   // batch_size x seqlen_q x num
     at::cuda::CUDAGuard device_guard{q.device()};
 
     auto opts = q.options();
+    if (flash::is_gfx1x_arch()) {
+        flash::check_gfx1x_bwd_supported(deterministic);
+    }
     auto softmax_d = torch::empty({batch_size, num_heads, seqlen_q}, opts.dtype(at::kFloat));
     at::Tensor dq_accum = torch::zeros({batch_size, num_heads, nsplits, seqlen_q, head_size}, opts.dtype(at::kFloat));
 

csrc/flash_attn_ck/mha_fwd.cpp

@@ -5,8 +5,15 @@
 #include "flash_common.hpp"
 
 #include "fmha_fwd.hpp"
+#include "fmha_fwd_head_grouping.hpp"
 #include "mask.hpp"
 
+#include <optional>
+#include <string>
+#include <iostream>
+
+namespace head_grouping = fmha_fwd_head_grouping;
+
 fmha_fwd_traits get_ck_fmha_fwd_traits(const mask_info &mask,
                                        std::string dtype,
                                        int head_size,
@@ -119,6 +126,8 @@ fmha_fwd_args get_ck_fmha_fwd_args(bool has_lse,
                          d,             // hdim_v
                          h,             // nhead
                          h_k,           // nhead_k
+                         0,             // num_head_q_total
+                         0,             // head_start
                          softmax_scale, // scale_s
                          0.0f,          // logits_soft_cap
                          stride_q,
@@ -329,8 +338,81 @@ mha_fwd(at::Tensor &q,                            // batch_size x seqlen_q x num
                 softmax_scale,
                 p_dropout,
                 drop_seed_offset);
-
-        float t = fmha_fwd(traits, args, stream_config);
+        float t = -1.0f;
+        if(head_grouping::disabled_by_env())
+        {
+            if(head_grouping::log_enabled())
+                std::cout << "[LLC Head Grouping] disabled by env" << std::endl;
+        }
+        else
+        {
+            const auto group_size_opt = head_grouping::get_head_group_size(
+                num_heads,
+                num_heads_k,
+                batch_size,
+                seqlen_k,
+                head_size,
+                head_size,
+                k.element_size(),
+                v.element_size());
+            if(group_size_opt.has_value() && group_size_opt.value() < num_heads)
+            {
+                if(head_grouping::log_enabled())
+                {
+                    const std::string arch = ck_tile::get_device_name();
+                    const size_t llc_bytes = head_grouping::get_llc_cache_bytes(arch);
+                    const ck_tile::index_t gqa_ratio = (num_heads_k > 0 ? (num_heads / num_heads_k) : 1);
+                    const ck_tile::index_t group_sz  = group_size_opt.value();
+                    const ck_tile::index_t n_groups = ck_tile::integer_divide_ceil(num_heads, group_sz);
+                    std::cout << "[LLC Head Grouping] enabled" << std::endl;
+                    std::cout << "[LLC Head Grouping] arch=" << (arch.empty() ? "unknown" : arch)
+                              << " llc_mb=" << (llc_bytes / (1024ull * 1024ull))
+                              << " nhead_q=" << num_heads << " nhead_k=" << num_heads_k
+                              << " gqa_ratio=" << gqa_ratio << " group_size=" << group_sz
+                              << " groups=" << n_groups << std::endl;
+                }
+                const bool use_blockscale_qscale = traits.qscale_type == quant_scale_enum::blockscale;
+                auto dispatch_grouped_fwd = [&](auto type_config_tag) {
+                    using TypeConfig = decltype(type_config_tag);
+                    return head_grouping::run_fwd_head_grouped<typename TypeConfig::QDataType,
+                                                               typename TypeConfig::KDataType,
+                                                               typename TypeConfig::VDataType,
+                                                               typename TypeConfig::ODataType,
+                                                               float,
+                                                               typename TypeConfig::LSEDataType,
+                                                               typename TypeConfig::RandValOutputDataType>(
+                        stream_config,
+                        traits,
+                        args,
+                        num_heads,
+                        num_heads_k,
+                        group_size_opt.value(),
+                        use_blockscale_qscale,
+                        [&](const auto& grouped_traits,
+                            auto& grouped_args,
+                            const auto& grouped_sc) {
+                            return fmha_fwd(grouped_traits, grouped_args, grouped_sc);
+                        });
+                };
+                if(q_dtype == torch::kFloat16)
+                {
+                    t = dispatch_grouped_fwd(FmhaFwdTypeConfig<FmhaFwdFp16>{});
+                }
+                else if(q_dtype == torch::kBFloat16)
+                {
+                    t = dispatch_grouped_fwd(FmhaFwdTypeConfig<FmhaFwdBf16>{});
+                }
+            }
+            else if(head_grouping::log_enabled())
+            {
+                std::cout << "[LLC Head Grouping] skipped (group_size not set or >= nhead)"
+                          << std::endl;
+            }
+        }
+        if(t < 0.0f)
+        {
+            t = fmha_fwd(traits, args, stream_config);
+        }
         TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd");
     }
     else {

csrc/flash_attn_ck/mha_varlen_bwd.cpp

@@ -363,6 +363,9 @@ mha_varlen_bwd(const at::Tensor &dout,                   // total_q x num_heads
     at::cuda::CUDAGuard device_guard{q.device()};
 
     auto opts = q.options();
+    if (flash::is_gfx1x_arch()) {
+        flash::check_gfx1x_bwd_supported(deterministic);
+    }
     auto softmax_d = torch::empty({batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat));
     at::Tensor dq_accum  = torch::zeros({num_heads, nsplits, total_q, head_size}, opts.dtype(at::kFloat));
 
@@ -450,4 +453,4 @@ mha_varlen_bwd(const at::Tensor &dout,                   // total_q x num_heads
     }
 
     return { dq, dk, dv, softmax_d };
-}
+}

csrc/flash_attn_ck/mha_varlen_fwd.cpp

@@ -5,8 +5,15 @@
 #include "flash_common.hpp"
 
 #include "fmha_fwd.hpp"
+#include "fmha_fwd_head_grouping.hpp"
 #include "mask.hpp"
 
+#include <optional>
+#include <string>
+#include <iostream>
+
+namespace head_grouping = fmha_fwd_head_grouping;
+
 fmha_fwd_traits get_ck_fmha_varlen_fwd_traits(const mask_info &mask,
                                               std::string dtype,
                                               int head_size,
@@ -141,6 +148,8 @@ fmha_fwd_args get_ck_fmha_varlen_fwd_args(bool has_lse,
                          d,             // hdim_v
                          h,             // nhead
                          h_k,           // nhead_k
+                         0,             // num_head_q_total
+                         0,             // head_start
                          softmax_scale, // scale_s
                          0.0f,          // logits_soft_cap
                          stride_q,
@@ -571,8 +580,82 @@ mha_varlen_fwd(at::Tensor &q,                   // total_q x num_heads x head_si
                     softmax_scale,
                     p_dropout,
                     drop_seed_offset);
-
-            float t = fmha_fwd(traits, args, stream_config);
+            float t = -1.0f;
+            if(head_grouping::disabled_by_env())
+            {
+                if(head_grouping::log_enabled())
+                    std::cout << "[LLC Head Grouping] disabled by env" << std::endl;
+            }
+            else
+            {
+                const auto group_size_opt = head_grouping::get_head_group_size(
+                    num_heads,
+                    num_heads_k,
+                    batch_size,
+                    max_seqlen_k,
+                    head_size,
+                    head_size,
+                    k.element_size(),
+                    v.element_size());
+                if(group_size_opt.has_value() && group_size_opt.value() < num_heads)
+                {
+                    if(head_grouping::log_enabled())
+                    {
+                        const std::string arch = ck_tile::get_device_name();
+                        const size_t llc_bytes = head_grouping::get_llc_cache_bytes(arch);
+                        const ck_tile::index_t gqa_ratio = (num_heads_k > 0 ? (num_heads / num_heads_k) : 1);
+                        const ck_tile::index_t group_sz  = group_size_opt.value();
+                        const ck_tile::index_t n_groups = ck_tile::integer_divide_ceil(num_heads, group_sz);
+                        std::cout << "[LLC Head Grouping] enabled" << std::endl;
+                        std::cout << "[LLC Head Grouping] arch=" << (arch.empty() ? "unknown" : arch)
+                                  << " llc_mb=" << (llc_bytes / (1024ull * 1024ull))
+                                  << " nhead_q=" << num_heads << " nhead_k=" << num_heads_k
+                                  << " gqa_ratio=" << gqa_ratio << " group_size=" << group_sz
+                                  << " groups=" << n_groups << std::endl;
+                    }
+                    const bool use_blockscale_qscale =
+                        traits.qscale_type == quant_scale_enum::blockscale;
+                    auto dispatch_grouped_fwd = [&](auto type_config_tag) {
+                        using TypeConfig = decltype(type_config_tag);
+                        return head_grouping::run_fwd_head_grouped<typename TypeConfig::QDataType,
+                                                                   typename TypeConfig::KDataType,
+                                                                   typename TypeConfig::VDataType,
+                                                                   typename TypeConfig::ODataType,
+                                                                   float,
+                                                                   typename TypeConfig::LSEDataType,
+                                                                   typename TypeConfig::RandValOutputDataType>(
+                            stream_config,
+                            traits,
+                            args,
+                            num_heads,
+                            num_heads_k,
+                            group_size_opt.value(),
+                            use_blockscale_qscale,
+                            [&](const auto& grouped_traits,
+                                auto& grouped_args,
+                                const auto& grouped_sc) {
+                                return fmha_fwd(grouped_traits, grouped_args, grouped_sc);
+                            });
+                    };
+                    if(q_dtype == torch::kFloat16)
+                    {
+                        t = dispatch_grouped_fwd(FmhaFwdTypeConfig<FmhaFwdFp16>{});
+                    }
+                    else if(q_dtype == torch::kBFloat16)
+                    {
+                        t = dispatch_grouped_fwd(FmhaFwdTypeConfig<FmhaFwdBf16>{});
+                    }
+                }
+                else if(head_grouping::log_enabled())
+                {
+                    std::cout << "[LLC Head Grouping] skipped (group_size not set or >= nhead)"
+                              << std::endl;
+                }
+            }
+            if(t < 0.0f)
+            {
+                t = fmha_fwd(traits, args, stream_config);
+            }
             TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd");
         }
     }

setup.py

@@ -199,7 +199,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", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200", "gfx1201"]
 
     # Validate if each element in archs is in allowed_archs
     assert all(
@@ -397,10 +397,24 @@ def validate_and_update_archs(archs):
             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)
+        archs = [arch.lower() for arch in os.getenv("GPU_ARCHS", "native").split(";")]
+        validate_and_update_archs(archs)
+
+        if archs != ["native"]:
+            kernel_targets = archs
+        else:
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "GPU_ARCHS not provided and no ROCm device detected. "
+                    "Set GPU_ARCHS (e.g., gfx942) to target your GPU."
+                )
+            detected_arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
+            kernel_targets = [detected_arch.lower()]
+            validate_and_update_archs(kernel_targets)
+
+        targets_arg = ",".join(kernel_targets)
+        for direction in ["fwd", "fwd_appendkv", "fwd_splitkv", "bwd"]:
+            subprocess.run([sys.executable, f"{ck_dir}/example/ck_tile/01_fmha/generate.py", "-d", direction, "--output_dir", "build", "--receipt", "2", "--optdim", optdim, "--targets", targets_arg], check=True)
 
         # Check, if ATen/CUDAGeneratorImpl.h is found, otherwise use ATen/cuda/CUDAGeneratorImpl.h
         # See https://github.com/pytorch/pytorch/pull/70650
@@ -410,14 +424,7 @@ def validate_and_update_archs(archs):
             generator_flag = ["-DOLD_GENERATOR_PATH"]
 
         check_if_rocm_home_none("flash_attn")
-        archs = os.getenv("GPU_ARCHS", "native").split(";")
-        validate_and_update_archs(archs)
-
-        if archs != ['native']:
-            cc_flag = [f"--offload-arch={arch}" for arch in archs]
-        else:
-            arch = torch.cuda.get_device_properties("cuda").gcnArchName.split(":")[0]
-            cc_flag = [f"--offload-arch={arch}"]
+        cc_flag = [f"--offload-arch={arch}" for arch in kernel_targets]
 
         # HACK: The compiler flag -D_GLIBCXX_USE_CXX11_ABI is set to be the same as
         # torch._C._GLIBCXX_USE_CXX11_ABI
@@ -468,7 +475,11 @@ def validate_and_update_archs(archs):
                     # "-DFLASHATTENTION_DISABLE_BACKWARD",
                     "-D__HIP_PLATFORM_HCC__=1"]
 
-        cc_flag += [f"-DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT={os.environ.get('CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT', 3)}"]
+        ck_tile_float_to_bfloat16_default = os.environ.get("CK_TILE_FLOAT_TO_BFLOAT16_DEFAULT")
+        if ck_tile_float_to_bfloat16_default is None:
+            has_gfx11_target = any(arch.startswith("gfx11") for arch in kernel_targets)
+            ck_tile_float_to_bfloat16_default = "0" if has_gfx11_target else "3"
+        cc_flag += [f"-DCK_TILE_FLOAT_TO_BFLOAT16_DEFAULT={ck_tile_float_to_bfloat16_default}"]
 
         # Imitate https://github.com/ROCm/composable_kernel/blob/c8b6b64240e840a7decf76dfaa13c37da5294c4a/CMakeLists.txt#L190-L214
         hip_version = get_hip_version()