[ROCm][DSv4] Enable DSpark adaptive verification (#52362)

Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
Signed-off-by: larryli2-amd <larryli2@amd.com>
Co-authored-by: larryli2-amd <larryli2@amd.com>

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -188,6 +188,7 @@ def _ref_sparse_decode_ragged(
     extra_rows: list[list[int]] | None = None,
     main_use_fnuz: bool = False,
     extra_use_fnuz: bool = False,
+    extra_block_size: int | None = None,
 ) -> torch.Tensor:
     q_f32 = q.float()
     out = torch.empty_like(q_f32)
@@ -209,7 +210,10 @@ def _ref_sparse_decode_ragged(
             )
             row_kv.append(
                 _read_fp8_ds_mla_cache_rows(
-                    extra_cache, extra_slots, block_size, extra_use_fnuz
+                    extra_cache,
+                    extra_slots,
+                    extra_block_size or block_size,
+                    extra_use_fnuz,
                 )
             )
 
@@ -244,6 +248,12 @@ def _ragged_from_rows(
     )
 
 
+def _rows_from_ragged(indices: torch.Tensor, indptr: torch.Tensor) -> list[list[int]]:
+    ends = indptr.cpu().tolist()
+    values = indices[: ends[-1]].cpu().tolist()
+    return [values[start:end] for start, end in zip(ends, ends[1:])]
+
+
 def _launch_sparse_decode_reduce(
     part_m: torch.Tensor,
     part_l: torch.Tensor,
@@ -1192,6 +1202,254 @@ def test_sparse_attn_decode_gfx950_graph_replay(monkeypatch) -> None:
     torch.testing.assert_close(out, expected_long, atol=2e-2, rtol=2e-2)
 
 
+@requires_gfx950
+@torch.inference_mode()
+def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
+    """Replay the production ROCm decode path after device-only reallocation."""
+    from tests.v1.attention.utils import create_vllm_config
+    from vllm.models.deepseek_v4.amd.rocm import (
+        DeepseekV4ROCMAiterMLASparseMetadataBuilder,
+        DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+    )
+    from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
+    from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec
+
+    device = torch.device("cuda")
+    num_reqs = 3
+    upper_query_len = 8
+    graph_tokens = num_reqs * upper_query_len
+    block_size = 256
+    compressed_block_size = block_size // 128
+    window_size = 32
+
+    vllm_config = create_vllm_config(
+        model_name="facebook/opt-125m",
+        max_model_len=1024,
+        block_size=block_size,
+        max_num_seqs=num_reqs,
+        max_num_batched_tokens=graph_tokens,
+        hf_config_override={
+            "compress_ratios": [128],
+            "index_topk": 2048,
+            "sliding_window": window_size,
+        },
+    )
+    vllm_config.speculative_config = SimpleNamespace(
+        num_speculative_tokens=upper_query_len - 1,
+        parallel_drafting=False,
+        enable_adaptive_verification=True,
+        use_dspark=lambda: True,
+    )
+    mla_spec = MLAAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=HEAD_DIM,
+        dtype=torch.bfloat16,
+        tokens_per_state=128,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+    )
+    swa_spec = SlidingWindowMLASpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=HEAD_DIM,
+        dtype=torch.bfloat16,
+        sliding_window=window_size,
+        cache_dtype_str="fp8_ds_mla",
+        model_version="deepseek_v4",
+    )
+    mla_builder = DeepseekV4ROCMAiterMLASparseMetadataBuilder(
+        mla_spec, ["c128a"], vllm_config, device
+    )
+    swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
+        swa_spec, ["c128a"], vllm_config, device
+    )
+    assert (
+        mla_builder.get_cudagraph_support(vllm_config, mla_spec)
+        == AttentionCGSupport.ALWAYS
+    )
+    assert (
+        swa_builder.get_cudagraph_support(vllm_config, swa_spec)
+        == AttentionCGSupport.ALWAYS
+    )
+
+    seq_lens = torch.tensor([520, 528, 536], dtype=torch.int32, device=device)
+    seq_lens_cpu = seq_lens.cpu()
+    cpu_query_start_loc = torch.arange(
+        0,
+        graph_tokens + 1,
+        upper_query_len,
+        dtype=torch.int32,
+    )
+    max_blocks = (int(seq_lens_cpu.max()) + block_size - 1) // block_size
+    block_table = torch.arange(
+        num_reqs * max_blocks, dtype=torch.int32, device=device
+    ).view(num_reqs, max_blocks)
+
+    def build_metadata(query_lens: list[int]):
+        query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device)
+        query_start_loc = torch.cat(
+            [
+                torch.zeros(1, dtype=torch.int32, device=device),
+                query_lens_tensor.cumsum(0),
+            ]
+        )
+        active_tokens = sum(query_lens)
+        positions = torch.zeros(graph_tokens, dtype=torch.int64, device=device)
+        position_rows = [
+            torch.arange(
+                seq_len - query_len,
+                seq_len,
+                dtype=torch.int64,
+                device=device,
+            )
+            for seq_len, query_len in zip(seq_lens_cpu.tolist(), query_lens)
+        ]
+        positions[:active_tokens] = torch.cat(position_rows)
+        slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device)
+        slot_mapping[:active_tokens] = torch.arange(
+            active_tokens, dtype=torch.int64, device=device
+        )
+        common = CommonAttentionMetadata(
+            query_start_loc=query_start_loc,
+            query_start_loc_cpu=cpu_query_start_loc,
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=seq_lens_cpu,
+            num_reqs=num_reqs,
+            num_actual_tokens=graph_tokens,
+            max_query_len=upper_query_len,
+            max_seq_len=int(seq_lens_cpu.max()),
+            block_table_tensor=block_table,
+            slot_mapping=slot_mapping,
+            positions=positions,
+            causal=True,
+        )
+        return (
+            mla_builder.build_for_cudagraph_capture(common),
+            swa_builder.build_for_cudagraph_capture(common),
+        )
+
+    mla_metadata, swa_metadata = build_metadata([8, 8, 8])
+    assert mla_metadata.for_cudagraph_capture
+    metadata_ptrs = (
+        mla_metadata.c128a_decode_topk_ragged_indices.data_ptr(),
+        mla_metadata.c128a_decode_topk_ragged_indptr.data_ptr(),
+        swa_metadata.decode_swa_ragged_indices.data_ptr(),
+        swa_metadata.decode_swa_ragged_indptr.data_ptr(),
+    )
+
+    torch.manual_seed(19)
+    num_heads = 16
+    q = (
+        torch.randn(
+            graph_tokens,
+            num_heads,
+            HEAD_DIM,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        * 0.125
+    )
+    num_cache_blocks = num_reqs * max_blocks
+    swa_cache = _pack_fp8_ds_mla_cache(
+        torch.randn(
+            num_cache_blocks * block_size,
+            HEAD_DIM,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        * 0.125,
+        block_size,
+        use_fnuz=False,
+    )
+    compressed_cache = _pack_fp8_ds_mla_cache(
+        torch.randn(
+            num_cache_blocks * compressed_block_size,
+            HEAD_DIM,
+            dtype=torch.bfloat16,
+            device=device,
+        )
+        * 0.125,
+        compressed_block_size,
+        use_fnuz=False,
+    )
+    attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
+    out = torch.empty_like(q)
+    monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
+
+    def run_decode(mla_md, swa_md) -> None:
+        mod.rocm_sparse_attn_decode(
+            q=q,
+            kv_cache=compressed_cache,
+            swa_k_cache=swa_cache,
+            swa_only=False,
+            topk_indices=mla_md.c128a_global_decode_topk_indices,
+            topk_lens=mla_md.c128a_decode_topk_lens,
+            swa_indices=swa_md.decode_swa_indices,
+            swa_lens=swa_md.decode_swa_lens,
+            swa_ragged_indices=swa_md.decode_swa_ragged_indices,
+            swa_ragged_indptr=swa_md.decode_swa_ragged_indptr,
+            topk_ragged_indices=mla_md.c128a_decode_topk_ragged_indices,
+            topk_ragged_indptr=mla_md.c128a_decode_topk_ragged_indptr,
+            attn_sink=attn_sink,
+            scale=HEAD_DIM**-0.5,
+            head_dim=HEAD_DIM,
+            nope_head_dim=NOPE_HEAD_DIM,
+            rope_head_dim=ROPE_HEAD_DIM,
+            output=out,
+            extra_cache_nan_free=True,
+            adaptive_splits=True,
+        )
+
+    def reference(mla_md, swa_md) -> torch.Tensor:
+        return _ref_sparse_decode_ragged(
+            q=q,
+            main_cache=swa_cache,
+            main_rows=_rows_from_ragged(
+                swa_md.decode_swa_ragged_indices,
+                swa_md.decode_swa_ragged_indptr,
+            ),
+            scale=HEAD_DIM**-0.5,
+            attn_sink=attn_sink,
+            block_size=block_size,
+            extra_cache=compressed_cache,
+            extra_rows=_rows_from_ragged(
+                mla_md.c128a_decode_topk_ragged_indices,
+                mla_md.c128a_decode_topk_ragged_indptr,
+            ),
+            extra_block_size=compressed_block_size,
+        )
+
+    run_decode(mla_metadata, swa_metadata)
+    torch.accelerator.synchronize()
+    expected_full = reference(mla_metadata, swa_metadata)
+    torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2)
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        run_decode(mla_metadata, swa_metadata)
+    torch.accelerator.synchronize()
+    captured_full = out.clone()
+
+    reallocated_mla, reallocated_swa = build_metadata([3, 8, 5])
+    assert metadata_ptrs == (
+        reallocated_mla.c128a_decode_topk_ragged_indices.data_ptr(),
+        reallocated_mla.c128a_decode_topk_ragged_indptr.data_ptr(),
+        reallocated_swa.decode_swa_ragged_indices.data_ptr(),
+        reallocated_swa.decode_swa_ragged_indptr.data_ptr(),
+    )
+    expected_reallocated = reference(reallocated_mla, reallocated_swa)
+    run_decode(reallocated_mla, reallocated_swa)
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
+    assert not torch.equal(captured_full, expected_reallocated)
+
+    graph.replay()
+    torch.accelerator.synchronize()
+    torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
+
+
 # ---------------------------------------------------------------------------
 # o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
 # ---------------------------------------------------------------------------

tests/models/test_deepseek_v4_dspark_rocm.py

@@ -0,0 +1,130 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.amd import dspark as dspark_module
+from vllm.models.deepseek_v4.amd.dspark import DSparkDeepseekV4ForCausalLM
+
+
+def _make_uninitialized_model(confidence_head):
+    model = DSparkDeepseekV4ForCausalLM.__new__(DSparkDeepseekV4ForCausalLM)
+    object.__setattr__(
+        model,
+        "model",
+        SimpleNamespace(confidence_head=confidence_head),
+    )
+    return model
+
+
+def _prepare_loader_model(model, named_parameters):
+    object.__setattr__(
+        model,
+        "config",
+        SimpleNamespace(
+            n_routed_experts=1,
+            expert_dtype="fp4",
+            num_attention_heads=1,
+        ),
+    )
+    object.__setattr__(model, "named_parameters", lambda: named_parameters)
+
+
+def _disable_distributed_loader_paths(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        dspark_module,
+        "fused_moe_make_expert_params_mapping",
+        lambda *args, **kwargs: [],
+    )
+    monkeypatch.setattr(
+        dspark_module, "get_tensor_model_parallel_world_size", lambda: 1
+    )
+    monkeypatch.setattr(dspark_module, "get_tensor_model_parallel_rank", lambda: 0)
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_maps_enabled_confidence_head():
+    model = _make_uninitialized_model(object())
+
+    assert (
+        model._remap_dspark_name("mtp.2.confidence_head.proj.weight")
+        == "model.confidence_head.proj.weight"
+    )
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_skips_disabled_confidence_head():
+    model = _make_uninitialized_model(None)
+
+    assert model._remap_dspark_name("mtp.2.confidence_head.proj.weight") is None
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_disables_unloaded_confidence_head(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    model = _make_uninitialized_model(object())
+    _prepare_loader_model(model, [])
+    _disable_distributed_loader_paths(monkeypatch)
+
+    assert model.load_weights([]) == set()
+    assert model.model.confidence_head is None
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_loads_complete_confidence_head(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class FakeParameter:
+        loaded_weight = None
+
+        def weight_loader(self, param, loaded_weight):
+            assert param is self
+            self.loaded_weight = loaded_weight
+
+    confidence_head = object()
+    parameter = FakeParameter()
+    model = _make_uninitialized_model(confidence_head)
+    _prepare_loader_model(
+        model,
+        [("model.confidence_head.proj.weight", parameter)],
+    )
+    _disable_distributed_loader_paths(monkeypatch)
+    loaded_weight = torch.tensor([[1.0, 2.0]])
+
+    loaded = model.load_weights([("mtp.2.confidence_head.proj.weight", loaded_weight)])
+
+    assert loaded == {"model.confidence_head.proj.weight"}
+    assert parameter.loaded_weight is loaded_weight
+    assert model.model.confidence_head is confidence_head
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_confidence_is_probability():
+    class ConfidenceHead:
+        def __call__(self, head_hidden, markov_embed):
+            return (head_hidden[:, 0] + markov_embed[:, 0]).float()
+
+    model = _make_uninitialized_model(ConfidenceHead())
+    head_hidden = torch.tensor([[0.0], [1.0]], dtype=torch.bfloat16)
+    markov_embed = torch.tensor([[0.0], [-2.0]], dtype=torch.bfloat16)
+
+    confidence = model.compute_confidence(head_hidden, markov_embed)
+
+    torch.testing.assert_close(
+        confidence,
+        torch.sigmoid(torch.tensor([0.0, -1.0])),
+    )
+    assert torch.all((confidence >= 0) & (confidence <= 1))
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_dspark_confidence_requires_a_head():
+    model = _make_uninitialized_model(None)
+    empty = torch.zeros((1, 1))
+
+    with pytest.raises(RuntimeError, match="confidence_head"):
+        model.compute_confidence(empty, empty)

tests/v1/attention/test_deepseek_v4_rocm_adaptive.py

@@ -0,0 +1,249 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.models.deepseek_v4.amd.rocm import (
+    DeepseekV4ROCMAiterMLASparseMetadataBuilder,
+    DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+)
+from vllm.v1.attention.backend import AttentionCGSupport
+from vllm.v1.attention.backends.mla import indexer
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV4IndexerBackend,
+    DeepseekV32IndexerMetadataBuilder,
+    DeepseekV41IndexerBackend,
+)
+
+
+def _make_indexer_config(*, architecture: str, adaptive: bool):
+    return SimpleNamespace(
+        model_config=SimpleNamespace(architectures=[architecture]),
+        num_speculative_tokens=1,
+        speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive),
+    )
+
+
+def _mock_rocm_platform(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Platform predicates are mutually exclusive in production. Override both
+    # so these ROCm policy tests do not inherit CUDA capabilities from the CI
+    # host running them.
+    monkeypatch.setattr(indexer.current_platform, "is_cuda", lambda: False)
+    monkeypatch.setattr(indexer.current_platform, "is_rocm", lambda: True)
+
+
+def _make_indexer_builder(*, adaptive: bool, capacity: int = 12):
+    builder = DeepseekV32IndexerMetadataBuilder.__new__(
+        DeepseekV32IndexerMetadataBuilder
+    )
+    builder.vllm_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=adaptive)
+    )
+    builder.supports_varlen = False
+    builder.decode_seq_lens_buffer = torch.zeros(capacity, dtype=torch.int32)
+    builder.expanded_block_table_buffer = torch.zeros((capacity, 2), dtype=torch.int32)
+    builder.decode_lens_buffer = torch.zeros(capacity, dtype=torch.int32)
+    builder.arange_buffer = torch.arange(capacity, dtype=torch.int32)
+    return builder
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs():
+    adaptive_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=True)
+    )
+    fixed_config = SimpleNamespace(
+        speculative_config=SimpleNamespace(enable_adaptive_verification=False)
+    )
+
+    for builder_cls in (
+        DeepseekV4ROCMAiterMLASparseMetadataBuilder,
+        DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
+    ):
+        assert (
+            builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace())
+            == AttentionCGSupport.ALWAYS
+        )
+        assert (
+            builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace())
+            == AttentionCGSupport.UNIFORM_BATCH
+        )
+
+
+@pytest.mark.cpu_test
+def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch):
+    _mock_rocm_platform(monkeypatch)
+    adaptive_config = _make_indexer_config(
+        architecture="DeepseekV4ForCausalLM", adaptive=True
+    )
+    fixed_config = _make_indexer_config(
+        architecture="DeepseekV4ForCausalLM", adaptive=False
+    )
+
+    assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch()
+    assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
+    assert (
+        DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder
+    )
+    assert indexer._use_flattening(adaptive_config)
+    assert (
+        DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
+            adaptive_config, SimpleNamespace()
+        )
+        == AttentionCGSupport.ALWAYS
+    )
+    assert not indexer._use_flattening(fixed_config)
+    assert (
+        DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
+            fixed_config, SimpleNamespace()
+        )
+        == AttentionCGSupport.UNIFORM_BATCH
+    )
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    "architecture",
+    ["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"],
+)
+def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4(
+    monkeypatch: pytest.MonkeyPatch,
+    architecture: str,
+):
+    _mock_rocm_platform(monkeypatch)
+    adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True)
+
+    assert not indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
+    assert not indexer._use_flattening(adaptive_config)
+    assert (
+        DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
+            adaptive_config, SimpleNamespace()
+        )
+        == AttentionCGSupport.UNIFORM_BATCH
+    )
+
+    if architecture == "DeepseekV41ForCausalLM":
+        assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch()
+
+
+@pytest.mark.cpu_test
+def test_rocm_adaptive_indexer_preserves_single_request_uniform_path(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    builder = _make_indexer_builder(adaptive=True, capacity=8)
+
+    class FakeUniformKernel:
+        called = False
+
+        def __call__(
+            self,
+            seq_lens,
+            decode_seq_lens,
+            block_table,
+            expanded_block_table,
+            decode_lens,
+            num_decode_tokens,
+            max_decode_len,
+        ):
+            self.called = True
+            assert num_decode_tokens == 4
+            assert block_table.stride(0) == expanded_block_table.stride(0) == 2
+            decode_seq_lens[:max_decode_len] = torch.arange(
+                seq_lens[0] - max_decode_len + 1,
+                seq_lens[0] + 1,
+                dtype=torch.int32,
+            )
+            expanded_block_table[:max_decode_len] = block_table[0]
+            decode_lens[:max_decode_len] = 1
+
+    fake_kernel = FakeUniformKernel()
+    monkeypatch.setattr(indexer, "_PREPARE_UNIFORM_DECODE_KERNEL", fake_kernel)
+
+    seq_lens, block_table, decode_lens, batch_size, requires_padding = (
+        builder._prepare_decode_tensors(
+            seq_lens=torch.tensor([10], dtype=torch.int32),
+            block_table=torch.tensor([[1, 2]], dtype=torch.int32),
+            decode_lens=torch.tensor([4], dtype=torch.int32),
+            decode_lens_cpu=torch.tensor([4], dtype=torch.int32),
+            query_start_loc=torch.tensor([0], dtype=torch.int32),
+            num_decodes=1,
+            num_decode_tokens=4,
+            use_native=False,
+            next_n=8,
+            max_decode_len=4,
+        )
+    )
+
+    assert fake_kernel.called
+    torch.testing.assert_close(seq_lens, torch.tensor([7, 8, 9, 10], dtype=torch.int32))
+    torch.testing.assert_close(
+        block_table,
+        torch.tensor([[1, 2], [1, 2], [1, 2], [1, 2]], dtype=torch.int32),
+    )
+    torch.testing.assert_close(decode_lens, torch.ones(4, dtype=torch.int32))
+    assert batch_size == 4
+    assert not requires_padding
+
+
+@pytest.mark.cpu_test
+def test_rocm_adaptive_indexer_replays_changed_allocations_with_stable_buffers():
+    builder = _make_indexer_builder(adaptive=True)
+    buffer_ptrs = (
+        builder.decode_seq_lens_buffer.data_ptr(),
+        builder.expanded_block_table_buffer.data_ptr(),
+        builder.decode_lens_buffer.data_ptr(),
+    )
+    block_table = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
+    decode_lens_cpu = torch.tensor([4, 4, 0], dtype=torch.int32)
+
+    first = builder._prepare_decode_tensors(
+        seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32),
+        block_table=block_table,
+        decode_lens=torch.tensor([7, 1, 0], dtype=torch.int32),
+        decode_lens_cpu=decode_lens_cpu,
+        query_start_loc=torch.tensor([0, 7, 8], dtype=torch.int32),
+        num_decodes=3,
+        num_decode_tokens=10,
+        use_native=False,
+        next_n=8,
+        max_decode_len=4,
+    )
+    torch.testing.assert_close(
+        first[0],
+        torch.tensor([14, 15, 16, 17, 18, 19, 20, 20, 0, 0], dtype=torch.int32),
+    )
+    torch.testing.assert_close(
+        first[1][:, 0],
+        torch.tensor([1, 1, 1, 1, 1, 1, 1, 3, 0, 0], dtype=torch.int32),
+    )
+
+    second = builder._prepare_decode_tensors(
+        seq_lens=torch.tensor([20, 20, 0], dtype=torch.int32),
+        block_table=block_table,
+        decode_lens=torch.tensor([1, 7, 0], dtype=torch.int32),
+        decode_lens_cpu=decode_lens_cpu,
+        query_start_loc=torch.tensor([0, 1, 8], dtype=torch.int32),
+        num_decodes=3,
+        num_decode_tokens=10,
+        use_native=False,
+        next_n=8,
+        max_decode_len=4,
+    )
+    torch.testing.assert_close(
+        second[0],
+        torch.tensor([20, 14, 15, 16, 17, 18, 19, 20, 0, 0], dtype=torch.int32),
+    )
+    torch.testing.assert_close(
+        second[1][:, 0],
+        torch.tensor([1, 3, 3, 3, 3, 3, 3, 3, 0, 0], dtype=torch.int32),
+    )
+    torch.testing.assert_close(second[2], torch.ones(10, dtype=torch.int32))
+    assert second[3:] == (10, False)
+    assert buffer_ptrs == (
+        builder.decode_seq_lens_buffer.data_ptr(),
+        builder.expanded_block_table_buffer.data_ptr(),
+        builder.decode_lens_buffer.data_ptr(),
+    )

