[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"