[ROCm][MLA][DCP] Support causal multi-token verification by YukioZzz · Pull Request #51705 · vllm-project/vllm

YukioZzz · 2026-08-10T16:40:18Z

Summary

Enable ROCm AITER MLA decode context parallelism for causal multi-token target verification. Each verification token gets its correct per-rank causal KV window, and rank-local partial outputs are merged through the existing DCP attention merge.

This PR is scoped to the ROCm AITER target backend. Hybrid prefix-cache geometry and external connector integration remain separate changes.

Implementation

Compute each verification row's global causal boundary and map it to the rank-local KV length.
Expand physical KV blocks into the bounded subpages consumed by segmented AITER MLA.
Run segmented MLA with unreduced partial output, then merge its base-2 statistics into BF16 output plus natural-log LSE for the cross-rank merge.
Make the metadata builder the single owner of decode routing. Presence of an AiterDCPVerifyView selects segmented DCP verification; the implementation consumes that decision without recomputing it.
Allocate stable row metadata and use a maximum DCP-local KV bound during FULL graph capture. Replay still masks inactive segments using actual row lengths.
Keep single-token decode on the existing Gluon/ASM selection and keep DCP multi-token verification out of those paths.
Allow ROCm DCP configurations to retain FULL CUDA graphs; PCP remains on the existing PIECEWISE fallback.

Validation

Focused ROCm tests

Revision: e1843114b7c233a9c71ad44b28bf63426ad64836

540 passed, 14 warnings in 31.56s

The suite includes an actual FP8 segmented-kernel correctness test with DCP=2, qlen=3, and 96 gathered query heads. It builds each rank's causal row view through AiterMLAMetadataBuilder, computes both rank-local partials through AiterMLAImpl.forward_mqa(), merges them using their returned LSE, and compares the result against global causal attention.

It also covers:

per-row causal DCP lengths and cross-rank coverage;
block-to-subpage expansion and bounds;
empty-row and natural-log LSE reduction;
static FULL-graph metadata bounds;
FP8 single-token decode routing;
DCP multi-token routing to segmented MLA rather than Gluon/ASM.

Negative control: the causal-reference test was rerun after changing only the row-boundary calculation back to the previous committed-prefix-only form. The fixed implementation passed. The old calculation failed deterministically with 2,671/147,456 output elements outside tolerance and a maximum absolute difference of 0.1311. This test-only mutation was not committed.

Full-model regression

Configuration:

Image: vllm/vllm-openai-rocm:nightly
PR revision: e1843114b7c233a9c71ad44b28bf63426ad64836
Base: 2c7d7dd64a2eaba0feedf42cab2f527486d7479c
Real Kimi-K3 target weights
TP8 / DCP8, A2A DCP, KV interleave 1
FP8 KV cache
maximum model length 1,048,576
FULL_AND_PIECEWISE graphs, capture sizes 1 through 32
maximum batch size 32
prefix caching enabled

Full lm-eval GSM8K, 5-shot, 1,319 samples, concurrency 64, default generation length (max_length=2047):

strict-match exact_match:     0.9613343442 +/- 0.0053105832
flexible-extract exact_match: 0.9613343442 +/- 0.0053105832
evaluation time:              335 seconds
server errors:                0
server tracebacks:            0
HSA memory-access faults:     0

The server selected ROCM_AITER_MLA, completed both PIECEWISE and FULL graph capture, exposed 27,012,229 KV-cache tokens, and remained healthy after the evaluation. The staged runtime files were checked against a SHA256 manifest in the container, and the server log recorded the exact PR revision above.

This target-only run validates the DCP8/FP8-KV full-model accuracy baseline. The causal multi-token target path is covered directly by the focused causal-reference test.

Stacked target-verification run

PR 54546 at 8474084088b12b7bc3ca35cf6d74f87cb7161f92 was cleanly stacked on this PR head to enable Triton MLA for the non-causal draft group. The target retained the same TP8/DCP8, FP8-KV, 1M, FULL graph, and maximum-batch-32 configuration. The speculative configuration used two draft tokens, probabilistic draft sampling, block rejection sampling, and Triton MLA for the draft.

combined focused tests:       543 passed, 14 warnings in 31.75s
GSM8K samples:                1319/1319 at concurrency 64
strict-match exact_match:     0.9628506444 +/- 0.0052095163
flexible-extract exact_match: 0.9636087945 +/- 0.0051581135
evaluation time:              353 seconds
mean acceptance length:       2.67-2.74 in the final 10 windows
server errors/tracebacks:     0/0
HSA memory-access faults:     0

This run exercises the ROCm AITER segmented causal target-verification path from this PR together with PR 54546's non-causal draft capability. The server completed FULL graph capture and remained healthy after evaluation.

AI assistance

This change was prepared with AI assistance (Codex). The human submitter has reviewed every changed line and is responsible for defending it.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

shen-shanshan

AMD CI passed (http://buildkite.com/vllm/amd-ci/builds/12452/canvas), let's run the full CI tests.

shen-shanshan · 2026-08-31T02:18:18Z

/ci run

github-actions · 2026-08-31T02:18:43Z

✅ Triggered Buildkite CI #86273 for commit cabd41db68c1.

shen-shanshan · 2026-08-31T03:20:25Z

/ci retry

github-actions · 2026-08-31T03:20:42Z

✅ Queued 1 failed job(s) for retry in Buildkite CI #86273.

billishyahao · 2026-08-31T08:04:21Z

CI passed https://buildkite.com/vllm/ci/builds/86273/list

billishyahao

Please address the comments. Thank you!

billishyahao · 2026-08-31T08:11:32Z

@@ -0,0 +1,128 @@
+# SPDX-License-Identifier: Apache-2.0


can we fold this one segment triton kernel into common dcp specific util file? e.g. https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/ops/dcp.py

billishyahao · 2026-08-31T08:15:37Z

    # Non-causal DSpark block is flattened to one decode row per query token in
    # forward_mqa, so no intra-block causal masking is required.
    supports_non_causal_multi_token_decode: ClassVar[bool] = True
+    supports_non_causal_multi_token_dcp: ClassVar[bool] = True


Will this new flag bring unexpected regression to other platform? Please double check

billishyahao · 2026-08-31T08:20:06Z

+def _segmented_mla_page_size(block_size: int) -> int:
+    """Largest supported power-of-two subpage dividing a physical KV block."""
+    assert block_size > 0
+    return min(128, largest_power_of_2_divisor(block_size))


Why limit this page size up to 128?

billishyahao · 2026-08-31T08:21:19Z

+def _get_segmented_mla_decode():
+    """Load AITER's segmented MLA decode with unreduced partial output."""
+    from aiter.ops.triton.attention.mla import mla_decode_fwd
+


eliminate blank row

Checked. In order to maintain the same format pattern, a blank row will remain after the inner import.

billishyahao · 2026-08-31T08:21:26Z

+@functools.lru_cache(maxsize=1)
+def _get_aiter_mla_decode():
+    from aiter.mla import mla_decode_fwd
+


billishyahao · 2026-08-31T08:25:36Z

+        kv_cache_dtype: str,
+        dcp_world_size: int = 1,
+    ) -> bool:
+        """Whether non-DCP multi-token verification uses Gluon."""


restore function docstring:

""" Whether a small-head multi-token verify uses native Gluon MTP. bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through ``mla_gluon``'s 4-D MTP entry (``q`` shaped ``[batch, qlen, nhead, dim]``) with ``use_2d_view=False``. fp8 has one via the q-row fold and must not come here: the MTP path hands Gluon the batch size its fp8 regime asserts against. A predicate rather than inline in forward_mqa so the builder sees the same answer the impl acts on. """

Done. Restored and updated.

billishyahao · 2026-08-31T08:30:30Z

+        skip_paged_kv_expand = use_segmented_dcp_verify
+
+        if not skip_paged_kv_expand:


why not use use_segmented_dcp_verify instead?

Removed. Redundant code introduced from reconstruction.

billishyahao · 2026-08-31T08:43:55Z

+        pages_per_block: int,
+        max_local_pages: int,
+    ) -> None:
+        if block_table.is_cuda:


is this always true?

Keep only one path.

YukioZzz · 2026-08-31T11:57:36Z

/amd-ci run

github-actions · 2026-08-31T11:57:54Z

✅ Triggered Buildkite AMD CI #12473 for commit e1843114b7c2.

YukioZzz · 2026-08-31T12:17:01Z

/ci run

github-actions · 2026-08-31T12:17:20Z

✅ Triggered Buildkite CI #86360 for commit e1843114b7c2.

shen-shanshan · 2026-08-31T13:16:46Z

/ci retry

github-actions · 2026-08-31T13:17:04Z

✅ Queued 1 failed job(s) for retry in Buildkite CI #86360.

Co-authored-by: andyluo7 <andy.luo@amd.com> Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

Signed-off-by: billishyahao <bill.he@amd.com>

Split the Triton MLA non-causal multi-token DCP capability into its own PR: it is the only part of this change that touches a backend shared with CUDA and XPU. This PR keeps only the ROCm AITER causal target verification path. Give every decode route a single owner. The builder decides, the decode metadata carries the decision, and the impl acts on it: - segmented DCP verify is recorded as an AiterDCPVerifyView, whose four tensors were four independent optional fields that were always set or cleared together; its presence is the routing decision, so five is-not-None assertions and the impl-side predicate all go away; - use_gluon_verify joins use_gluon_decode on the metadata instead of being recomputed in the impl from a different dtype source; - configuration support for the segmented route is one predicate evaluated once per builder, with the interleave restriction inside it rather than only in the constructor, replacing a query-length probe that passed a placeholder 2. Collapse the DCP verify page-table fill onto one implementation. It was split into a Triton kernel and a torch fallback selected by block_table.is_cuda, so unit tests exercised the fallback while production ran the kernel, and the kernel did not bound its block-table load. The torch expansion runs on both devices and states the block count invariant it relies on. Stop aliasing the query as the segmented kernel's unused output pointer, so a build that does write it fails loudly instead of scribbling over q. Build the verify row lengths through dcp_local_verify_row_lens instead of a second inline copy, and state on AiterMLAImpl that its decode LSE is natural-log: aiter mla_decode_fwd(return_lse=True) matches logsumexp on gfx950, and the segment merge converts AITER's base-2 statistics. Rename that merge to merge_mla_segments_triton, in module rocm_aiter_mla_merge, so it is not read as a collective reduce. Also revert refactors unrelated to DCP verify (the use_gluon_decode rewrite and the num_kernel_reqs rename), restore the comments and docstrings they dropped, inline skip_paged_kv_expand, and explain the segmented subpage size bound. Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Assisted-by: Codex

