[ROCm][MLA] Mask the AITER MLA small-head verify flatten causally (#50476)
Signed-off-by: yudigege86 <naqin@amd.com> Co-authored-by: yudigege86 <naqin@amd.com>
Matt Qin committed
Aug 1, 2026 at 02:17 UTC
77469c9057bec3212a64877dbbf3b9c48c22d786
2 files changed
+284
-10
tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
new
+257
@@ -0,0 +1,257 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Regression test for causal masking in the AITER MLA verify flatten.
4
+
5
+When ``num_heads < 16`` the ROCm AITER MLA backend cannot use the ASM decode
6
+kernel for a multi-token block, so ``AiterMLAImpl.forward_mqa`` flattens a
7
+``1 + num_speculative_tokens`` verify block into that many single-token Gluon
8
+decodes. Every resulting row used to be handed its request's entire paged-KV
9
+range. ``paged_kv_indptr`` is the cumulative sum of ``seq_lens``, which counts
10
+the tokens scheduled in the current step, so that range already spans the
11
+verify block: each verify position could attend to the draft tokens after it,
12
+which are the tokens it is supposed to be checking. Nothing crashes and nothing
13
+warns, so the only way to catch it is to look at the KV window per row.
14
+
15
+The test drives the real builder and the real ``forward_mqa`` over a
16
+multi-token decode batch, intercepts the ``(page_table, seq_info)`` actually
17
+handed to the Gluon kernel, and pins verify row ``t`` of request ``r`` to the
18
+``context_r + t + 1`` entries of that request's ascending slice.
19
+"""
20
+
21
+import types
22
+from unittest.mock import patch
23
+
24
+import pytest
25
+import torch
26
+
27
+from vllm._aiter_ops import is_aiter_found
28
+from vllm.platforms import current_platform
29
+
30
+
31
+def _on_rocm_with_aiter() -> bool:
32
+ return current_platform.is_rocm() and is_aiter_found()
33
+
34
+
35
+# Unlike the fp8 persistent-decode metadata this file's neighbour guards, the
36
+# verify flatten is not gfx950-only: it is selected by head count alone and is
37
+# reached on any ROCm AITER MLA deployment with fewer than 16 heads per rank.
38
+# The Gluon kernel is replaced by a spy here, so nothing below needs it to be
39
+# runnable either.
40
+pytestmark = pytest.mark.skipif(
41
+ not _on_rocm_with_aiter(),
42
+ reason="ROCM_AITER_MLA verify flatten requires ROCm and AITER",
43
+)
44
+
45
+# Below the 16-head threshold that selects the flatten. DeepSeek-V3 / R1 at
46
+# TP=16 and Kimi-K3 at TP=8 both land here.
47
+NUM_QUERY_HEADS = 12
48
+KV_LORA_RANK = 512
49
+QK_NOPE_HEAD_DIM = 128
50
+QK_ROPE_HEAD_DIM = 64
51
+V_HEAD_DIM = 128
52
+HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM
53
+
54
+# 1 + num_speculative_tokens, i.e. the verify block length.
55
+QLEN = 8
56
+# Committed context per request, including a fresh request at zero context
57
+# where the whole KV range is the verify block itself.
58
+CONTEXT_LENS = [0, 1, 37, 512]
59
+# A cudagraph padding request: seq_len 0, so every one of its rows must clamp
60
+# to an empty window.
61
+PADDING_ROWS = 1
62
+
63
+# The Gluon path flattens the KV cache to one page per token.
64
+PAGE_SIZE = 1
65
+MAX_MODEL_LEN = 1024
66
+MAX_NUM_SEQS = 8
67
+
68
+
69
+def _seq_lens() -> list[int]:
70
+ return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS
71
+
72
+
73
+def _expected_row_lens() -> list[int]:
74
+ """Causal row lengths: row r*QLEN + t sees seq_len_r - (QLEN - 1) + t."""
75
+ return [max(0, s - (QLEN - 1) + t) for s in _seq_lens() for t in range(QLEN)]
76
+
77
+
78
+def _run_verify_block():
79
+ """Drive the real builder + forward_mqa over one multi-token verify block.
80
+
81
+ Returns ``(metadata, captured)`` where ``captured`` holds the ``page_table``,
82
+ ``seq_info`` and ``min_kv_seq_len`` the backend passed to the Gluon kernel.
83
+ """
84
+ from tests.v1.attention.utils import (
85
+ BatchSpec,
86
+ create_common_attn_metadata,
87
+ create_vllm_config,
88
+ )
89
+ from vllm.config import SpeculativeConfig
90
+ from vllm.config.vllm import set_current_vllm_config
91
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
92
+ from vllm.v1.kv_cache_interface import MLAAttentionSpec
93
+ from vllm.v1.worker.workspace import init_workspace_manager
94
+
95
+ device = torch.device("cuda:0")
96
+ seq_lens = _seq_lens()
97
+ num_reqs = len(seq_lens)
98
+ # One flat page per token, arange block indices, plus room to spare.
99
+ num_gpu_blocks = num_reqs * max(seq_lens) + 256
100
+
101
+ vllm_config = create_vllm_config(
102
+ model_name="deepseek-ai/DeepSeek-R1",
103
+ max_model_len=MAX_MODEL_LEN,
104
+ num_gpu_blocks=num_gpu_blocks,
105
+ block_size=PAGE_SIZE,
106
+ max_num_seqs=MAX_NUM_SEQS,
107
+ max_num_batched_tokens=8192,
108
+ hf_config_override={"num_attention_heads": NUM_QUERY_HEADS},
109
+ )
110
+ # What raises reorder_batch_threshold to QLEN, so a QLEN-token block is
111
+ # classified as a decode instead of a prefill. ngram needs no draft model.
112
+ vllm_config.speculative_config = SpeculativeConfig(
113
+ method="ngram", num_speculative_tokens=QLEN - 1
114
+ )
115
+
116
+ spec = MLAAttentionSpec(
117
+ block_size=PAGE_SIZE,
118
+ num_kv_heads=1,
119
+ head_size=vllm_config.model_config.get_head_size(),
120
+ dtype=vllm_config.model_config.dtype,
121
+ cache_dtype_str="auto",
122
+ )
123
+
124
+ backend_cls = AttentionBackendEnum.ROCM_AITER_MLA.get_class()
125
+ builder_cls = backend_cls.get_builder_cls()
126
+ impl_cls = backend_cls.get_impl_cls()
127
+
128
+ # The builder reads layer.prefill_backend from static_forward_context; a
129
+ # stub with the attribute is enough for metadata construction.
130
+ layer_name = "placeholder"
131
+ vllm_config.compilation_config.static_forward_context[layer_name] = (
132
+ types.SimpleNamespace(prefill_backend=torch.empty((1,)))
133
+ )
134
+
135
+ init_workspace_manager(device)
136
+
137
+ batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=[QLEN] * num_reqs)
138
+
139
+ captured: dict = {}
140
+
141
+ def spy(**kwargs):
142
+ captured["page_table"] = kwargs["page_table"].detach().clone()
143
+ captured["seq_info"] = kwargs["seq_info"].detach().clone()
144
+ captured["min_kv_seq_len"] = kwargs["min_kv_seq_len"]
145
+
146
+ with set_current_vllm_config(vllm_config):
147
+ builder = builder_cls(spec, [layer_name], vllm_config, device)
148
+ common_attn_metadata = create_common_attn_metadata(
149
+ batch_spec, PAGE_SIZE, device, arange_block_indices=True
150
+ )
151
+ metadata = builder.build(
152
+ common_prefix_len=0, common_attn_metadata=common_attn_metadata
153
+ )
154
+
155
+ impl = impl_cls(
156
+ num_heads=NUM_QUERY_HEADS,
157
+ head_size=HEAD_SIZE,
158
+ scale=(QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM) ** -0.5,
159
+ num_kv_heads=1,
160
+ alibi_slopes=None,
161
+ sliding_window=None,
162
+ kv_cache_dtype="auto",
163
+ logits_soft_cap=None,
164
+ attn_type="decoder",
165
+ kv_sharing_target_layer_name=None,
166
+ q_lora_rank=None,
167
+ kv_lora_rank=KV_LORA_RANK,
168
+ qk_nope_head_dim=QK_NOPE_HEAD_DIM,
169
+ qk_rope_head_dim=QK_ROPE_HEAD_DIM,
170
+ qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM,
171
+ v_head_dim=V_HEAD_DIM,
172
+ kv_b_proj=None,
173
+ )
174
+
175
+ num_rows = num_reqs * QLEN
176
+ dtype = torch.bfloat16
177
+ q_nope = torch.zeros(
178
+ num_rows, NUM_QUERY_HEADS, KV_LORA_RANK, dtype=dtype, device=device
179
+ )
180
+ q_pe = torch.zeros(
181
+ num_rows, NUM_QUERY_HEADS, QK_ROPE_HEAD_DIM, dtype=dtype, device=device
182
+ )
183
+ kv_cache = torch.zeros(
184
+ num_gpu_blocks, PAGE_SIZE, HEAD_SIZE, dtype=dtype, device=device
185
+ )
186
+
187
+ # The Gluon kernel only reads the metadata this test is about, so a spy in
188
+ # its place keeps the assertions independent of the AITER build.
189
+ with patch(
190
+ "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
191
+ lambda: spy,
192
+ ):
193
+ impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None)
194
+
195
+ return metadata, captured
196
+
197
+
198
+def test_verify_flatten_rows_are_causal():
199
+ """Verify row t must see the committed prefix plus tokens 0..t, and no more.
200
+
201
+ Regression guard: before the fix every row of a request got that request's
202
+ whole paged-KV range, so verify position t attended to the draft tokens
203
+ after it. Fails on unmodified upstream.
204
+ """
205
+ metadata, captured = _run_verify_block()
206
+
207
+ decode = metadata.decode
208
+ assert decode is not None, "batch was not classified as a decode"
209
+ assert decode.max_qo_len == QLEN, (
210
+ f"expected a {QLEN}-token verify block, got max_qo_len="
211
+ f"{decode.max_qo_len}; the flatten under test was not reached"
212
+ )
213
+ assert captured, "forward_mqa did not reach the Gluon kernel"
214
+
215
+ seq_lens = _seq_lens()
216
+ indptr = captured["seq_info"].tolist()
217
+ page_table = captured["page_table"]
218
+ got_row_lens = [indptr[i + 1] - indptr[i] for i in range(len(indptr) - 1)]
219
+ want_row_lens = _expected_row_lens()
220
+
221
+ assert got_row_lens == want_row_lens, (
222
+ "AITER MLA verify flatten is not causal: verify row r*qlen+t must get "
223
+ f"seq_len_r - {QLEN - 1} + t KV entries.\n"
224
+ f" seq_lens {seq_lens}\n"
225
+ f" got {got_row_lens}\n"
226
+ f" expected {want_row_lens}\n"
227
+ "Rows longer than expected let a verify position attend to the draft "
228
+ "tokens it is supposed to be checking."
229
+ )
230
+
231
+ # arange_block_indices lays request r's pages out ascending from
232
+ # r * max_blocks, so each causal window must be that slice's prefix.
233
+ max_blocks = max(seq_lens)
234
+ for r, seq_len in enumerate(seq_lens):
235
+ for t in range(QLEN):
236
+ row = r * QLEN + t
237
+ window = page_table[indptr[row] : indptr[row + 1]].tolist()
238
+ n = max(0, seq_len - (QLEN - 1) + t)
239
+ want = [r * max_blocks + j for j in range(n)]
240
+ assert window == want, (
241
+ f"request {r} (seq_len {seq_len}) verify token {t} reads KV "
242
+ f"pages {window}, expected the ascending causal prefix {want}"
243
+ )
244
+
245
+ padding_rows = got_row_lens[len(CONTEXT_LENS) * QLEN :]
246
+ assert padding_rows == [0] * (PADDING_ROWS * QLEN), (
247
+ f"cudagraph padding requests (seq_len 0) must clamp to empty windows, "
248
+ f"got {padding_rows}"
249
+ )
250
+
251
+ # min_kv_seq_len tells Gluon how short the shortest row is, so it must be
252
+ # the minimum over the causal rows actually submitted, not over the
253
+ # per-request lengths they were cut from.
254
+ assert captured["min_kv_seq_len"] == min(want_row_lens), (
255
+ f"min_kv_seq_len={captured['min_kv_seq_len']} does not match the "
256
+ f"shortest submitted row ({min(want_row_lens)})"
257
+ )
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+27
-10
@@ -1106,11 +1106,13 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1106
)
1107
return o, None
1108
1109
- # 12-head (<16) non-causal multi-token verify (DSpark): the asm path has
1110
- # no gqa<16, qseqlen>1 kernel. Flatten each verify token to a qseqlen=1
1111
- # gluon decode over the same committed prefix (non-causal), mirroring the
1112
- # TRITON_MLA / sparse-backend flatten but on the fast gluon kernel. Each
1113
- # request's KV metadata is repeated max_qo_len times.
1109
+ # 12-head (<16) multi-token verify (DSpark): the asm path has no
1110
+ # gqa<16, qseqlen>1 kernel. Flatten each verify token to its own
1111
+ # qseqlen=1 gluon decode, mirroring the TRITON_MLA / sparse-backend
1112
+ # flatten but on the fast gluon kernel. The block is causal -- the
1113
+ # target is checking draft tokens, so position t must not see t+1 --
1114
+ # and attention rows are independent, so giving row t the KV range
1115
+ # [0, context + t] is exactly causal multi-token attention.
1116
if (
1117
self.num_heads < AiterMLAHelper._AITER_MIN_MLA_HEADS
1118
and int(decode.max_qo_len) > 1
@@ -1131,16 +1133,31 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1133
device=q_nope.device,
1134
)
1135
kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])
1134
- # Expand per-request paged-KV to per-verify-token: each request's KV
1135
- # block is repeated qlen times (row r*qlen+t attends request r's
1136
- # committed prefix). Fully vectorized (no host loop).
1136
+ # Expand per-request paged-KV to per-verify-token. Row r*qlen+t is
1137
+ # request r's verify token t, and seq_lens counts the tokens
1138
+ # scheduled in this step, so a request's KV range already spans its
1139
+ # whole verify block and context_r = seq_len_r - qlen. Token t may
1140
+ # attend to [0, context_r + t], i.e. seq_len_r - (qlen - 1) + t
1141
+ # entries. paged_kv_indices lists a request's pages in ascending
1142
+ # position order, so each row's causal window is a prefix of that
1143
+ # request's slice and only the row length changes. Rows clamp to
1144
+ # zero for cudagraph padding requests, whose seq_len is 0. Fully
1145
+ # vectorized (no host loop).
1146
old_indptr = decode.paged_kv_indptr
1147
per_req_len = old_indptr[1:] - old_indptr[:-1]
1148
dev = q_nope.device
1149
row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave(
1150
qlen
1151
)
1143
- row_len = per_req_len[row_req]
1152
+ row_len = (
1153
+ (
1154
+ per_req_len.unsqueeze(1)
1155
+ - (qlen - 1)
1156
+ + torch.arange(qlen, device=dev, dtype=per_req_len.dtype)
1157
+ )
1158
+ .clamp_(min=0)
1159
+ .flatten()
1160
+ )
1161
new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to(
1162
torch.int32
1163
)
@@ -1165,7 +1182,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1182
kv_pe_offset=self.kv_lora_rank,
1183
use_2d_view=False,
1184
kv_scale=1.0,
1168
- min_kv_seq_len=int(per_req_len.min()),
1185
+ min_kv_seq_len=int(row_len.min()),
1186
)
1187
return o, None
1188