[Bugfix] add missing function params to rocm_aiter_mla.py

Looks like #17864 had an outdated
branch. So its [merge commit][1] caused `qo_indptr` and `max_seqlen_qo` to go
into the function signature of `aiter_mla_decode_fwd()` where they're not used
and into the body of `mla_decode_fwd_impl()` where they aren't defined.

This PR fixes the discrepancies and call-sites.

Signed-off-by: David Xia <[email protected]>

[1]: 9f64e93#diff-88fd09f50e8cfc77678ade87483ab9a89ce58904203578f8816882763bd577c2

Author: davidxia

vllm/attention/backends/rocm_aiter_mla.py

@@ -425,11 +425,14 @@ def _forward_decode(
 
         kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
 
-        aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,
+        aiter_mla_decode_fwd(q,
+                             kv_buffer,
+                             o,
                              attn_metadata.qo_indptr,
                              attn_metadata.max_query_len,
                              attn_metadata.paged_kv_indptr,
                              attn_metadata.paged_kv_indices,
-                             attn_metadata.paged_kv_last_page_lens)
+                             attn_metadata.paged_kv_last_page_lens,
+                             sm_scale=self.scale)
 
         return self._v_up_proj(o)

vllm/attention/ops/rocm_aiter_mla.py

@@ -28,19 +28,21 @@ def aiter_mla_decode_fwd(
     q: torch.Tensor,
     kv_buffer: torch.Tensor,
     o: torch.Tensor,
-    sm_scale: float,
     qo_indptr: torch.Tensor,
     max_seqlen_qo: int,
     kv_indptr: Optional[torch.Tensor] = None,
     kv_indices: Optional[torch.Tensor] = None,
     kv_last_page_lens: Optional[torch.Tensor] = None,
+    sm_scale: float = 1.0,
     logit_cap: float = 0.0,
 ):
 
     torch.ops.vllm.rocm_aiter_mla_decode_fwd(q,
                                              kv_buffer.view(
                                                  -1, 1, 1, q.shape[-1]),
                                              o,
+                                             qo_indptr,
+                                             max_seqlen_qo,
                                              kv_indptr,
                                              kv_indices,
                                              kv_last_page_lens,
@@ -52,6 +54,8 @@ def mla_decode_fwd_impl(
     q: torch.Tensor,
     kv_buffer: torch.Tensor,
     o: torch.Tensor,
+    qo_indptr: torch.Tensor,
+    max_seqlen_qo: int,
     kv_indptr: Optional[torch.Tensor] = None,
     kv_indices: Optional[torch.Tensor] = None,
     kv_last_page_lens: Optional[torch.Tensor] = None,
@@ -76,6 +80,8 @@ def mla_decode_fwd_fake(
     q: torch.Tensor,
     kv_buffer: torch.Tensor,
     o: torch.Tensor,
+    qo_indptr: torch.Tensor,
+    max_seqlen_qo: int,
     kv_indptr: Optional[torch.Tensor] = None,
     kv_indices: Optional[torch.Tensor] = None,
     kv_last_page_lens: Optional[torch.Tensor] = None,

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -186,9 +186,14 @@ def _forward_decode(
 
         kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
 
-        aiter_mla_decode_fwd(q, kv_buffer, o, self.scale,
+        aiter_mla_decode_fwd(q,
+                             kv_buffer,
+                             o,
+                             attn_metadata.qo_indptr,
+                             attn_metadata.max_query_len,
                              attn_metadata.decode.paged_kv_indptr,
                              attn_metadata.decode.paged_kv_indices,
-                             attn_metadata.decode.paged_kv_last_page_len)
+                             attn_metadata.decode.paged_kv_last_page_len,
+                             sm_scale=self.scale)
 
         return self._v_up_proj(o)