[ROCm][MLA][DCP] Support causal multi-token verification (#51705)

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

tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py

@@ -58,6 +58,13 @@ PAGE_SIZE = 1
 MAX_MODEL_LEN = 1024
 MAX_NUM_SEQS = 8
 
+# The builder resolves the verify routing and records it on the metadata, so
+# the gfx950 feature probe has to be forced for the build as well as for the
+# forward; ROCm CI also runs gfx942, where the real probe is False.
+_GLUON_SUPPORTED_TARGET = (
+    "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported"
+)
+
 
 def _seq_lens() -> list[int]:
     return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS
@@ -99,6 +106,10 @@ def _run_verify_block():
     vllm_config.speculative_config = SpeculativeConfig(
         method="ngram", num_speculative_tokens=QLEN - 1
     )
+    vllm_config.model_config.get_num_attention_heads = types.MethodType(
+        lambda self, parallel_config, arch_config=None: NUM_QUERY_HEADS,
+        vllm_config.model_config,
+    )
 
     spec = MLAAttentionSpec(
         block_size=PAGE_SIZE,
@@ -138,7 +149,10 @@ def _run_verify_block():
         captured["page_table"] = kwargs["page_table"].detach().clone()
         captured["seq_info"] = kwargs["seq_info"].detach().clone()
 
-    with set_current_vllm_config(vllm_config):
+    with (
+        patch(_GLUON_SUPPORTED_TARGET, lambda: True),
+        set_current_vllm_config(vllm_config),
+    ):
         builder = builder_cls(spec, [layer_name], vllm_config, device)
         common_attn_metadata = create_common_attn_metadata(
             batch_spec, PAGE_SIZE, device, arange_block_indices=True
@@ -182,10 +196,7 @@ def _run_verify_block():
     # The Gluon kernel only reads the metadata this test is about, so a spy in
     # its place keeps the assertions independent of the AITER build.
     with (
-        patch(
-            "vllm.v1.attention.backends.mla.rocm_aiter_mla._gluon_mla_decode_supported",
-            lambda: True,
-        ),
+        patch(_GLUON_SUPPORTED_TARGET, lambda: True),
         patch(
             "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
             lambda: spy,
@@ -211,6 +222,10 @@ def test_verify_mtp_uses_native_4d_gluon_entry():
         f"expected a {QLEN}-token verify block, got max_qo_len="
         f"{decode.max_qo_len}; the MTP path under test was not reached"
     )
+    assert decode.use_gluon_verify, (
+        "the builder owns verify routing, so the metadata -- not the impl -- "
+        "has to select the Gluon MTP entry for a small-head bf16 verify block"
+    )
     assert captured, "forward_mqa did not reach the Gluon kernel"
 
     seq_lens = _seq_lens()

tests/kernels/attention/test_rocm_aiter_mla_head_padding.py

@@ -324,7 +324,9 @@ def test_h12_aiter_mla_decode_matches_reference():
 
     impl = object.__new__(AiterMLAImpl)
     impl.num_heads = NUM_HEADS
+    impl.dcp_world_size = 1
     impl.kv_lora_rank = KV_LORA_RANK
+    impl.qk_rope_head_dim = QK_ROPE_HEAD_DIM
     impl.scale = SCALE
     impl.kv_cache_dtype = "auto"
 

tests/v1/attention/test_rocm_aiter_mla_fp8_decode_routing.py

@@ -8,6 +8,10 @@ batch_size == 1, so it cannot serve a decode batch. Every fp8 shape therefore
 has to land on the asm kernels, which ship real fp8 variants for gqa=16. These
 tests pin that down at the predicates, since the failure it prevents is either a
 batch assertion or -- worse -- a silently wrong result.
+
+Gluon is still the route for supported non-DCP decode shapes. DCP multi-token
+verification uses segmented MLA instead, so it depends on neither Gluon's
+query-head nor its long-context pipeline limits.
 """
 
 import pytest
@@ -41,8 +45,8 @@ def gluon_available(monkeypatch):
 @pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES)
 @pytest.mark.parametrize("num_heads", [1, 2, 4, 5, 6, 8, 12, 16, 32, 128])
 @pytest.mark.parametrize("max_qo_len", [1, 2, 4, 5, 8, 15])
-def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
-    """No fp8 shape reaches either Gluon entry point.
+def test_non_dcp_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
+    """No non-DCP fp8 shape reaches either Gluon entry point.
 
     The head count is deliberately swept across divisors of 16 as well as
     non-divisors: the divisor case (e.g. 8 heads at TP8) is the one that stays
@@ -58,7 +62,7 @@ def test_fp8_never_routes_to_gluon(kv_cache_dtype, num_heads, max_qo_len):
 def test_fp8_never_routes_to_gluon_under_any_mode(
     monkeypatch, kv_cache_dtype, num_heads, mode
 ):
-    """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force an fp8 cache onto Gluon.
+    """VLLM_ROCM_AITER_MLA_ASM_PADDING cannot force a non-DCP fp8 cache onto Gluon.
 
     The dtype guard deliberately precedes the mode knob: honouring an explicit
     "gluon" request under fp8 would hand Gluon the batch it asserts against, so
@@ -80,6 +84,26 @@ def test_large_head_counts_never_use_gluon(kv_cache_dtype, num_heads, max_qo_len
     assert not AiterMLAHelper.use_gluon_verify(num_heads, max_qo_len, kv_cache_dtype)
 
 
+@pytest.mark.parametrize("kv_cache_dtype", FP8_DTYPES + UNQUANTIZED_DTYPES)
+@pytest.mark.parametrize("num_heads", [8, 16, 32, 96])
+@pytest.mark.parametrize("max_qo_len", [2, 3, 8])
+def test_dcp_multitoken_verify_never_uses_gluon(
+    gluon_available, kv_cache_dtype, num_heads, max_qo_len
+):
+    assert not AiterMLAHelper.use_gluon_verify(
+        num_heads, max_qo_len, kv_cache_dtype, dcp_world_size=8
+    )
+
+
+def test_segmented_dcp_verify_does_not_depend_on_gluon(monkeypatch):
+    monkeypatch.setattr(rocm_aiter_mla, "_gluon_mla_decode_supported", lambda: False)
+    monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
+
+    assert rocm_aiter_mla._segmented_dcp_verify_supported(8, 1)
+    # Round-robin interleaving other than 1 is not served by this route.
+    assert not rocm_aiter_mla._segmented_dcp_verify_supported(8, 4)
+
+
 @pytest.mark.parametrize("kv_cache_dtype", UNQUANTIZED_DTYPES)
 @pytest.mark.parametrize("num_heads", [1, 2, 4, 8])
 def test_unquantized_divisor_heads_keep_gluon_decode(

tests/v1/attention/test_rocm_aiter_mla_mtp_split.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
 import sys
 from types import SimpleNamespace
 from unittest import mock
@@ -16,8 +17,12 @@ if not current_platform.is_rocm():
 from vllm.v1.attention.backends.mla import rocm_aiter_mla  # noqa: E402
 from vllm.v1.attention.backends.mla.rocm_aiter_mla import (  # noqa: E402
     AiterMLAHelper,
+    AiterMLAImpl,
     AiterMLAMetadataBuilder,
 )
+from vllm.v1.attention.ops.rocm_aiter_mla_merge import (  # noqa: E402
+    merge_mla_segments_triton,
+)
 
 
 class _NoOpTritonKernel:
@@ -29,6 +34,19 @@ class _NoOpTritonKernel:
         pass
 
 
+def _lse_combine_natural(
+    out_a: torch.Tensor,
+    lse_a: torch.Tensor,
+    out_b: torch.Tensor,
+    lse_b: torch.Tensor,
+    out_dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    normalizer = torch.logaddexp(lse_a, lse_b)
+    out = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1)
+    out += out_b * torch.exp(lse_b - normalizer).unsqueeze(-1)
+    return out.to(out_dtype), normalizer
+
+
 class _ExpandPageIndicesKernel:
     def __getitem__(self, grid):
         self.grid = grid
@@ -67,10 +85,27 @@ def _builder(
     max_decode_rows: int = 32,
     num_heads: int = 16,
     kv_cache_dtype: str = "auto",
+    dcp_world_size: int = 1,
+    dcp_rank: int = 0,
 ):
-    return SimpleNamespace(
+    stub = SimpleNamespace(
         device=torch.device("cpu"),
         num_heads=num_heads,
+        # DCP gathers every rank's head shard before decode, and the routing
+        # predicates read the gathered count.
+        _decode_num_heads=num_heads * dcp_world_size,
+        dcp_world_size=dcp_world_size,
+        dcp_rank=dcp_rank,
+        cp_kv_cache_interleave_size=1,
+        # Mirrors the production constructor: configuration decides whether the
+        # segmented route is available, query length decides per batch.
+        _supports_segmented_dcp_verify=rocm_aiter_mla._segmented_dcp_verify_supported(
+            dcp_world_size, 1
+        ),
+        # Derived once in the real constructor, so derive it once here too.
+        _segmented_page_size=rocm_aiter_mla._segmented_mla_page_size(kernel_block_size),
+        _dcp_verify_buffers=None,
+        _graph_seq_lens=None,
         _kv_cache_dtype_str=kv_cache_dtype,
         paged_kv_last_page_len=torch.ones(max_decode_rows, dtype=torch.int32),
         paged_kv_indices=torch.empty(1024, dtype=torch.int32),
@@ -96,6 +131,15 @@ def _builder(
         _mla_kv_dtype=torch.bfloat16,
         decode_attn_out_dtype=torch.bfloat16,
     )
+    # Bound to the stub rather than faked: the verify flatten's per-row view is
+    # part of what _build_decode is being tested for.
+    stub._build_dcp_verify_row_view = (
+        AiterMLAMetadataBuilder._build_dcp_verify_row_view.__get__(stub)
+    )
+    stub._fill_dcp_verify_page_table = (
+        AiterMLAMetadataBuilder._fill_dcp_verify_page_table.__get__(stub)
+    )
+    return stub
 
 
 def test_backend_declares_uniform_batch_support():
@@ -111,6 +155,324 @@ def test_backend_declares_uniform_batch_support():
     )
 
 
+def test_dcp_verify_row_view_is_causal_per_row():
+    """Row t must cover this rank's share of positions [0, seq_len - qlen + t].
+
+    Regression guard: giving every row of a request the committed prefix only
+    drops the whole verify block, including each row's own token, since nothing
+    else adds it back on the target path.
+    """
+    qlen = 4
+    builder = _builder(mtp_decode_qlen=qlen, dcp_world_size=2, kernel_block_size=2)
+    block_table = torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32)
+
+    view = builder._build_dcp_verify_row_view(
+        qlen,
+        block_table,
+        torch.tensor([10, 12], dtype=torch.int32),
+    )
+
+    # seq_len 10 -> rows see 7, 8, 9, 10 global positions, of which rank 0 of 2
+    # holds ceil(n / 2); seq_len 12 -> 9, 10, 11, 12.
+    assert view.row_lens.tolist() == [4, 4, 5, 5, 5, 5, 6, 6]
+    assert view.qo_indptr.tolist() == list(range(2 * qlen + 1))
+    # One query per row, so the page table is the request's ascending shard and
+    # only the row length selects the causal prefix.
+    assert view.block_table.tolist() == [[7, 8, 9]] * qlen + [[10, 11, 12]] * qlen
+    assert view.max_kv_seq_len == 6
+
+
+def test_dcp_verify_rows_sum_to_the_global_window_across_ranks():
+    """Every rank's rows together must cover each row's window exactly once."""
+    qlen = 4
+    dcp_world_size = 2
+    seq_lens = torch.tensor([10, 12], dtype=torch.int32)
+    per_rank = []
+    for dcp_rank in range(dcp_world_size):
+        builder = _builder(
+            mtp_decode_qlen=qlen,
+            dcp_world_size=dcp_world_size,
+            dcp_rank=dcp_rank,
+            kernel_block_size=2,
+        )
+        view = builder._build_dcp_verify_row_view(
+            qlen,
+            torch.tensor([[7, 8, 9], [10, 11, 12]], dtype=torch.int32),
+            seq_lens,
+        )
+        per_rank.append(view.row_lens.tolist())
+
+    expected = [
+        int(seq_len) - qlen + t + 1
+        for seq_len in seq_lens.tolist()
+        for t in range(qlen)
+    ]
+    assert [sum(lens) for lens in zip(*per_rank)] == expected
+
+
+def test_dcp_verify_row_view_uses_static_graph_bound():
+    """Under full graphs the bound comes from the buffers, not the batch."""
+    qlen = 3
+    builder = _builder(
+        mtp_decode_qlen=qlen,
+        dcp_world_size=8,
+        has_full_cudagraphs=True,
+        kernel_block_size=1536,
+    )
+    builder._dcp_verify_buffers = rocm_aiter_mla.AiterMLADCPVerifyMetadata(
+        row_lens=torch.zeros(qlen, dtype=torch.int32),
+        block_table=torch.zeros((qlen, 24), dtype=torch.int32),
+        qo_indptr=torch.zeros(qlen + 1, dtype=torch.int32),
+        page_size=128,
+        max_kv_seq_len=3072,
+    )
+
+    view = builder._build_dcp_verify_row_view(
+        qlen,
+        torch.arange(2, dtype=torch.int32).view(1, 2),
+        torch.tensor([1000], dtype=torch.int32),
+    )
+
+    buffers = builder._dcp_verify_buffers
+    assert view.row_lens.tolist() == [125] * qlen
+    assert view.block_table.data_ptr() == buffers.block_table.data_ptr()
+    assert view.qo_indptr.data_ptr() == buffers.qo_indptr.data_ptr()
+    assert view.block_table[0].tolist() == list(range(24))
+    assert view.max_kv_seq_len == 3072
+
+
+def test_dcp_fp8_verify_build_uses_segmented(monkeypatch):
+    qlen = 4
+    monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
+
+    metadata = AiterMLAMetadataBuilder._build_decode(
+        _builder(
+            mtp_decode_qlen=qlen,
+            dcp_world_size=2,
+            kv_cache_dtype="fp8",
+            kernel_block_size=2,
+        ),
+        block_table_tensor=torch.tensor([[0, 1, 2], [10, 11, 12]], dtype=torch.int32),
+        seq_lens_device=torch.tensor([5, 6], dtype=torch.int32),
+        max_seq_len=6,
+        query_start_loc_cpu=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
+        query_start_loc_device=torch.tensor([0, qlen, 2 * qlen], dtype=torch.int32),
+        num_decode_tokens=2 * qlen,
+        dcp_tot_seq_lens_device=torch.tensor([10, 12], dtype=torch.int32),
+    )
+
+    assert metadata.dcp_verify is not None
+
+
+def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch):
+    """Single-token DCP decode must come back with an LSE the merge can use.
+
+    This is the path every DCP step takes outside verification. The vLLM custom
+    op drops aiter's LSE, so the impl calls aiter directly; without an LSE the
+    MLA layer's cross-rank combine asserts. The head count is deliberately a
+    non-divisor of 16 so both the output and the LSE have to be unpadded back
+    to the gathered head count.
+    """
+    num_heads, dcp_world_size = 6, 2
+    decode_heads = num_heads * dcp_world_size  # 12, padded to 16 by aiter
+    num_tokens, head_dim = 3, 576
+    captured = {}
+
+    def fake_aiter_decode(q, kv_buffer, out, *args, **kwargs):
+        captured["q_heads"] = q.shape[1]
+        captured["return_lse"] = kwargs.get("return_lse")
+        return None, torch.zeros(num_tokens, q.shape[1])
+
+    monkeypatch.setattr(
+        rocm_aiter_mla, "_get_aiter_mla_decode", lambda: fake_aiter_decode
+    )
+
+    impl = object.__new__(AiterMLAImpl)
+    impl.num_heads = num_heads
+    impl.dcp_world_size = dcp_world_size
+    impl.kv_cache_dtype = "auto"
+    impl.kv_lora_rank = 512
+    impl.qk_rope_head_dim = 64
+    impl.scale = head_dim**-0.5
+    decode = SimpleNamespace(
+        max_qo_len=1,
+        qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32),
+        paged_kv_indptr=torch.zeros(num_tokens + 1, dtype=torch.int32),
+        paged_kv_indices=torch.zeros(1, dtype=torch.int32),
+        paged_kv_last_page_len=torch.ones(num_tokens, dtype=torch.int32),
+        use_gluon_decode=False,
+        use_gluon_verify=False,
+        dcp_verify=None,
+        has_persistent_metadata=False,
+        attn_out_dtype=torch.bfloat16,
+    )
+    attn_metadata = SimpleNamespace(decode=decode, causal=True, work_meta_data=None)
+    layer = SimpleNamespace(_q_scale=torch.tensor(1.0), _k_scale=torch.tensor(1.0))
+
+    q = torch.zeros(num_tokens, decode_heads, head_dim, dtype=torch.bfloat16)
+    output, lse = impl.forward_mqa(q, torch.zeros(1, 1, head_dim), attn_metadata, layer)
+
+    assert captured["return_lse"] is True
+    # aiter is handed the padded head count, but both results come back at the
+    # gathered count the cross-rank merge expects.
+    assert captured["q_heads"] == 16
+    assert lse is not None
+    assert output.shape[1] == decode_heads
+    assert lse.shape == (num_tokens, decode_heads)
+
+
+def test_verify_partial_attention_merge():
+    device = torch.device("cuda")
+    out_a = torch.tensor([[[1.0, 2.0]], [[3.0, 4.0]]], device=device)
+    out_b = torch.tensor([[[5.0, 6.0]], [[7.0, 8.0]]], device=device)
+    lse_a = torch.tensor([[0.0], [1.0]], device=device)
+    lse_b = torch.tensor([[1.0], [0.0]], device=device)
+
+    out, lse = _lse_combine_natural(out_a, lse_a, out_b, lse_b, torch.float32)
+    normalizer = torch.logaddexp(lse_a, lse_b)
+    expected = out_a * torch.exp(lse_a - normalizer).unsqueeze(-1) + out_b * torch.exp(
+        lse_b - normalizer
+    ).unsqueeze(-1)
+
+    torch.testing.assert_close(out, expected)
+    torch.testing.assert_close(lse, normalizer)
+
+
+def test_segmented_verify_reduce_returns_natural_lse_and_masks_empty_rows():
+    segment_output = torch.tensor(
+        [1.0, 3.0, 99.0, 99.0],
+        device="cuda",
+    ).reshape(2, 1, 2, 1)
+    segment_max = torch.tensor(
+        [[[0.0, 1.0]], [[99.0, 99.0]]],
+        device="cuda",
+    )
+    segment_expsum = torch.ones_like(segment_max)
+
+    output, lse = merge_mla_segments_triton(
+        segment_output,
+        segment_max,
+        segment_expsum,
+        torch.tensor([2, 0], dtype=torch.int32, device="cuda"),
+        tile_size=1,
+        out_dtype=torch.float32,
+    )
+
+    torch.testing.assert_close(output[0], torch.tensor([[7.0 / 3]], device="cuda"))
+    assert output[1].item() == 0
+    torch.testing.assert_close(lse[0], torch.tensor([math.log(3.0)], device="cuda"))
+    assert lse[1].item() == float("-inf")
+
+
+def test_segmented_dcp_verify_matches_causal_attention(monkeypatch):
+    """Exercise the target qlen>1 path and merge two rank-local results."""
+    monkeypatch.setattr(rocm_aiter_mla, "_segmented_mla_decode_supported", lambda: True)
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    dcp_world_size = 2
+    qlen = 3
+    num_heads = 64
+    kv_lora_rank = 512
+    rope_dim = 64
+    head_dim = kv_lora_rank + rope_dim
+    global_seq_len = 259
+    block_size = 128
+    kv_scale = 0.02
+    sm_scale = head_dim**-0.5
+
+    q_nope = torch.randn(
+        qlen, num_heads, kv_lora_rank, dtype=torch.bfloat16, device=device
+    )
+    q_pe = torch.randn(qlen, num_heads, rope_dim, dtype=torch.bfloat16, device=device)
+    kv_source = torch.randn(
+        global_seq_len, head_dim, dtype=torch.float32, device=device
+    )
+    kv_fp8 = (kv_source / kv_scale).to(torch.float8_e4m3fn)
+    kv_dequant = kv_fp8.float() * kv_scale
+
+    partials = []
+    for dcp_rank in range(dcp_world_size):
+        local_kv = kv_fp8[dcp_rank::dcp_world_size]
+        num_blocks = math.ceil(local_kv.shape[0] / block_size)
+        kv_cache = torch.zeros(
+            num_blocks,
+            block_size,
+            head_dim,
+            dtype=torch.float8_e4m3fn,
+            device=device,
+        )
+        kv_cache.view(-1, head_dim)[: local_kv.shape[0]].copy_(local_kv)
+
+        builder = _builder(
+            mtp_decode_qlen=qlen,
+            kernel_block_size=block_size,
+            num_heads=num_heads // dcp_world_size,
+            kv_cache_dtype="fp8",
+            dcp_world_size=dcp_world_size,
+            dcp_rank=dcp_rank,
+        )
+        block_table = torch.arange(
+            num_blocks, dtype=torch.int32, device=device
+        ).unsqueeze(0)
+        view = builder._build_dcp_verify_row_view(
+            qlen,
+            block_table,
+            torch.tensor([global_seq_len], dtype=torch.int32, device=device),
+        )
+        decode = SimpleNamespace(
+            max_qo_len=qlen,
+            paged_kv_indptr=torch.tensor(
+                [0, local_kv.shape[0]], dtype=torch.int32, device=device
+            ),
+            paged_kv_indices=torch.arange(
+                local_kv.shape[0], dtype=torch.int32, device=device
+            ),
+            use_gluon_decode=False,
+            use_gluon_verify=False,
+            dcp_verify=view,
+            attn_out_dtype=torch.bfloat16,
+        )
+        attn_metadata = SimpleNamespace(decode=decode, causal=True)
+        impl = object.__new__(AiterMLAImpl)
+        impl.num_heads = num_heads // dcp_world_size
+        impl.dcp_world_size = dcp_world_size
+        impl.kv_cache_dtype = "fp8"
+        impl.kv_lora_rank = kv_lora_rank
+        impl.qk_rope_head_dim = rope_dim
+        impl.scale = sm_scale
+        layer = SimpleNamespace(_k_scale=torch.tensor(kv_scale, device=device))
+
+        partials.append(
+            impl.forward_mqa((q_nope, q_pe), kv_cache, attn_metadata, layer)
+        )
+
+    output, lse = partials[0]
+    assert lse is not None
+    for rank_output, rank_lse in partials[1:]:
+        assert rank_lse is not None
+        output, lse = _lse_combine_natural(
+            output.float(), lse, rank_output.float(), rank_lse, torch.float32
+        )
+
+    reference = torch.empty_like(output)
+    q_nope_fp32 = q_nope.float()
+    q_pe_fp32 = q_pe.float()
+    for query_pos in range(qlen):
+        visible = global_seq_len - qlen + query_pos + 1
+        keys = kv_dequant[:visible]
+        scores = torch.einsum(
+            "hd,nd->hn", q_nope_fp32[query_pos], keys[:, :kv_lora_rank]
+        )
+        scores += torch.einsum(
+            "hd,nd->hn", q_pe_fp32[query_pos], keys[:, kv_lora_rank:]
+        )
+        probs = torch.softmax(scores * sm_scale, dim=-1)
+        reference[query_pos] = probs @ keys[:, :kv_lora_rank]
+
+    torch.testing.assert_close(output, reference, rtol=3e-2, atol=3e-2)
+
+
 @pytest.mark.parametrize("num_heads", [8, 16, 24, 32, 64, 128])
 @pytest.mark.parametrize(
     "spec_method, parallel_drafting",
@@ -120,7 +482,7 @@ def test_backend_declares_uniform_batch_support():
         # and a parallel one, so the threshold is 1 + 2 * num_spec rather than
         # 1 + num_spec. Sizing the metadata off a method name instead leaves
         # these at qlen=1 while the router still admits the full range.
-        ("dspark", True),
+        ("custom", True),
         ("eagle", False),
     ],
 )
@@ -161,6 +523,8 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
 
     def init_common_builder(self, *args, **kwargs):
         self.num_heads = num_heads
+        self.dcp_world_size = 1
+        self.cp_kv_cache_interleave_size = 1
         # Mirror what _init_reorder_batch_threshold would have left behind: the
         # metadata is sized from the routing threshold, so a stub that skips it
         # would size for qlen=1 and hide the very mismatch this test covers.
@@ -190,12 +554,23 @@ def test_mtp_builder_init_sizes_native_fp8_metadata(
             num_speculative_tokens=3,
             parallel_drafting=parallel_drafting,
         ),
-        parallel_config=SimpleNamespace(tensor_parallel_size=8),
-        model_config=SimpleNamespace(max_model_len=16, dtype=torch.bfloat16),
+        parallel_config=SimpleNamespace(
+            tensor_parallel_size=8,
+            decode_context_parallel_size=1,
+            cp_kv_cache_interleave_size=1,
+        ),
+        model_config=SimpleNamespace(
+            max_model_len=16,
+            dtype=torch.bfloat16,
+            get_num_attention_heads=lambda parallel_config: num_heads,
+        ),
         scheduler_config=SimpleNamespace(max_num_seqs=2),
         cache_config=SimpleNamespace(cache_dtype="fp8_e4m3"),
         compilation_config=SimpleNamespace(
-            cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False)
+            cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False),
+            # Empty: the per-layer head-count probe finds no attention layer and
+            # falls back to the model config above.
+            static_forward_context={},
         ),
     )
     builder = AiterMLAMetadataBuilder(

vllm/platforms/rocm.py

@@ -900,23 +900,17 @@ class RocmPlatform(Platform):
         compilation_config = vllm_config.compilation_config
         parallel_config = vllm_config.parallel_config
 
-        if compilation_config.cudagraph_mode.has_full_cudagraphs():
-            # decode context parallel does not support full cudagraphs
-            if parallel_config.decode_context_parallel_size > 1:
-                logger.warning_once(
-                    "Decode context parallel (DCP) is enabled, which is "
-                    "incompatible with full CUDA graphs. "
-                    "Overriding cudagraph_mode to PIECEWISE."
-                )
-                compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
+        if (
+            compilation_config.cudagraph_mode.has_full_cudagraphs()
+            and parallel_config.prefill_context_parallel_size > 1
+        ):
             # prefill context parallel do not support full cudagraphs
-            elif parallel_config.prefill_context_parallel_size > 1:
-                logger.warning_once(
-                    "Prefill context parallel (PCP) is enabled, which is "
-                    "incompatible with full CUDA graphs. "
-                    "Overriding cudagraph_mode to PIECEWISE."
-                )
-                compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
+            logger.warning_once(
+                "Prefill context parallel (PCP) is enabled, which is "
+                "incompatible with full CUDA graphs. "
+                "Overriding cudagraph_mode to PIECEWISE."
+            )
+            compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
 
         if parallel_config.worker_cls == "auto":
             parallel_config.worker_cls = "vllm.v1.worker.gpu_worker.Worker"

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

@@ -11,6 +11,7 @@ import torch
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
+from vllm.distributed import get_dcp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention.mla_attention import (
     MLACommonBackend,
@@ -21,17 +22,39 @@ from vllm.model_executor.layers.attention.mla_attention import (
     QueryLenSupport,
 )
 from vllm.triton_utils import tl, triton
+from vllm.utils.math_utils import cdiv, largest_power_of_2_divisor
 from vllm.v1.attention.backend import (
     AttentionCGSupport,
     AttentionLayer,
     CommonAttentionMetadata,
     MultipleOf,
 )
+from vllm.v1.attention.backends.utils import (
+    get_dcp_local_seq_lens,
+)
+from vllm.v1.attention.ops.rocm_aiter_mla_merge import (
+    merge_mla_segments_triton,
+)
 from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache
 
 logger = init_logger(__name__)
 
 
+def _segmented_mla_page_size(block_size: int) -> int:
+    """Largest supported power-of-two subpage dividing a physical KV block.
+
+    The subpage size becomes the segmented kernel's ``TILE_SIZE``: the cache is
+    reinterpreted as ``(-1, page_size, 1, head_dim)`` before the call. It must
+    divide the physical block and be a power of two, since the kernel walks a
+    tile with ``tl.arange(0, TILE_SIZE)``. 128 is the largest tile that tiling
+    supports -- a tile holds ``TILE_SIZE x kv_lora_rank`` keys and its scores
+    are ``BLOCK_M x TILE_SIZE`` -- so a larger manager block is split into
+    several subpages instead of widening the tile.
+    """
+    assert block_size > 0
+    return min(128, largest_power_of_2_divisor(block_size))
+
+
 @functools.lru_cache(maxsize=1)
 def _get_mla_gluon():
     """Load the small-head Gluon MLA entry point."""
@@ -137,6 +160,43 @@ def _gluon_mla_decode_supported() -> bool:
     return on_gfx950()
 
 
+@functools.lru_cache(maxsize=1)
+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
+
+    return mla_decode_fwd
+
+
+@functools.lru_cache(maxsize=1)
+def _get_aiter_mla_decode():
+    from aiter.mla import mla_decode_fwd
+
+    return mla_decode_fwd
+
+
+@functools.lru_cache(maxsize=1)
+def _segmented_mla_decode_supported() -> bool:
+    """Whether AITER exposes the segmented MLA decode used by DCP verify."""
+    try:
+        _get_segmented_mla_decode()
+    except Exception:  # noqa: BLE001
+        return False
+    return True
+
+
+def _segmented_dcp_verify_supported(dcp_world_size: int, cp_interleave: int) -> bool:
+    """Whether this configuration can serve DCP verify on segmented MLA.
+
+    Configuration only -- whether a given batch takes the route additionally
+    depends on its query length. Round-robin interleaving other than 1 is
+    excluded because the per-row causal window has not been validated there.
+    """
+    return (
+        dcp_world_size > 1 and cp_interleave == 1 and _segmented_mla_decode_supported()
+    )
+
+
 def _aiter_mla_small_head_mode() -> str:
     """Small-head (<16) MLA decode kernel selection.
 
@@ -200,6 +260,31 @@ class AiterMLABackend(MLACommonBackend):
         return AiterMLAMetadataBuilder
 
 
+@dataclass
+class AiterMLADCPVerifyMetadata:
+    """One paged-KV row per verify token, for segmented DCP verification.
+
+    These are produced together or not at all, so they travel as one value: its
+    presence on the decode metadata *is* the routing decision the builder made,
+    and the impl does not re-derive it.
+    """
+
+    # Local KV length of each verify row, already causally bounded.
+    row_lens: torch.Tensor
+    # Subpage IDs each row reads, shape [num_rows, max_local_pages].
+    block_table: torch.Tensor
+    # One query per row, so this is an arange.
+    qo_indptr: torch.Tensor
+    # Subpage size the block table is expressed in, and the segmented kernel's
+    # TILE_SIZE. Carried so the impl reinterprets the cache exactly the way the
+    # block table was built, instead of re-deriving it from the cache shape.
+    page_size: int
+    # Kernel-visible KV bound. Whenever full graphs are enabled this is the
+    # configuration's maximum for every batch, not just during capture, so the
+    # page table keeps one shape across replays.
+    max_kv_seq_len: int
+
+
 @dataclass
 class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
     # The indptr of the paged kv cache, shape: [batch_size + 1]
@@ -217,8 +302,12 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata):
     max_qo_len: int | None = None
     # Minimum KV length used by Gluon to choose a safe split count.
     min_kv_seq_len: int = 1
+    # Set exactly when this batch routes to segmented DCP verification.
+    dcp_verify: AiterMLADCPVerifyMetadata | None = None
     # Small-head decode uses Gluon (avoids padding to 16).
     use_gluon_decode: bool = False
+    # Small-head non-DCP multi-token verify uses Gluon's native MTP entry.
+    use_gluon_verify: bool = False
     # Whether persistent MLA metadata was computed
     has_persistent_metadata: bool = False
 
@@ -305,13 +394,27 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         vllm_config: VllmConfig,
         device: torch.device,
     ):
-        super().__init__(
-            kv_cache_spec, layer_names, vllm_config, device, AiterMLAMetadata
+        parallel_config = vllm_config.parallel_config
+        supports_segmented_dcp_verify = _segmented_dcp_verify_supported(
+            parallel_config.decode_context_parallel_size,
+            parallel_config.cp_kv_cache_interleave_size,
         )
+        super().__init__(
+            kv_cache_spec,
+            layer_names,
+            vllm_config,
+            device,
+            AiterMLAMetadata,
+            supports_dcp_with_varlen=supports_segmented_dcp_verify,
+        )
+        self._supports_segmented_dcp_verify = supports_segmented_dcp_verify
 
         self.compilation_config = vllm_config.compilation_config
         self.decode_attn_out_dtype = vllm_config.model_config.dtype
 
+        # Needed to place a verify row's causal window on this rank's KV shard.
+        self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
+
         # reorder_batch_threshold is the largest query length decode can be
         # handed, and already accounts for the drafting scheme. A method-name
         # whitelist sizes unlisted drafters for qlen=1, which closes the
@@ -323,6 +426,7 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         # (e.g. 16 with Eagle3), we expand block-level indices into per-token
         # flat indices since the aiter kernel always uses page_size=1 internally.
         self.kernel_block_size = kv_cache_spec.block_size
+        self._segmented_page_size = _segmented_mla_page_size(self.kernel_block_size)
 
         # In the flat view (.view(-1,1,1,H)), each token is its own page,
         # so max_num_pages_per_req = max_model_len regardless of
@@ -351,10 +455,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
 
         from aiter import dtypes, get_mla_metadata_info_v1
 
+        # Decode kernels consume the DCP-gathered query heads.
+        self._decode_num_heads = self.num_heads * self.dcp_world_size
         # Keep metadata sizing consistent with the padded tensor shape passed
-        # to mla_decode_fwd.
+        # to mla_decode_fwd, including native 24-head AITER builds.
         self._num_attention_heads = AiterMLAHelper.get_actual_mla_num_heads(
-            self.num_heads
+            self._decode_num_heads
         )
         kv_cache_dtype_str = getattr(vllm_config.cache_config, "cache_dtype", "auto")
         if kv_cache_dtype_str in ("fp8", "fp8_e4m3", "fp8_e5m2"):
@@ -437,6 +543,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 device,
             )
 
+        # Persistent buffers for segmented DCP verification. Captured graphs
+        # require stable addresses while row lengths vary between replays.
+        self._dcp_verify_buffers: AiterMLADCPVerifyMetadata | None = None
+        self._graph_seq_lens: torch.Tensor | None = None
         if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
             self.paged_kv_indptr = torch.zeros(
                 max_num_reqs + 1, dtype=torch.int32, device=device
@@ -446,6 +556,39 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 max_num_reqs + 1, dtype=torch.int32, device=device
             )
 
+            # Full graphs require a stable address after uniform-MTP padding.
+            self._graph_seq_lens = torch.zeros(
+                max_num_reqs, dtype=torch.int32, device=device
+            )
+
+            if self._supports_segmented_dcp_verify and self._mtp_decode_qlen > 1:
+                # A DCP rank's shard of the longest sequence bounds every verify
+                # row, and full graphs need that bound to be constant.
+                num_dcp_partitions = (
+                    self.dcp_world_size * self.cp_kv_cache_interleave_size
+                )
+                graph_max_kv_seq_len = (
+                    cdiv(vllm_config.model_config.max_model_len, num_dcp_partitions)
+                    * self.cp_kv_cache_interleave_size
+                )
+                max_verify_rows = max_num_reqs * self._mtp_decode_qlen
+                max_local_pages = cdiv(graph_max_kv_seq_len, self._segmented_page_size)
+                self._dcp_verify_buffers = AiterMLADCPVerifyMetadata(
+                    row_lens=torch.zeros(
+                        max_verify_rows, dtype=torch.int32, device=device
+                    ),
+                    block_table=torch.zeros(
+                        (max_verify_rows, max_local_pages),
+                        dtype=torch.int32,
+                        device=device,
+                    ),
+                    qo_indptr=torch.arange(
+                        max_verify_rows + 1, dtype=torch.int32, device=device
+                    ),
+                    page_size=self._segmented_page_size,
+                    max_kv_seq_len=graph_max_kv_seq_len,
+                )
+
     def _init_fp8_prefill_ps_buffers(
         self,
         max_num_reqs: int,
@@ -667,6 +810,99 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         metadata.fp8_prefill_max_q_len = prefill.max_query_len
         metadata.fp8_prefill_num_partial_tiles = num_partial_tiles
 
+    def _build_dcp_verify_row_view(
+        self,
+        qlen: int,
+        block_table: torch.Tensor,
+        dcp_tot_seq_lens: torch.Tensor,
+    ) -> AiterMLADCPVerifyMetadata:
+        """Build one paged-KV row per verify token for segmented DCP verification.
+
+        Every row is a single query (``qo_indptr`` is an arange), so the whole
+        causal structure is carried by ``dcp_local_verify_row_lens`` and the
+        kernel applies no tail of its own.
+        """
+        assert self.dcp_world_size > 1
+        num_reqs = dcp_tot_seq_lens.numel()
+        row_lens = AiterMLAHelper.dcp_local_verify_row_lens(
+            dcp_tot_seq_lens,
+            qlen,
+            self.dcp_world_size,
+            self.dcp_rank,
+            self.cp_kv_cache_interleave_size,
+        )
+        num_rows = row_lens.numel()
+        buffers = self._dcp_verify_buffers
+        # Persistent buffers exist exactly when full graphs are on, and they
+        # already carry the static bound they were sized for.
+        max_kv_seq_len = (
+            buffers.max_kv_seq_len
+            if buffers is not None
+            else max(1, int(row_lens.max().item()))
+        )
+        page_size = self._segmented_page_size
+        max_local_pages = cdiv(max_kv_seq_len, page_size)
+        if buffers is not None:
+            row_block_table = buffers.block_table[:num_rows]
+            buffers.row_lens[:num_rows].copy_(row_lens, non_blocking=True)
+            row_lens = buffers.row_lens[:num_rows]
+            qo_indptr = buffers.qo_indptr[: num_rows + 1]
+        else:
+            row_block_table = torch.empty(
+                (num_rows, max_local_pages),
+                dtype=torch.int32,
+                device=block_table.device,
+            )
+            qo_indptr = torch.arange(
+                num_rows + 1,
+                dtype=torch.int32,
+                device=block_table.device,
+            )
+        self._fill_dcp_verify_page_table(row_block_table, block_table, num_reqs, qlen)
+        return AiterMLADCPVerifyMetadata(
+            row_lens=row_lens,
+            block_table=row_block_table,
+            qo_indptr=qo_indptr,
+            page_size=page_size,
+            max_kv_seq_len=max_kv_seq_len,
+        )
+
+    def _fill_dcp_verify_page_table(
+        self,
+        row_block_table: torch.Tensor,
+        block_table: torch.Tensor,
+        num_reqs: int,
+        qlen: int,
+    ) -> None:
+        """Expand each request's blocks into the subpages one verify row reads.
+
+        Every row of a request shares the request's shard, so the page list is
+        built once per request and repeated; only ``row_lens`` distinguishes the
+        rows. A DCP rank holds ``1/dcp_world_size`` of the sequence, so the
+        request's block table is always wider than the pages a row can reach.
+        """
+        pages_per_block = self.kernel_block_size // self._segmented_page_size
+        max_local_pages = row_block_table.shape[1]
+        max_local_blocks = cdiv(max_local_pages, pages_per_block)
+        assert max_local_blocks <= block_table.shape[1], (
+            f"DCP verify needs {max_local_blocks} blocks per request but the "
+            f"block table only has {block_table.shape[1]}"
+        )
+        subpage_offsets = torch.arange(
+            pages_per_block,
+            dtype=block_table.dtype,
+            device=block_table.device,
+        )
+        per_req_page_table = (
+            block_table[:num_reqs, :max_local_blocks, None] * pages_per_block
+            + subpage_offsets
+        ).flatten(1)[:, :max_local_pages]
+        # Broadcast the request's page list across its rows instead of
+        # materializing a repeat_interleave copy.
+        row_block_table.unflatten(0, (num_reqs, qlen)).copy_(
+            per_req_page_table.unsqueeze(1)
+        )
+
     def _build_decode(
         self,
         block_table_tensor: torch.Tensor,
@@ -701,6 +937,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
                 seq_lens_for_kernel.new_full((), max_qo_len),
             )
 
+        if self._graph_seq_lens is not None:
+            self._graph_seq_lens[:num_kernel_reqs].copy_(
+                seq_lens_for_kernel, non_blocking=True
+            )
+            seq_lens_for_kernel = self._graph_seq_lens[:num_kernel_reqs]
+
         # The aiter kernel always operates with page_size=1 (the wrapper
         # flattens kv_buffer). last_page_len is always 1.
         paged_kv_last_page_len = self.paged_kv_last_page_len[:num_kernel_reqs]
@@ -713,26 +955,43 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             ]
         )
         use_gluon_decode = AiterMLAHelper.use_gluon_decode(
-            self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
+            self._decode_num_heads,
+            int(max_qo_len),
+            self._kv_cache_dtype_str,
+        )
+        use_gluon_verify = AiterMLAHelper.use_gluon_verify(
+            self._decode_num_heads,
+            int(max_qo_len),
+            self._kv_cache_dtype_str,
+            self.dcp_world_size,
+        )
+        use_segmented_dcp_verify = (
+            self._supports_segmented_dcp_verify and max_qo_len > 1
         )
 
-        if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
-            self.paged_kv_indices.fill_(-1)
+        # Segmented DCP verify carries its own per-row subpage table, so the
+        # flat per-token view is dead work for it. Leave the buffer alone and
+        # hand the metadata None, so a future reader cannot pick up whatever
+        # the previous batch left behind.
+        paged_kv_indices = None
+        if not use_segmented_dcp_verify:
+            if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
+                self.paged_kv_indices.fill_(-1)
 
-        # Expand block_table entries into per-token flat indices.
-        # When kernel_block_size=1, this degrades to a direct copy (identical
-        # to the original _copy_page_indices_kernel).
-        # When kernel_block_size=K>1, block_table entry b covering K tokens
-        # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
-        _expand_page_indices_kernel[(num_reqs,)](
-            self.paged_kv_indices,
-            block_table_tensor,
-            block_table_tensor.stride(0),
-            paged_kv_indptr,
-            KERNEL_BLOCK_SIZE=self.kernel_block_size,
-            BLOCK_SIZE=1024,
-        )
-        paged_kv_indices = self.paged_kv_indices
+            # Expand block_table entries into per-token flat indices.
+            # When kernel_block_size=1, this degrades to a direct copy (identical
+            # to the original _copy_page_indices_kernel).
+            # When kernel_block_size=K>1, block_table entry b covering K tokens
+            # gets expanded to flat indices b*K, b*K+1, ..., b*K+(K-1).
+            _expand_page_indices_kernel[(num_reqs,)](
+                self.paged_kv_indices,
+                block_table_tensor,
+                block_table_tensor.stride(0),
+                paged_kv_indptr,
+                KERNEL_BLOCK_SIZE=self.kernel_block_size,
+                BLOCK_SIZE=1024,
+            )
+            paged_kv_indices = self.paged_kv_indices
 
         if self.compilation_config.cudagraph_mode.has_full_cudagraphs():
             self.paged_kv_indptr[: 1 + num_kernel_reqs].copy_(
@@ -745,12 +1004,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             # (set above), so no copy needed - buffer is always 1s.
 
             if pad_uniform_mtp:
-                qo_indptr_src = torch.arange(
-                    0,
-                    (num_kernel_reqs + 1) * max_qo_len,
-                    step=max_qo_len,
-                    dtype=torch.int32,
-                    device=device,
+                qo_indptr_src = AiterMLAHelper.qo_indptr_for_uniform_qlen(
+                    num_kernel_reqs, int(max_qo_len), device
                 )
             else:
                 qo_indptr_src = query_start_loc_device[: 1 + num_kernel_reqs]
@@ -762,42 +1017,32 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
 
         else:
             if max_qo_len == 1:
-                qo_indptr = torch.arange(
-                    0,
-                    num_kernel_reqs + 1,
-                    step=1,
-                    dtype=torch.int32,
-                    device=device,
+                qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
+                    num_kernel_reqs, 1, device
                 )
             else:
                 if pad_uniform_mtp:
-                    qo_indptr = torch.arange(
-                        0,
-                        (num_kernel_reqs + 1) * max_qo_len,
-                        step=max_qo_len,
-                        dtype=torch.int32,
-                        device=device,
+                    qo_indptr = AiterMLAHelper.qo_indptr_for_uniform_qlen(
+                        num_kernel_reqs, int(max_qo_len), device
                     )
                 else:
                     qo_indptr = query_start_loc_device[: 1 + num_kernel_reqs]
 
+        has_persistent_metadata = False
         # Only the asm decode consumes the schedule, so gate on the routing
         # rather than on num_heads >= 16, which denies it to a padded rank
-        # running the same asm kernels. The two predicates are disjoint --
-        # decode is qlen==1, verify is qlen>1 -- and cover both Gluon entries.
-        has_persistent_metadata = False
+        # running the same asm kernels. The predicates are disjoint -- decode
+        # is qlen==1, verify is qlen>1 -- and cover both Gluon entries plus the
+        # segmented DCP verify.
         use_persistent_metadata = (
-            not AiterMLAHelper.use_gluon_decode(
-                self.num_heads, max_qo_len, self._kv_cache_dtype_str
-            )
-            and not AiterMLAHelper.use_gluon_verify(
-                self.num_heads, max_qo_len, self._kv_cache_dtype_str
-            )
+            not use_gluon_decode
+            and not use_gluon_verify
+            and not use_segmented_dcp_verify
             # A padded rank has no bf16 persistent kernel past qlen 4 where the
             # gfx950 fold is absent; the non-persistent entry covers it. fp8
             # keeps the schedule -- its fold rejects non-persistent outright.
             and (
-                self.num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
+                self._decode_num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS
                 or max_qo_len <= AiterMLAHelper._ASM_PADDED_MAX_PS_QLEN
                 or is_quantized_kv_cache(self._kv_cache_dtype_str)
             )
@@ -840,11 +1085,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
         # active requests, not cudagraph padding rows pinned to max_qo_len.
         # The .item() below runs in the builder, outside the captured region,
         # so it does not abort HIP graph capture the way the per-layer syncs
-        # in forward_mqa did.
+        # in forward_mqa did. DCP verify never reaches here -- it carries its
+        # own per-row lengths instead of a single split bound.
         min_kv_seq_len = 1
-        if AiterMLAHelper.use_gluon_verify(
-            self.num_heads, int(max_qo_len), self._kv_cache_dtype_str
-        ):
+        if use_gluon_verify:
             per_req_len = paged_kv_indptr[1:] - paged_kv_indptr[:-1]
             if pad_uniform_mtp:
                 active = qo_lens_device > 0
@@ -853,6 +1097,15 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             else:
                 min_kv_seq_len = int(per_req_len.min().item())
 
+        dcp_verify = None
+        if use_segmented_dcp_verify:
+            assert dcp_tot_seq_lens_device is not None
+            dcp_verify = self._build_dcp_verify_row_view(
+                int(max_qo_len),
+                block_table_tensor,
+                dcp_tot_seq_lens_device,
+            )
+
         attn_metadata = AiterMLADecodeMetadata(
             block_table=block_table_tensor,
             seq_lens=seq_lens_for_kernel,
@@ -863,7 +1116,9 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]):
             dcp_tot_seq_lens=dcp_tot_seq_lens_device,
             max_qo_len=max_qo_len,
             min_kv_seq_len=min_kv_seq_len,
+            dcp_verify=dcp_verify,
             use_gluon_decode=use_gluon_decode,
+            use_gluon_verify=use_gluon_verify,
             attn_out_dtype=self.decode_attn_out_dtype,
             has_persistent_metadata=has_persistent_metadata,
         )
@@ -966,6 +1221,22 @@ class AiterMLAHelper:
     _ASM_PADDED_MAX_PS_QLEN: Final = 4
     _AITER_UNSUPPORTED_HEADS: ClassVar[tuple[int, ...]] = ()
 
+    @staticmethod
+    def qo_indptr_for_uniform_qlen(
+        num_reqs: int,
+        qlen: int,
+        device: torch.device,
+        dtype: torch.dtype = torch.int32,
+    ) -> torch.Tensor:
+        """Build ``[0, qlen, 2*qlen, ..., num_reqs*qlen]``."""
+        return torch.arange(
+            0,
+            (num_reqs + 1) * qlen,
+            step=qlen,
+            dtype=dtype,
+            device=device,
+        )
+
     @staticmethod
     def check_num_heads_validity(num_heads: int):
         assert AiterMLAHelper.is_valid_num_heads(num_heads), (
@@ -1042,14 +1313,22 @@ class AiterMLAHelper:
 
     @staticmethod
     def get_mla_unpadded_o(num_heads: int, o: torch.Tensor) -> torch.Tensor:
+        return AiterMLAHelper._get_mla_unpadded_heads(num_heads, o)
+
+    @staticmethod
+    def _get_mla_unpadded_heads(num_heads: int, tensor: torch.Tensor) -> torch.Tensor:
         m = AiterMLAHelper.get_actual_mla_num_heads(num_heads)
         if num_heads == m:
-            return o
+            return tensor
         if m % num_heads == 0:
-            return o[:, :: m // num_heads, :]
+            return tensor[:, :: m // num_heads, ...]
         # Undo the tile-padding from get_mla_padded_q: the real heads are the
         # first num_heads.
-        return o[:, :num_heads, :]
+        return tensor[:, :num_heads, ...]
+
+    @staticmethod
+    def get_mla_unpadded_lse(num_heads: int, lse: torch.Tensor) -> torch.Tensor:
+        return AiterMLAHelper._get_mla_unpadded_heads(num_heads, lse)
 
     @staticmethod
     def use_gluon_decode(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
@@ -1074,7 +1353,12 @@ class AiterMLAHelper:
         return m % num_heads == 0 and gluon_supported
 
     @staticmethod
-    def use_gluon_verify(num_heads: int, max_qo_len: int, kv_cache_dtype: str) -> bool:
+    def use_gluon_verify(
+        num_heads: int,
+        max_qo_len: int,
+        kv_cache_dtype: str,
+        dcp_world_size: int = 1,
+    ) -> bool:
         """Whether a small-head multi-token verify uses native Gluon MTP.
 
         bf16 has no gqa<16, qseqlen>1 asm kernel, so verify goes through
@@ -1083,16 +1367,67 @@ class AiterMLAHelper:
         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.
+
+        DCP verify is excluded: its per-row causal windows are served by the
+        segmented path, which Gluon's MTP entry cannot express.
         """
-        if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS or max_qo_len <= 1:
+        if max_qo_len <= 1 or dcp_world_size > 1:
             return False
         if is_quantized_kv_cache(kv_cache_dtype):
             return False
+        if num_heads >= AiterMLAHelper._AITER_MIN_MLA_HEADS:
+            return False
+        if not _gluon_mla_decode_supported():
+            return False
         # Same arch and mode gating as use_gluon_decode.
-        return _aiter_mla_small_head_mode() != "asm" and _gluon_mla_decode_supported()
+        return _aiter_mla_small_head_mode() != "asm"
+
+    @staticmethod
+    def dcp_local_verify_row_lens(
+        tot_seq_lens: torch.Tensor,
+        qlen: int,
+        dcp_world_size: int,
+        dcp_rank: int,
+        cp_interleave: int,
+    ) -> torch.Tensor:
+        """Local KV length of every verify row, in global-causal order.
+
+        Row ``t`` of a qlen-token verify attends global positions
+        ``[0, seq_len - qlen + t]``, its own token included, so its local length
+        is the round-robin count evaluated at that bound. Truncating causally in
+        global coordinates has to happen before the shard mapping: the tokens a
+        row drops sit on ``qlen - 1 - t`` different ranks, while subtracting the
+        offset from a request's local length takes one from every rank.
+        """
+        offsets = torch.arange(
+            1,
+            qlen + 1,
+            device=tot_seq_lens.device,
+            dtype=tot_seq_lens.dtype,
+        )
+        visible = (tot_seq_lens.unsqueeze(1) - qlen + offsets).clamp_(min=0)
+        return get_dcp_local_seq_lens(
+            visible,
+            dcp_world_size,
+            dcp_rank,
+            cp_interleave,
+        ).flatten()
 
 
 class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
+    # DCP decode paths return natural-log softmax LSE for the cross-rank merge.
+    can_return_lse_for_decode: bool = True
+    # Measured on gfx950: aiter.mla.mla_decode_fwd(return_lse=True) matches
+    # logsumexp to fp32 exactly, and merge_mla_segments_triton converts AITER's
+    # base-2 segment statistics with LOGE2. Stated rather than inherited because
+    # the DCP combine silently corrupts the softmax denominator if it disagrees.
+    lse_base_on_e: bool = True
+
+    @property
+    def _decode_num_heads(self) -> int:
+        """Return the query-head count after DCP gathering."""
+        return self.num_heads * self.dcp_world_size
+
     def __init__(
         self,
         num_heads: int,
@@ -1122,6 +1457,7 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
             **mla_args,
         )
         AiterMLAHelper.check_num_heads_validity(num_heads)
+        AiterMLAHelper.check_num_heads_validity(self._decode_num_heads)
 
         unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]
         if any(unsupported_features):
@@ -1363,6 +1699,59 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
 
         self._mla_fp8_prefill_attn(q, k, v, attn_metadata, output)
 
+    def _forward_segmented_dcp_verify(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        verify: AiterMLADCPVerifyMetadata,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        layer: AttentionLayer,
+        out_dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Run segmented attention over this rank's shard of every verify row.
+
+        Each row's length already covers the current tokens this rank holds, so
+        the cross-rank LSE merge in the MLA layer completes the causal block.
+        """
+        q_mla = torch.cat([q_nope, q_pe], dim=-1)
+        # skip_reduce with NUM_SEGMENTS>1 returns the partials and never writes
+        # `out`. Pass None rather than aliasing q_mla: if that ever stops
+        # holding, this fails loudly instead of scribbling over the query.
+        segment_partials = _get_segmented_mla_decode()(
+            q_mla,
+            kv_c_and_k_pe_cache.view(
+                -1,
+                verify.page_size,
+                1,
+                kv_c_and_k_pe_cache.shape[-1],
+            ),
+            None,
+            verify.qo_indptr,
+            verify.row_lens,
+            verify.max_kv_seq_len,
+            verify.block_table,
+            self.scale,
+            self.kv_lora_rank,
+            self.qk_rope_head_dim,
+            causal=True,
+            q_descale=None,
+            kv_descale=layer._k_scale,
+            skip_reduce=True,
+        )
+        assert isinstance(segment_partials, tuple) and len(segment_partials) == 3, (
+            "AITER segmented MLA verify must return segment partials "
+            "when skip_reduce=True."
+        )
+        segm_output, segm_max, segm_expsum = segment_partials
+        return merge_mla_segments_triton(
+            segm_output,
+            segm_max,
+            segm_expsum,
+            verify.row_lens,
+            verify.page_size,
+            out_dtype,
+        )
+
     def forward_mqa(
         self,
         q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
@@ -1375,9 +1764,9 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
 
         decode = attn_metadata.decode
         assert decode.max_qo_len is not None
-        assert decode.paged_kv_indptr is not None
-        assert decode.paged_kv_indices is not None
         if decode.use_gluon_decode:
+            assert decode.paged_kv_indptr is not None
+            assert decode.paged_kv_indices is not None
             if type(q) is tuple:
                 q_nope, q_pe = q
             else:
@@ -1416,12 +1805,11 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
         # The block is causal -- the target is checking draft tokens, so
         # position t must not see t+1 -- and the kernel bounds each query
         # position's scores itself.
-        # Arch, mode and dtype gating all live in use_gluon_verify, so that the
-        # builder -- which has to know whether the asm decode will run -- sees
-        # the same answer as this branch.
-        if AiterMLAHelper.use_gluon_verify(
-            self.num_heads, int(decode.max_qo_len), self.kv_cache_dtype
-        ):
+        # Arch, mode and dtype gating all live in use_gluon_verify, which the
+        # builder evaluates and records on the metadata, so this branch acts on
+        # the same answer the builder used to decide whether the asm decode
+        # would run. DCP verify is routed to the segmented path below instead.
+        if decode.use_gluon_verify:
             if type(q) is tuple:
                 q_nope, q_pe = q
             else:
@@ -1473,14 +1861,48 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
             )
             return o, None
 
+        verify = decode.dcp_verify
+        if verify is not None:
+            if type(q) is tuple:
+                q_nope, q_pe = q
+            else:
+                q_nope, q_pe = torch.split(
+                    q, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
+                )
+            if (
+                is_quantized_kv_cache(self.kv_cache_dtype)
+                and q_nope.dtype != torch.bfloat16
+            ):
+                q_nope = q_nope.to(torch.bfloat16) * layer._q_scale
+                q_pe = q_pe.to(torch.bfloat16) * layer._q_scale
+            return self._forward_segmented_dcp_verify(
+                q_nope,
+                q_pe,
+                verify,
+                kv_c_and_k_pe_cache,
+                layer,
+                decode.attn_out_dtype,
+            )
+
+        if self.dcp_world_size > 1 and int(decode.max_qo_len) > 1:
+            raise RuntimeError(
+                "ROCM_AITER_MLA DCP multi-token verify requires segmented MLA."
+            )
+
         if type(q) is tuple:
             q = torch.cat(q, dim=-1)
 
         assert isinstance(q, torch.Tensor)
+        assert decode.paged_kv_indptr is not None
+        assert decode.paged_kv_indices is not None
         B = q.shape[0]
 
-        mla_padded_q = AiterMLAHelper.get_mla_padded_q(self.num_heads, q)
-        mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self.num_heads)
+        assert q.shape[1] == self._decode_num_heads, (
+            "ROCM_AITER_MLA decode expected the DCP-gathered query head count "
+            f"{self._decode_num_heads}, got {q.shape[1]}"
+        )
+        mla_padded_q = AiterMLAHelper.get_mla_padded_q(self._decode_num_heads, q)
+        mla_num_heads = AiterMLAHelper.get_actual_mla_num_heads(self._decode_num_heads)
         o = torch.empty(
             B,
             mla_num_heads,
@@ -1512,17 +1934,43 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):
                 reduce_partial_map=attn_metadata.reduce_partial_map,
             )
 
-        rocm_aiter_ops.mla_decode_fwd(
-            mla_padded_q,
-            kv_buffer,
-            o,
-            self.scale,
-            decode.qo_indptr,
-            decode.max_qo_len,
-            decode.paged_kv_indptr,
-            decode.paged_kv_indices,
-            decode.paged_kv_last_page_len,
-            **mla_kwargs,
-        )
+        lse = None
+        if self.dcp_world_size > 1:
+            # The vLLM custom-op wrapper exposes only the in-place output and
+            # drops aiter's final LSE, which the cross-shard merge needs, so go
+            # through aiter's native entry point on the DCP path.
+            _, lse = _get_aiter_mla_decode()(
+                mla_padded_q,
+                kv_buffer.view(-1, 1, 1, mla_padded_q.shape[-1]),
+                o,
+                decode.qo_indptr,
+                decode.paged_kv_indptr,
+                decode.paged_kv_indices,
+                decode.paged_kv_last_page_len,
+                decode.max_qo_len,
+                sm_scale=self.scale,
+                return_lse=True,
+                **mla_kwargs,
+            )
+            assert lse is not None, (
+                "aiter mla_decode_fwd(return_lse=True) returned no LSE; upgrade "
+                "aiter to a build with decode LSE support."
+            )
+        else:
+            rocm_aiter_ops.mla_decode_fwd(
+                mla_padded_q,
+                kv_buffer,
+                o,
+                self.scale,
+                decode.qo_indptr,
+                decode.max_qo_len,
+                decode.paged_kv_indptr,
+                decode.paged_kv_indices,
+                decode.paged_kv_last_page_len,
+                **mla_kwargs,
+            )
 
-        return AiterMLAHelper.get_mla_unpadded_o(self.num_heads, o), None
+        output = AiterMLAHelper.get_mla_unpadded_o(self._decode_num_heads, o)
+        if lse is not None:
+            lse = AiterMLAHelper.get_mla_unpadded_lse(self._decode_num_heads, lse)
+        return output, lse

vllm/v1/attention/ops/rocm_aiter_mla_merge.py

@@ -0,0 +1,136 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Merge AITER segmented MLA split-K partials into output plus natural-log LSE.
+
+This mirrors AITER's own segment reduction, including its
+``tiles_per_segment = cdiv(seq_len, NUM_SEGMENTS * TILE_SIZE)`` partitioning, so
+it has to move together with the ``skip_reduce=True`` call in the AITER MLA
+backend. It is a rank-local split-K merge and unrelated to any collective
+reduce; the natural-log LSE it returns is what the cross-rank DCP merge
+consumes.
+"""
+
+import torch
+
+from vllm.triton_utils import LOGE2, tl, triton
+
+
+@triton.jit
+def _merge_mla_segments_kernel(
+    out_ptr,
+    lse_ptr,
+    segm_output_ptr,
+    segm_max_ptr,
+    segm_expsum_ptr,
+    seq_lens_ptr,
+    num_query_heads: tl.constexpr,
+    out_stride0: tl.int64,
+    out_stride1: tl.int64,
+    lse_stride0: tl.int64,
+    TILE_SIZE: tl.constexpr,
+    KV_LORA_RANK: tl.constexpr,
+    NUM_SEGMENTS_PER_SEQ: tl.constexpr,
+    LOGE2: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    head_idx = tl.program_id(1)
+
+    seq_len = tl.load(seq_lens_ptr + token_idx)
+    tiles_per_segment = tl.maximum(
+        1, tl.cdiv(seq_len, NUM_SEGMENTS_PER_SEQ * TILE_SIZE)
+    )
+    active_segments = tl.where(
+        seq_len > 0,
+        tl.cdiv(seq_len, tiles_per_segment * TILE_SIZE),
+        0,
+    )
+    segment_offsets = tl.arange(0, NUM_SEGMENTS_PER_SEQ)
+    segment_mask = segment_offsets < active_segments
+
+    segment_base = (
+        token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ)
+        + head_idx * NUM_SEGMENTS_PER_SEQ
+    )
+    segment_indices = segment_base + segment_offsets
+    segment_max = tl.load(
+        segm_max_ptr + segment_indices,
+        mask=segment_mask,
+        other=float("-inf"),
+    )
+    overall_max = tl.max(segment_max)
+
+    segment_expsum = tl.load(
+        segm_expsum_ptr + segment_indices,
+        mask=segment_mask,
+        other=0.0,
+    )
+    segment_weight = tl.where(
+        segment_mask,
+        tl.exp2(segment_max - overall_max),
+        0.0,
+    )
+    overall_expsum = tl.sum(segment_expsum * segment_weight)
+
+    output_offsets = (
+        token_idx.to(tl.int64) * (num_query_heads * NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK)
+        + head_idx * (NUM_SEGMENTS_PER_SEQ * KV_LORA_RANK)
+        + segment_offsets[:, None] * KV_LORA_RANK
+        + tl.arange(0, KV_LORA_RANK)[None, :]
+    )
+    segment_output = tl.load(
+        segm_output_ptr + output_offsets,
+        mask=segment_mask[:, None],
+        other=0.0,
+    )
+    accumulator = tl.sum(segment_output * segment_weight[:, None], axis=0)
+    accumulator = tl.where(overall_expsum == 0.0, 0.0, accumulator / overall_expsum)
+    lse = tl.where(
+        overall_expsum == 0.0,
+        float("-inf"),
+        (overall_max + tl.log2(overall_expsum)) * LOGE2,
+    )
+
+    output_indices = (
+        token_idx * out_stride0 + head_idx * out_stride1 + tl.arange(0, KV_LORA_RANK)
+    )
+    tl.store(out_ptr + output_indices, accumulator)
+    tl.store(lse_ptr + token_idx * lse_stride0 + head_idx, lse)
+
+
+def merge_mla_segments_triton(
+    segm_output: torch.Tensor,
+    segm_max: torch.Tensor,
+    segm_expsum: torch.Tensor,
+    seq_lens: torch.Tensor,
+    tile_size: int,
+    out_dtype: torch.dtype,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Merge AITER base-2 segment partials into output and natural-log LSE."""
+    num_tokens, num_heads, num_segments, kv_lora_rank = segm_output.shape
+    output = torch.empty(
+        (num_tokens, num_heads, kv_lora_rank),
+        dtype=out_dtype,
+        device=segm_output.device,
+    )
+    lse = torch.empty(
+        (num_tokens, num_heads),
+        dtype=torch.float32,
+        device=segm_output.device,
+    )
+    _merge_mla_segments_kernel[(num_tokens, num_heads)](
+        output,
+        lse,
+        segm_output,
+        segm_max,
+        segm_expsum,
+        seq_lens,
+        num_query_heads=num_heads,
+        out_stride0=output.stride(0),
+        out_stride1=output.stride(1),
+        lse_stride0=lse.stride(0),
+        TILE_SIZE=tile_size,
+        KV_LORA_RANK=kv_lora_rank,
+        NUM_SEGMENTS_PER_SEQ=num_segments,
+        LOGE2=LOGE2,
+    )
+    return output, lse