[ROCm][MLA] Cut ~21x host dispatch from the MLA metadata build (#58381)

Signed-off-by: lifulu <fululi12@amd.com>
Signed-off-by: lifulu <lifulu@users.noreply.github.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: lifulu <lifulu@users.noreply.github.com>

tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Differential test for the AITER MLA uniform-MTP padding decision.
+
+`_uniform_padded_mtp_qo_len` decides whether a full-CUDA-graph decode batch may
+synthesize dummy query rows. It used to reach that decision through eight tensor
+dispatches -- sum/item, boolean-mask index, all, eq, nonzero, item, any -- on a
+CPU tensor holding one entry per request. It now reads the tensor out once and
+decides on Python ints.
+
+That rewrite also **reordered** the checks: the `num_decode_tokens <= sum(...)`
+early-out moved below the divisibility and `uniform_qo_len <= 1` guards. Every
+one of those paths returns 0, so the result should be unchanged -- but "should
+be" is exactly the sort of claim worth checking exhaustively, so both
+implementations are carried here and compared over the whole small input space.
+
+Pure host arithmetic; no GPU required.
+"""
+
+import itertools
+
+import pytest
+import torch
+
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import AiterMLAMetadataBuilder
+
+
+def _reference(qo_len: torch.Tensor, max_qo_len: int, num_decode_tokens: int) -> int:
+    """The pre-#58381 implementation, verbatim."""
+    num_reqs = qo_len.numel()
+    if num_reqs == 0 or num_decode_tokens <= 0:
+        return 0
+    if num_decode_tokens <= int(qo_len.sum().item()):
+        return 0
+    if num_decode_tokens % num_reqs != 0:
+        return 0
+
+    uniform_qo_len = num_decode_tokens // num_reqs
+    if uniform_qo_len <= 1:
+        return 0
+
+    positive_qo_len = qo_len[qo_len > 0]
+    if positive_qo_len.numel() == qo_len.numel():
+        return 0
+    if positive_qo_len.numel() > 0:
+        if max_qo_len != uniform_qo_len:
+            return 0
+        if not torch.all(positive_qo_len == uniform_qo_len):
+            return 0
+
+    zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten()
+    if zero_positions.numel() > 0:
+        first_zero = int(zero_positions[0].item())
+        if torch.any(qo_len[first_zero:] > 0):
+            return 0
+
+    return uniform_qo_len
+
+
+_ACTUAL = AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len
+
+
+@pytest.mark.parametrize("num_reqs", [1, 2, 3, 4])
+def test_matches_reference_exhaustively(num_reqs):
+    """Every qo_len vector x decode-token count x max_qo_len in a small space."""
+    checked = 0
+    for lens in itertools.product(range(0, 4), repeat=num_reqs):
+        qo = torch.tensor(lens, dtype=torch.int32)
+        for num_decode_tokens in range(0, 13):
+            for max_qo_len in range(0, 5):
+                got = _ACTUAL(qo, max_qo_len, num_decode_tokens)
+                want = _reference(qo, max_qo_len, num_decode_tokens)
+                assert got == want, (
+                    f"lens={lens} num_decode_tokens={num_decode_tokens} "
+                    f"max_qo_len={max_qo_len}: got {got}, reference {want}"
+                )
+                checked += 1
+    assert checked > 0
+
+
+def test_empty_batch():
+    assert _ACTUAL(torch.zeros(0, dtype=torch.int32), 4, 8) == 0
+    assert _reference(torch.zeros(0, dtype=torch.int32), 4, 8) == 0
+
+
+def test_the_padded_mtp_case_that_should_succeed():
+    """A live prefix at the uniform width, then zero-padded dummies."""
+    qo = torch.tensor([4, 4, 0, 0], dtype=torch.int32)
+    # 4 requests x 4 tokens = 16 captured, 8 live -> dummies are synthesizable.
+    assert _ACTUAL(qo, 4, 16) == 4
+    assert _reference(qo, 4, 16) == 4
+
+
+def test_positive_after_zero_is_rejected():
+    """A live request after a padding hole breaks the contiguous-prefix layout."""
+    qo = torch.tensor([4, 0, 4, 0], dtype=torch.int32)
+    assert _ACTUAL(qo, 4, 16) == 0
+    assert _reference(qo, 4, 16) == 0
+
+
+def test_ragged_live_lengths_are_rejected():
+    qo = torch.tensor([4, 2, 0, 0], dtype=torch.int32)
+    assert _ACTUAL(qo, 4, 16) == 0
+    assert _reference(qo, 4, 16) == 0

vllm/model_executor/layers/attention/mla_attention.py

@@ -2540,16 +2540,26 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                     "Non-causal multi-token MLA requires an explicitly supported "
                     "attention group."
                 )
