@tduncan / Nvidia-Optimization-Stack-VLLM / commits / dbf662c9e8

[ROCm][MLA] Reach FULL cudagraphs for AITER MLA speculative decoding (#51171)

Signed-off-by: yudigege86 <naqin@amd.com> Co-authored-by: yudigege86 <naqin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

Matt Qin committed Aug 29, 2026 at 21:49 UTC dbf662c9e810ef4c3f466c7c4b2b3282ba889622

4 files changed +181 -127

tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py

+48 -71

@@ -1,21 +1,17 @@

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.

3 +"""Regression test for causal masking in the AITER MLA verify MTP path.

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.

6 +kernel for a multi-token block, so ``AiterMLAImpl.forward_mqa`` routes verify

7 +through ``mla_gluon``'s native 4-D MTP entry (``q`` shaped

8 +``[batch, qlen, nhead, dim]`` with ``use_2d_view=False``). aiter applies the

9 +causal tail internally: position ``t`` attends KV ``[0, seq_len - qlen + t]``.

10 +

11 +The test drives the real builder and ``forward_mqa`` over a multi-token decode

12 +batch, intercepts the arguments actually handed to the Gluon kernel, and pins

13 +that the 4-D MTP contract is used with per-request ``block_table`` and

14 +``cache_seqlens`` rather than a flattened per-token paged-KV view.

15 """

17 import types

@@ -32,15 +28,15 @@ def _on_rocm_with_aiter() -> bool:

28 return current_platform.is_rocm() and is_aiter_found()

35 -# The production verify flatten is gfx950-only. This test replaces the Gluon

31 +# The production verify MTP path is gfx950-only. This test replaces the Gluon

32 # kernel with a spy and forces the architecture feature probe so the metadata

33 # regression remains testable on every ROCm AITER CI runner.

34 pytestmark = pytest.mark.skipif(

35 not _on_rocm_with_aiter(),

40 - reason="ROCM_AITER_MLA verify flatten requires ROCm and AITER",

36 + reason="ROCM_AITER_MLA verify MTP requires ROCm and AITER",

37 )

43 -# Below the 16-head threshold that selects the flatten. DeepSeek-V3 / R1 at

39 +# Below the 16-head threshold that selects the MTP path. DeepSeek-V3 / R1 at

40 # TP=16 and Kimi-K3 at TP=8 both land here.

41 NUM_QUERY_HEADS = 12

42 KV_LORA_RANK = 512

@@ -54,8 +50,7 @@ QLEN = 8

50 # Committed context per request, including a fresh request at zero context

51 # where the whole KV range is the verify block itself.

52 CONTEXT_LENS = [0, 1, 37, 512]

57 -# A cudagraph padding request: seq_len 0, so every one of its rows must clamp

58 -# to an empty window.

53 +# A cudagraph padding request: seq_len 0.

54 PADDING_ROWS = 1

56 # The Gluon path flattens the KV cache to one page per token.

@@ -68,16 +63,10 @@ def _seq_lens() -> list[int]:

63 return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS

71 -def _expected_row_lens() -> list[int]:

72 - """Causal row lengths: row r*QLEN + t sees seq_len_r - (QLEN - 1) + t."""

73 - return [max(0, s - (QLEN - 1) + t) for s in _seq_lens() for t in range(QLEN)]

74 -

75 -

66 def _run_verify_block():

67 """Drive the real builder + forward_mqa over one multi-token verify block.

79 - Returns ``(metadata, captured)`` where ``captured`` holds the ``page_table``,

80 - ``seq_info`` and ``min_kv_seq_len`` the backend passed to the Gluon kernel.

69 + Returns ``(metadata, captured)`` where ``captured`` holds the Gluon kwargs.

70 """

71 from tests.v1.attention.utils import (

72 BatchSpec,

@@ -144,9 +133,10 @@ def _run_verify_block():

133 captured: dict = {}

134

135 def spy(**kwargs):

136 + captured.update(kwargs)

137 + captured["q_nope"] = kwargs["q_nope"].detach().clone()

138 captured["page_table"] = kwargs["page_table"].detach().clone()

139 captured["seq_info"] = kwargs["seq_info"].detach().clone()

149 - captured["min_kv_seq_len"] = kwargs["min_kv_seq_len"]

140

141 with set_current_vllm_config(vllm_config):

142 builder = builder_cls(spec, [layer_name], vllm_config, device)

@@ -206,12 +196,12 @@ def _run_verify_block():

196 return metadata, captured

197

198

209 -def test_verify_flatten_rows_are_causal():

210 - """Verify row t must see the committed prefix plus tokens 0..t, and no more.

199 +def test_verify_mtp_uses_native_4d_gluon_entry():

200 + """Verify uses 4-D q with per-request paged_kv metadata.

201

212 - Regression guard: before the fix every row of a request got that request's

213 - whole paged-KV range, so verify position t attended to the draft tokens

214 - after it. Fails on unmodified upstream.

202 + Regression guard: before the fix every flattened row of a request got that

203 + request's whole paged-KV range, so verify position t attended to the draft

204 + tokens after it. The native MTP path delegates causal masking to aiter.

205 """

206 metadata, captured = _run_verify_block()

207

@@ -219,50 +209,37 @@ def test_verify_flatten_rows_are_causal():

209 assert decode is not None, "batch was not classified as a decode"

210 assert decode.max_qo_len == QLEN, (

211 f"expected a {QLEN}-token verify block, got max_qo_len="

222 - f"{decode.max_qo_len}; the flatten under test was not reached"

212 + f"{decode.max_qo_len}; the MTP path under test was not reached"

213 )

214 assert captured, "forward_mqa did not reach the Gluon kernel"

215

216 seq_lens = _seq_lens()

227 - indptr = captured["seq_info"].tolist()

228 - page_table = captured["page_table"]

229 - got_row_lens = [indptr[i + 1] - indptr[i] for i in range(len(indptr) - 1)]

230 - want_row_lens = _expected_row_lens()

231 -

232 - assert got_row_lens == want_row_lens, (

233 - "AITER MLA verify flatten is not causal: verify row r*qlen+t must get "

234 - f"seq_len_r - {QLEN - 1} + t KV entries.\n"

235 - f" seq_lens {seq_lens}\n"

236 - f" got {got_row_lens}\n"

237 - f" expected {want_row_lens}\n"

238 - "Rows longer than expected let a verify position attend to the draft "

239 - "tokens it is supposed to be checking."

240 - )

217 + num_reqs = len(seq_lens)

218

242 - # arange_block_indices lays request r's pages out ascending from

243 - # r * max_blocks, so each causal window must be that slice's prefix.

244 - max_blocks = max(seq_lens)

245 - for r, seq_len in enumerate(seq_lens):

246 - for t in range(QLEN):

247 - row = r * QLEN + t

248 - window = page_table[indptr[row] : indptr[row + 1]].tolist()

249 - n = max(0, seq_len - (QLEN - 1) + t)

250 - want = [r * max_blocks + j for j in range(n)]

251 - assert window == want, (

252 - f"request {r} (seq_len {seq_len}) verify token {t} reads KV "

253 - f"pages {window}, expected the ascending causal prefix {want}"

254 - )

255 -

256 - padding_rows = got_row_lens[len(CONTEXT_LENS) * QLEN :]

257 - assert padding_rows == [0] * (PADDING_ROWS * QLEN), (

258 - f"cudagraph padding requests (seq_len 0) must clamp to empty windows, "

259 - f"got {padding_rows}"

219 + q_nope = captured["q_nope"]

220 + assert q_nope.dim() == 4, (

221 + f"verify must pass 4-D q to mla_gluon, got {q_nope.dim()}-D shape "

222 + f"{tuple(q_nope.shape)}"

223 )

224 + assert q_nope.shape[:2] == (num_reqs, QLEN), (

225 + f"expected q shape [num_reqs={num_reqs}, qlen={QLEN}, ...], "

226 + f"got {tuple(q_nope.shape)}"

227 + )

228 + assert q_nope.shape[2] == NUM_QUERY_HEADS

229

262 - # min_kv_seq_len tells Gluon how short the shortest row is, so it must be

263 - # the minimum over the causal rows actually submitted, not over the

264 - # per-request lengths they were cut from.

265 - assert captured["min_kv_seq_len"] == min(want_row_lens), (

266 - f"min_kv_seq_len={captured['min_kv_seq_len']} does not match the "

267 - f"shortest submitted row ({min(want_row_lens)})"

230 + assert captured["use_2d_view"] is False

231 + page_table = captured["page_table"]

232 + seq_info = captured["seq_info"]

233 + assert page_table.dim() == 1, (

234 + f"verify MTP passes the 1-D paged_kv_indices buffer, got {page_table.dim()}-D"

235 + )

236 + assert seq_info.dim() == 1 and seq_info.numel() == num_reqs + 1, (

237 + f"verify MTP passes paged_kv_indptr [num_reqs + 1], got shape "

238 + f"{tuple(seq_info.shape)}"

239 + )

240 + assert decode is not None and decode.paged_kv_indptr is not None

241 + assert torch.equal(seq_info, decode.paged_kv_indptr[: num_reqs + 1])

242 + assert torch.equal(

243 + page_table[: int(seq_info[-1].item())],

244 + decode.paged_kv_indices[: int(seq_info[-1].item())],

245 )

tests/v1/attention/test_rocm_aiter_mla_mtp_split.py

+48 -2

@@ -40,7 +40,6 @@ class _ExpandPageIndicesKernel:

40 block_table_tensor,

41 stride,

42 paged_kv_indptr,

43 - seq_lens_for_kernel,

43 *,

44 KERNEL_BLOCK_SIZE,

45 BLOCK_SIZE,

@@ -48,7 +47,9 @@ class _ExpandPageIndicesKernel:

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())

51 - seq_len = int(seq_lens_for_kernel[req_idx].item())

50 + seq_len = int(

51 + (paged_kv_indptr[req_idx + 1] - paged_kv_indptr[req_idx]).item()

52 + )

53 for token_idx in range(seq_len):

54 block_id = int(

55 block_table_tensor[req_idx, token_idx // KERNEL_BLOCK_SIZE].item()

@@ -249,6 +250,51 @@ def test_mtp_decode_qlen4_keeps_uniform_rows_with_metadata(monkeypatch):

250 assert get_mla_metadata_v1.call_args.kwargs["uni_seqlen_qo"] == 4

251

252

253 +def test_min_kv_seq_len_ignores_cudagraph_padding_rows(monkeypatch):

254 + """min_kv_seq_len must not be driven by padded dummy requests.

255 +

256 + Full-CG verify pins zero-qo rows to max_qo_len for paged_kv metadata; taking

257 + per_req_len.min() over all rows then reports qlen instead of a real KV length

258 + and collapses Gluon KV-split parallelism on single-request decode.

259 + """

260 + monkeypatch.setitem(

261 + sys.modules,

262 + "aiter",

263 + SimpleNamespace(get_mla_metadata_v1=mock.MagicMock()),

264 + )

265 + monkeypatch.setattr(

266 + rocm_aiter_mla, "_expand_page_indices_kernel", _NoOpTritonKernel()

267 + )

268 + monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: True)

269 + monkeypatch.setattr(rocm_aiter_mla, "_aiter_mla_small_head_mode", lambda: "auto")

270 +

271 + mtp_qlen = 8

272 + num_reqs = 8

273 + active_seq_len = 1032

274 + seq_lens = torch.tensor([active_seq_len] + [0] * (num_reqs - 1), dtype=torch.int32)

275 + query_start_loc = torch.tensor(

276 + [0, mtp_qlen] + [mtp_qlen] * (num_reqs - 1), dtype=torch.int32

277 + )

278 +

279 + metadata = AiterMLAMetadataBuilder._build_decode(

280 + _builder(

281 + mtp_decode_qlen=mtp_qlen,

282 + has_full_cudagraphs=True,

283 + max_decode_rows=num_reqs,

284 + num_heads=12,

285 + ),

286 + block_table_tensor=torch.zeros(num_reqs, active_seq_len, dtype=torch.int32),

287 + seq_lens_device=seq_lens,

288 + max_seq_len=active_seq_len,

289 + query_start_loc_cpu=query_start_loc,

290 + query_start_loc_device=query_start_loc,

291 + num_decode_tokens=num_reqs * mtp_qlen,

292 + dcp_tot_seq_lens_device=None,

293 + )

294 +

295 + assert metadata.min_kv_seq_len == active_seq_len

296 +

297 +

298 def test_full_cudagraph_padded_uniform_mtp_synthesizes_decode_indptr(

299 monkeypatch,

300 ):

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

+55 -54

@@ -729,7 +729,6 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

729 block_table_tensor,

730 block_table_tensor.stride(0),

731 paged_kv_indptr,

732 - seq_lens_for_kernel,

732 KERNEL_BLOCK_SIZE=self.kernel_block_size,

733 BLOCK_SIZE=1024,

734 )

@@ -834,6 +833,26 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

833 )

834 has_persistent_metadata = True

835

836 + # Small-head multi-token verify uses mla_gluon's 4-D MTP entry over the

837 + # ordinary per-request paged-KV view, so there is no expanded per-token

838 + # buffer to build here. mla_gluon still wants a lower bound on the KV

839 + # length it is asked to split; on the verify path that bound is over

840 + # active requests, not cudagraph padding rows pinned to max_qo_len.

841 + # The .item() below runs in the builder, outside the captured region,

842 + # so it does not abort HIP graph capture the way the per-layer syncs

843 + # in forward_mqa did.

844 + min_kv_seq_len = 1

845 + if AiterMLAHelper.use_gluon_verify(

846 + self.num_heads, int(max_qo_len), self._kv_cache_dtype_str

847 + ):

848 + per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1]

849 + if pad_uniform_mtp:

850 + active = qo_lens_device > 0

851 + if active.any():

852 + min_kv_seq_len = int(per_req_len[active].min().item())

853 + else:

854 + min_kv_seq_len = int(per_req_len.min().item())

855 +

856 attn_metadata = AiterMLADecodeMetadata(

857 block_table=block_table_tensor,

858 seq_lens=seq_lens_for_kernel,

@@ -843,6 +862,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):

862 qo_indptr=qo_indptr,

863 dcp_tot_seq_lens=dcp_tot_seq_lens_device,

864 max_qo_len=max_qo_len,

865 + min_kv_seq_len=min_kv_seq_len,

866 use_gluon_decode=use_gluon_decode,

867 attn_out_dtype=self.decode_attn_out_dtype,

868 has_persistent_metadata=has_persistent_metadata,

@@ -884,7 +904,6 @@ def _expand_page_indices_kernel(

904 block_table,

905 block_table_stride,

906 cu_num_tokens,

887 - seq_lens,

907 KERNEL_BLOCK_SIZE: tl.constexpr,

908 BLOCK_SIZE: tl.constexpr,

909 ):

@@ -904,7 +923,7 @@ def _expand_page_indices_kernel(

923 req_idx = tl.program_id(0)

924 row_ptr = block_table + req_idx * block_table_stride

925 start_idx = tl.load(cu_num_tokens + req_idx)

907 - num_tokens = tl.load(seq_lens + req_idx)

926 + num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx

927

928 offset = tl.arange(0, BLOCK_SIZE)

929 for i in tl.range(0, num_tokens, BLOCK_SIZE):

@@ -1056,11 +1075,12 @@ class AiterMLAHelper:

1075

1076 @staticmethod

1077 def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:

1059 - """Whether a small-head multi-token verify is flattened onto Gluon.

1078 + """Whether a small-head multi-token verify uses native Gluon MTP.

1079

1061 - bf16 has no gqa<16, qseqlen>1 asm kernel, so the verify is flattened

1062 - into per-token Gluon decodes. fp8 has one via the q-row fold and must

1063 - not come here: the flatten hands Gluon the batch size its fp8 regime

1080 + bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through

1081 + ``mla_gluon``'s 4-D MTP entry (``q`` shaped ``[batch, qlen, nhead, dim]``)

1082 + with ``use_2d_view=False``. fp8 has one via the q-row fold and must not

1083 + come here: the MTP path hands Gluon the batch size its fp8 regime

1084 asserts against. A predicate rather than inline in forward_mqa so the

1085 builder sees the same answer the impl acts on.

1086 """

@@ -1391,19 +1411,17 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1411 return o, None

1412

1413 # 12-head (<16) multi-token verify (DSpark): the asm path has no

1394 - # gqa<16, qseqlen>1 kernel. Flatten each verify token to its own

1395 - # qseqlen=1 gluon decode, mirroring the TRITON_MLA / sparse-backend

1396 - # flatten but on the fast gluon kernel. The block is causal -- the

1397 - # target is checking draft tokens, so position t must not see t+1 --

1398 - # and attention rows are independent, so giving row t the KV range

1399 - # [0, context + t] is exactly causal multi-token attention.

1414 + # gqa<16, qseqlen>1 kernel, so the block goes to the gluon kernel's 4-D

1415 + # MTP entry, which serves a whole (1 + num_spec) block in one launch.

1416 + # The block is causal -- the target is checking draft tokens, so

1417 + # position t must not see t+1 -- and the kernel bounds each query

1418 + # position's scores itself.

1419 # Arch, mode and dtype gating all live in use_gluon_verify, so that the

1420 # builder -- which has to know whether the asm decode will run -- sees

1421 # the same answer as this branch.

1422 if AiterMLAHelper.use_gluon_verify(

1423 self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype

1424 ):

1406 - qlen = int(decode.max_qo_len)

1425 if type(q) is tuple:

1426 q_nope, q_pe = q

1427 else:

@@ -1419,56 +1437,39 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1437 device=q_nope.device,

1438 )

1439 kv_buffer = kv_c_and_k_pe_cache.reshape(-1, kv_c_and_k_pe_cache.shape[-1])

1422 - # Expand per-request paged-KV to per-verify-token. Row r*qlen+t is

1423 - # request r's verify token t, and seq_lens counts the tokens

1424 - # scheduled in this step, so a request's KV range already spans its

1425 - # whole verify block and context_r = seq_len_r - qlen. Token t may

1426 - # attend to [0, context_r + t], i.e. seq_len_r - (qlen - 1) + t

1427 - # entries. paged_kv_indices lists a request's pages in ascending

1428 - # position order, so each row's causal window is a prefix of that

1429 - # request's slice and only the row length changes. Rows clamp to

1430 - # zero for cudagraph padding requests, whose seq_len is 0. Fully

1431 - # vectorized (no host loop).

1432 - old_indptr = decode.paged_kv_indptr

1433 - per_req_len = old_indptr[1:] - old_indptr[:-1]

1434 - dev = q_nope.device

1435 - row_req = torch.arange(per_req_len.shape[0], device=dev).repeat_interleave(

1436 - qlen

1440 + assert attn_metadata.causal, (

1441 + "AITER MLA small-head verify MTP is causal-only"

1442 )

1438 - row_len = (

1439 - (

1440 - per_req_len.unsqueeze(1)

1441 - - (qlen - 1)

1442 - + torch.arange(qlen, device=dev, dtype=per_req_len.dtype)

1443 + # Hand mla_gluon its 4-D MTP entry instead of an expanded

1444 + # per-verify-token paged-KV view. The flat query layout is already

1445 + # row-major (r * qlen + t), so unflatten is a free view. mla_gluon

1446 + # applies the per-position causal bound

1447 + # score_end = min(split_kv_end, seq_len - qlen + q_pos + 1)

1448 + # in-kernel, and seq_lens already spans the verify block, so that

1449 + # bound is context_r + q_pos + 1 -- the same window the expanded

1450 + # view supplied by truncating each row's page list.

1451 + qlen = int(decode.max_qo_len)

1452 + num_reqs = B // qlen

1453 + if num_reqs * qlen != B:

1454 + raise ValueError(

1455 + f"verify block {B} rows is not a multiple of qlen {qlen}"

1456 )

1444 - .clamp_(min=0)

1445 - .flatten()

1446 - )

1447 - new_indptr = torch.cat([old_indptr.new_zeros(1), row_len.cumsum(0)]).to(

1448 - torch.int32

1449 - )

1450 - total = int(new_indptr[-1].item())

1451 - within = torch.arange(total, device=dev, dtype=torch.int64) - new_indptr[

1452 - :-1

1453 - ].to(torch.int64).repeat_interleave(row_len)

1454 - src = (

1455 - old_indptr[row_req].to(torch.int64).repeat_interleave(row_len) + within

1456 - )

1457 - new_indices = decode.paged_kv_indices[src]

1457 + assert decode.paged_kv_indptr is not None

1458 + assert decode.paged_kv_indices is not None

1459 mla_gluon = _get_mla_gluon()

1460 mla_gluon(

1460 - q_nope=q_nope,

1461 - q_pe=q_pe,

1461 + q_nope=q_nope.unflatten(0, (num_reqs, qlen)),

1462 + q_pe=q_pe.unflatten(0, (num_reqs, qlen)),

1463 kv_c=kv_buffer,

1463 - o=o,

1464 - page_table=new_indices,

1465 - seq_info=new_indptr,

1464 + o=o.unflatten(0, (num_reqs, qlen)),

1465 + page_table=decode.paged_kv_indices,

1466 + seq_info=decode.paged_kv_indptr,

1467 sm_scale=self.scale,

1468 k_pe=None,

1469 kv_pe_offset=self.kv_lora_rank,

1470 use_2d_view=False,

1471 kv_scale=1.0,

1471 - min_kv_seq_len=int(row_len.min()),

1472 + min_kv_seq_len=decode.min_kv_seq_len,

1473 )

1474 return o, None

1475

vllm/v1/attention/backends/mla/triton_mla.py

+30

@@ -6,6 +6,7 @@ from typing import ClassVar

6 import torch

8 import vllm.envs as envs

9 +from vllm.config import VllmConfig

10 from vllm.config.cache import CacheDType

11 from vllm.logger import init_logger

12 from vllm.model_executor.layers.attention.mla_attention import (

@@ -25,6 +26,7 @@ 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

30 from vllm.v1.worker.workspace import (

31 current_workspace_manager,

32 is_workspace_manager_initialized,

@@ -55,6 +57,34 @@ class TritonMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

57 # forward_mqa, so no intra-block causal masking is required.

58 supports_non_causal_multi_token_decode: ClassVar[bool] = True

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 +

88 def __init__(self, kv_cache_spec, layer_names, vllm_config, device):

89 super().__init__(kv_cache_spec, layer_names, vllm_config, device)

90 # DCP local sequence lengths are not advanced between draft steps.