vllm/models/deepseek_v4/amd/dspark.py

@@ -44,6 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
 )
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.qwen3_dspark import (
+    DSparkConfidenceHead,
     DSparkMarkovHead,
 )
 from vllm.model_executor.models.utils import maybe_prefix
@@ -102,7 +103,7 @@ class DSparkDeepseekV4Model(nn.Module):
             ]
         )
 
-        # Heads: final norm + hc_head, and the Markov head
+        # Heads: final norm + hc_head, and the Markov + confidence heads
         # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         hc_dim = self.hc_mult * config.hidden_size
@@ -125,6 +126,12 @@ class DSparkDeepseekV4Model(nn.Module):
             config.dspark_markov_rank,
             prefix=maybe_prefix(prefix, "markov_head"),
         )
+        self.confidence_head: DSparkConfidenceHead | None = None
+        if getattr(config, "enable_confidence_head", True):
+            self.confidence_head = DSparkConfidenceHead(
+                config.hidden_size + config.dspark_markov_rank,
+                prefix=maybe_prefix(prefix, "confidence_head"),
+            )
 
         # MHC head CustomOp dispatcher (aiter / tilelang / triton / torch),
         # replacing the direct nvidia tilelang kernel call.
@@ -358,6 +365,17 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
     def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
         return self.model.markov_head.bias(markov_embed, self.logits_processor)
 
+    def compute_confidence(
+        self, head_hidden: torch.Tensor, markov_embed: torch.Tensor
+    ) -> torch.Tensor:
+        """Per-position acceptance probability for each drafted token."""
+        if self.model.confidence_head is None:
+            raise RuntimeError(
+                "compute_confidence() requires a confidence head, but the "
+                "checkpoint did not provide confidence_head weights."
+            )
+        return torch.sigmoid(self.model.confidence_head(head_hidden, markov_embed))
+
     # --- Weight loading ----------------------------------------------------
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
@@ -391,6 +409,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
+        loaded_confidence_head = False
 
         tp_size = get_tensor_model_parallel_world_size()
         tp_rank = get_tensor_model_parallel_rank()
@@ -403,6 +422,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
             if mapped is None:
                 continue
             name = mapped
+            if "confidence_head." in name:
+                loaded_confidence_head = True
 
             # ``.scale`` -> per-method scale suffix.
             if name.endswith(".scale"):
@@ -439,8 +460,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
                 continue
 
             # Stacked rules only apply to decoder-layer weights. Head-stack params
-            # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g.
-            # "markov_w1" would collide with the "w1" shard rule.
+            # (main_proj/norm/hc_head/markov_head/confidence_head) load directly —
+            # otherwise e.g. "markov_w1" would collide with the "w1" shard rule.
             is_layer_param = name.startswith("model.layers.")
             for param_name, weight_name, stacked_shard_id in stacked_params_mapping:
                 if not is_layer_param or weight_name not in name:
