fix: expand build combinations and include all files

build.toml

@@ -33,99 +33,98 @@ src = [
   "flash_attn/src/static_switch.h",
   "flash_attn/src/utils.h",
 
-  ## TODO: include bwd kernels
+  ## bwd kernels
 
-  # "flash_attn/src/flash_bwd_hdim128_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim128_bf16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim128_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim128_fp16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim160_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim160_bf16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim160_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim160_fp16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim192_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim192_bf16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim192_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim192_fp16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim256_bf16_causal_sm80.cu",
-  # "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_hdim128_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim128_bf16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim128_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim128_fp16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim160_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim160_bf16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim160_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim160_fp16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim192_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim192_bf16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim192_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim192_fp16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim256_bf16_causal_sm80.cu",
+  "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_hdim64_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim64_bf16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim64_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim64_fp16_sm80.cu",
-  # "flash_attn/src/flash_bwd_hdim96_bf16_causal_sm80.cu",
-  # "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_hdim64_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim64_bf16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim64_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim64_fp16_sm80.cu",
+  "flash_attn/src/flash_bwd_hdim96_bf16_causal_sm80.cu",
+  "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",
 
-  ## TODO: include fwd kernels
-
-  # "flash_attn/src/flash_fwd_hdim128_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim128_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim128_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim128_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim160_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim160_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim160_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim160_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim192_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim192_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim192_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim192_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim256_bf16_causal_sm80.cu",
-  # "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",
+  ## fwd kernels
+  "flash_attn/src/flash_fwd_hdim128_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim128_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim128_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim128_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim160_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim160_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim160_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim160_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim192_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim192_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim192_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim192_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim256_bf16_causal_sm80.cu",
+  "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_hdim64_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim64_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim64_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_hdim64_fp16_sm80.cu",
-  # "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_hdim64_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim64_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim64_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_hdim64_fp16_sm80.cu",
+  "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_kernel.h",
   "flash_attn/src/flash_fwd_launch_template.h",
-  # "flash_attn/src/flash_fwd_split_hdim128_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim128_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim128_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim160_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim160_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim160_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim160_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim192_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim192_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim192_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim192_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim256_bf16_causal_sm80.cu",
-  # "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_hdim128_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim128_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim128_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim128_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim160_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim160_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim160_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim160_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim192_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim192_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim192_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim192_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim256_bf16_causal_sm80.cu",
+  "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_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",
-  # "flash_attn/src/flash_fwd_split_hdim64_fp16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim96_bf16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim96_bf16_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim96_fp16_causal_sm80.cu",
-  # "flash_attn/src/flash_fwd_split_hdim96_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",
+  "flash_attn/src/flash_fwd_split_hdim64_fp16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim96_bf16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim96_bf16_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim96_fp16_causal_sm80.cu",
+  "flash_attn/src/flash_fwd_split_hdim96_fp16_sm80.cu",
 ]
 depends = ["torch", "cutlass_3_6"]

flash_attn/flash_api.cpp

