Refactor dispatch

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer.h

@@ -8,43 +8,44 @@
 
 #include "ck_tiled_fmha_batched_infer_dispatch.h"
 #include "ck_tiled_fmha_batched_infer_splitkv_dispatch.h"
+#include "ck_tiled_fmha_dispatch_tags.h"
 #include "ck_tiled_fmha_seqlen_q_switch.h"
 
 template <
     typename ScalarType,
-    bool kHasMask,
-    bool kHasBias,
-    bool kHasDropout,
-    ck_tile::index_t MaxK>
+    typename HasMask,
+    typename HasBias,
+    typename HasDropout,
+    typename MaxHeadDimension>
 void run_batched_infer_mask_bias_dropout_dispatch(
     BatchedForwardParams& param,
     hipStream_t stream) {
   // currently split-kv implementation does not support dropout
-  if constexpr (!kHasDropout) {
+  if constexpr (!HasDropout::value) {
 #ifndef FMHA_FWD_SPLITKV_NOT_USED
     if (param.use_split_kv) {
-      FMHA_FWD_SEQLEN_Q_SWITCH(param.M, MaxSeqlenQ, [&] {
+      FMHA_FWD_SEQLEN_Q_SWITCH(param.M, kMaxSeqlenQ, [&] {
         batched_infer_splitkv_mask_bias_dropout_dispatch<
             ScalarType,
-            kHasMask,
-            kHasBias,
-            MaxK,
-            MaxSeqlenQ>::Run(param, stream);
+            HasMask,
+            HasBias,
+            MaxHeadDimension,
+            max_query_seqlen_t<kMaxSeqlenQ>>::Run(param, stream);
       });
     } else
 #endif
       batched_infer_mask_bias_dropout_dispatch<
           ScalarType,
-          kHasMask,
-          kHasBias,
-          kHasDropout,
-          MaxK>::Run(param, stream);
+          HasMask,
+          HasBias,
+          HasDropout,
+          MaxHeadDimension>::Run(param, stream);
   } else {
     batched_infer_mask_bias_dropout_dispatch<
         ScalarType,
-        kHasMask,
-        kHasBias,
-        kHasDropout,
-        MaxK>::Run(param, stream);
+        HasMask,
+        HasBias,
+        HasDropout,
+        MaxHeadDimension>::Run(param, stream);
   }
 };

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_bf16.cpp

@@ -15,23 +15,23 @@
 void batched_infer_bf16(BatchedForwardParams& param, hipStream_t stream) {
   const bool has_dropout = (param.dropout_prob > 0.0f);
   BOOL_SWITCH_2(param.has_attn_bias, kHasBias, has_dropout, kHasDropout, [&] {
-    FMHA_FWD_HEADDIM_SWITCH(param.K, param.Kv, MaxK, [&] {
+    FMHA_FWD_HEADDIM_SWITCH(param.K, param.Kv, kMaxHeadDimension, [&] {
       if (param.custom_mask_type == 0 && param.window_size <= 0)
         run_batched_infer_mask_bias_dropout_dispatch<
             ck_tile::bf16_t,
-            false,
-            kHasBias,
-            kHasDropout,
-            MaxK>(param, stream);
+            has_mask_t<false>,
+            has_bias_t<kHasBias>,
+            has_dropout_t<kHasDropout>,
+            max_head_dimension_t<kMaxHeadDimension>>(param, stream);
       else if (
           param.custom_mask_type == 1 || param.custom_mask_type == 2 ||
           param.window_size > 0)
         run_batched_infer_mask_bias_dropout_dispatch<
             ck_tile::bf16_t,
-            true,
-            kHasBias,
-            kHasDropout,
-            MaxK>(param, stream);
+            has_mask_t<true>,
+            has_bias_t<kHasBias>,
+            has_dropout_t<kHasDropout>,
+            max_head_dimension_t<kMaxHeadDimension>>(param, stream);
       else
         throw std::runtime_error("Invalid custom_mask_type value");
     });

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_dispatch.h

@@ -19,10 +19,10 @@
 
 template <
     typename ScalarType,
-    bool kHasMask,
-    bool kHasBias,
-    bool kHasDropout,
-    ck_tile::index_t MaxK>
+    typename HasMask,
+    typename HasBias,
+    typename HasDropout,
+    typename MaxHeadDimension>
 struct batched_infer_mask_bias_dropout_dispatch {
   template <typename FmhaTraits, typename FmhaMask>
   using FmhaPipelineProblemTemp = ck_tile::BlockFmhaPipelineProblem<
@@ -37,20 +37,21 @@ struct batched_infer_mask_bias_dropout_dispatch {
       typename FmhaFwdTypeConfig<ScalarType>::PDataType,
       typename FmhaFwdTypeConfig<ScalarType>::OaccDataType,
       typename FmhaFwdTypeConfig<ScalarType>::ODataType,
-      FmhaFwdShape<MaxK>,
+      FmhaFwdShape<MaxHeadDimension::value>,
       false, // kIsGroupMode
       FmhaMask,
       FmhaTraits>;
 
   static void Run(BatchedForwardParams& param, hipStream_t stream) {
-    using FmhaMask = ck_tile::SimplifiedGenericAttentionMask<kHasMask>;
+    using FmhaMask = ck_tile::SimplifiedGenericAttentionMask<HasMask::value>;
 
-    using FmhaShape = FmhaFwdShape<MaxK>;
+    using FmhaShape = FmhaFwdShape<MaxHeadDimension::value>;
     using FmhaTilePartitioner = ck_tile::FmhaFwdTilePartitioner<FmhaShape>;
-    constexpr ck_tile::index_t occupancy =
-        (MaxK == 64) ? 3 : ((MaxK == 256) ? 1 : 2);
+    constexpr ck_tile::index_t occupancy = (MaxHeadDimension::value == 64)
+        ? 3
+        : ((MaxHeadDimension::value == 256) ? 1 : 2);
 
-    constexpr auto kBiasEnum = kHasBias
+    constexpr auto kBiasEnum = HasBias::value
         ? ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS
         : ck_tile::BlockAttentionBiasEnum::NO_BIAS;
 
@@ -65,8 +66,8 @@ struct batched_infer_mask_bias_dropout_dispatch {
     const bool pad_headdim = (pad_headdim_q || pad_headdim_v);
 
     const bool use_async_pipeline =
-        (!kHasBias && (param.K % 8 == 0) && (param.Kv % 8 == 0) &&
-         (MaxK <= 128));
+        (!HasBias::value && (param.K % 8 == 0) && (param.Kv % 8 == 0) &&
+         (MaxHeadDimension::value <= 128));
 
     if (!use_async_pipeline) {
       BOOL_SWITCH_3(
@@ -85,7 +86,7 @@ struct batched_infer_mask_bias_dropout_dispatch {
                 kBiasEnum,
                 false, // kHasBiasGrad place-holder
                 false, // kStoreLSE
-                kHasDropout,
+                HasDropout::value,
                 false, // kDoFp8StaticQuant place-holder
                 occupancy>;
 
@@ -117,7 +118,7 @@ struct batched_infer_mask_bias_dropout_dispatch {
             kBiasEnum,
             false, // kHasBiasGrad place-holder
             false, // kStoreLSE
-            kHasDropout,
+            HasDropout::value,
             false, // kDoFp8StaticQuant place-holder
             occupancy>;
 

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_fp16.cpp

@@ -15,23 +15,23 @@
 void batched_infer_fp16(BatchedForwardParams& param, hipStream_t stream) {
   const bool has_dropout = (param.dropout_prob > 0.0f);
   BOOL_SWITCH_2(param.has_attn_bias, kHasBias, has_dropout, kHasDropout, [&] {
-    FMHA_FWD_HEADDIM_SWITCH(param.K, param.Kv, MaxK, [&] {
+    FMHA_FWD_HEADDIM_SWITCH(param.K, param.Kv, kMaxHeadDimension, [&] {
       if (param.custom_mask_type == 0 && param.window_size <= 0)
         run_batched_infer_mask_bias_dropout_dispatch<
             ck_tile::fp16_t,
-            false,
-            kHasBias,
-            kHasDropout,
-            MaxK>(param, stream);
+            has_mask_t<false>,
+            has_bias_t<kHasBias>,
+            has_dropout_t<kHasDropout>,
+            max_head_dimension_t<kMaxHeadDimension>>(param, stream);
       else if (
           param.custom_mask_type == 1 || param.custom_mask_type == 2 ||
           param.window_size > 0)
         run_batched_infer_mask_bias_dropout_dispatch<
             ck_tile::fp16_t,
-            true,
-            kHasBias,
-            kHasDropout,
-            MaxK>(param, stream);
+            has_mask_t<true>,
+            has_bias_t<kHasBias>,
+            has_dropout_t<kHasDropout>,
+            max_head_dimension_t<kMaxHeadDimension>>(param, stream);
       else
         throw std::runtime_error("Invalid custom_mask_type value");
     });

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_batched_infer_splitkv_dispatch.h

@@ -19,10 +19,10 @@
 
 template <
     typename ScalarType,
-    bool kHasMask,
-    bool kHasBias,
-    ck_tile::index_t MaxK,
-    ck_tile::index_t MaxSeqlenQ>
+    typename HasMask,
+    typename HasBias,
+    typename MaxHeadDimension,
+    typename MaxSeqlenQ>
 struct batched_infer_splitkv_mask_bias_dropout_dispatch {
   template <
       typename FmhaFwdSplitKVTraits,
@@ -40,7 +40,9 @@ struct batched_infer_splitkv_mask_bias_dropout_dispatch {
           typename FmhaFwdTypeConfig<ScalarType>::PDataType,
           typename FmhaFwdTypeConfig<ScalarType>::OaccDataType,
           ODataType,
-          typename FmhaFwdSplitKVShape<MaxK, MaxSeqlenQ>::Type,
+          typename FmhaFwdSplitKVShape<
+              MaxHeadDimension::value,
+              MaxSeqlenQ::value>::Type,
           false, // kIsGroupMode
           FmhaMask,
           FmhaFwdSplitKVTraits>;
@@ -54,23 +56,24 @@ struct batched_infer_splitkv_mask_bias_dropout_dispatch {
           typename FmhaFwdTypeConfig<ScalarType>::LSEDataType,
           typename FmhaFwdTypeConfig<ScalarType>::OaccDataType,
           typename FmhaFwdTypeConfig<ScalarType>::ODataType,
-          MaxK, // headdim_v
+          MaxHeadDimension::value, // headdim_v
           kM0,
           kN1,
           false, // kIsGroupMode
           FmhaSplitKVCombineTraits>;
 
   static void Run(BatchedForwardParams& param, hipStream_t stream) {
     {
-      using FmhaMask = ck_tile::SimplifiedGenericAttentionMask<kHasMask>;
+      using FmhaMask = ck_tile::SimplifiedGenericAttentionMask<HasMask::value>;
 
-      using FmhaTileShape =
-          typename FmhaFwdSplitKVShape<MaxK, MaxSeqlenQ>::Type;
+      using FmhaTileShape = typename FmhaFwdSplitKVShape<
+          MaxHeadDimension::value,
+          MaxSeqlenQ::value>::Type;
       using FmhaTilePartitioner =
           ck_tile::FmhaFwdSplitKVTilePartitioner<FmhaTileShape>;
       constexpr ck_tile::index_t occupancy = -1;
 
-      constexpr auto kBiasEnum = kHasBias
+      constexpr auto kBiasEnum = HasBias::value
           ? ck_tile::BlockAttentionBiasEnum::ELEMENTWISE_BIAS
           : ck_tile::BlockAttentionBiasEnum::NO_BIAS;
 
@@ -174,8 +177,9 @@ struct batched_infer_splitkv_mask_bias_dropout_dispatch {
     };
 
     if (param.num_kv_splits > 1) {
-      using FmhaTileShape =
-          typename FmhaFwdSplitKVShape<MaxK, MaxSeqlenQ>::Type;
+      using FmhaTileShape = typename FmhaFwdSplitKVShape<
+          MaxHeadDimension::value,
+          MaxSeqlenQ::value>::Type;
 
       constexpr ck_tile::index_t kM0 = FmhaTileShape::kM0 / 2;
       constexpr ck_tile::index_t kN1 = FmhaTileShape::kN1 / 2;

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_dispatch_tags.h

@@ -0,0 +1,25 @@
+/*
+ * Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+#pragma once
+
+#include "ck_tile/core/numeric/integral_constant.hpp"
+
+template <bool v>
+struct has_mask_t : ck_tile::bool_constant<v> {};
+
+template <bool v>
+struct has_bias_t : ck_tile::bool_constant<v> {};
+
+template <bool v>
+struct has_dropout_t : ck_tile::bool_constant<v> {};
+
+template <ck_tile::index_t v>
+struct max_head_dimension_t : ck_tile::integral_constant<ck_tile::index_t, v> {
+};
+
+template <ck_tile::index_t v>
+struct max_query_seqlen_t : ck_tile::integral_constant<ck_tile::index_t, v> {};

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_infer.h

@@ -6,45 +6,46 @@
  */
 #pragma once
 
+#include "ck_tiled_fmha_dispatch_tags.h"
 #include "ck_tiled_fmha_grouped_infer_dispatch.h"
 #include "ck_tiled_fmha_grouped_infer_splitkv_dispatch.h"
 #include "ck_tiled_fmha_seqlen_q_switch.h"
 
 template <
     typename ScalarType,
-    bool kHasMask,
-    bool kHasBias,
-    bool kHasDropout,
-    ck_tile::index_t MaxK>
+    typename HasMask,
+    typename HasBias,
+    typename HasDropout,
+    typename MaxHeadDimension>
 void run_grouped_infer_mask_bias_dropout_dispatch(
     GroupedForwardParams& param,
     hipStream_t stream) {
   // currently split-kv implementation does not support dropout
-  if constexpr (!kHasDropout) {
+  if constexpr (!HasDropout::value) {
 #ifndef FMHA_FWD_SPLITKV_NOT_USED
     if (param.use_split_kv) {
-      FMHA_FWD_SEQLEN_Q_SWITCH(param.max_seqlen_q, MaxSeqlenQ, [&] {
+      FMHA_FWD_SEQLEN_Q_SWITCH(param.max_seqlen_q, kMaxSeqlenQ, [&] {
         grouped_infer_splitkv_mask_bias_dropout_dispatch<
             ScalarType,
-            kHasMask,
-            kHasBias,
-            MaxK,
-            MaxSeqlenQ>::Run(param, stream);
+            HasMask,
+            HasBias,
+            MaxHeadDimension,
+            max_query_seqlen_t<kMaxSeqlenQ>>::Run(param, stream);
       });
     } else
 #endif
       grouped_infer_mask_bias_dropout_dispatch<
           ScalarType,
-          kHasMask,
-          kHasBias,
-          kHasDropout,
-          MaxK>::Run(param, stream);
+          HasMask,
+          HasBias,
+          HasDropout,
+          MaxHeadDimension>::Run(param, stream);
   } else {
     grouped_infer_mask_bias_dropout_dispatch<
         ScalarType,
-        kHasMask,
-        kHasBias,
-        kHasDropout,
-        MaxK>::Run(param, stream);
+        HasMask,
+        HasBias,
+        HasDropout,
+        MaxHeadDimension>::Run(param, stream);
   }
 };

xformers/csrc/attention/hip_fmha/ck_tiled_fmha_grouped_infer_bf16.cpp

@@ -15,23 +15,23 @@
 void grouped_infer_bf16(GroupedForwardParams& param, hipStream_t stream) {
   const bool has_dropout = (param.dropout_prob > 0.0f);
   BOOL_SWITCH_2(param.has_attn_bias, kHasBias, has_dropout, kHasDropout, [&] {
-    FMHA_FWD_HEADDIM_SWITCH(param.K, param.Kv, MaxK, [&] {
+    FMHA_FWD_HEADDIM_SWITCH(param.K, param.Kv, kMaxHeadDimension, [&] {
       if (param.custom_mask_type == 0 && param.window_size <= 0)
         run_grouped_infer_mask_bias_dropout_dispatch<
             ck_tile::bf16_t,
-            false,
-            kHasBias,
-            kHasDropout,
-            MaxK>(param, stream);
+            has_mask_t<false>,
+            has_bias_t<kHasBias>,
+            has_dropout_t<kHasDropout>,
+            max_head_dimension_t<kMaxHeadDimension>>(param, stream);
       else if (
           param.custom_mask_type == 1 || param.custom_mask_type == 2 ||
           param.window_size > 0)
         run_grouped_infer_mask_bias_dropout_dispatch<
             ck_tile::bf16_t,
-            true,
-            kHasBias,
-            kHasDropout,
-            MaxK>(param, stream);
+            has_mask_t<true>,
+            has_bias_t<kHasBias>,
+            has_dropout_t<kHasDropout>,
+            max_head_dimension_t<kMaxHeadDimension>>(param, stream);
       else
         throw std::runtime_error("Invalid custom_mask_type value");
     });