[ROCm][Perf] Parallelise AITER MLA page-index expansion over token chunks (#57978)

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

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)

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
 
 

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: