diff --git a/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py b/tests/v1/attention/test_aiter_mla_uniform_mtp_qo_len.py new file mode 100644 index 0000000000..6c5416d1f4 --- /dev/null +++ b/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 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 861c894c38..db9e8dbc88 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/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. diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 32154105a9..1115da630b 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/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