fix: revise bindings and wrapper typing

flash_attn/flash_api.cpp

@@ -1476,10 +1476,10 @@ mha_fwd_kvcache(at::Tensor &q,                 // batch_size x seqlen_q x num_he
 } // namespace FLASH_NAMESPACE
 
 // 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)
+std::vector<torch::Tensor>
+mha_fwd(const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
+        const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
+        const torch::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,
@@ -1511,17 +1511,17 @@ mha_fwd(const at::Tensor &q,                               // batch_size x seqle
     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(at::Tensor &q,  // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
-               const at::Tensor &k,  // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_>
-               const at::Tensor &v,  // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_>
-               const std::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
-               const at::Tensor &cu_seqlens_q,  // b+1
-               const at::Tensor &cu_seqlens_k,  // b+1
-               const std::optional<at::Tensor> &seqused_k_, // b. If given, only this many elements of each batch element's keys are used.
-               const std::optional<const at::Tensor> &leftpad_k_, // batch_size
-               const std::optional<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
-               const std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
+std::vector<torch::Tensor>
+mha_varlen_fwd(const torch::Tensor &q,                            // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+               const torch::Tensor &k,                            // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_>
+               const torch::Tensor &v,                            // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_>
+               const c10::optional<torch::Tensor> &out_,          // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+               const torch::Tensor &cu_seqlens_q,                 // b+1
+               const torch::Tensor &cu_seqlens_k,                 // b+1
+               const c10::optional<torch::Tensor> &seqused_k_,    // b. If given, only this many elements of each batch element's keys are used.
+               const c10::optional<torch::Tensor> &leftpad_k_,    // batch_size
+               const c10::optional<torch::Tensor> &block_table_,  // batch_size x max_num_blocks_per_seq
+               const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or b x num_heads
                const int64_t max_seqlen_q,
                const int64_t max_seqlen_k,
                const double p_dropout,
@@ -1532,7 +1532,8 @@ mha_varlen_fwd(at::Tensor &q,  // total_q x num_heads x head_size, total_q := \s
                const int64_t window_size_right,
                const double softcap,
                const bool return_softmax,
-               const std::optional<at::Generator> gen_) {
+               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;
@@ -1564,26 +1565,26 @@ mha_varlen_fwd(at::Tensor &q,  // total_q x num_heads x head_size, total_q := \s
         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
+std::vector<torch::Tensor>
+mha_bwd(const torch::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
+        const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size
+        const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size
+        const torch::Tensor &softmax_lse,                  // b x h x seqlen_q
+        const c10::optional<torch::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size
+        const c10::optional<torch::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const c10::optional<torch::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const c10::optional<torch::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) {
+        c10::optional<torch::Generator> gen_,
+        const c10::optional<torch::Tensor> &rng_state) {
 
     auto gen = gen_.value_or(at::cuda::detail::getDefaultCUDAGenerator());
     
@@ -1600,6 +1601,13 @@ mha_bwd(const at::Tensor &dout,                         // batch_size x seqlen_q
     int window_size_left_int = static_cast<int>(window_size_left);
     int window_size_right_int = static_cast<int>(window_size_right);
 
+    // TODO: avoid copying rng_state if possible
+    // Create a non-const copy of rng_state
+    std::optional<at::Tensor> rng_state_copy;
+    if (rng_state.has_value()) {
+        rng_state_copy = rng_state.value().clone();
+    }
+
     return FLASH_NAMESPACE::mha_bwd(
         const_cast<at::Tensor &>(dout), 
         q, k, v, out, softmax_lse, 
@@ -1608,23 +1616,23 @@ mha_bwd(const at::Tensor &dout,                         // batch_size x seqlen_q
         is_causal, 
         window_size_left_int, window_size_right_int, 
         softcap_float, deterministic, 
-        gen, rng_state);
+        gen, rng_state_copy);
 }
 
 
-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 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 at::Tensor &cu_seqlens_q,          // batch_size + 1
-               const at::Tensor &cu_seqlens_k,          // batch_size + 1
-               const std::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
+std::vector<torch::Tensor>
+mha_varlen_bwd(const torch::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
+               const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size
+               const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size
+               const torch::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size
+               const torch::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size
+               const torch::Tensor &softmax_lse,                  // b x h x seqlen_q
+               const c10::optional<torch::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size
+               const c10::optional<torch::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size
+               const c10::optional<torch::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size
+               const torch::Tensor &cu_seqlens_q,                 // batch_size + 1
+               const torch::Tensor &cu_seqlens_k,                 // batch_size + 1
+               const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or b x num_heads
                const int64_t max_seqlen_q,
                const int64_t max_seqlen_k,
                const double p_dropout,
@@ -1635,8 +1643,8 @@ mha_varlen_bwd(const at::Tensor &dout,                  // batch_size x seqlen_q
                const int64_t window_size_right,
                const double softcap,
                const bool deterministic,
-               std::optional<at::Generator> gen_,
-               std::optional<at::Tensor> &rng_state) {
+               c10::optional<torch::Generator> gen_,
+               const c10::optional<torch::Tensor> &rng_state) {
     
     auto gen = gen_.value_or(at::cuda::detail::getDefaultCUDAGenerator());
     
@@ -1655,6 +1663,14 @@ mha_varlen_bwd(const at::Tensor &dout,                  // batch_size x seqlen_q
     int window_size_left_int = static_cast<int>(window_size_left);
     int window_size_right_int = static_cast<int>(window_size_right);
 
+
+    // TODO: avoid copying rng_state if possible
+    // Create a non-const copy of rng_state
+    std::optional<at::Tensor> rng_state_copy;
+    if (rng_state.has_value()) {
+        rng_state_copy = rng_state.value().clone();
+    }
+    
     return FLASH_NAMESPACE::mha_varlen_bwd(
         const_cast<at::Tensor &>(dout), 
         q, k, v, out, softmax_lse, 
@@ -1666,31 +1682,30 @@ mha_varlen_bwd(const at::Tensor &dout,                  // batch_size x seqlen_q
         zero_tensors, is_causal, 
         window_size_left_int, window_size_right_int, 
         softcap_float, deterministic, 
-        gen, rng_state);
+        gen, rng_state_copy);
 }
 
-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
+std::vector<torch::Tensor>
+mha_fwd_kvcache(const torch::Tensor &q,                                // batch_size x seqlen_q x num_heads x head_size
+                const torch::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 torch::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<torch::Tensor> &block_table_,      // batch_size x max_num_blocks_per_seq
+                const c10::optional<torch::Tensor> &alibi_slopes_,     // num_heads or batch_size x num_heads
+                const c10::optional<torch::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
-                ) {
+                const int64_t num_splits) {
     
     // Prepare the optional arguments as const references where needed
     std::optional<const at::Tensor> k = k_.has_value() ? std::optional<const at::Tensor>(k_.value()) : std::nullopt;

torch-ext/torch_binding.cpp

@@ -17,7 +17,7 @@ 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.def("mha_varlen_fwd(Tensor! q, Tensor! k, Tensor! v, Tensor? out_, Tensor cu_seqlens_q, Tensor cu_seqlens_k, Tensor? seqused_k_, Tensor? leftpad_k_, Tensor? block_table_, Tensor? alibi_slopes_, int max_seqlen_q, int max_seqlen_k, float p_dropout, float softmax_scale, bool zero_tensors, 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);
 
   ops.def("mha_bwd(Tensor! dout, Tensor! q, Tensor! k, Tensor! v, Tensor! out, Tensor! softmax_lse, Tensor? dq_, Tensor? dk_, Tensor? dv_, Tensor? alibi_slopes_, float p_dropout, float softmax_scale, bool is_causal, int window_size_left, int window_size_right, float softcap, bool deterministic, Generator? gen_, Tensor? rng_state) -> Tensor[]");

torch-ext/torch_binding.h

@@ -2,11 +2,11 @@
 
 #include <torch/torch.h>
 
-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)
+std::vector<torch::Tensor>
+mha_fwd(const torch::Tensor &q,                               // batch_size x seqlen_q x num_heads x round_multiple(head_size, 8)
+        const torch::Tensor &k,                               // batch_size x seqlen_k x num_heads_k x round_multiple(head_size, 8)
+        const torch::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,
@@ -15,4 +15,101 @@ mha_fwd(const at::Tensor &q,         // batch_size x seqlen_q x num_heads x roun
         const int64_t window_size_right,
         const double softcap,
         const bool return_softmax,
-        const c10::optional<at::Generator> gen_);
+        const c10::optional<at::Generator> gen_);
+
+std::vector<torch::Tensor>
+mha_varlen_fwd(
+        const torch::Tensor &q,                            // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
+        const torch::Tensor &k,                            // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_>
+        const torch::Tensor &v,                            // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_>
+        const c10::optional<torch::Tensor> &out_,          // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
+        const torch::Tensor &cu_seqlens_q,                 // b+1
+        const torch::Tensor &cu_seqlens_k,                 // b+1
+        const c10::optional<torch::Tensor> &seqused_k_,    // b. If given, only this many elements of each batch element's keys are used.
+        const c10::optional<torch::Tensor> &leftpad_k_,    // batch_size
+        const c10::optional<torch::Tensor> &block_table_,  // batch_size x max_num_blocks_per_seq
+        const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or b x num_heads
+        const int64_t max_seqlen_q,
+        const int64_t max_seqlen_k,
+        const double p_dropout,
+        const double softmax_scale,
+        const bool zero_tensors,
+        const 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_);
+
+
+std::vector<torch::Tensor>
+mha_bwd(const torch::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
+        const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size
+        const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size
+        const torch::Tensor &softmax_lse,                  // b x h x seqlen_q
+        const c10::optional<torch::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size
+        const c10::optional<torch::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const c10::optional<torch::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const c10::optional<torch::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,
+        c10::optional<at::Generator> gen_,
+        const c10::optional<torch::Tensor> &rng_state);
+
+
+std::vector<torch::Tensor>
+mha_varlen_bwd(
+        const torch::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
+        const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size
+        const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size
+        const torch::Tensor &softmax_lse,                  // b x h x seqlen_q
+        const c10::optional<torch::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size
+        const c10::optional<torch::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const c10::optional<torch::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size
+        const torch::Tensor &cu_seqlens_q,                 // batch_size + 1
+        const torch::Tensor &cu_seqlens_k,                 // batch_size + 1
+        const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or b x num_heads
+        const int64_t max_seqlen_q,
+        const int64_t max_seqlen_k,
+        const double p_dropout,
+        const double softmax_scale,
+        const bool zero_tensors,
+        const bool is_causal,
+        const int64_t window_size_left,
+        const int64_t window_size_right,
+        const double softcap,
+        const bool deterministic,
+        c10::optional<at::Generator> gen_,
+        const c10::optional<torch::Tensor> &rng_state);
+
+std::vector<torch::Tensor>
+mha_fwd_kvcache(
+        const torch::Tensor &q,                                // batch_size x seqlen_q x num_heads x head_size
+        const torch::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 torch::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<torch::Tensor> &block_table_,      // batch_size x max_num_blocks_per_seq
+        const c10::optional<torch::Tensor> &alibi_slopes_,     // num_heads or batch_size x num_heads
+        const c10::optional<torch::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,
+        const int64_t num_splits);