[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>
fululi12 committed
Oct 1, 2026 at 10:00 UTC
26bfdfb6bc5ef85e2128b9ddcf27c369356edd3b
3 files changed
+221
-26
tests/v1/attention/test_expand_page_indices.py
new
+159
@@ -0,0 +1,159 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Equivalence tests for the AITER MLA page-index expansion kernel.
4
+
5
+`_expand_page_indices_kernel` turns a block table into the per-token flat page
6
+indices the AITER MLA decode kernel consumes (it always runs with page_size=1
7
+internally, so `kv_buffer` is flattened and each token needs its own index).
8
+
9
+The kernel was parallelised over token chunks as well as requests; at
10
+concurrency 1 the previous `(num_reqs,)` grid put a single workgroup in a serial
11
+loop over the whole sequence. These tests pin the result to a straightforward
12
+PyTorch expansion so the parallelisation cannot change what is produced.
13
+
14
+Both block-table strides are passed by the caller, so a non-contiguous view is
15
+covered explicitly rather than assumed away.
16
+"""
17
+
18
+import pytest
19
+import torch
20
+
21
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
22
+ _expand_page_indices_kernel,
23
+)
24
+
25
+
26
+def _reference(block_table, cu_num_tokens, kernel_block_size, out_len, device):
27
+ """Per-token flat index, written the obvious way."""
28
+ out = torch.full((out_len,), -1, dtype=torch.int64, device=device)
29
+ num_reqs = cu_num_tokens.numel() - 1
30
+ for r in range(num_reqs):
31
+ start = int(cu_num_tokens[r])
32
+ end = int(cu_num_tokens[r + 1])
33
+ n = end - start
34
+ if n <= 0:
35
+ continue
36
+ tok = torch.arange(n, device=device)
37
+ blk = tok // kernel_block_size
38
+ off = tok % kernel_block_size
39
+ out[start:end] = block_table[r, blk].to(torch.int64) * kernel_block_size + off
40
+ return out
41
+
42
+
43
+def _run(block_table, cu_num_tokens, kernel_block_size, out_len, device):
44
+ out = torch.full((out_len,), -1, dtype=torch.int64, device=device)
45
+ num_reqs = cu_num_tokens.numel() - 1
46
+ max_tokens_per_req = block_table.shape[1] * kernel_block_size
47
+ num_chunks = max(1, -(-max_tokens_per_req // 1024))
48
+ _expand_page_indices_kernel[(num_reqs, num_chunks)](
49
+ out,
50
+ block_table,
51
+ block_table.stride(0),
52
+ block_table.stride(1),
53
+ cu_num_tokens,
54
+ KERNEL_BLOCK_SIZE=kernel_block_size,
55
+ BLOCK_SIZE=1024,
56
+ )
57
+ return out
58
+
59
+
60
+def _make(seq_lens, kernel_block_size, device, noncontig=False):
61
+ num_reqs = len(seq_lens)
62
+ max_blocks = (
63
+ max(1, max((s + kernel_block_size - 1) // kernel_block_size for s in seq_lens))
64
+ + 2
65
+ )
66
+ width = max_blocks * 2 if noncontig else max_blocks
67
+ torch.manual_seed(0)
68
+ bt = torch.randint(0, 10_000, (num_reqs, width), dtype=torch.int64, device=device)
69
+ if noncontig:
70
+ bt = bt[:, ::2]
71
+ assert not bt.is_contiguous()
72
+ cu = torch.zeros(num_reqs + 1, dtype=torch.int64, device=device)
73
+ cu[1:] = torch.tensor(seq_lens, dtype=torch.int64, device=device).cumsum(0)
74
+ return bt, cu, int(cu[-1])
75
+
76
+
77
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
78
+@pytest.mark.parametrize(
79
+ "seq_lens", [[1], [128], [4096], [1, 1], [4096, 7], [37, 4096, 1], [1024] * 8]
80
+)
81
+@pytest.mark.parametrize("kernel_block_size", [1, 16, 64])
82
+def test_matches_reference(seq_lens, kernel_block_size):
83
+ bt, cu, total = _make(seq_lens, kernel_block_size, "cuda")
84
+ got = _run(bt, cu, kernel_block_size, total, "cuda")
85
+ want = _reference(bt, cu, kernel_block_size, total, "cuda")
86
+ torch.testing.assert_close(got, want, atol=0, rtol=0)
87
+
88
+
89
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
90
+@pytest.mark.parametrize("kernel_block_size", [1, 16])
91
+def test_non_contiguous_block_table(kernel_block_size):
92
+ """Column stride != 1 must be honoured, not assumed away."""
93
+ bt, cu, total = _make([4096, 37], kernel_block_size, "cuda", noncontig=True)
94
+ got = _run(bt, cu, kernel_block_size, total, "cuda")
95
+ want = _reference(bt, cu, kernel_block_size, total, "cuda")
96
+ torch.testing.assert_close(got, want, atol=0, rtol=0)
97
+
98
+
99
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
100
+def test_ragged_batch_leaves_other_rows_untouched():
101
+ """A short request must not have its tail written by a longer one's chunks."""
102
+ kbs = 16
103
+ bt, cu, total = _make([4096, 1, 512], kbs, "cuda")
104
+ out = torch.full((total + 8,), -999, dtype=torch.int64, device="cuda")
105
+ num_chunks = max(1, -(-(bt.shape[1] * kbs) // 1024))
106
+ _expand_page_indices_kernel[(3, num_chunks)](
107
+ out,
108
+ bt,
109
+ bt.stride(0),
110
+ bt.stride(1),
111
+ cu,
112
+ KERNEL_BLOCK_SIZE=kbs,
113
+ BLOCK_SIZE=1024,
114
+ )
115
+ want = _reference(bt, cu, kbs, total, "cuda")
116
+ torch.testing.assert_close(out[:total], want, atol=0, rtol=0)
117
+ # The slack past the last request must be untouched.
118
+ assert torch.all(out[total:] == -999)
119
+
120
+
121
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
122
+def test_zero_length_request_writes_nothing():
123
+ """A zero-token request in a padded batch must produce no stores."""
124
+ kbs = 16
125
+ bt, cu, total = _make([256, 0, 128], kbs, "cuda")
126
+ got = _run(bt, cu, kbs, total, "cuda")
127
+ want = _reference(bt, cu, kbs, total, "cuda")
128
+ torch.testing.assert_close(got, want, atol=0, rtol=0)
129
+ assert int(cu[2] - cu[1]) == 0
130
+
131
+
132
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU")
133
+def test_over_provisioned_chunks_do_not_change_result():
134
+ """Extra chunks must be inert.
135
+
136
+ Each program covers `chunk_idx * BLOCK_SIZE + [0, BLOCK_SIZE)`, so the grid
137
+ has to span the request: 8192 tokens need at least ceil(8192/1024) == 8
138
+ chunks, and fewer would simply leave the tail unwritten. The caller sizes
139
+ the grid from `block_table.shape[1] * kernel_block_size`, an upper bound
140
+ over the batch, so in a ragged batch most requests get more chunks than
141
+ they need. That over-provisioning is the case worth pinning: the surplus
142
+ programs must return early and change nothing.
143
+ """
144
+ kbs = 1
145
+ bt, cu, total = _make([8192], kbs, "cuda")
146
+ ref = _reference(bt, cu, kbs, total, "cuda")
147
+ required = -(-total // 1024)
148
+ for num_chunks in (required, required * 2, required * 8):
149
+ out = torch.full((total,), -1, dtype=torch.int64, device="cuda")
150
+ _expand_page_indices_kernel[(1, num_chunks)](
151
+ out,
152
+ bt,
153
+ bt.stride(0),
154
+ bt.stride(1),
155
+ cu,
156
+ KERNEL_BLOCK_SIZE=kbs,
157
+ BLOCK_SIZE=1024,
158
+ )
159
+ torch.testing.assert_close(out, ref, atol=0, rtol=0)
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
+18
-2
@@ -56,12 +56,21 @@ class _ExpandPageIndicesKernel:
56
self,
57
page_indices,
58
block_table_tensor,
59
- stride,
59
+ stride_0,
60
+ stride_1,
61
paged_kv_indptr,
62
*,
63
KERNEL_BLOCK_SIZE,
64
BLOCK_SIZE,
65
):
66
+ # Both block-table strides are passed explicitly: the kernel no longer
67
+ # assumes the column stride is 1. This stand-in indexes the tensor
68
+ # directly so torch applies the strides for it, but the values are
69
+ # checked against the tensor actually handed in -- a caller passing the
70
+ # wrong stride is exactly the bug the parameter exists to prevent, and
71
+ # silently ignoring them here would hide it.
72
+ assert stride_0 == block_table_tensor.stride(0)
73
+ assert stride_1 == block_table_tensor.stride(1)
74
self.kernel_block_size = KERNEL_BLOCK_SIZE
75
for req_idx in range(self.grid[0]):
76
out_start = int(paged_kv_indptr[req_idx].item())
@@ -1052,7 +1061,14 @@ def test_decode_expands_kernel_block_page_indices(monkeypatch):
1061
metadata.paged_kv_indices[: expected_indices.numel()],
1062
expected_indices,
1063
)
1055
- assert expand_kernel.grid == (seq_lens.numel(),)
1064
+ # The launch now carries a second grid dimension over token chunks, so
1065
+ # the expansion is spread across the context instead of one workgroup
1066
+ # walking it serially. The width comes from the block-table upper bound
1067
+ # on tokens per request, which is available host-side and so costs no
1068
+ # device sync; derived here rather than hard-coded so the assertion
1069
+ # still describes the launch if the chunk size changes.
1070
+ expected_chunks = max(1, -(-block_table.shape[1] * kernel_block_size // 1024))
1071
+ assert expand_kernel.grid == (seq_lens.numel(), expected_chunks)
1072
assert expand_kernel.kernel_block_size == kernel_block_size
1073
1074
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+44
-24
@@ -1331,10 +1331,17 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1331
# to the original _copy_page_indices_kernel).
1332
# When kernel_block_size=K>1, block_table entry b covering K tokens
1333
# gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
1334
- _expand_page_indices_kernel[(num_reqs,)](
1334
+ # Chunk count comes from the block table width, an upper bound on
1335
+ # tokens per request that is available host-side, so this adds no
1336
+ # device synchronisation. Programs whose chunk lies past
1337
+ # num_tokens mask out entirely.
1338
+ max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size
1339
+ num_chunks = max(1, cdiv(max_tokens_per_req, 1024))
1340
+ _expand_page_indices_kernel[(num_reqs, num_chunks)](
1341
self.paged_kv_indices,
1342
block_table_tensor,
1343
block_table_tensor.stride(0),
1344
+ block_table_tensor.stride(1),
1345
paged_kv_indptr,
1346
KERNEL_BLOCK_SIZE=self.kernel_block_size,
1347
BLOCK_SIZE=1024,
@@ -1573,7 +1580,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
1580
def _expand_page_indices_kernel(
1581
page_indices,
1582
block_table,
1576
- block_table_stride,
1583
+ block_table_stride_0,
1584
+ block_table_stride_1,
1585
cu_num_tokens,
1586
KERNEL_BLOCK_SIZE: tl.constexpr,
1587
BLOCK_SIZE: tl.constexpr,
@@ -1591,32 +1599,44 @@ def _expand_page_indices_kernel(
1599
When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens)
1600
is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
1601
"""
1602
+ # One program per (request, token-chunk). Parallelising over requests
1603
+ # alone degenerates at low concurrency: with num_reqs == 1 a single
1604
+ # workgroup walked the whole sequence in a serial loop.
1605
req_idx = tl.program_id(0)
1595
- row_ptr = block_table + req_idx * block_table_stride
1606
+ chunk_idx = tl.program_id(1)
1607
+ row_ptr = block_table + req_idx * block_table_stride_0
1608
start_idx = tl.load(cu_num_tokens + req_idx)
1609
num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx
1610
1599
- offset = tl.arange(0, BLOCK_SIZE)
1600
- for i in tl.range(0, num_tokens, BLOCK_SIZE):
1601
- token_offsets = i + offset
1602
- mask = token_offsets < num_tokens
1603
-
1604
- # Which block in the block table does this token belong to?
1605
- block_idx = token_offsets // KERNEL_BLOCK_SIZE
1606
- # Offset within that block
1607
- offset_in_block = token_offsets % KERNEL_BLOCK_SIZE
1608
-
1609
- # Load the block ID from the block table
1610
- block_ids = tl.load(row_ptr + block_idx, mask=mask)
1611
-
1612
- # Compute flat index in the flattened kv_buffer
1613
- flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block
1614
-
1615
- tl.store(
1616
- page_indices + start_idx + token_offsets,
1617
- flat_indices,
1618
- mask=mask,
1619
- )
1611
+ # The grid is sized from the block table width, an upper bound over all
1612
+ # requests, so a ragged batch launches chunks past a short request's end.
1613
+ # Returning here keeps those programs from issuing masked-out loads and
1614
+ # stores at all.
1615
+ chunk_start = chunk_idx * BLOCK_SIZE
1616
+ if chunk_start >= num_tokens:
1617
+ return
1618
+
1619
+ token_offsets = chunk_start + tl.arange(0, BLOCK_SIZE)
1620
+ mask = token_offsets < num_tokens
1621
+
1622
+ # Which block in the block table does this token belong to?
1623
+ block_idx = token_offsets // KERNEL_BLOCK_SIZE
1624
+ # Offset within that block
1625
+ offset_in_block = token_offsets % KERNEL_BLOCK_SIZE
1626
+
1627
+ # Load the block ID from the block table
1628
+ # Both strides are taken from the caller: the block table is a view owned
1629
+ # elsewhere, so a unit column stride must not be assumed.
1630
+ block_ids = tl.load(row_ptr + block_idx * block_table_stride_1, mask=mask)
1631
+
1632
+ # Compute flat index in the flattened kv_buffer
1633
+ flat_indices = block_ids * KERNEL_BLOCK_SIZE + offset_in_block
1634
+
1635
+ tl.store(
1636
+ page_indices + start_idx + token_offsets,
1637
+ flat_indices,
1638
+ mask=mask,
1639
+ )
1640
1641
1642
class AiterMLAHelper: