fix: remove unused trailing param

flash_mla/flash_mla_api.cu

@@ -70,10 +70,10 @@ mha_fwd_kvcache_mla(
     const double softmax_scale,
     const bool is_causal_,
     const at::Tensor &tile_scheduler_metadata,      // num_sm_parts x TileSchedulerMetaDataSize
-    const at::Tensor &num_splits,                   // batch_size + 1
+    const at::Tensor &num_splits                    // batch_size + 1
 
     // TODO: remove this once determined why build is adding this parameter
-    const int64_t unknown_param
+    // const int64_t unknown_param
 ) {
     auto dprops = at::cuda::getCurrentDeviceProperties();
     bool is_sm90 = dprops->major == 9 && dprops->minor == 0;

torch-ext/flash_mla/__init__.py

@@ -19,8 +19,6 @@ def mha_fwd_kvcache_mla(
     tile_scheduler_metadata: torch.Tensor,
     num_splits: torch.Tensor,
 ) -> torch.Tensor:
-    # TODO: remove when resolved
-    unknown_param = 0
     return ops.mha_fwd_kvcache_mla(
         q,
         kcache,
@@ -31,6 +29,5 @@ def mha_fwd_kvcache_mla(
         softmax_scale,
         is_causal_,
         tile_scheduler_metadata,
-        num_splits,
-        unknown_param,
+        num_splits
     )

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, int unknown_param) -> 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

@@ -29,8 +29,5 @@ mha_fwd_kvcache_mla(
     const bool is_causal_,
 
     const torch::Tensor &tile_scheduler_metadata,
-    const torch::Tensor &num_splits,
-
-    // TODO: remove when resolved
-    const int64_t unknown_param = 0
+    const torch::Tensor &num_splits
 );