-            query_lens = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
-            num_active_reqs = int(torch.count_nonzero(query_lens > 0))
-            uniform_active_queries = num_active_reqs > 0 and bool(
-                torch.all(query_lens[:num_active_reqs] == query_lens[0])
+            # query_start_loc_cpu is host-side and holds one entry per
+            # request, so this check costs more in tensor dispatch than in
+            # the work it does. Read it out once and validate on Python ints:
+            # the block runs on every build, between decode steps.
+            qsl = query_start_loc_cpu.tolist()
+            query_lens = [qsl[i + 1] - qsl[i] for i in range(len(qsl) - 1)]
+            num_active_reqs = 0
+            for qlen in query_lens:
+                if qlen > 0:
+                    num_active_reqs += 1
+            uniform_active_queries = num_active_reqs > 0 and all(
+                qlen == query_lens[0] for qlen in query_lens[:num_active_reqs]
+            )
+            trailing_graph_padding = all(
+                qlen == 0 for qlen in query_lens[num_active_reqs:]
             )
-            trailing_graph_padding = bool(torch.all(query_lens[num_active_reqs:] == 0))
             if not (uniform_active_queries and trailing_graph_padding):
                 raise ValueError(
                     "Non-causal MLA requires a uniform query block; got query "
-                    f"lengths {query_lens.tolist()}."
+                    f"lengths {query_lens}."
                 )
             # Use exact GPU sequence lengths instead of the prefill path's CPU
             # context-length upper bounds.

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

@@ -563,8 +563,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         # query_start_loc flat for dummy requests. Only synthesize dummy rows
         # when every padded request maps to the same qlen and the q buffer has
         # exactly that many rows.
-        if num_decode_tokens <= int(qo_len.sum().item()):
-            return 0
         if num_decode_tokens % num_reqs != 0:
             return 0
 
@@ -572,19 +570,34 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         if uniform_qo_len <= 1:
             return 0
 
-        positive_qo_len = qo_len[qo_len > 0]
-        if positive_qo_len.numel() == qo_len.numel():
+        # qo_len is derived from query_start_loc_cpu, so it is already
+        # host-side and tiny (one entry per request). Reading it out once and
+        # deciding on Python ints avoids eight separate tensor dispatches --
+        # sum/item, boolean-mask index, all, eq, nonzero, item, any -- on the
+        # metadata-build path, which runs between decode steps while the GPUs
+        # are idle.
+        lens = qo_len.tolist()
+        if num_decode_tokens <= sum(lens):
             return 0
-        if positive_qo_len.numel() > 0:
+
+        first_zero = -1
+        num_positive = 0
+        for i, qlen in enumerate(lens):
+            if qlen > 0:
+                num_positive += 1
+                # A positive entry after a zero breaks the contiguous live
+                # prefix that the padded layout requires.
+                if first_zero >= 0:
+                    return 0
+            elif first_zero < 0:
+                first_zero = i
+
+        if num_positive == num_reqs:
+            return 0
+        if num_positive > 0:
             if max_qo_len != uniform_qo_len:
                 return 0
-            if not torch.all(positive_qo_len == uniform_qo_len):
-                return 0
-
-        zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten()
-        if zero_positions.numel() > 0:
-            first_zero = int(zero_positions[0].item())
-            if torch.any(qo_len[first_zero:] > 0):
+            if any(qlen != uniform_qo_len for qlen in lens if qlen > 0):
                 return 0
 
         return uniform_qo_len