The small-head verify routing decision now lives in the metadata builder instead of AiterMLAImpl.forward_mqa, so _gluon_mla_decode_supported is probed during build(). The test patched it only around forward_mqa, so on gfx942 the builder recorded use_gluon_verify=False and forward_mqa fell through to the asm decode path, which dereferences layer._q_scale on the layer=None handed in by the test. Hoist the probe patch so the build and the forward see the same answer, and assert the metadata flag directly so the build-time contract is pinned rather than incidentally satisfied. Forcing the probe also makes the build architecture independent: use_persistent_metadata is False on gfx942 as well now, matching gfx950. Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

YukioZzz · 2026-08-31T15:38:33Z

/ci run

github-actions · 2026-08-31T15:38:53Z

✅ Triggered Buildkite CI #86381 for commit 6d1bcafa7000.

Pin the small-head Gluon verify test to the intended 12-head metadata shape instead of relying on a partial HF config override that still leaves the builder at the full model head count. Use 64 query heads in the segmented DCP verify causal-reference test. The test still exercises DCP > 1 and qlen > 1 against causal attention, while avoiding a gfx942 Triton launch that requires 128 KiB of shared memory on hardware capped at 64 KiB. Co-authored-by: andyluo7 <andy.luo@amd.com> Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

YukioZzz · 2026-08-31T17:20:50Z

/ci run

github-actions · 2026-08-31T17:21:11Z

✅ Triggered Buildkite CI #86397 for commit a6615cd07d50.

GirasoleY

Thanks for the iteration! LGTM

claude Bot reviewed Aug 10, 2026

View reviewed changes

mergify Bot added kimi k3 rocm Related to AMD ROCm mrv2 Model Runner V2 specific labels Aug 10, 2026

github-project-automation Bot added this to AMD Aug 10, 2026

github-project-automation Bot moved this to Todo in AMD Aug 10, 2026

shen-shanshan approved these changes Aug 31, 2026

View reviewed changes

billishyahao suggested changes Aug 31, 2026

View reviewed changes

YukioZzz mentioned this pull request Aug 31, 2026

[ROCm][MLA][DCP] Advertise Triton MLA non-causal multi-token DCP #54546

Open

YukioZzz and others added 5 commits August 31, 2026 15:27

[ROCm][MLA][DCP] Support causal multi-token verification

13b7b40

Co-authored-by: andyluo7 <andy.luo@amd.com> Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

[MLA][DCP] Advertise Triton non-causal multi-token support

659ff0d

Co-authored-by: andyluo7 <andy.luo@amd.com> Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>

fix test failure

58ec0ae

Signed-off-by: billishyahao <bill.he@amd.com>

GirasoleY approved these changes Aug 31, 2026

View reviewed changes

ywang96 approved these changes Aug 31, 2026

View reviewed changes

okorzh-amd mentioned this pull request Aug 31, 2026

[AMD][DCP] Close two decode-LSE holes left by AITER DCP support #54639

Closed

4 tasks

This was referenced Sep 14, 2026

[ROCm][MLA] Add an AITER ASM round-robin decode route for DCP multi-token verify #56861

Open

[Bugfix][Spec Decode][DCP] Align ranks before DFlash speculator graph capture #56869

Draft

seungrokj mentioned this pull request Sep 15, 2026

[ROCm][MLA][DCP] Enable ROCM_AITER_MLA non-causal DSpark draft under DCP #57085

Open

		skip_paged_kv_expand = use_segmented_dcp_verify

		if not skip_paged_kv_expand:

Uh oh!

Uh oh!

Conversation

YukioZzz commented Aug 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation

Validation

Focused ROCm tests

Full-model regression

Stacked target-verification run

AI assistance

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

shen-shanshan left a comment

Choose a reason for hiding this comment

Uh oh!

shen-shanshan commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

shen-shanshan commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

billishyahao commented Aug 31, 2026

Uh oh!

billishyahao left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

YukioZzz commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

YukioZzz commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

shen-shanshan commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

YukioZzz commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

YukioZzz commented Aug 31, 2026

Uh oh!

github-actions Bot commented Aug 31, 2026

Uh oh!

GirasoleY left a comment

YukioZzz commented Aug 10, 2026 •

edited

Loading