fix: limit build head size and cuda caps for dev

build.toml

@@ -5,7 +5,11 @@ name = "flash_attn"
 src = ["torch-ext/torch_binding.cpp", "torch-ext/torch_binding.h"]
 
 [kernel.flash_attn]
-cuda-capabilities = ["7.0", "7.2", "7.5", "8.0", "8.6", "8.7", "8.9", "9.0"]
+cuda-capabilities = [
+  # "7.0", "7.2", "7.5", "8.0", "8.6", "8.7", 
+  "8.9",
+  # "9.0",
+]
 src = [
   "flash_attn/flash_api.cpp",
   "flash_attn/src/philox_unpack.cuh",
@@ -47,10 +51,10 @@ src = [
   # "flash_attn/src/flash_bwd_hdim256_bf16_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim256_fp16_causal_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim256_fp16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim32_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim32_bf16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim32_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim32_fp16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim32_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim32_bf16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim32_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim32_fp16_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim64_bf16_causal_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim64_bf16_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim64_fp16_causal_sm80.cu",
@@ -59,9 +63,9 @@ src = [
   # "flash_attn/src/flash_bwd_hdim96_bf16_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim96_fp16_causal_sm80.cu",
   # "flash_attn/src/flash_bwd_hdim96_fp16_sm80.cu",
-  # "flash_attn/src/flash_bwd_kernel.h",
-  # "flash_attn/src/flash_bwd_launch_template.h",
-  # "flash_attn/src/flash_bwd_preprocess_kernel.h",
+  "flash_attn/src/flash_bwd_kernel.h",
+  "flash_attn/src/flash_bwd_launch_template.h",
+  "flash_attn/src/flash_bwd_preprocess_kernel.h",
 
   ## TODO: include fwd kernels
 
@@ -81,10 +85,10 @@ src = [
   # "flash_attn/src/flash_fwd_hdim256_bf16_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim256_fp16_causal_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim256_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim32_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim32_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim32_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim32_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim32_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim32_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim32_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim32_fp16_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim64_bf16_causal_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim64_bf16_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim64_fp16_causal_sm80.cu",
@@ -92,7 +96,7 @@ src = [
   # "flash_attn/src/flash_fwd_hdim96_bf16_causal_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim96_bf16_sm80.cu",
   # "flash_attn/src/flash_fwd_hdim96_fp16_causal_sm80.cu",
-  "flash_attn/src/flash_fwd_hdim96_fp16_sm80.cu",
+  # "flash_attn/src/flash_fwd_hdim96_fp16_sm80.cu",
   "flash_attn/src/flash_fwd_kernel.h",
   "flash_attn/src/flash_fwd_launch_template.h",
   # "flash_attn/src/flash_fwd_split_hdim128_bf16_causal_sm80.cu",
@@ -111,10 +115,10 @@ src = [
   # "flash_attn/src/flash_fwd_split_hdim256_bf16_sm80.cu",
   # "flash_attn/src/flash_fwd_split_hdim256_fp16_causal_sm80.cu",
   # "flash_attn/src/flash_fwd_split_hdim256_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim32_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim32_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim32_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim32_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim32_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim32_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim32_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim32_fp16_sm80.cu",
   # "flash_attn/src/flash_fwd_split_hdim64_bf16_causal_sm80.cu",
   # "flash_attn/src/flash_fwd_split_hdim64_bf16_sm80.cu",
   # "flash_attn/src/flash_fwd_split_hdim64_fp16_causal_sm80.cu",

flake.lock

@@ -41,11 +41,11 @@
         "rocm-nix": "rocm-nix"
       },
       "locked": {
-        "lastModified": 1742905006,
-        "narHash": "sha256-SCi1f5Lti4AM0kNPlAidcgN/5YM4HgJP4KwCsMrB0IE=",
+        "lastModified": 1742916494,
+        "narHash": "sha256-crH7vQjPJZ1yrS0GA/waYvsRLLUN4PIj85L+Rpy0Q+U=",
         "ref": "refs/heads/main",
-        "rev": "517a2bf2d0a3f1faf058ab995b6ca280b0999e7c",
-        "revCount": 105,
+        "rev": "faf433757fefae660ada7d003e394b1939989a5c",
+        "revCount": 106,
         "type": "git",
         "url": "ssh://git@github.com/huggingface/kernel-builder"
       },

flake.nix

@@ -2,7 +2,7 @@
   description = "Flake for ReLU kernel";
 
   inputs = {
-    kernel-builder.url = "git+ssh://git@github.com/huggingface/kernel-builder";
+    kernel-builder.url = "github:huggingface/kernel-builder";
   };
 
   outputs =

flash_attn/src/static_switch.h

@@ -87,28 +87,33 @@
     }                                        \
   }()
 
+// #define HEADDIM_SWITCH(HEADDIM, ...)   \
+//   [&] {                                    \
+//     if (HEADDIM <= 32) {                   \
+//       constexpr static int kHeadDim = 32;  \
+//       return __VA_ARGS__();                \
+//     } else if (HEADDIM <= 64) {            \
+//       constexpr static int kHeadDim = 64;  \
+//       return __VA_ARGS__();                \
+//     } else if (HEADDIM <= 96) {            \
+//       constexpr static int kHeadDim = 96;  \
+//       return __VA_ARGS__();                \
+//     } else if (HEADDIM <= 128) {           \
+//       constexpr static int kHeadDim = 128; \
+//       return __VA_ARGS__();                \
+//     } else if (HEADDIM <= 160) {           \
+//       constexpr static int kHeadDim = 160; \
+//       return __VA_ARGS__();                \
+//     } else if (HEADDIM <= 192) {           \
+//       constexpr static int kHeadDim = 192; \
+//       return __VA_ARGS__();                \
+//     } else if (HEADDIM <= 256) {           \
+//       constexpr static int kHeadDim = 256; \
+//       return __VA_ARGS__();                \
+//     }                                      \
+//   }()
 #define HEADDIM_SWITCH(HEADDIM, ...)   \
-  [&] {                                    \
-    if (HEADDIM <= 32) {                   \
-      constexpr static int kHeadDim = 32;  \
-      return __VA_ARGS__();                \
-    } else if (HEADDIM <= 64) {            \
-      constexpr static int kHeadDim = 64;  \
-      return __VA_ARGS__();                \
-    } else if (HEADDIM <= 96) {            \
-      constexpr static int kHeadDim = 96;  \
-      return __VA_ARGS__();                \
-    } else if (HEADDIM <= 128) {           \
-      constexpr static int kHeadDim = 128; \
-      return __VA_ARGS__();                \
-    } else if (HEADDIM <= 160) {           \
-      constexpr static int kHeadDim = 160; \
-      return __VA_ARGS__();                \
-    } else if (HEADDIM <= 192) {           \
-      constexpr static int kHeadDim = 192; \
-      return __VA_ARGS__();                \
-    } else if (HEADDIM <= 256) {           \
-      constexpr static int kHeadDim = 256; \
-      return __VA_ARGS__();                \
-    }                                      \
+  [&] {                                  \
+    constexpr static int kHeadDim = 32;  \
+    return __VA_ARGS__();                \
   }()