kernel: generate smaller kernel instantiations (#395)

guocuimi · web-flow · commit 29a9b316c322 · 2025-02-04T13:45:02.000-08:00
diff --git a/src/kernels/attention/CMakeLists.txt b/src/kernels/attention/CMakeLists.txt
@@ -16,7 +16,7 @@ cc_library(
     mha_tile.h
     mha_traits_sm80.h
     mha_kernel_sm80.cuh
-    mha_launch_sm80.cuh
+    mha_dispatch_sm80.cuh
   DEPS
     cutlass
 )
diff --git a/src/kernels/attention/attn_api.cpp b/src/kernels/attention/attn_api.cpp
@@ -3,16 +3,13 @@
 #include <ATen/cuda/CUDAContext.h>
 
 #include "cute/layout.hpp"
+#include "mha_dispatch_sm80.cuh"
 #include "mha_params.h"
 #include "static_dispatch.h"
 
 namespace llm {
 using namespace cute;
 
-// forward declaration
-template <typename Dtype, int HEAD_DIM, typename Params>
-void run_mha_kernel_sm80(Params& params, cudaStream_t stream);
-
 void paged_kv_varlen_mha(
     torch::Tensor& out,                // [n_tokens, n_heads, head_dim]
     const torch::Tensor& query,        // [n_tokens, n_heads, head_dim]
diff --git a/src/kernels/attention/generate_instantiation_cu.py b/src/kernels/attention/generate_instantiation_cu.py
@@ -7,51 +7,119 @@
 from pathlib import Path
 from typing import Iterator
 
+# map from python to c++ types
 DTYPE_MAP = {
     "fp16": "cute::half_t",
     "bf16": "cute::bfloat16_t",
 }
 
-HEAD_DIMENSIONS = [64, 96, 128, 256]
+BOOL_MAP = {
+    False: "false",
+    True: "true",
+}
+
 
-PAGEDKV_KERNEL_IMPL_TEMPLATE = """
-#include "mha_launch_sm80.cuh" // IWYU pragma: keep
+MHA_KERNEL_TEMPLATE = """
+#include "mha_kernel_sm80.cuh"  // IWYU pragma: export
+#include "mha_params.h"         // IWYU pragma: export
+#include "mha_traits_sm80.h"    // IWYU pragma: export
 
 namespace llm {{
 
+using Traits = MHATraitsSM80<{DTYPE}, {HEAD_DIM}, {BLK_M}, {BLK_N}, {BLK_K}>;
 using Params = MHAPagedKVParams;
-template void run_mha_kernel_sm80<{DTYPE}, {HEAD_DIM}, Params>(
-    Params& params, cudaStream_t stream);
 
+template void launch_mha_kernel_sm80<Traits,
+                                     Params,
+                                     /*EVEN_K=*/{EVEN_K},
+                                     /*ALIBI=*/{ALIBI},
+                                     /*SOFT_CAP=*/{SOFT_CAP},
+                                     /*LOCAL=*/{LOCAL}>(const Params& params, 
+                                                        cudaStream_t stream);
 }}  // namespace llm
 """
 
+
 @dataclass
-class Kernel:
+class MHAKernel:
     dtype: str
     head_dim: int
+    blk_m: int
+    blk_n: int
+    blk_k: int
+    even_k: bool
+    alibi: bool
+    soft_cap: bool
+    local: bool
 
     @property
     def template(self) -> str:
-        return PAGEDKV_KERNEL_IMPL_TEMPLATE.format(
-            DTYPE=DTYPE_MAP[self.dtype], HEAD_DIM=self.head_dim
+        assert self.head_dim % self.blk_k == 0
+
+        return MHA_KERNEL_TEMPLATE.format(
+            DTYPE=DTYPE_MAP[self.dtype],
+            HEAD_DIM=self.head_dim,
+            BLK_M=self.blk_m,
+            BLK_N=self.blk_n,
+            BLK_K=self.blk_k,
+            EVEN_K=BOOL_MAP[self.even_k],
+            ALIBI=BOOL_MAP[self.alibi],
+            SOFT_CAP=BOOL_MAP[self.soft_cap],
+            LOCAL=BOOL_MAP[self.local],
         )
 
     @property
     def filename(self) -> str:
-        return f"mha_{self.dtype}_hd{self.head_dim}_sm80.cu"
+        def to_str(val: bool) -> str:
+            return "1" if val else "0"
 
+        return f"mha_{self.dtype}_hd{self.head_dim}_m{self.blk_m}_n{self.blk_n}_k{self.blk_k}_ek{to_str(self.even_k)}_al{to_str(self.alibi)}_sc{to_str(self.soft_cap)}_lc{to_str(self.local)}.cu"
 
-def get_all_kernels() -> Iterator[Kernel]:
-    for dtype, head_dim in itertools.product(DTYPE_MAP.keys(), HEAD_DIMENSIONS):
-        yield Kernel(dtype=dtype, head_dim=head_dim)
+
+def gen_all_kernels() -> Iterator[MHAKernel]:
+    # mha kernel instantiations
+    for (
+        dtype,
+        head_dim,
+        blk_m,
+        blk_n,
+        blk_k,
+        even_k,
+        alibi,
+        soft_cap,
+        local,
+    ) in itertools.product(
+        ["fp16", "bf16"],  # dtype
+        [64, 96, 128, 256],  # head_dim
+        [64],  # blk_m
+        [64],  # blk_n
+        [32, 64],  # blk_k
+        [False, True],  # even_k
+        [False, True],  # alibi
+        [False, True],  # soft_cap
+        [False, True],  # local
+    ):
+        # skip invalid configurations
+        if head_dim % blk_k != 0:
+            continue
+        yield MHAKernel(
+            dtype=dtype,
+            head_dim=head_dim,
+            blk_m=blk_m,
+            blk_n=blk_n,
+            blk_k=blk_k,
+            even_k=even_k,
+            alibi=alibi,
+            soft_cap=soft_cap,
+            local=local,
+        )
 
 
 if __name__ == "__main__":
     output_dir = Path.cwd() / "generated"
     shutil.rmtree(output_dir, ignore_errors=True)
     output_dir.mkdir(parents=True, exist_ok=True)
-    
+
     # written to several files to speed up compilation
-    for kernel in get_all_kernels():
-        (output_dir / kernel.filename).write_text(kernel.template)
+    for kernel in gen_all_kernels():
+        (output_dir / kernel.filename).write_text(kernel.template)
diff --git a/src/kernels/attention/mha_dispatch_sm80.cuh b/src/kernels/attention/mha_dispatch_sm80.cuh
@@ -3,45 +3,34 @@
 #include <cute/int_tuple.hpp>
 #include <cute/layout.hpp>
 
-#include "mha_kernel_sm80.cuh"
 #include "mha_traits_sm80.h"
 #include "static_dispatch.h"
 
 namespace llm {
-namespace detail {
+// forward declaration
 template <typename Traits,
           typename Params,
           bool EVEN_K,
           bool ALIBI,
           bool SOFT_CAP,
           bool LOCAL>
-void launch_mha_kernel(const Params& params, cudaStream_t stream) {
-  const auto batch_size = params.batch_size;
-  const auto n_kv_heads = params.n_kv_heads;
-  const auto max_q_packed_len = params.max_q_len * params.group_size;
+void launch_mha_kernel_sm80(const Params& params, cudaStream_t stream);
 
-  const auto smem_size = Traits::kSmemSize;
-  auto mha_kernel =
-      mha_kernel_sm80<Traits, Params, EVEN_K, ALIBI, SOFT_CAP, LOCAL>;
-  cudaFuncSetAttribute(
-      mha_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
-  // TODO: support persistent kernels
-  dim3 grid(cute::ceil_div(max_q_packed_len, Traits::kBlockM),
-            batch_size,
-            n_kv_heads);
-  dim3 block = Traits::kThreadNum;
-  mha_kernel<<<grid, block, smem_size, stream>>>(params);
-}
+namespace detail {
 
 template <typename Traits, typename Params>
-void run_mha_kernel(const Params& params, cudaStream_t stream) {
+void dispatch_mha_kernel_sm80(const Params& params, cudaStream_t stream) {
   // dispatch to proper kernel instantiation based on params
   DISPATCH_BOOL(params.head_dim == Traits::kHeadDim, EVEN_K, [&] {
     DISPATCH_BOOL(params.alibi_slopes_ptr != nullptr, ALIBI, [&] {
       DISPATCH_BOOL(params.logits_soft_cap > 0, SOFT_CAP, [&] {
         DISPATCH_BOOL(params.sliding_window >= 0, LOCAL, [&] {
-          launch_mha_kernel<Traits, Params, EVEN_K, ALIBI, SOFT_CAP, LOCAL>(
-              params, stream);
+          launch_mha_kernel_sm80<Traits,
+                                 Params,
+                                 EVEN_K,
+                                 ALIBI,
+                                 SOFT_CAP,
+                                 LOCAL>(params, stream);
         });
       });
     });
@@ -63,36 +52,36 @@ void run_mha_kernel_sm80(Params& params, cudaStream_t stream = nullptr) {
                                  /*BLK_M=*/64,
                                  /*BLK_N=*/64,
                                  /*BLK_K=*/64>;
-    detail::run_mha_kernel<Traits>(params, stream);
+    detail::dispatch_mha_kernel_sm80<Traits>(params, stream);
   } else if constexpr (HEAD_DIM == 96) {
     using Traits = MHATraitsSM80<Dtype,
                                  HEAD_DIM,
                                  /*BLK_M=*/64,
                                  /*BLK_N=*/64,
                                  /*BLK_K=*/32>;
-    detail::run_mha_kernel<Traits>(params, stream);
+    detail::dispatch_mha_kernel_sm80<Traits>(params, stream);
   } else if constexpr (HEAD_DIM == 128) {
     using Traits = MHATraitsSM80<Dtype,
                                  HEAD_DIM,
                                  /*BLK_M=*/64,
                                  /*BLK_N=*/64,
                                  /*BLK_K=*/64>;
-    detail::run_mha_kernel<Traits>(params, stream);
+    detail::dispatch_mha_kernel_sm80<Traits>(params, stream);
   } else if constexpr (HEAD_DIM == 256) {
     using Traits = MHATraitsSM80<Dtype,
                                  HEAD_DIM,
                                  /*BLK_M=*/64,
                                  /*BLK_N=*/64,
                                  /*BLK_K=*/64>;
-    detail::run_mha_kernel<Traits>(params, stream);
+    detail::dispatch_mha_kernel_sm80<Traits>(params, stream);
   } else {
     // use the default block size
     using Traits = MHATraitsSM80<Dtype,
                                  HEAD_DIM,
                                  /*BLK_M=*/64,
                                  /*BLK_N=*/64,
                                  /*BLK_K=*/64>;
-    detail::run_mha_kernel<Traits>(params, stream);
+    detail::dispatch_mha_kernel_sm80<Traits>(params, stream);
   }
 }
 
diff --git a/src/kernels/attention/mha_kernel_sm80.cuh b/src/kernels/attention/mha_kernel_sm80.cuh
@@ -432,4 +432,28 @@ __global__ void mha_kernel_sm80(__grid_constant__ const Params params) {
   epilogue(tOrAccO);
 }
 
+template <typename Traits,
+          typename Params,
+          bool EVEN_K,
+          bool ALIBI,
+          bool SOFT_CAP,
+          bool LOCAL>
+void launch_mha_kernel_sm80(const Params& params, cudaStream_t stream) {
+  const auto batch_size = params.batch_size;
+  const auto n_kv_heads = params.n_kv_heads;
+  const auto max_q_packed_len = params.max_q_len * params.group_size;
+
+  const auto smem_size = Traits::kSmemSize;
+  auto mha_kernel =
+      mha_kernel_sm80<Traits, Params, EVEN_K, ALIBI, SOFT_CAP, LOCAL>;
+  cudaFuncSetAttribute(
+      mha_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
+  // TODO: support persistent kernels
+  dim3 grid(cute::ceil_div(max_q_packed_len, Traits::kBlockM),
+            batch_size,
+            n_kv_heads);
+  dim3 block = Traits::kThreadNum;
+  mha_kernel<<<grid, block, smem_size, stream>>>(params);
+}
+
 }  // namespace llm
diff --git a/src/kernels/attention/mha_kernel_sm80_pagedkv_test.cu b/src/kernels/attention/mha_kernel_sm80_pagedkv_test.cu
@@ -3,7 +3,8 @@
 #include <torch/torch.h>
 
 #include "cute/layout.hpp"
-#include "mha_launch_sm80.cuh"
+#include "mha_dispatch_sm80.cuh"
+#include "mha_kernel_sm80.cuh"  // IWYU pragma: keep
 #include "mha_params.h"
 #include "mha_ref.h"
 
diff --git a/src/kernels/attention/mha_kernel_sm80_test.cu b/src/kernels/attention/mha_kernel_sm80_test.cu
@@ -4,7 +4,8 @@
 #include <cstdint>
 
 #include "cute/layout.hpp"
-#include "mha_launch_sm80.cuh"
+#include "mha_dispatch_sm80.cuh"
+#include "mha_kernel_sm80.cuh"  // IWYU pragma: keep
 #include "mha_params.h"
 #include "mha_ref.h"
 
diff --git a/src/kernels/attention/mha_kernel_sm80_varlen_test.cu b/src/kernels/attention/mha_kernel_sm80_varlen_test.cu
@@ -3,7 +3,8 @@
 #include <torch/torch.h>
 
 #include "cute/layout.hpp"
-#include "mha_launch_sm80.cuh"
+#include "mha_dispatch_sm80.cuh"
+#include "mha_kernel_sm80.cuh"  // IWYU pragma: keep
 #include "mha_params.h"
 #include "mha_ref.h"
 
diff --git a/src/kernels/attention/mha_sm80_bench.cu b/src/kernels/attention/mha_sm80_bench.cu
@@ -4,7 +4,8 @@
 #include <cuda/std/chrono>
 #include <nvbench/nvbench.cuh>
 
-#include "mha_launch_sm80.cuh"
+#include "mha_dispatch_sm80.cuh"
+#include "mha_kernel_sm80.cuh"  // IWYU pragma: keep
 #include "mha_params.h"
 
 using namespace llm;
diff --git a/src/kernels/attention/mha_sm80_pagedkv_bench.cu b/src/kernels/attention/mha_sm80_pagedkv_bench.cu
@@ -5,7 +5,8 @@
 #include <cuda/std/chrono>
 #include <nvbench/nvbench.cuh>
 
-#include "mha_launch_sm80.cuh"
+#include "mha_dispatch_sm80.cuh"
+#include "mha_kernel_sm80.cuh"  // IWYU pragma: keep
 #include "mha_params.h"
 
 using namespace llm;
diff --git a/src/kernels/attention/mha_traits_sm80.h b/src/kernels/attention/mha_traits_sm80.h
@@ -40,6 +40,8 @@ struct MHATraitsSM80 {
   static constexpr int kBlockK = BLK_K;
   static constexpr int kRowsPerMMA = 2;
 
+  static_assert(kHeadDim % kBlockK == 0);
+
   using DType = DTYPE;
   using _BLK_M = Int<kBlockM>;
   using _BLK_N = Int<kBlockN>;

Original file line number	Diff line number	Diff line change
`@@ -16,7 +16,7 @@ cc_library(`
`16`	`16`	`mha_tile.h`
`17`	`17`	`mha_traits_sm80.h`
`18`	`18`	`mha_kernel_sm80.cuh`
`19`		`- mha_launch_sm80.cuh`
	`19`	`+ mha_dispatch_sm80.cuh`
`20`	`20`	`DEPS`
`21`	`21`	`cutlass`
`22`	`22`	`)`