diff --git a/tests/v1/attention/test_expand_page_indices.py b/tests/v1/attention/test_expand_page_indices.py new file mode 100644 index 0000000000..73163fe13d --- /dev/null +++ b/tests/v1/attention/test_expand_page_indices.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Equivalence tests for the AITER MLA page-index expansion kernel. + +`_expand_page_indices_kernel` turns a block table into the per-token flat page +indices the AITER MLA decode kernel consumes (it always runs with page_size=1 +internally, so `kv_buffer` is flattened and each token needs its own index). + +The kernel was parallelised over token chunks as well as requests; at +concurrency 1 the previous `(num_reqs,)` grid put a single workgroup in a serial +loop over the whole sequence. These tests pin the result to a straightforward +PyTorch expansion so the parallelisation cannot change what is produced. + +Both block-table strides are passed by the caller, so a non-contiguous view is +covered explicitly rather than assumed away. +""" + +import pytest +import torch + +from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( + _expand_page_indices_kernel, +) + + +def _reference(block_table, cu_num_tokens, kernel_block_size, out_len, device): + """Per-token flat index, written the obvious way.""" + out = torch.full((out_len,), -1, dtype=torch.int64, device=device) + num_reqs = cu_num_tokens.numel() - 1 + for r in range(num_reqs): + start = int(cu_num_tokens[r]) + end = int(cu_num_tokens[r + 1]) + n = end - start + if n <= 0: + continue + tok = torch.arange(n, device=device) + blk = tok // kernel_block_size + off = tok % kernel_block_size + out[start:end] = block_table[r, blk].to(torch.int64) * kernel_block_size + off + return out + + +def _run(block_table, cu_num_tokens, kernel_block_size, out_len, device): + out = torch.full((out_len,), -1, dtype=torch.int64, device=device) + num_reqs = cu_num_tokens.numel() - 1 + max_tokens_per_req = block_table.shape[1] * kernel_block_size + num_chunks = max(1, -(-max_tokens_per_req // 1024)) + _expand_page_indices_kernel[(num_reqs, num_chunks)]( + out, + block_table, + block_table.stride(0), + block_table.stride(1), + cu_num_tokens, + KERNEL_BLOCK_SIZE=kernel_block_size, + BLOCK_SIZE=1024, + ) + return out + + +def _make(seq_lens, kernel_block_size, device, noncontig=False): + num_reqs = len(seq_lens) + max_blocks = ( + max(1, max((s + kernel_block_size - 1) // kernel_block_size for s in seq_lens)) + + 2 + ) + width = max_blocks * 2 if noncontig else max_blocks + torch.manual_seed(0) + bt = torch.randint(0, 10_000, (num_reqs, width), dtype=torch.int64, device=device) + if noncontig: + bt = bt[:, ::2] + assert not bt.is_contiguous() + cu = torch.zeros(num_reqs + 1, dtype=torch.int64, device=device) + cu[1:] = torch.tensor(seq_lens, dtype=torch.int64, device=device).cumsum(0) + return bt, cu, int(cu[-1]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize( + "seq_lens", [[1], [128], [4096], [1, 1], [4096, 7], [37, 4096, 1], [1024] * 8] +) +@pytest.mark.parametrize("kernel_block_size", [1, 16, 64]) +def test_matches_reference(seq_lens, kernel_block_size): + bt, cu, total = _make(seq_lens, kernel_block_size, "cuda") + got = _run(bt, cu, kernel_block_size, total, "cuda") + want = _reference(bt, cu, kernel_block_size, total, "cuda") + torch.testing.assert_close(got, want, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +@pytest.mark.parametrize("kernel_block_size", [1, 16]) +def test_non_contiguous_block_table(kernel_block_size): + """Column stride != 1 must be honoured, not assumed away.""" + bt, cu, total = _make([4096, 37], kernel_block_size, "cuda", noncontig=True) + got = _run(bt, cu, kernel_block_size, total, "cuda") + want = _reference(bt, cu, kernel_block_size, total, "cuda") + torch.testing.assert_close(got, want, atol=0, rtol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_ragged_batch_leaves_other_rows_untouched(): + """A short request must not have its tail written by a longer one's chunks.""" + kbs = 16 + bt, cu, total = _make([4096, 1, 512], kbs, "cuda") + out = torch.full((total + 8,), -999, dtype=torch.int64, device="cuda") + num_chunks = max(1, -(-(bt.shape[1] * kbs) // 1024)) + _expand_page_indices_kernel[(3, num_chunks)]( + out, + bt, + bt.stride(0), + bt.stride(1), + cu, + KERNEL_BLOCK_SIZE=kbs, + BLOCK_SIZE=1024, + ) + want = _reference(bt, cu, kbs, total, "cuda") + torch.testing.assert_close(out[:total], want, atol=0, rtol=0) + # The slack past the last request must be untouched. + assert torch.all(out[total:] == -999) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_zero_length_request_writes_nothing(): + """A zero-token request in a padded batch must produce no stores.""" + kbs = 16 + bt, cu, total = _make([256, 0, 128], kbs, "cuda") + got = _run(bt, cu, kbs, total, "cuda") + want = _reference(bt, cu, kbs, total, "cuda") + torch.testing.assert_close(got, want, atol=0, rtol=0) + assert int(cu[2] - cu[1]) == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_over_provisioned_chunks_do_not_change_result(): + """Extra chunks must be inert. + + Each program covers `chunk_idx * BLOCK_SIZE + [0, BLOCK_SIZE)`, so the grid + has to span the request: 8192 tokens need at least ceil(8192/1024) == 8 + chunks, and fewer would simply leave the tail unwritten. The caller sizes + the grid from `block_table.shape[1] * kernel_block_size`, an upper bound + over the batch, so in a ragged batch most requests get more chunks than + they need. That over-provisioning is the case worth pinning: the surplus + programs must return early and change nothing. + """ + kbs = 1 + bt, cu, total = _make([8192], kbs, "cuda") + ref = _reference(bt, cu, kbs, total, "cuda") + required = -(-total // 1024) + for num_chunks in (required, required * 2, required * 8): + out = torch.full((total,), -1, dtype=torch.int64, device="cuda") + _expand_page_indices_kernel[(1, num_chunks)]( + out, + bt, + bt.stride(0), + bt.stride(1), + cu, + KERNEL_BLOCK_SIZE=kbs, + BLOCK_SIZE=1024, + ) + torch.testing.assert_close(out, ref, atol=0, rtol=0) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 309bfd609c..da141acf68 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -56,12 +56,21 @@ class _ExpandPageIndicesKernel: self, page_indices, block_table_tensor, - stride, + stride_0, + stride_1, paged_kv_indptr, *, KERNEL_BLOCK_SIZE, BLOCK_SIZE, ): + # Both block-table strides are passed explicitly: the kernel no longer + # assumes the column stride is 1. This stand-in indexes the tensor + # directly so torch applies the strides for it, but the values are + # checked against the tensor actually handed in -- a caller passing the + # wrong stride is exactly the bug the parameter exists to prevent, and + # silently ignoring them here would hide it. + assert stride_0 == block_table_tensor.stride(0) + assert stride_1 == block_table_tensor.stride(1) self.kernel_block_size = KERNEL_BLOCK_SIZE for req_idx in range(self.grid[0]): out_start = int(paged_kv_indptr[req_idx].item()) @@ -1052,7 +1061,14 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch): metadata.paged_kv_indices[: expected_indices.numel()], expected_indices, ) - assert expand_kernel.grid == (seq_lens.numel(),) + # The launch now carries a second grid dimension over token chunks, so + # the expansion is spread across the context instead of one workgroup + # walking it serially. The width comes from the block-table upper bound + # on tokens per request, which is available host-side and so costs no + # device sync; derived here rather than hard-coded so the assertion + # still describes the launch if the chunk size changes. + expected_chunks = max(1, -(-block_table.shape[1] * kernel_block_size // 1024)) + assert expand_kernel.grid == (seq_lens.numel(), expected_chunks) assert expand_kernel.kernel_block_size == kernel_block_size diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index ab6c11d41a..c833a56a83 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -1331,10 +1331,17 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # to the original _copy_page_indices_kernel). # When kernel_block_size=K>1, block_table entry b covering K tokens # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). - _expand_page_indices_kernel[(num_reqs,)]( + # Chunk count comes from the block table width, an upper bound on + # tokens per request that is available host-side, so this adds no + # device synchronisation. Programs whose chunk lies past + # num_tokens mask out entirely. + max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size + num_chunks = max(1, cdiv(max_tokens_per_req, 1024)) + _expand_page_indices_kernel[(num_reqs, num_chunks)]( self.paged_kv_indices, block_table_tensor, block_table_tensor.stride(0), + block_table_tensor.stride(1), paged_kv_indptr, KERNEL_BLOCK_SIZE=self.kernel_block_size, BLOCK_SIZE=1024, @@ -1573,7 +1580,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): def _expand_page_indices_kernel( page_indices, block_table, - block_table_stride, + block_table_stride_0, + block_table_stride_1, cu_num_tokens, KERNEL_BLOCK_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, @@ -1591,32 +1599,44 @@ def _expand_page_indices_kernel( When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). """ + # One program per (request, token-chunk). Parallelising over requests + # alone degenerates at low concurrency: with num_reqs == 1 a single + # workgroup walked the whole sequence in a serial loop. req_idx = tl.program_id(0) - row_ptr = block_table + req_idx * block_table_stride + chunk_idx = tl.program_id(1) + row_ptr = block_table + req_idx * block_table_stride_0 start_idx = tl.load(cu_num_tokens + req_idx) num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx - offset = tl.arange(0, BLOCK_SIZE) - for i in tl.range(0, num_tokens, BLOCK_SIZE): - token_offsets = i + offset - mask = token_offsets < num_tokens + # The grid is sized from the block table width, an upper bound over all + # requests, so a ragged batch launches chunks past a short request's end. + # Returning here keeps those programs from issuing masked-out loads and + # stores at all. + chunk_start = chunk_idx * BLOCK_SIZE + if chunk_start >= num_tokens: + return - # Which block in the block table does this token belong to? - block_idx = token_offsets // KERNEL_BLOCK_SIZE - # Offset within that block - offset_in_block = token_offsets % KERNEL_BLOCK_SIZE + token_offsets = chunk_start + tl.arange(0, BLOCK_SIZE) + mask = token_offsets < num_tokens - # Load the block ID from the block table - block_ids = tl.load(row_ptr + block_idx, mask=mask) + # Which block in the block table does this token belong to? + block_idx = token_offsets // KERNEL_BLOCK_SIZE + # Offset within that block + offset_in_block = token_offsets % KERNEL_BLOCK_SIZE - # Compute flat index in the flattened kv_buffer - flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block + # Load the block ID from the block table + # Both strides are taken from the caller: the block table is a view owned + # elsewhere, so a unit column stride must not be assumed. + block_ids = tl.load(row_ptr + block_idx * block_table_stride_1, mask=mask) - tl.store( - page_indices + start_idx + token_offsets, - flat_indices, - mask=mask, - ) + # Compute flat index in the flattened kv_buffer + flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block + + tl.store( + page_indices + start_idx + token_offsets, + flat_indices, + mask=mask, + ) class AiterMLAHelper: