fix: adjust sig types

flash_mla/flash_mla_api.cu

@@ -53,40 +53,26 @@ get_mla_metadata(
     return {tile_scheduler_metadata, num_splits};
 }
 
+// note doubles and longs are used in place of floats and ints
+// https://github.com/pytorch/pytorch/blob/338ed67a1e7aa98dd849f297533c5a71bea4b661/aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h#L211
 std::vector<at::Tensor>
 mha_fwd_kvcache_mla(
     at::Tensor &q,                               // batch_size x seqlen_q x num_heads x head_size
     const at::Tensor &kcache,                    // num_blocks x page_block_size x num_heads_k x head_size
-    
-    // TODO: fix for optional
-    // std::optional<const at::Tensor> &vcache_,    // num_blocks x page_block_size x num_heads_k x head_size_v
-    const at::Tensor &vcache_,                      // num_blocks x page_block_size x num_heads_k x head_size_v
-
+    const c10::optional<torch::Tensor> &vcache_,    // num_blocks x page_block_size x num_heads_k x head_size_v
     const int64_t head_size_v,
     const at::Tensor &seqlens_k,                    // batch_size
     const at::Tensor &block_table,                  // batch_size x max_num_blocks_per_seq
-    
-    // TODO: should be float
     const double softmax_scale,
-    const bool is_causal_,
+    bool is_causal,
     const at::Tensor &tile_scheduler_metadata,      // num_sm_parts x TileSchedulerMetaDataSize
     const at::Tensor &num_splits                    // batch_size + 1
-
-    // TODO: remove this once determined why build is adding this parameter
-    // const int64_t unknown_param
 ) {
     auto dprops = at::cuda::getCurrentDeviceProperties();
     bool is_sm90 = dprops->major == 9 && dprops->minor == 0;
     TORCH_CHECK(is_sm90);
 
-    // TODO: fix for mutable bool
-    bool is_causal = is_causal_;
-
-
-    // TODO: fix for optional
-    // at::Tensor vcache = vcache_.has_value() ? vcache_.value() : kcache;
-    at::Tensor vcache = vcache_;
-
+    at::Tensor vcache = vcache_.has_value() ? vcache_.value() : kcache;
     auto q_dtype = q.dtype();
     TORCH_CHECK(kcache.dtype() == q_dtype, "query and key must have the same dtype");
 

torch-ext/torch_binding.cpp

@@ -8,7 +8,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.impl("get_mla_metadata", torch::kCUDA, &get_mla_metadata);
 
   // TOOD: remove last unknown_param when resolved
-  ops.def("mha_fwd_kvcache_mla(Tensor! q, Tensor! kcache, Tensor! vcache_, int head_size_v, Tensor! seqlens_k, Tensor! block_table, float softmax_scale, bool is_causal_, Tensor! tile_scheduler_metadata, Tensor! num_splits) -> Tensor[]");
+  ops.def("mha_fwd_kvcache_mla(Tensor! q, Tensor! kcache, Tensor? vcache_, int head_size_v, Tensor! seqlens_k, Tensor! block_table, float softmax_scale, bool is_causal_, Tensor! tile_scheduler_metadata, Tensor! num_splits) -> Tensor[]");
   ops.impl("mha_fwd_kvcache_mla", torch::kCUDA, &mha_fwd_kvcache_mla);
 }
 

torch-ext/torch_binding.h

@@ -13,21 +13,13 @@ std::vector<torch::Tensor>
 mha_fwd_kvcache_mla(
     torch::Tensor &q,
     const torch::Tensor &kcache,
-
-    // TODO: fix for optional
-    // std::optional<torch::Tensor> &vcache_,
-    
-    const torch::Tensor &vcache_,
+    const c10::optional<torch::Tensor> &vcache_,
     const int64_t head_size_v,
     const torch::Tensor &seqlens_k,
     const torch::Tensor &block_table,
-
     // TODO:should be float
-    const double softmax_scale,
-
-    // TODO: fix for mutable bool
-    const bool is_causal_,
-
+    const torch::kFloat softmax_scale,
+    bool is_causal,
     const torch::Tensor &tile_scheduler_metadata,
     const torch::Tensor &num_splits
 );