[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,46 @@ 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
+    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};
+#else
+    return "";
+#endif
+}
+
+inline bool is_gfx11_arch() {
+    const std::string arch = get_gcn_arch_name();
+    return !arch.empty() && arch.rfind("gfx11", 0) == 0;
+}
+
+inline bool is_gfx12_arch() {
+    const std::string arch = get_gcn_arch_name();
+    return !arch.empty() && arch.rfind("gfx12", 0) == 0;
+}
+
+inline bool is_gfx1x_arch() {
+    return is_gfx11_arch() || is_gfx12_arch();
+}
+
+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

@@ -2,11 +2,13 @@
  * Copyright (c) 2024, Tri Dao.
  ******************************************************************************/
 
-#include "flash_common.hpp"
+#include "mha_fwd_head_grouping_utils.hpp"
 
-#include "fmha_fwd.hpp"
 #include "mask.hpp"
 
+#include <optional>
+#include <string>
+
 fmha_fwd_traits get_ck_fmha_fwd_traits(const mask_info &mask,
                                        std::string dtype,
                                        int head_size,
@@ -119,6 +121,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,
@@ -330,7 +334,27 @@ mha_fwd(at::Tensor &q,                            // batch_size x seqlen_q x num
                 p_dropout,
                 drop_seed_offset);
 
-        float t = fmha_fwd(traits, args, stream_config);
+        float t =
+            flash::maybe_dispatch_head_grouped_fwd(
+                stream_config,
+                traits,
+                args,
+                num_heads,
+                num_heads_k,
+                batch_size,
+                seqlen_k,
+                head_size,
+                head_size,
+                k.element_size(),
+                v.element_size(),
+                q.scalar_type(),
+                [&](const auto& grouped_traits, auto& grouped_args, const auto& grouped_sc) {
+                    return fmha_fwd(grouped_traits, grouped_args, grouped_sc);
+                });
+
+        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_fwd_head_grouping_utils.hpp

@@ -0,0 +1,101 @@
+/******************************************************************************
+ * Copyright (c) 2024, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include "flash_common.hpp"
+
+#include "fmha_fwd.hpp"
+#include "fmha_fwd_head_grouping.hpp"
+
+#include <iostream>
+
+namespace flash {
+
+template <typename FmhaFwdTraits, typename FmhaFwdArgs, typename FmhaFwdFn>
+inline float maybe_dispatch_head_grouped_fwd(const ck_tile::stream_config& stream_config,
+                                             const FmhaFwdTraits& traits,
+                                             const FmhaFwdArgs& args,
+                                             int num_heads,
+                                             int num_heads_k,
+                                             int batch_size,
+                                             int seqlen_k,
+                                             int head_size_q,
+                                             int head_size_v,
+                                             size_t elem_bytes_k,
+                                             size_t elem_bytes_v,
+                                             at::ScalarType q_dtype,
+                                             FmhaFwdFn&& fmha_fwd_fn)
+{
+    namespace head_grouping = fmha_fwd_head_grouping;
+
+    if (head_grouping::disabled_by_env()) {
+        if (head_grouping::log_enabled()) {
+            std::cout << "[LLC Head Grouping] disabled by env" << std::endl;
+        }
+        return -1.0f;
+    }
+
+    const auto group_size_opt = head_grouping::get_head_group_size(num_heads,
+                                                                   num_heads_k,
+                                                                   batch_size,
+                                                                   seqlen_k,
+                                                                   head_size_q,
+                                                                   head_size_v,
+                                                                   elem_bytes_k,
+                                                                   elem_bytes_v);
+    if (!group_size_opt.has_value() || group_size_opt.value() >= num_heads) {
+        if (head_grouping::log_enabled()) {
+            std::cout << "[LLC Head Grouping] skipped (group_size not set or >= nhead)"
+                      << std::endl;
+        }
+        return -1.0f;
+    }
+
+    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"
+                  << " 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_fn(grouped_traits, grouped_args, grouped_sc);
+            });
+    };
+
+    if (q_dtype == torch::kFloat16) {
+        return dispatch_grouped_fwd(FmhaFwdTypeConfig<FmhaFwdFp16>{});
+    }
+    if (q_dtype == torch::kBFloat16) {
+        return dispatch_grouped_fwd(FmhaFwdTypeConfig<FmhaFwdBf16>{});
+    }
+    return -1.0f;
+}
+
+} // namespace flash

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

@@ -2,11 +2,13 @@
  * Copyright (c) 2024, Tri Dao.
  ******************************************************************************/
 
-#include "flash_common.hpp"
+#include "mha_fwd_head_grouping_utils.hpp"
 
-#include "fmha_fwd.hpp"
 #include "mask.hpp"
 
+#include <optional>
+#include <string>
+
 fmha_fwd_traits get_ck_fmha_varlen_fwd_traits(const mask_info &mask,
                                               std::string dtype,
                                               int head_size,
@@ -141,6 +143,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,
@@ -572,7 +576,27 @@ mha_varlen_fwd(at::Tensor &q,                   // total_q x num_heads x head_si
                     p_dropout,
                     drop_seed_offset);
 
-            float t = fmha_fwd(traits, args, stream_config);
+            float t =
+                flash::maybe_dispatch_head_grouped_fwd(
+                    stream_config,
+                    traits,
+                    args,
+                    num_heads,
+                    num_heads_k,
+                    batch_size,
+                    max_seqlen_k,
+                    head_size,
+                    head_size,
+                    k.element_size(),
+                    v.element_size(),
+                    q.scalar_type(),
+                    [&](const auto& grouped_traits, auto& grouped_args, const auto& grouped_sc) {
+                        return fmha_fwd(grouped_traits, grouped_args, grouped_sc);
+                    });
+
+            if (t < 0.0f) {
+                t = fmha_fwd(traits, args, stream_config);
+            }
             TORCH_CHECK(t >= 0, "invalid argument for fmha_fwd");
         }
     }