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

tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py

@@ -58,6 +58,13 @@ def _on_rocm_with_aiter() -> bool:
 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 spy(**kwargs):
         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 spy(**kwargs):
     # 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 @@
 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(