get_mla_decoding_metadata_dense_fp8和社区保持一致

csrc/extension/flash_api.h

@@ -265,9 +265,7 @@ std::vector<at::Tensor>
 get_mla_decoding_metadata_dense_fp8(
     at::Tensor &seqlens_k,
     const int num_heads_per_head_k,
-    const int num_heads_k,
-    const std::optional<int> h_q
-) {
+    const int num_heads_k) {
     // This should match the logic in the MLA kernel.
     int block_size_m = 16;
     static constexpr int block_size_n = 64;

flash_mla/flash_mla_interface.py

@@ -213,9 +213,8 @@ def flash_mla_sparse_fwd(
 def get_mla_decoding_metadata_dense_fp8(
     cache_seqlens: torch.Tensor,
     num_heads_per_head_k: int,
-    num_heads_k: int,
-    num_heads_q : int = 16,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+    num_heads_k: int
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Arguments:
         cache_seqlens: (batch_size), dtype torch.int32.
@@ -226,7 +225,7 @@ def get_mla_decoding_metadata_dense_fp8(
         tile_scheduler_metadata: (num_sm_parts, TileSchedulerMetaDataSize), dtype torch.int32.
         num_splits: (batch_size + 1), dtype torch.int32.
     """
-    return flash_mla_cuda.get_mla_decoding_metadata_dense_fp8(cache_seqlens, num_heads_per_head_k, num_heads_k, num_heads_q)
+    return flash_mla_cuda.get_mla_decoding_metadata_dense_fp8(cache_seqlens, num_heads_per_head_k, num_heads_k)
 
 
 

tests/test_flash_mla_fp8.py

@@ -79,7 +79,7 @@ def test_flash_mla_fp8_e5m2(b, s_q, mean_sk, h_q, h_kv, d, dv, causal, varlen, i
     blocked_v = blocked_k[..., :dv]
 
     tile_scheduler_metadata, num_splits = get_mla_decoding_metadata_dense_fp8(
-        cache_seqlens, s_q * h_q // h_kv, h_kv, h_q
+        cache_seqlens, s_q * h_q // h_kv, h_kv
     )
 
     # print("q:", q.shape, q.dtype, q)

tests/test_flash_mla_qkvfp8.py

@@ -88,7 +88,7 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, causal, varlen, is_prof=Fa
     blocked_v = blocked_k[..., :dv]
 
     tile_scheduler_metadata, num_splits = get_mla_decoding_metadata_dense_fp8(
-        cache_seqlens, s_q * h_q // h_kv, h_kv, h_q
+        cache_seqlens, s_q * h_q // h_kv, h_kv
     )
 
     init_dtype = q.dtype