@@ -469,6 +490,8 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
                 weight_loader(param, loaded_weight)
                 loaded_params.add(name)
 
+        if self.model.confidence_head is not None and not loaded_confidence_head:
+            self.model.confidence_head = None
         logger.info_once("DSpark draft model loaded: %d params", len(loaded_params))
         return loaded_params
 
@@ -482,8 +505,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
             return None
         stage = int(m.group(1))
         rest = m.group(2)
-        # The confidence head is not wired into inference yet; drop its weights.
-        if rest.startswith("confidence_head."):
+        if rest.startswith("confidence_head.") and self.model.confidence_head is None:
             return None
         # Head-stack params live at model level (mtp.last), context combiner at
         # model level (mtp.0); everything else is a per-layer decoder block.
@@ -493,6 +515,7 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
             "hc_head_base",
             "hc_head_scale",
             "markov_head.",
+            "confidence_head.",
         )
         if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith(
             head_prefixes

vllm/models/deepseek_v4/amd/rocm.py

@@ -8,6 +8,7 @@ from typing import cast
 import torch
 
 from vllm import envs
+from vllm.config import VllmConfig
 from vllm.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
@@ -26,6 +27,7 @@ from vllm.platforms.rocm import _ON_GFX950
 from vllm.triton_utils import tl, triton
 from vllm.utils.multi_stream_utils import execute_in_parallel
 from vllm.v1.attention.backend import (
+    AttentionCGSupport,
     CommonAttentionMetadata,
 )
 from vllm.v1.attention.backends.mla.sparse_swa import (
@@ -38,6 +40,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
     rocm_sparse_attn_decode,
     rocm_sparse_attn_prefill,
 )
+from vllm.v1.kv_cache_interface import KVCacheSpec
 from vllm.v1.worker.workspace import current_workspace_manager
 
 logger = init_logger(__name__)
@@ -445,6 +448,20 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
 
 
 class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> AttentionCGSupport:
+        spec_config = vllm_config.speculative_config
+        if spec_config is not None and spec_config.enable_adaptive_verification:
+            # All per-token metadata is built from device query boundaries into
+            # persistent buffers, so adaptive verification can replay varlen
+            # FULL decode graphs after reallocating drafts across requests.
+            return AttentionCGSupport.ALWAYS
+        return super().get_cudagraph_support(vllm_config, kv_cache_spec)
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.c128a_decode_topk_ragged_indices_buffer: torch.Tensor | None = None
@@ -512,6 +529,20 @@ class DeepseekV4ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBui
 
 
 class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekSparseSWAMetadataBuilder):
+    @classmethod
+    def get_cudagraph_support(
+        cls,
+        vllm_config: VllmConfig,
+        kv_cache_spec: KVCacheSpec,
+    ) -> AttentionCGSupport:
+        spec_config = vllm_config.speculative_config
+        if spec_config is not None and spec_config.enable_adaptive_verification:
+            # SWA indices, lengths, and token-to-request mappings are built from
+            # device boundaries into persistent buffers, so adaptive verification
+            # can replay varlen FULL decode graphs safely.
+            return AttentionCGSupport.ALWAYS
+        return super().get_cudagraph_support(vllm_config, kv_cache_spec)
+
     # Keep fused multi-step decode disabled until update_draft_decode_metadata()
     # also refreshes the ROCm-specific ragged SWA indices and indptrs.
     supports_draft_decode_metadata_update = False

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

@@ -237,6 +237,15 @@ class KpoolTailBackend(DeepseekV32IndexerBackend):
 
 
 class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        # ROCm runs adaptive verification through the per-token flattened
+        # indexer path, which derives row ownership from device decode lengths.
+        return (
+            current_platform.is_rocm()
+            or super().supports_device_cpu_query_lens_mismatch()
+        )
+
     @staticmethod
     def get_name() -> str:
         return "DEEPSEEK_V4_INDEXER"
@@ -254,6 +263,12 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
 
 
 class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        # The ROCm flattened-query support above is validated for the
+        # DeepSeek-V4 adaptive DSpark path only.
+        return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch()
+
     @staticmethod
     def get_name() -> str:
         return "DEEPSEEK_V41_INDEXER"
@@ -802,6 +817,15 @@ def _supports_flattened_device_query_lens() -> bool:
     )
 
 
+def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool:
+    model_config = vllm_config.model_config
+    return (
+        current_platform.is_rocm()
+        and model_config is not None
+        and "DeepseekV4ForCausalLM" in model_config.architectures
+    )
+
+
 def _supports_native_decode(next_n: int) -> bool:
     """Whether decode can pass `next_n` Q rows per request to the kernel
     instead of flattening to one single-token row per query, which re-reads
@@ -822,7 +846,10 @@ def _use_flattening(vllm_config: VllmConfig) -> bool:
     return not _supports_native_decode(next_n) or (
         speculative_config is not None
         and speculative_config.enable_adaptive_verification
-        and _supports_flattened_device_query_lens()
+        and (
+            _supports_flattened_device_query_lens()
+            or _rocm_supports_flattened_device_query_lens(vllm_config)
+        )
     )