[Bugfix][MLA] TritonMLA: fix illegal memory access on causal multi-token decode (#51065)
Signed-off-by: Oleksandr Kachur <oleksandr.kachur@amd.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
olka-amd committed
Sep 18, 2026 at 11:49 UTC
32636580a6f1c3bc41deefc4bf7800f850c24034
2 files changed
+252
-54
tests/kernels/attention/test_triton_mla_causal_verify_flatten.py
new
+215
@@ -0,0 +1,215 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the multi-token decode flatten in ``TritonMLAImpl.forward_mqa``.
4
+
5
+``decode_attention_fwd`` has no causal flag: it launches one program per row of
6
+``q`` and reads that row's KV extent from ``B_Seqlen``. So a ``query_len``-token
7
+decode block is flattened to one row per query token, and intra-block causality
8
+has to be expressed through the per-row sequence lengths -- row ``t`` gets
9
+``seq_len - (query_len - 1) + t``, i.e. the committed prefix plus block tokens
10
+``0..t``. A non-causal draft block instead gives every row the same full extent.
11
+
12
+The kernel is replaced by a spy, so these tests assert on the ``block_table``
13
+and ``seq_lens`` the backend actually submits and need no GPU or Triton runtime.
14
+"""
15
+
16
+from types import SimpleNamespace
17
+from unittest.mock import patch
18
+
19
+import pytest
20
+import torch
21
+
22
+from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport
23
+from vllm.v1.attention.backend import AttentionCGSupport
24
+from vllm.v1.attention.backends.mla.triton_mla import (
25
+ TritonMLAImpl,
26
+ TritonMLAMetadataBuilder,
27
+)
28
+
29
+NUM_HEADS = 16
30
+KV_LORA_RANK = 512
31
+QK_ROPE_HEAD_DIM = 64
32
+HEAD_SIZE = KV_LORA_RANK + QK_ROPE_HEAD_DIM
33
+PAGE_SIZE = 16
34
+
35
+# Committed context per request, including a fresh request at zero context
36
+# where the whole KV extent is the block itself. A decode block's seq_len
37
+# already counts the block, so seq_len = context + query_len.
38
+CONTEXT_LENS = [0, 1, 37, 512]
39
+
40
+
41
+def _seq_lens(query_len: int) -> list[int]:
42
+ return [c + query_len for c in CONTEXT_LENS]
43
+
44
+
45
+def _make_impl(dcp_world_size: int = 1) -> TritonMLAImpl:
46
+ """Only the attributes ``forward_mqa`` reads; __init__ wants a full config."""
47
+ impl = object.__new__(TritonMLAImpl)
48
+ impl.kv_lora_rank = KV_LORA_RANK
49
+ impl.scale = HEAD_SIZE**-0.5
50
+ impl._sm_count = 304
51
+ impl.dcp_world_size = dcp_world_size
52
+ return impl
53
+
54
+
55
+def _make_metadata(seq_lens: list[int], query_len: int, causal: bool):
56
+ num_decodes = len(seq_lens)
57
+ return SimpleNamespace(
58
+ num_decodes=num_decodes,
59
+ num_decode_tokens=num_decodes * query_len,
60
+ max_seq_len=max(seq_lens),
61
+ causal=causal,
62
+ decode=SimpleNamespace(
63
+ # arange page ids so each request's window is an identifiable slice.
64
+ block_table=torch.arange(num_decodes * 64, dtype=torch.int32).reshape(
65
+ num_decodes, 64
66
+ ),
67
+ seq_lens=torch.tensor(seq_lens, dtype=torch.int32),
68
+ ),
69
+ )
70
+
71
+
72
+def _run_forward_mqa(
73
+ seq_lens: list[int],
74
+ query_len: int,
75
+ causal: bool,
76
+ q_rows=None,
77
+ dcp_world_size: int = 1,
78
+):
79
+ """Drive forward_mqa with the decode kernel spied out.
80
+
81
+ Returns the ``(block_table, seq_lens)`` handed to ``decode_attention_fwd``.
82
+ """
83
+ captured: dict = {}
84
+
85
+ def spy(q, kv_cache, kv_c_cache, o, lse, block_table, b_seq_len, *args, **kwargs):
86
+ captured["block_table"] = block_table.detach().clone()
87
+ captured["seq_lens"] = b_seq_len.detach().clone()
88
+
89
+ metadata = _make_metadata(seq_lens, query_len, causal)
90
+ rows = metadata.num_decode_tokens if q_rows is None else q_rows
91
+ q = torch.zeros(rows, NUM_HEADS, HEAD_SIZE, dtype=torch.bfloat16)
92
+ kv_cache = torch.zeros(
93
+ len(seq_lens) * 64, PAGE_SIZE, HEAD_SIZE, dtype=torch.bfloat16
94
+ )
95
+ layer = SimpleNamespace(_k_scale=torch.ones(1), layer_name="test")
96
+
97
+ with patch("vllm.v1.attention.backends.mla.triton_mla.decode_attention_fwd", spy):
98
+ _make_impl(dcp_world_size).forward_mqa(q, kv_cache, metadata, layer)
99
+
100
+ assert captured, "forward_mqa did not reach the decode kernel"
101
+ return captured
102
+
103
+
104
+def test_multi_token_decode_flags_are_declared_together():
105
+ """Capture probes max_query_len > 1, which only clears the assert in
106
+ build_for_cudagraph_capture once query_len_support raised the threshold."""
107
+ assert (
108
+ TritonMLAMetadataBuilder._cudagraph_support == AttentionCGSupport.UNIFORM_BATCH
109
+ )
110
+ assert TritonMLAMetadataBuilder.query_len_support == QueryLenSupport.UNIFORM
111
+
112
+
113
+@pytest.mark.parametrize("query_len", [2, 3, 4, 5, 8])
114
+def test_causal_block_rows_see_prefix_plus_own_position(query_len):
115
+ """Causal verify: row t sees seq_len - (query_len - 1) + t KV entries.
116
+
117
+ Regression guard: the causal branch used to be skipped entirely, so a verify
118
+ token could attend the draft siblings it is supposed to be checking. Fails
119
+ on unmodified upstream, where no flatten happens at all for causal blocks.
120
+ """
121
+ seq_lens = _seq_lens(query_len)
122
+ captured = _run_forward_mqa(seq_lens, query_len, causal=True)
123
+
124
+ got = captured["seq_lens"].tolist()
125
+ want = [s - (query_len - 1) + t for s in seq_lens for t in range(query_len)]
126
+ assert got == want, (
127
+ "causal flatten is not causal: verify row r*query_len+t must get "
128
+ f"seq_len_r - {query_len - 1} + t entries.\n"
129
+ f" seq_lens {seq_lens}\n got {got}\n expected {want}"
130
+ )
131
+
132
+ # Per-request invariants the arithmetic must hold regardless of query_len.
133
+ for r, seq_len in enumerate(seq_lens):
134
+ rows = got[r * query_len : (r + 1) * query_len]
135
+ assert rows == sorted(rows) and len(set(rows)) == query_len, (
136
+ f"request {r} rows must strictly increase, got {rows}"
137
+ )
138
+ assert rows[-1] == seq_len, "last row must see the full sequence"
139
+ assert rows[0] == seq_len - (query_len - 1)
140
+ assert all(n > 0 for n in rows), f"non-positive KV extent in {rows}"
141
+
142
+
143
+def test_cudagraph_padding_rows_present_no_kv_extent():
144
+ """Causal flattening pins cudagraph padding rows (seq_len 0) to exactly
145
+ zero KV extent, keeping them in the kernel's skipped branch."""
146
+ query_len = 4
147
+ seq_lens = _seq_lens(query_len) + [0]
148
+ captured = _run_forward_mqa(seq_lens, query_len, causal=True)
149
+
150
+ padding = captured["seq_lens"].tolist()[-query_len:]
151
+ assert padding == [0] * query_len, (
152
+ f"padding rows must present exactly zero KV extent, got {padding}"
153
+ )
154
+
155
+
156
+@pytest.mark.parametrize("query_len", [2, 4])
157
+def test_non_causal_block_rows_all_see_the_full_prefix(query_len):
158
+ """The DSpark draft block is generated in one pass, so no row is masked."""
159
+ seq_lens = _seq_lens(query_len)
160
+ captured = _run_forward_mqa(seq_lens, query_len, causal=False)
161
+
162
+ want = [s for s in seq_lens for _ in range(query_len)]
163
+ assert captured["seq_lens"].tolist() == want
164
+
165
+
166
+@pytest.mark.parametrize("causal", [True, False])
167
+def test_block_table_is_expanded_to_match_q_rows(causal):
168
+ """The kernel indexes block_table by program_id, so it needs one row per
169
+ query token -- a short block_table is read out of bounds."""
170
+ query_len = 4
171
+ captured = _run_forward_mqa(_seq_lens(query_len), query_len, causal=causal)
172
+
173
+ num_rows = len(CONTEXT_LENS) * query_len
174
+ assert captured["block_table"].shape[0] == num_rows
175
+ assert captured["seq_lens"].shape[0] == num_rows
176
+ # repeat_interleave, not repeat: a request's rows must stay adjacent.
177
+ for r in range(len(CONTEXT_LENS)):
178
+ block = captured["block_table"][r * query_len : (r + 1) * query_len]
179
+ assert torch.equal(block, block[:1].expand_as(block))
180
+
181
+
182
+@pytest.mark.parametrize("causal", [True, False])
183
+def test_single_token_decode_is_untouched(causal):
184
+ """query_len == 1 must short-circuit: ordinary decode sees no expansion."""
185
+ seq_lens = _seq_lens(1)
186
+ captured = _run_forward_mqa(seq_lens, query_len=1, causal=causal)
187
+
188
+ assert captured["seq_lens"].tolist() == seq_lens
189
+ assert captured["block_table"].shape[0] == len(seq_lens)
190
+
191
+
192
+def test_expansion_factor_follows_the_q_row_count():
193
+ """The kernel indexes block_table and seq_lens by program_id, so the
194
+ expansion must follow q rows, not the token/request counts."""
195
+ num_rows = len(CONTEXT_LENS) * 4
196
+ captured = _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=num_rows)
197
+
198
+ assert captured["block_table"].shape[0] == num_rows
199
+ assert captured["seq_lens"].shape[0] == num_rows
200
+
201
+
202
+def test_non_uniform_block_is_rejected():
203
+ """Q rows must divide evenly across the decode requests."""
204
+ with pytest.raises(AssertionError, match="non-uniform decode block"):
205
+ _run_forward_mqa(_seq_lens(1), query_len=1, causal=True, q_rows=6)
206
+
207
+
208
+def test_causal_multi_token_decode_is_rejected_under_dcp():
209
+ """Per-row extents offset the global seq_len, but DCP passes the rank-local
210
+ slice, so the causal arithmetic would silently address the wrong KV."""
211
+ with pytest.raises(AssertionError, match="not supported with DCP"):
212
+ _run_forward_mqa(_seq_lens(4), query_len=4, causal=True, dcp_world_size=2)
213
+
214
+ # Non-causal rows all take the same extent, which stays correct when local.
215
+ _run_forward_mqa(_seq_lens(4), query_len=4, causal=False, dcp_world_size=2)
vllm/v1/attention/backends/mla/triton_mla.py
+37
-54
@@ -6,7 +6,6 @@ from typing import ClassVar
6
import torch
7
8
import vllm.envs as envs
9
-from vllm.config import VllmConfig
9
from vllm.config.cache import CacheDType
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.attention.mla_attention import (
@@ -14,6 +13,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
13
MLACommonImpl,
14
MLACommonMetadata,
15
MLACommonMetadataBuilder,
16
+ QueryLenSupport,
17
)
18
from vllm.platforms import current_platform
19
from vllm.platforms.interface import DeviceCapability
@@ -26,7 +26,6 @@ from vllm.v1.attention.backend import (
26
MultipleOf,
27
)
28
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
29
-from vllm.v1.kv_cache_interface import KVCacheSpec
29
from vllm.v1.worker.workspace import (
30
current_workspace_manager,
31
is_workspace_manager_initialized,
@@ -50,50 +49,16 @@ def _compute_num_kv_splits(max_seq_len: int, sm_count: int) -> int:
49
50
51
class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
53
- _cudagraph_support: ClassVar[AttentionCGSupport] = (
54
- AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
55
- )
56
- # Non-causal DSpark block is flattened to one decode row per query token in
57
- # forward_mqa, so no intra-block causal masking is required.
52
+ # forward_mqa flattens a uniform multi-token block to one decode row per
53
+ # query token, so causal and non-causal blocks both take the decode path.
54
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
55
+ query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
56
supports_non_causal_multi_token_decode: ClassVar[bool] = True
57
60
- @classmethod
61
- def get_cudagraph_support(
62
- cls,
63
- vllm_config: VllmConfig,
64
- kv_cache_spec: KVCacheSpec,
65
- ) -> AttentionCGSupport:
66
- """Report UNIFORM_BATCH where a non-causal multi-token block is served.
67
-
68
- ``_cudagraph_support`` is a class constant, so serving the DSpark
69
- draft's (1 + num_spec) block through the decode path reports
70
- UNIFORM_SINGLE_TOKEN_DECODE and, because the engine takes the minimum
71
- over all attention groups, downgrades the *whole* engine off full
72
- cudagraphs. ``forward_mqa`` flattens that block with
73
- ``repeat_interleave`` on a Python int and performs no device->host
74
- sync, so it does satisfy the UNIFORM_BATCH contract.
75
-
76
- ``non_causal_multi_token_decode`` is a KV-cache-group property, not a
77
- per-layer one: ``MLAAttentionSpec.merge`` ORs it over every layer in
78
- the group, so a group holding both a draft and its target reports it
79
- for both. That is the same predicate ``__init__`` below already uses to
80
- raise ``reorder_batch_threshold``, so the two stay consistent, but it
81
- does mean this lifts a causal target sharing the draft's KV cache group
82
- as well.
83
- """
84
- if getattr(kv_cache_spec, "non_causal_multi_token_decode", False):
85
- return AttentionCGSupport.UNIFORM_BATCH
86
- return cls._cudagraph_support
87
-
58
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
59
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
60
# DCP local sequence lengths are not advanced between draft steps.
61
self.supports_draft_decode_metadata_update = self.dcp_world_size == 1
92
- # Only the non-causal DSpark draft group serves multi-token blocks via
93
- # the decode path; raise its reorder threshold to the spec block length
94
- # so full-cudagraph capture admits it. Causal usage stays single-token.
95
- if getattr(self, "non_causal_multi_token_decode", False):
96
- self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
62
self._reserve_attn_logits_workspace()
63
64
def update_draft_decode_metadata(self, _metadata: MLACommonMetadata) -> None:
@@ -109,12 +74,12 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
74
"""
75
if not is_workspace_manager_initialized():
76
return
112
- # Decode reorder threshold is 1, so decode tokens <= max_num_seqs.
113
- B = self.vllm_config.scheduler_config.max_num_seqs
114
- # Non-causal DSpark draft flattens each request's block to query_len
115
- # decode rows; cover max_num_seqs * block_len rows.
116
- if getattr(self, "non_causal_multi_token_decode", False):
117
- B *= self.reorder_batch_threshold
77
+ # forward_mqa flattens each request's block to query_len decode rows,
78
+ # and query_len is bounded by the reorder threshold.
79
+ B = (
80
+ self.vllm_config.scheduler_config.max_num_seqs
81
+ * self.reorder_batch_threshold
82
+ )
83
# DCP all-gathers the query heads before forward_mqa.
84
q_num_heads = self.num_heads * self.dcp_world_size
85
max_splits = _compute_num_kv_splits(
@@ -313,14 +278,32 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
278
279
block_table = attn_metadata.decode.block_table
280
seq_lens = attn_metadata.decode.seq_lens
316
- if not attn_metadata.causal:
317
- # Non-causal DSpark block: flatten to one decode row per query token.
318
- # Each row attends to the same committed KV prefix (per-row seq_lens)
319
- # and never to sibling block tokens = non-causal block semantics.
320
- # Mirrors FlashInferMLA's non-causal path.
321
- query_len = attn_metadata.num_decode_tokens // attn_metadata.num_decodes
322
- if query_len > 1:
323
- block_table = block_table.repeat_interleave(query_len, dim=0)
281
+ # decode_attention_fwd has no causal flag: it launches one program per
282
+ # row of q and reads that row's KV extent from seq_lens, so intra-block
283
+ # causality is expressed as per-row extents. Deriving query_len from the
284
+ # tensors the kernel indexes keeps the three row counts in step.
285
+ num_decodes = seq_lens.shape[0]
286
+ query_len, remainder = divmod(B, num_decodes) if num_decodes else (1, 0)
287
+ assert remainder == 0, (
288
+ f"non-uniform decode block: {B} query rows over {num_decodes} requests"
289
+ )
290
+ if query_len > 1:
291
+ block_table = block_table.repeat_interleave(query_len, dim=0)
292
+ if attn_metadata.causal:
293
+ # Per-row extents are offsets off the global sequence length;
294
+ # under DCP seq_lens is this rank's local slice instead.
295
+ assert self.dcp_world_size == 1, (
296
+ "causal multi-token decode is not supported with DCP"
297
+ )
298
+ # Row t attends the prefix plus block tokens 0..t. Clamp holds
299
+ # padding rows at extent 0; the kernel skips them either way
300
+ # (split_kv_end == split_kv_start).
301
+ offsets = torch.arange(
302
+ 1 - query_len, 1, device=seq_lens.device, dtype=seq_lens.dtype
303
+ )
304
+ seq_lens = (seq_lens.unsqueeze(1) + offsets).flatten().clamp(min=0)
305
+ else:
306
+ # Non-causal draft block: every row sees the same prefix.
307
seq_lens = seq_lens.repeat_interleave(query_len)
308
309
# Run MQA — always pass layer scales. When KV cache is