@@ -1477,10 +1477,10 @@ mha_fwd_kvcache(at::Tensor &q,                 // batch_size x seqlen_q x num_he
 
 // NOTE: wrap the namespaced functions so all types are doubles and longs
 std::vector<at::Tensor>
-mha_fwd(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
-        const at::Tensor &k,         // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
-        const at::Tensor &v,         // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
-        const c10::optional<torch::Tensor> &out_,             // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
+mha_fwd(const at::Tensor &q,                               // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
+        const at::Tensor &k,                               // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
+        const at::Tensor &v,                               // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
+        const c10::optional<torch::Tensor> &out_,          // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
         const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
         const double p_dropout,
         const double softmax_scale,
@@ -1509,4 +1509,159 @@ mha_fwd(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x roun
     int window_size_right_int = static_cast<int>(window_size_right);
     
     return FLASH_NAMESPACE::mha_fwd(const_cast<at::Tensor &>(q), k, v, out, alibi_slopes, p_dropout_float, softmax_scale_float, is_causal, window_size_left_int, window_size_right_int, softcap_float, return_softmax, gen);
+}
+
+std::vector<at::Tensor>
+mha_varlen_fwd(const at::Tensor &q,                        // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
+               const at::Tensor &k,                        // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
+               const at::Tensor &v,                        // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
+               const c10::optional<torch::Tensor> &out_,   // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
+               const at::Tensor &cu_seqlens_q,             // batch_size + 1
+               const at::Tensor &cu_seqlens_k,             // batch_size + 1
+               const int64_t max_seqlen_q,
+               const int64_t max_seqlen_k,
+               const double p_dropout,
+               const double softmax_scale,
+               bool is_causal,
+               const int64_t window_size_left,
+               const int64_t window_size_right,
+               const double softcap,
+               const bool return_softmax,
+               const c10::optional<at::Generator> gen_) {
+
+    auto gen = gen_.value_or(at::cuda::detail::getDefaultCUDAGenerator());
+
+    // Prepare the optional arguments as non-const references.
+    std::optional<at::Tensor> out = out_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(out_.value())) : std::nullopt;
+
+    if (!out.has_value()){
+        out = torch::empty_like(q);
+    }
+
+    // Convert double to float and int64_t to int.
+    float p_dropout_float = static_cast<float>(p_dropout);
+    float softmax_scale_float = static_cast<float>(softmax_scale);
+    float softcap_float = static_cast<float>(softcap);
+    int window_size_left_int = static_cast<int>(window_size_left);
+    int window_size_right_int = static_cast<int>(window_size_right);
+
+    return FLASH_NAMESPACE::mha_varlen_fwd(const_cast<at::Tensor &>(q), k, v, out, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, p_dropout_float, softmax_scale_float, is_causal, window_size_left_int, window_size_right_int, softcap_float, return_softmax, gen);
+}
+
+std::vector<at::Tensor>
+mha_bwd(const at::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
+        const at::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size
+        const at::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const at::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const at::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size
+        const at::Tensor &softmax_lse,                  // b x h x seqlen_q
+        const std::optional<at::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size
+        const std::optional<at::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const std::optional<at::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const std::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
+        const double p_dropout,                         // probability to drop
+        const double softmax_scale,
+        const bool is_causal,
+        const int64_t window_size_left,
+        const int64_t window_size_right,
+        const double softcap,
+        const bool deterministic,
+        std::optional<at::Generator> gen_,
+        std::optional<at::Tensor> &rng_state) {
+
+    auto gen = gen_.value_or(at::cuda::detail::getDefaultCUDAGenerator());
+
+    // Prepare the optional arguments as non-const references.
+    std::optional<at::Tensor> dq = dq_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(dq_.value())) : std::nullopt;
+    std::optional<at::Tensor> dk = dk_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(dk_.value())) : std::nullopt;
+    std::optional<at::Tensor> dv = dv_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(dv_.value())) : std::nullopt;
+    std::optional<at::Tensor> alibi_slopes = alibi_slopes_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(alibi_slopes_.value())) : std::nullopt;
+    
+    // Convert double to float and int64_t to int.
+    float p_dropout_float = static_cast<float>(p_dropout);
+    float softmax_scale_float = static_cast<float>(softmax_scale);
+    float softcap_float = static_cast<float>(softcap);
+    int window_size_left_int = static_cast<int>(window_size_left);
+    int window_size_right_int = static_cast<int>(window_size_right);
+
+    return FLASH_NAMESPACE::mha_bwd(const_cast<at::Tensor &>(dout), q, k, v, out, softmax_lse, dq, dk, dv, alibi_slopes, p_dropout_float, softmax_scale_float, is_causal, window_size_left_int, window_size_right_int, softcap_float, deterministic, gen, rng_state);
+}
+
+
+std::vector<at::Tensor>
+mha_varlen_bwd(const at::Tensor &dout,                  // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
+               const at::Tensor &q,                     // batch_size x seqlen_q x num_heads x head_size
+               const at::Tensor &k,                     // batch_size x seqlen_k x num_heads_k x head_size
+               const at::Tensor &v,                     // batch_size x seqlen_k x num_heads_k x head_size
+               const at::Tensor &out,                   // batch_size x seqlen_q x num_heads x head_size
+               const at::Tensor &softmax_lse,           // b x h x seqlen_q
+               const at::Tensor &cu_seqlens_q,          // batch_size + 1
+               const at::Tensor &cu_seqlens_k,          // batch_size + 1
+               const int64_t max_seqlen_q,
+               const int64_t max_seqlen_k,
+               const double p_dropout,
+               const double softmax_scale,
+               const bool is_causal,
+               const int64_t window_size_left,
+               const int64_t window_size_right,
+               const double softcap,
+               const bool deterministic,
+               std::optional<at::Generator> gen_,
+               std::optional<at::Tensor> &rng_state) {
+
+    auto gen = gen_.value_or(at::cuda::detail::getDefaultCUDAGenerator());
+
+    // Convert double to float and int64_t to int.
+    float p_dropout_float = static_cast<float>(p_dropout);
+    float softmax_scale_float = static_cast<float>(softmax_scale);
+    float softcap_float = static_cast<float>(softcap);
+    int window_size_left_int = static_cast<int>(window_size_left);
+    int window_size_right_int = static_cast<int>(window_size_right);
+
+    return FLASH_NAMESPACE::mha_varlen_bwd(const_cast<at::Tensor &>(dout), q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, p_dropout_float, softmax_scale_float, is_causal, window_size_left_int, window_size_right_int, softcap_float, deterministic, gen, rng_state);
+}
+
+std::vector<at::Tensor>
+mha_fwd_kvcache(const at::Tensor &q,                                    // batch_size x seqlen_q x num_heads x head_size
+                const at::Tensor &kcache,                               // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
+                const at::Tensor &vcache,                               // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
+                const c10::optional<torch::Tensor> &k_,                 // batch_size x seqlen_knew x num_heads_k x head_size
+                const c10::optional<torch::Tensor> &v_,                 // batch_size x seqlen_knew x num_heads_k x head_size
+                const c10::optional<torch::Tensor> &seqlens_k_,         // batch_size
+                const c10::optional<torch::Tensor> &rotary_cos_,        // seqlen_ro x (rotary_dim / 2)
+                const c10::optional<torch::Tensor> &rotary_sin_,        // seqlen_ro x (rotary_dim / 2)
+                const c10::optional<torch::Tensor> &cache_batch_idx_,   // indices to index into the KV cache
+                const c10::optional<torch::Tensor> &leftpad_k_,         // batch_size
+                const c10::optional<at::Tensor> &block_table_,          // batch_size x max_num_blocks_per_seq
+                const c10::optional<at::Tensor> &alibi_slopes_,         // num_heads or batch_size x num_heads
+                const c10::optional<at::Tensor> &out_,                  // batch_size x seqlen_q x num_heads x head_size
+                const double softmax_scale,
+                bool is_causal,
+                const int64_t window_size_left,
+                const int64_t window_size_right,
+                const double softcap,
+                bool is_rotary_interleaved,   // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
+                const int64_t num_splits
+                ) {
+
+    // Prepare the optional arguments as non-const references.
+    std::optional<at::Tensor> k = k_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(k_.value())) : std::nullopt;
+    std::optional<at::Tensor> v = v_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(v_.value())) : std::nullopt;
+    std::optional<at::Tensor> seqlens_k = seqlens_k_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(seqlens_k_.value())) : std::nullopt;
+    std::optional<at::Tensor> rotary_cos = rotary_cos_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(rotary_cos_.value())) : std::nullopt;
+    std::optional<at::Tensor> rotary_sin = rotary_sin_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(rotary_sin_.value())) : std::nullopt;
+    std::optional<at::Tensor> cache_batch_idx = cache_batch_idx_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(cache_batch_idx_.value())) : std::nullopt;
+    std::optional<at::Tensor> leftpad_k = leftpad_k_.has_value() ? std::optional<at::Tensor>(const_cast<at::at::Tensor &>(leftpad_k_.value())) : std::nullopt;
+    std::optional<at::Tensor> block_table = block_table_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(block_table_.value())) : std::nullopt;
+    std::optional<at::Tensor> alibi_slopes = alibi_slopes_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(alibi_slopes_.value())) : std::nullopt;
+    std::optional<at::Tensor> out = out_.has_value() ? std::optional<at::Tensor>(const_cast<at::Tensor &>(out_.value())) : std::nullopt;
+
+    // Convert double to float and int64_t to int.
+    float softmax_scale_float = static_cast<float>(softmax_scale);
+    float softcap_float = static_cast<float>(softcap);
+    int window_size_left_int = static_cast<int>(window_size_left);
+    int window_size_right_int = static_cast<int>(window_size_right);
+    int num_splits_int = static_cast<int>(num_splits);
+
+    return FLASH_NAMESPACE::mha_fwd_kvcache(const_cast<at::Tensor &>(q), kcache, vcache, k, v, seqlens_k, rotary_cos, rotary_sin, cache_batch_idx, leftpad_k, block_table, alibi_slopes, out, softmax_scale_float, is_causal, window_size_left_int, window_size_right_int, softcap_float, is_rotary_interleaved, num_splits_int);
 }

flash_attn/src/static_switch.h

@@ -87,33 +87,28 @@
     }                                        \
   }()
 
-// #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, ...)   \
-  [&] {                                  \
-    constexpr static int kHeadDim = 32;  \
-    return __VA_ARGS__();                \
+  [&] {                                    \
+    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__();                \
+    }                                      \
   }()

torch-ext/torch_binding.cpp

@@ -16,6 +16,9 @@
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("mha_fwd(Tensor! q, Tensor! k, Tensor! v, Tensor? out_, Tensor? alibi_slopes_, float p_dropout, float softmax_scale, bool is_causal, int window_size_left, int window_size_right, float softcap, bool return_softmax, Generator? gen_) -> Tensor[]");
   ops.impl("mha_fwd", torch::kCUDA, &mha_fwd);
+
+  ops.def("mha_varlen_fwd(Tensor! q, Tensor! k, Tensor! v, Tensor? out_, Tensor cu_seqlens_q, Tensor cu_seqlens_k, int max_seqlen_q, int max_seqlen_k, float p_dropout, float softmax_scale, bool is_causal, int window_size_left, int window_size_right, float softcap, bool return_softmax, Generator? gen_) -> Tensor[]");
+  ops.impl("mha_varlen_fwd", torch::kCUDA, &mha_varlen_fwd);
 }
 
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)