[ROCm][Perf] Drop two dead per-step ops from the AITER sparse MLA metadata build (#58354)

Signed-off-by: ZhengGong-amd <zgong@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/kernels/attention/test_rocm_aiter_mla_sparse_metadata_sync.py

@@ -223,7 +223,7 @@ def _patch_build_deps(monkeypatch, events=None):
     on CPU."""
 
     def fake_generate_sparse_seqlen_triton(
-        query_lens, seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len
+        seq_lens, cu_query_lens, topk_token, num_tokens, max_query_len
     ):
         return torch.zeros(num_tokens, dtype=torch.int32, device="cpu")
 
@@ -276,6 +276,25 @@ def test_build_populates_mixed_split_fields(monkeypatch):
     assert md.prefill is None
 
 
+def test_build_fills_paged_kv_indptr_tail(monkeypatch):
+    """Rows past num_tokens must read as empty: the tail repeats the total."""
+    builder = _make_builder()
+    _patch_build_deps(monkeypatch)
+    monkeypatch.setattr(
+        sparse_mod,
+        "generate_sparse_seqlen_triton",
+        lambda *args, **kwargs: torch.tensor([3, 5], dtype=torch.int32),
+    )
+    builder.paged_kv_indptr[1:].fill_(-1)
+
+    md = builder.build(
+        common_prefix_len=0, common_attn_metadata=_make_common_metadata()
+    )
+
+    assert md.paged_kv_indptr.tolist() == [0, 3, 8]
+    assert builder.paged_kv_indptr.tolist() == [0, 3, 8, 8, 8, 8, 8, 8, 8]
+
+
 def test_sink_build_skips_persistent_metadata(monkeypatch):
     builder = _make_builder()
     builder._use_persistent_metadata = False
@@ -385,10 +404,10 @@ def test_dense_prefill_builds_sparse_metadata_for_decode_subset(monkeypatch):
     sparse_call: dict[str, torch.Tensor] = {}
 
     def fake_generate_sparse_seqlen(
-        query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q
+        sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q
     ):
         sparse_call.update(
-            query_lens=query_lens.clone(),
+            query_lens=torch.diff(sparse_query_start_loc),
             seq_lens=sparse_seq_lens.clone(),
             query_start_loc=sparse_query_start_loc.clone(),
             num_tokens=num_tokens,
@@ -471,10 +490,10 @@ def test_long_prefill_shares_full_sparse_metadata_with_decode(monkeypatch):
     sparse_call: dict[str, torch.Tensor] = {}
 
     def fake_generate_sparse_seqlen(
-        query_lens, sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q
+        sparse_seq_lens, sparse_query_start_loc, topk, num_tokens, max_q
     ):
         sparse_call.update(
-            query_lens=query_lens.clone(),
+            query_lens=torch.diff(sparse_query_start_loc),
             seq_lens=sparse_seq_lens.clone(),
             query_start_loc=sparse_query_start_loc.clone(),
             num_tokens=num_tokens,
@@ -699,7 +718,7 @@ def test_sparse_metadata_workload_length_dtype_matrix(
     monkeypatch.setattr(
         sparse_mod,
         "generate_sparse_seqlen_triton",
-        lambda *args, **kwargs: torch.ones(args[4], dtype=torch.int32),
+        lambda *args, **kwargs: torch.ones(args[3], dtype=torch.int32),
     )
     get_persistent_metadata = Mock()
     fake_aiter = _FakeAiter("aiter")

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

@@ -322,16 +322,15 @@ def generate_sparse_seqlen_kernel(
 
 
 def generate_sparse_seqlen_triton(
-    query_lens: torch.Tensor,
     seq_lens: torch.Tensor,
     cu_query_lens: torch.Tensor,
     topk_token: int,
     num_tokens: int,
     max_query_len: int,
 ):
-    num_seqs = query_lens.size(0)
+    num_seqs = cu_query_lens.size(0) - 1
     # zero initialize the tensor to make sure invalid positions will be zero
-    out = torch.zeros([num_tokens], dtype=torch.int32, device=query_lens.device)
+    out = torch.zeros([num_tokens], dtype=torch.int32, device=cu_query_lens.device)
     block_size = 64
     num_block_per_row = triton.cdiv(max_query_len, block_size)
     grid = (
@@ -627,10 +626,8 @@ class ROCMAiterMLASparseMetadataBuilder(
             np_to_pinned_tensor(req_id_per_token), non_blocking=True
         )
         query_start_loc = common_attn_metadata.query_start_loc[: num_reqs + 1]
-        query_lens = query_start_loc[1:] - query_start_loc[:-1]
         seq_lens = common_attn_metadata.seq_lens[:num_reqs]
         sparse_seqlen = generate_sparse_seqlen_triton(
-            query_lens,
             seq_lens,
             query_start_loc,
             self.topk_tokens,
@@ -639,7 +636,9 @@ class ROCMAiterMLASparseMetadataBuilder(
         )
 
         torch.cumsum(sparse_seqlen, dim=0, out=self.paged_kv_indptr[1 : num_tokens + 1])
-        self.paged_kv_indptr[num_tokens + 1 :].fill_(self.paged_kv_indptr[num_tokens])
+        # fill_() with a value that aliases the destination clones it first.
+        indptr_tail = self.paged_kv_indptr[num_tokens + 1 :]
+        indptr_tail.copy_(self.paged_kv_indptr[num_tokens].expand_as(indptr_tail))
 
         req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
         qo_indptr = self.qo_indptr[: num_tokens + 1]