[Bugfix][ROCm] AITER MLA: size MTP verification decode metadata for real qlen/dtype (#45227)
Signed-off-by: Edwin Lim <edwin.lim@mangoboost.io> Signed-off-by: Chaemin Lim <chaemin.lim@mangoboost.io> Signed-off-by: QinPR <1905873179@qq.com> Co-authored-by: Edwin Lim <edwin.lim@mangoboost.io> Co-authored-by: limeward <32970461+edwinlim0919@users.noreply.github.com> Co-authored-by: QinPR <1905873179@qq.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
chaeminlim-mb committed
Jul 31, 2026 at 10:50 UTC
f1899b2ffdf87598350c416a074c8c291c845902
2 files changed
+555
-36
tests/v1/attention/test_rocm_aiter_mla_mtp_split.py
new
+408
@@ -0,0 +1,408 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import sys
5
+from types import SimpleNamespace
6
+from unittest import mock
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.platforms import current_platform
12
+
13
+if not current_platform.is_rocm():
14
+ pytest.skip("ROCm AITER MLA tests", allow_module_level=True)
15
+
16
+from vllm.v1.attention.backends.mla import rocm_aiter_mla # noqa: E402
17
+from vllm.v1.attention.backends.mla.rocm_aiter_mla import ( # noqa: E402
18
+ AiterMLAMetadataBuilder,
19
+)
20
+
21
+
22
+class _NoOpTritonKernel:
23
+ def __getitem__(self, grid):
24
+ self.grid = grid
25
+ return self
26
+
27
+ def __call__(self, *args, **kwargs):
28
+ pass
29
+
30
+
31
+class _ExpandPageIndicesKernel:
32
+ def __getitem__(self, grid):
33
+ self.grid = grid
34
+ return self
35
+
36
+ def __call__(
37
+ self,
38
+ page_indices,
39
+ block_table_tensor,
40
+ stride,
41
+ paged_kv_indptr,
42
+ seq_lens_for_kernel,
43
+ *,
44
+ KERNEL_BLOCK_SIZE,
45
+ BLOCK_SIZE,
46
+ ):
47
+ self.kernel_block_size = KERNEL_BLOCK_SIZE
48
+ for req_idx in range(self.grid[0]):
49
+ out_start = int(paged_kv_indptr[req_idx].item())
50
+ seq_len = int(seq_lens_for_kernel[req_idx].item())
51
+ for token_idx in range(seq_len):
52
+ block_id = int(
53
+ block_table_tensor[req_idx, token_idx // KERNEL_BLOCK_SIZE].item()
54
+ )
55
+ page_indices[out_start + token_idx] = (
56
+ block_id * KERNEL_BLOCK_SIZE + token_idx % KERNEL_BLOCK_SIZE
57
+ )
58
+
59
+
60
+def _builder(
61
+ *,
62
+ mtp_decode_qlen: int,
63
+ has_full_cudagraphs: bool = False,
64
+ kernel_block_size: int = 1,
65
+ max_decode_rows: int = 32,
66
+ num_heads: int = 16,
67
+):
68
+ return SimpleNamespace(
69
+ device=torch.device("cpu"),
70
+ num_heads=num_heads,
71
+ paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
72
+ paged_kv_indices=torch.empty(1024, dtype=torch.int32),
73
+ paged_kv_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
74
+ qo_indptr=torch.empty(max_decode_rows + 1, dtype=torch.int32),
75
+ compilation_config=SimpleNamespace(
76
+ cudagraph_mode=SimpleNamespace(
77
+ has_full_cudagraphs=lambda: has_full_cudagraphs
78
+ )
79
+ ),
80
+ _mtp_decode_qlen=mtp_decode_qlen,
81
+ _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len),
82
+ _use_persistent_metadata=False,
83
+ kernel_block_size=kernel_block_size,
84
+ _num_attention_heads=16,
85
+ _mla_work_meta_data=torch.empty(1, dtype=torch.int32),
86
+ _mla_work_info_set=torch.empty(1, dtype=torch.int32),
87
+ _mla_work_indptr=torch.empty(1, dtype=torch.int32),
88
+ _mla_reduce_indptr=torch.empty(1, dtype=torch.int32),
89
+ _mla_reduce_final_map=torch.empty(1, dtype=torch.int32),
90
+ _mla_reduce_partial_map=torch.empty(1, dtype=torch.int32),
91
+ _mla_q_dtype=torch.bfloat16,
92
+ _mla_kv_dtype=torch.bfloat16,
93
+ decode_attn_out_dtype=torch.bfloat16,
94
+ )
95
+
96
+
97
+def test_backend_declares_uniform_batch_support():
98
+ # UNIFORM/UNIFORM_BATCH is unconditional: MTP yields uniform qlen>1 and
99
+ # non-MTP yields qlen==1, both uniform batches.
100
+ assert (
101
+ AiterMLAMetadataBuilder.query_len_support
102
+ == rocm_aiter_mla.QueryLenSupport.UNIFORM
103
+ )
104
+ assert (
105
+ AiterMLAMetadataBuilder._cudagraph_support
106
+ == rocm_aiter_mla.AttentionCGSupport.UNIFORM_BATCH
107
+ )
108
+
109
+
110
+@pytest.mark.parametrize("num_heads", [8, 16, 32, 64, 128])
111
+def test_mtp_builder_init_sizes_native_fp8_metadata(monkeypatch, num_heads):
112
+ """Aiter init passes real MTP qlen/dtypes to rocm_aiter_mla.py:349-355.
113
+
114
+ Sweeping num_heads asserts the max(16, num_heads) clamp is what sizes the
115
+ metadata, covering the fp8 nhead=32 (TP4) fold path.
116
+ """
117
+
118
+ dtypes = SimpleNamespace(fp8="fp8", fp16="fp16", bf16="bf16")
119
+ info_calls = []
120
+
121
+ def get_mla_metadata_info_v1(
122
+ max_batch_size,
123
+ max_qo_len,
124
+ num_attention_heads,
125
+ q_dtype,
126
+ kv_dtype,
127
+ *,
128
+ is_sparse,
129
+ fast_mode,
130
+ ):
131
+ info_calls.append(
132
+ {
133
+ "max_batch_size": max_batch_size,
134
+ "max_qo_len": max_qo_len,
135
+ "num_attention_heads": num_attention_heads,
136
+ "q_dtype": q_dtype,
137
+ "kv_dtype": kv_dtype,
138
+ "is_sparse": is_sparse,
139
+ "fast_mode": fast_mode,
140
+ }
141
+ )
142
+ return tuple((1, torch.int32) for _ in range(6))
143
+
144
+ def init_common_builder(self, *args, **kwargs):
145
+ self.num_heads = num_heads
146
+
147
+ monkeypatch.setitem(
148
+ sys.modules,
149
+ "aiter",
150
+ SimpleNamespace(
151
+ dtypes=dtypes,
152
+ get_mla_metadata_info_v1=get_mla_metadata_info_v1,
153
+ ),
154
+ )
155
+ monkeypatch.setattr(
156
+ rocm_aiter_mla.MLACommonMetadataBuilder,
157
+ "__init__",
158
+ init_common_builder,
159
+ )
160
+ monkeypatch.setattr(rocm_aiter_mla, "_fp8_mla_prefill_supported", lambda: False)
161
+
162
+ config = SimpleNamespace(
163
+ speculative_config=SimpleNamespace(
164
+ method="deepseek_mtp",
165
+ num_speculative_tokens=3,
166
+ ),
167
+ parallel_config=SimpleNamespace(tensor_parallel_size=8),
168
+ model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16),
169
+ scheduler_config=SimpleNamespace(max_num_seqs=2),
170
+ cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"),
171
+ compilation_config=SimpleNamespace(
172
+ cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False)
173
+ ),
174
+ )
175
+ builder = AiterMLAMetadataBuilder(
176
+ kv_cache_spec=SimpleNamespace(block_size=1, dtype=torch.bfloat16),
177
+ layer_names=["layer.0"],
178
+ vllm_config=config,
179
+ device=torch.device("cpu"),
180
+ )
181
+
182
+ assert info_calls == [
183
+ {
184
+ "max_batch_size": config.scheduler_config.max_num_seqs,
185
+ "max_qo_len": config.speculative_config.num_speculative_tokens + 1,
186
+ "num_attention_heads": max(16, num_heads),
187
+ "q_dtype": dtypes.fp8,
188
+ "kv_dtype": dtypes.fp8,
189
+ "is_sparse": False,
190
+ "fast_mode": True,
191
+ }
192
+ ]
193
+ assert builder._mla_q_dtype == dtypes.fp8
194
+ assert builder._mla_kv_dtype == dtypes.fp8
195
+
196
+
197
+def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch):
198
+ get_mla_metadata_v1 = mock.MagicMock()
199
+ monkeypatch.setitem(
200
+ sys.modules,
201
+ "aiter",
202
+ SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
203
+ )
204
+ monkeypatch.setattr(
205
+ rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
206
+ )
207
+
208
+ metadata = AiterMLAMetadataBuilder._build_decode(
209
+ _builder(mtp_decode_qlen=4),
210
+ block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
211
+ seq_lens_device=torch.tensor([7, 5], dtype=torch.int32),
212
+ max_seq_len=7,
213
+ query_start_loc_cpu=torch.tensor([0, 4, 8], dtype=torch.int32),
214
+ query_start_loc_device=torch.tensor([0, 4, 8], dtype=torch.int32),
215
+ num_decode_tokens=8,
216
+ dcp_tot_seq_lens_device=None,
217
+ )
218
+
219
+ assert metadata.max_qo_len == 4
220
+ assert torch.equal(metadata.seq_lens, torch.tensor([7, 5], dtype=torch.int32))
221
+ assert torch.equal(metadata.qo_indptr, torch.tensor([0, 4, 8], dtype=torch.int32))
222
+ assert metadata.has_persistent_metadata
223
+ assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == 4
224
+ assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == 4
225
+
226
+
227
+def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(
228
+ monkeypatch,
229
+):
230
+ """Full-CG zero-qo rows follow rocm_aiter_mla.py:608-657,717-759."""
231
+
232
+ get_mla_metadata_v1 = mock.MagicMock()
233
+ monkeypatch.setitem(
234
+ sys.modules,
235
+ "aiter",
236
+ SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
237
+ )
238
+ monkeypatch.setattr(
239
+ rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
240
+ )
241
+
242
+ mtp_qlen = 4
243
+ seq_lens = torch.tensor([7, 0], dtype=torch.int32)
244
+ qo_lens = torch.tensor([mtp_qlen, 0], dtype=torch.int32)
245
+ expected_seq_lens = torch.where(qo_lens > 0, seq_lens, mtp_qlen)
246
+ expected_paged_kv_indptr = torch.cat(
247
+ [
248
+ torch.zeros(1, dtype=torch.int32),
249
+ expected_seq_lens.cumsum(dim=0, dtype=torch.int32),
250
+ ]
251
+ )
252
+ expected_qo_indptr = torch.arange(
253
+ 0,
254
+ (seq_lens.numel() + 1) * mtp_qlen,
255
+ step=mtp_qlen,
256
+ dtype=torch.int32,
257
+ )
258
+
259
+ builder = _builder(
260
+ mtp_decode_qlen=mtp_qlen,
261
+ has_full_cudagraphs=True,
262
+ max_decode_rows=4,
263
+ )
264
+ metadata = AiterMLAMetadataBuilder._build_decode(
265
+ builder,
266
+ block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
267
+ seq_lens_device=seq_lens,
268
+ max_seq_len=int(seq_lens.max().item()),
269
+ query_start_loc_cpu=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),
270
+ query_start_loc_device=torch.tensor([0, mtp_qlen, mtp_qlen], dtype=torch.int32),
271
+ num_decode_tokens=seq_lens.numel() * mtp_qlen,
272
+ dcp_tot_seq_lens_device=None,
273
+ )
274
+
275
+ assert metadata.max_qo_len == mtp_qlen
276
+ assert torch.equal(metadata.seq_lens, expected_seq_lens)
277
+ assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr)
278
+ assert torch.equal(metadata.qo_indptr, expected_qo_indptr)
279
+ assert torch.all(
280
+ builder.paged_kv_indptr[expected_paged_kv_indptr.numel() :]
281
+ == expected_paged_kv_indptr[-1]
282
+ )
283
+ assert torch.all(
284
+ builder.qo_indptr[expected_qo_indptr.numel() :] == expected_qo_indptr[-1]
285
+ )
286
+ assert metadata.has_persistent_metadata
287
+ assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == mtp_qlen
288
+ assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == mtp_qlen
289
+
290
+
291
+def test_decode_expands_kernel_block_page_indices(monkeypatch):
292
+ """kernel_block_size>1 expands b -> b*K+offset at rocm_aiter_mla.py:696-704."""
293
+
294
+ expand_kernel = _ExpandPageIndicesKernel()
295
+ monkeypatch.setattr(rocm_aiter_mla, "_expand_page_indices_kernel", expand_kernel)
296
+ # qlen==1 now takes the persistent-metadata path, which imports
297
+ # get_mla_metadata_v1 from aiter; mock it so this test needs no real kernel.
298
+ monkeypatch.setitem(
299
+ sys.modules,
300
+ "aiter",
301
+ SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()),
302
+ )
303
+
304
+ kernel_block_size = 2
305
+ block_table = torch.tensor(
306
+ [
307
+ [10, 11, 99],
308
+ [20, 21, 22],
309
+ ],
310
+ dtype=torch.int32,
311
+ )
312
+ seq_lens = torch.tensor([3, 5], dtype=torch.int32)
313
+ expected_paged_kv_indptr = torch.cat(
314
+ [
315
+ torch.zeros(1, dtype=torch.int32),
316
+ seq_lens.cumsum(dim=0, dtype=torch.int32),
317
+ ]
318
+ )
319
+ expected_indices = torch.tensor(
320
+ [
321
+ int(block_table[req_idx, token_idx // kernel_block_size].item())
322
+ * kernel_block_size
323
+ + token_idx % kernel_block_size
324
+ for req_idx, seq_len in enumerate(seq_lens.tolist())
325
+ for token_idx in range(seq_len)
326
+ ],
327
+ dtype=torch.int32,
328
+ )
329
+
330
+ metadata = AiterMLAMetadataBuilder._build_decode(
331
+ _builder(
332
+ mtp_decode_qlen=1,
333
+ kernel_block_size=kernel_block_size,
334
+ ),
335
+ block_table_tensor=block_table,
336
+ seq_lens_device=seq_lens,
337
+ max_seq_len=int(seq_lens.max().item()),
338
+ query_start_loc_cpu=torch.tensor([0, 1, 2], dtype=torch.int32),
339
+ query_start_loc_device=torch.tensor([0, 1, 2], dtype=torch.int32),
340
+ num_decode_tokens=seq_lens.numel(),
341
+ dcp_tot_seq_lens_device=None,
342
+ )
343
+
344
+ assert metadata.max_qo_len == 1
345
+ assert torch.equal(metadata.paged_kv_indptr, expected_paged_kv_indptr)
346
+ assert torch.equal(
347
+ metadata.paged_kv_indices[: expected_indices.numel()],
348
+ expected_indices,
349
+ )
350
+ assert expand_kernel.grid == (seq_lens.numel(),)
351
+ assert expand_kernel.kernel_block_size == kernel_block_size
352
+
353
+
354
+@pytest.mark.parametrize(
355
+ "mtp_decode_qlen, qo_len, num_heads, expect_persistent",
356
+ [
357
+ (1, 1, 16, True), # non-MTP decode
358
+ (4, 2, 16, True), # MTP deployment, in-range step
359
+ (4, 4, 16, True), # MTP deployment, full-qlen verification step
360
+ (2, 4, 16, False), # step demand exceeds provisioned K -> fallback
361
+ (1, 1, 8, False), # small head count -> Gluon decode owns qlen==1
362
+ (4, 4, 8, False), # small head count -> Gluon flatten owns qlen>1
363
+ ],
364
+)
365
+def test_persistent_metadata_gate(
366
+ monkeypatch, mtp_decode_qlen, qo_len, num_heads, expect_persistent
367
+):
368
+ """Persistent metadata is passed iff num_heads >= 16 and 1 <= max_qo_len <= K.
369
+
370
+ K = _mtp_decode_qlen sizes the metadata buffers at init; a decode step gets
371
+ the pre-built schedule only when its qlen fits those buffers, otherwise it
372
+ falls back to the kernel computing its own. qlen==1 (non-MTP) must stay
373
+ in-range -- dropping it is the regression this guards. Fewer than 16 heads
374
+ is served by the Gluon decode paths, which never read this schedule.
375
+ """
376
+ get_mla_metadata_v1 = mock.MagicMock()
377
+ monkeypatch.setitem(
378
+ sys.modules,
379
+ "aiter",
380
+ SimpleNamespace(get_mla_metadata_v1=get_mla_metadata_v1),
381
+ )
382
+ monkeypatch.setattr(
383
+ rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()
384
+ )
385
+
386
+ # Uniform, non-CUDA-graph batch: every request has exactly qo_len tokens, so
387
+ # num_decode_tokens == sum(qo_len) and no dummy-row padding kicks in.
388
+ num_reqs = 2
389
+ query_start_loc = torch.arange(
390
+ 0, (num_reqs + 1) * qo_len, step=qo_len, dtype=torch.int32
391
+ )
392
+ metadata = AiterMLAMetadataBuilder._build_decode(
393
+ _builder(mtp_decode_qlen=mtp_decode_qlen, num_heads=num_heads),
394
+ block_table_tensor=torch.arange(16, dtype=torch.int32).view(2, 8),
395
+ seq_lens_device=torch.tensor([8, 8], dtype=torch.int32),
396
+ max_seq_len=8,
397
+ query_start_loc_cpu=query_start_loc,
398
+ query_start_loc_device=query_start_loc,
399
+ num_decode_tokens=num_reqs * qo_len,
400
+ dcp_tot_seq_lens_device=None,
401
+ )
402
+
403
+ assert metadata.max_qo_len == qo_len
404
+ assert metadata.has_persistent_metadata is expect_persistent
405
+ assert get_mla_metadata_v1.called is expect_persistent
406
+ if expect_persistent:
407
+ assert get_mla_metadata_v1.call_args.kwargs["max_seqlen_qo"] == qo_len
408
+ assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == qo_len
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+147
-36
@@ -132,7 +132,7 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
132
min_kv_seq_len: int = 1
133
# Small-head decode uses Gluon (avoids padding to 16).
134
use_gluon_decode: bool = False
135
- # Whether persistent MLA metadata was computed (only for qseqlen=1)
135
+ # Whether persistent MLA metadata was computed
136
has_persistent_metadata: bool = False
137
138
@@ -171,6 +171,46 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
171
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
172
query_len_support: ClassVar[QueryLenSupport] = QueryLenSupport.UNIFORM
173
174
+ @staticmethod
175
+ def _uniform_padded_mtp_qo_len(
176
+ qo_len: torch.Tensor,
177
+ max_qo_len: int,
178
+ num_decode_tokens: int,
179
+ ) -> int:
180
+ num_reqs = qo_len.numel()
181
+ if num_reqs == 0 or num_decode_tokens <= 0:
182
+ return 0
183
+
184
+ # Full-CG pads q to a captured token count while leaving
185
+ # query_start_loc flat for dummy requests. Only synthesize dummy rows
186
+ # when every padded request maps to the same qlen and the q buffer has
187
+ # exactly that many rows.
188
+ if num_decode_tokens <= int(qo_len.sum().item()):
189
+ return 0
190
+ if num_decode_tokens % num_reqs != 0:
191
+ return 0
192
+
193
+ uniform_qo_len = num_decode_tokens // num_reqs
194
+ if uniform_qo_len <= 1:
195
+ return 0
196
+
197
+ positive_qo_len = qo_len[qo_len > 0]
198
+ if positive_qo_len.numel() == qo_len.numel():
199
+ return 0
200
+ if positive_qo_len.numel() > 0:
201
+ if max_qo_len != uniform_qo_len:
202
+ return 0
203
+ if not torch.all(positive_qo_len == uniform_qo_len):
204
+ return 0
205
+
206
+ zero_positions = torch.nonzero(qo_len == 0, as_tuple=False).flatten()
207
+ if zero_positions.numel() > 0:
208
+ first_zero = int(zero_positions[0].item())
209
+ if torch.any(qo_len[first_zero:] > 0):
210
+ return 0
211
+
212
+ return uniform_qo_len
213
+
214
def __init__(
215
self,
216
kv_cache_spec: AttentionSpec,
@@ -185,6 +225,18 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
225
self.compilation_config = vllm_config.compilation_config
226
self.decode_attn_out_dtype = vllm_config.model_config.dtype
227
228
+ # MTP/deepseek_mtp verification runs decode with qlen = num_spec + 1;
229
+ # any other config (including no spec) stays at single-token decode.
230
+ speculative_config = vllm_config.speculative_config
231
+ if (
232
+ speculative_config is not None
233
+ and speculative_config.method in ("mtp", "deepseek_mtp")
234
+ and speculative_config.num_speculative_tokens is not None
235
+ ):
236
+ self._mtp_decode_qlen = int(speculative_config.num_speculative_tokens) + 1
237
+ else:
238
+ self._mtp_decode_qlen = 1
239
+
240
# Store the kernel block size from the spec. When kernel_block_size=1
241
# (no spec-dec), behavior is identical to the original. When > 1
242
# (e.g. 16 with Eagle3), we expand block-level indices into per-token
@@ -222,14 +274,20 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
274
# make sure get_mla_metadata_info_v1 / get_mla_metadata_v1 are consistent
275
# with the actual tensor shape passed to mla_decode_fwd.
276
self._num_attention_heads = max(16, self.num_heads)
225
- q_dtype = self.decode_attn_out_dtype
277
kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
278
if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
279
kv_cache_dtype_str = "fp8"
229
- q_dtype = dtypes.fp8
280
+ kv_dtype = dtypes.fp8
281
else:
231
- kv_cache_dtype_str = "bf16"
232
- kv_dtype = dtypes.d_dtypes.get(kv_cache_dtype_str, dtypes.bf16)
282
+ kv_dtype = {
283
+ torch.float16: dtypes.fp16,
284
+ torch.bfloat16: dtypes.bf16,
285
+ }[kv_cache_spec.dtype]
286
+ # MLAAttention quantizes decode Q to FP8 before calling this backend
287
+ # whenever the KV cache is FP8 and supports_quant_query_input is true.
288
+ q_dtype = (
289
+ dtypes.fp8 if kv_cache_dtype_str == "fp8" else self.decode_attn_out_dtype
290
+ )
291
# Persist for get_mla_metadata_v1 (decode build): omitting these causes
292
# wrong split/reduce metadata for the gfx950 fp8 nhead=32 fold path.
293
self._mla_q_dtype = q_dtype
@@ -243,7 +301,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
301
(reduce_partial_map_size, reduce_partial_map_type),
302
) = get_mla_metadata_info_v1(
303
max_num_reqs,
246
- 1,
304
+ self._mtp_decode_qlen,
305
self._num_attention_heads,
306
q_dtype,
307
kv_dtype,
@@ -488,20 +546,39 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
546
) -> AiterMLADecodeMetadata:
547
device = self.device
548
num_reqs = seq_lens_device.size(0)
549
+ qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
550
+ max_qo_len = qo_len.max().item()
551
+ padded_mtp_qo_len = self._uniform_padded_mtp_qo_len(
552
+ qo_len, max_qo_len, num_decode_tokens
553
+ )
554
+ if padded_mtp_qo_len > 0:
555
+ max_qo_len = padded_mtp_qo_len
556
+ pad_uniform_mtp = padded_mtp_qo_len > 0
557
+
558
+ seq_lens_for_kernel = seq_lens_device
559
+ num_kernel_reqs = num_reqs
560
+ if pad_uniform_mtp:
561
+ qo_lens_device = (
562
+ query_start_loc_device[1 : num_reqs + 1]
563
+ - query_start_loc_device[:num_reqs]
564
+ ).to(torch.int32)
565
+ seq_lens_for_kernel = torch.where(
566
+ qo_lens_device > 0,
567
+ seq_lens_for_kernel,
568
+ seq_lens_for_kernel.new_full((), max_qo_len),
569
+ )
570
571
# The aiter kernel always operates with page_size=1 (the wrapper
572
# flattens kv_buffer). last_page_len is always 1.
494
- paged_kv_last_page_len = self.paged_kv_last_page_len[:num_reqs]
573
+ paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs]
574
575
# indptr: cumsum of seq_lens (one page per token in the flat view)
576
paged_kv_indptr = torch.cat(
577
[
499
- torch.zeros(1, dtype=seq_lens_device.dtype, device=device),
500
- seq_lens_device.cumsum(dim=0, dtype=torch.int32),
578
+ torch.zeros(1, dtype=torch.int32, device=device),
579
+ seq_lens_for_kernel.cumsum(dim=0, dtype=torch.int32),
580
]
581
)
503
- qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
504
- max_qo_len = qo_len.max().item()
582
use_gluon_decode = AiterMLAHelper.use_gluon_decode(
583
self.num_heads, int(max_qo_len)
584
)
@@ -519,44 +596,75 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
596
block_table_tensor,
597
block_table_tensor.stride(0),
598
paged_kv_indptr,
522
- seq_lens_device,
599
+ seq_lens_for_kernel,
600
KERNEL_BLOCK_SIZE=self.kernel_block_size,
601
BLOCK_SIZE=1024,
602
)
603
paged_kv_indices = self.paged_kv_indices
604
605
if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
529
- self.paged_kv_indptr[: 1 + num_reqs].copy_(
606
+ self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_(
607
paged_kv_indptr, non_blocking=True
608
)
532
- self.paged_kv_indptr[1 + num_reqs :].fill_(paged_kv_indptr[-1])
533
- paged_kv_indptr = self.paged_kv_indptr[: 1 + num_reqs]
609
+ self.paged_kv_indptr[1 + num_kernel_reqs :].fill_(paged_kv_indptr[-1])
610
+ paged_kv_indptr = self.paged_kv_indptr[: 1 + num_kernel_reqs]
611
612
# paged_kv_last_page_len already uses the pre-initialized buffer slice
613
# (set above), so no copy needed - buffer is always 1s.
614
538
- self.qo_indptr[: 1 + num_reqs].copy_(
539
- query_start_loc_device, non_blocking=True
615
+ if pad_uniform_mtp:
616
+ qo_indptr_src = torch.arange(
617
+ 0,
618
+ (num_kernel_reqs + 1) * max_qo_len,
619
+ step=max_qo_len,
620
+ dtype=torch.int32,
621
+ device=device,
622
+ )
623
+ else:
624
+ qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs]
625
+ self.qo_indptr[: 1 + num_kernel_reqs].copy_(
626
+ qo_indptr_src, non_blocking=True
627
)
541
- self.qo_indptr[1 + num_reqs :] = query_start_loc_device[-1]
542
- qo_indptr = self.qo_indptr[: 1 + num_reqs]
628
+ self.qo_indptr[1 + num_kernel_reqs :] = qo_indptr_src[-1]
629
+ qo_indptr = self.qo_indptr[: 1 + num_kernel_reqs]
630
631
else:
545
- qo_indptr = torch.arange(
546
- 0, num_reqs + 1, step=1, dtype=torch.int32, device=device
547
- )
548
-
549
- # The aiter MLA ASM kernel only supports qseqlen=1 (single-token
550
- # decode). With speculative decoding, the verification step has
551
- # qseqlen > 1 (e.g. 8 for spec7). get_mla_metadata_v1 calls
552
- # get_heuristic_kernel_mla which fails for qseqlen > 1.
553
- # We track whether persistent metadata was successfully computed
554
- # so forward_mqa can skip passing it (falling back to the kernel
555
- # computing its own metadata internally, like v0.18.0).
632
+ if max_qo_len == 1:
633
+ qo_indptr = torch.arange(
634
+ 0,
635
+ num_kernel_reqs + 1,
636
+ step=1,
637
+ dtype=torch.int32,
638
+ device=device,
639
+ )
640
+ else:
641
+ if pad_uniform_mtp:
642
+ qo_indptr = torch.arange(
643
+ 0,
644
+ (num_kernel_reqs + 1) * max_qo_len,
645
+ step=max_qo_len,
646
+ dtype=torch.int32,
647
+ device=device,
648
+ )
649
+ else:
650
+ qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
651
+
652
+ # Pass persistent metadata for every uniform decode we sized buffers for
653
+ # (normal qlen==1 through MTP verification qlen==K): the fp8 nhead=32 fold
654
+ # path breaks without it. qlen>K falls back to kernel-internal metadata.
655
+ # Small-head (<16) decode takes the Gluon paths and never consumes it.
656
has_persistent_metadata = False
557
- if max_qo_len == 1 and not use_gluon_decode:
657
+ use_persistent_metadata = (
658
+ self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
659
+ and max_qo_len >= 1
660
+ and max_qo_len <= self._mtp_decode_qlen
661
+ )
662
+ if use_persistent_metadata:
663
from aiter import get_mla_metadata_v1
664
665
+ uni_qo_len = (
666
+ max_qo_len if pad_uniform_mtp or torch.all(qo_len == max_qo_len) else -1
667
+ )
668
get_mla_metadata_v1(
669
qo_indptr,
670
paged_kv_indptr,
@@ -573,7 +681,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
681
page_size=1,
682
kv_granularity=16,
683
max_seqlen_qo=max_qo_len,
576
- uni_seqlen_qo=max_qo_len,
684
+ uni_seqlen_qo=uni_qo_len,
685
fast_mode=True,
686
dtype_q=self._mla_q_dtype,
687
dtype_kv=self._mla_kv_dtype,
@@ -582,7 +690,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
690
691
attn_metadata = AiterMLADecodeMetadata(
692
block_table=block_table_tensor,
585
- seq_lens=seq_lens_device,
693
+ seq_lens=seq_lens_for_kernel,
694
paged_kv_indptr=paged_kv_indptr,
695
paged_kv_indices=paged_kv_indices,
696
paged_kv_last_page_len=paged_kv_last_page_len,
@@ -1076,13 +1184,16 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
1184
dtype=attn_metadata.decode.attn_out_dtype,
1185
device=q.device,
1186
)
1187
+ if decode.max_qo_len > 1 and not decode.has_persistent_metadata:
1188
+ # MTP verification can call the AITER MLA decode kernel with
1189
+ # qlen > 1. If that path is running without persistent metadata,
1190
+ # zero-fill so unwritten lanes cannot leak into logits.
1191
+ o.zero_()
1192
1193
kv_buffer = kv_c_and_k_pe_cache.unsqueeze(2)
1194
1195
# Build kwargs for mla_decode_fwd. Pass persistent metadata only
1083
- # when it was successfully computed (qseqlen=1 decode steps).
1084
- # For multi-token verification steps (spec-dec), the kernel falls
1085
- # back to computing metadata internally.
1196
+ # when it was successfully computed.
1197
mla_kwargs = dict(
1198
q_scale=layer._q_scale,
1199
kv_scale=layer._k_scale,