[Spec Decode] Enable adaptive DSpark on SM100 sparse MLA (#52783)

Signed-off-by: mgoin <mgoin64@gmail.com>

tests/v1/attention/test_dspark_noncausal_sparse_mla.py

@@ -55,6 +55,7 @@ if not current_platform.is_cuda():
     )
 
 from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseTRTLLMBackend,
 )
@@ -118,6 +119,9 @@ def _run_sparse_backend_vs_sdpa(
     num_heads: int,
     device: torch.device,
     force_future_dominance: bool = False,
+    qk_nope_head_dim: int = 128,
+    v_head_dim: int = 128,
+    stale_cpu_query_lens: list[int] | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Run a sparse-MLA backend with the given per-token indices and compute a
     dense per-token SDPA reference over the SAME indices.
@@ -145,9 +149,7 @@ def _run_sparse_backend_vs_sdpa(
     use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
 
     kv_lora_rank = 512
-    qk_nope_head_dim = 128
     qk_rope_head_dim = 64
-    v_head_dim = 128
     head_size = kv_lora_rank + qk_rope_head_dim
 
     max_seqlen = max(seq_lens)
@@ -290,6 +292,16 @@ def _run_sparse_backend_vs_sdpa(
     common_attn_metadata = create_common_attn_metadata(
         batch_spec, block_size, device, arange_block_indices=True
     )
+    if stale_cpu_query_lens is not None:
+        # Adaptive verification updates only device boundaries; the stale CPU
+        # copy proves token_to_req_indices consumes the device layout.
+        assert len(stale_cpu_query_lens) == len(query_lens)
+        cpu_query_start_loc = [0]
+        for query_len in stale_cpu_query_lens:
+            cpu_query_start_loc.append(cpu_query_start_loc[-1] + query_len)
+        common_attn_metadata.query_start_loc_cpu = torch.tensor(
+            cpu_query_start_loc, dtype=torch.int32
+        )
     kv_cache = create_and_prepopulate_kv_cache(
         kv_c_contexts=kv_c_contexts,
         k_pe_contexts=k_pe_contexts,
@@ -409,6 +421,52 @@ def _skip_if_backend_unavailable(backend_cls, kv_cache_dtype: str, block_size: i
             pytest.skip("FlashInferMLASparseTRTLLMBackend requires SM 10.x capability")
 
 
+def test_flashinfer_sparse_mla_adaptive_varlen_matches_sdpa(
+    default_vllm_config,
+    dist_init,
+    workspace_init,
+):
+    """Adaptive request boundaries must drive SM100 sparse index conversion."""
+    backend_cls = FlashInferMLASparseTRTLLMBackend
+    _skip_if_backend_unavailable(backend_cls, "fp8", 64)
+    assert (
+        backend_cls.get_builder_cls().get_cudagraph_support(None, None)
+        == AttentionCGSupport.ALWAYS
+    )
+
+    device = torch.device(DEVICE_TYPE)
+    seq_lens = [257, 270, 265, 276]
+    query_lens = [1, 7, 3, 5]
+    sparse_indices = _build_dspark_noncausal_indices(
+        seq_lens,
+        query_lens,
+        window=128,
+        topk_width=256,
+        device=device,
+    )
+
+    backend_output, sdpa_reference, _ = _run_sparse_backend_vs_sdpa(
+        backend_cls,
+        seq_lens,
+        query_lens,
+        sparse_indices,
+        "fp8",
+        64,
+        16,
+        device,
+        qk_nope_head_dim=192,
+        v_head_dim=256,
+        stale_cpu_query_lens=[4, 4, 4, 4],
+    )
+
+    torch.testing.assert_close(
+        backend_output,
+        sdpa_reference,
+        rtol=0.065,
+        atol=0.05,
+    )
+
+
 @pytest.mark.parametrize(
     "backend_cls",
     [FlashMLASparseBackend, FlashInferMLASparseTRTLLMBackend],

tests/v1/spec_decode/test_adaptive_verification.py

@@ -5,9 +5,13 @@ from types import SimpleNamespace
 
 import numpy as np
 
+from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.worker.gpu.async_utils import StepTimingSample
+from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
+from vllm.v1.worker.gpu.spec_decode import adaptive_verification as adaptive_module
 from vllm.v1.worker.gpu.spec_decode.adaptive_verification import (
     AdaptiveVerificationManager,
+    maybe_create_adaptive_verification_manager,
 )
 from vllm.v1.worker.gpu.structured_outputs import _build_grammar_mapping
 
@@ -31,6 +35,60 @@ def make_manager(
     return manager
 
 
+def test_manager_scopes_varlen_check_without_weakening_runner_cg_mode(monkeypatch):
+    class Backend:
+        @classmethod
+        def supports_device_cpu_query_lens_mismatch(cls):
+            return True
+
+    class Builder:
+        def __init__(self, support):
+            self.support = support
+
+        def get_cudagraph_support(self, *_args):
+            return self.support
+
+    def group(layer_name, support):
+        builder = Builder(support)
+        return SimpleNamespace(
+            layer_names=[layer_name],
+            backend=Backend,
+            kv_cache_spec=None,
+            get_metadata_builder=lambda _index: builder,
+        )
+
+    groups = [
+        [
+            group("target", AttentionCGSupport.ALWAYS),
+            group("draft", AttentionCGSupport.UNIFORM_BATCH),
+        ]
+    ]
+    runner_support = AttentionCGSupportInfo(
+        AttentionCGSupport.UNIFORM_BATCH, "DraftBackend"
+    )
+    created = object()
+    monkeypatch.setattr(
+        adaptive_module,
+        "AdaptiveVerificationManager",
+        lambda *_args, **_kwargs: created,
+    )
+
+    manager = maybe_create_adaptive_verification_manager(
+        enable_adaptive_verification=True,
+        attn_groups=groups,
+        attn_cg_support=runner_support,
+        req_states=object(),
+        query_start_loc=object(),
+        num_bonus_tokens=1,
+        max_total_logits=1,
+        vllm_config=None,
+        target_layer_names={"target"},
+    )
+
+    assert manager is created
+    assert runner_support.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
+
+
 def test_budget_stops_where_marginal_drafts_stop_paying_for_themselves():
     # Verification is cheap up to two extra tokens, then jumps 100x; only the
     # highest-confidence draft is worth the cheap slot.

tests/v1/worker/test_attn_utils.py

@@ -11,6 +11,7 @@ import pytest
 import torch
 
 from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
@@ -21,12 +22,102 @@ from vllm.v1.kv_cache_interface import (
     MLAAttentionSpec,
     compute_layout_strides,
 )
+from vllm.v1.worker.gpu.attn_utils import (
+    get_attn_cg_support,
+    get_query_lens_mismatch_unsupported_backend,
+)
 from vllm.v1.worker.utils import (
+    AttentionGroup,
     allocate_kv_cache,
     copy_kv_cache_blocks_inplace,
 )
 
 
+class _FakeMetadataBuilder:
+    def __init__(self, support: AttentionCGSupport):
+        self.support = support
+
+    def get_cudagraph_support(self, *_args):
+        return self.support
+
+
+class _TargetBackend:
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        return True
+
+
+class _DraftBackend:
+    @classmethod
+    def supports_device_cpu_query_lens_mismatch(cls) -> bool:
+        return False
+
+
+def test_attention_checks_preserve_global_and_target_scoped_support():
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=128,
+        dtype=torch.bfloat16,
+    )
+    target_group = AttentionGroup(
+        _TargetBackend,
+        ["target"],
+        spec,
+        0,  # type: ignore[arg-type]
+    )
+    target_group.metadata_builders = [
+        _FakeMetadataBuilder(AttentionCGSupport.ALWAYS)  # type: ignore[list-item]
+    ]
+    draft_group = AttentionGroup(
+        _DraftBackend,
+        ["draft"],
+        spec,
+        0,  # type: ignore[arg-type]
+    )
+    draft_group.metadata_builders = [
+        _FakeMetadataBuilder(AttentionCGSupport.UNIFORM_BATCH)  # type: ignore[list-item]
+    ]
+    groups = [[target_group, draft_group]]
+
+    # The runner-wide execution mode must still honor the drafter's limit.
+    unfiltered = get_attn_cg_support(groups, None)  # type: ignore[arg-type]
+    assert unfiltered.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
+    assert unfiltered.min_cg_attn_backend == "_DraftBackend"
+
+    # Adaptive verification validates only the target's varlen graphs.
+    target_only = get_attn_cg_support(
+        groups,
+        None,  # type: ignore[arg-type]
+        checked_layer_names={"target"},
+    )
+    assert target_only.min_cg_support == AttentionCGSupport.ALWAYS
+    assert target_only.min_cg_attn_backend is None
+    assert (
+        get_query_lens_mismatch_unsupported_backend(
+            groups,
+            checked_layer_names={"target"},
+        )
+        is None
+    )
+
+    # Shared target/draft groups still participate in target-scoped checks.
+    draft_group.layer_names.append("target")
+    target_with_shared_group = get_attn_cg_support(
+        groups,
+        None,  # type: ignore[arg-type]
+        checked_layer_names={"target"},
+    )
+    assert target_with_shared_group.min_cg_support == AttentionCGSupport.UNIFORM_BATCH
+    assert (
+        get_query_lens_mismatch_unsupported_backend(
+            groups,
+            checked_layer_names={"target"},
+        )
+        == "_DraftBackend"
+    )
+
+
 def test_reshape_padded_kv_cache_strides_by_padded_page():
     num_blocks = 3
     spec = FullAttentionSpec(

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

@@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import AttentionSpec
 
 if TYPE_CHECKING:
     from vllm.model_executor.models.deepseek_v2 import Indexer
+    from vllm.v1.attention.backend import CommonAttentionMetadata
 
 logger = init_logger(__name__)
 
@@ -83,6 +84,10 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
     def get_impl_cls() -> type[MLAAttentionImpl]:
         return FlashInferMLASparseImpl
 
+    @staticmethod
+    def get_builder_cls() -> type["FlashInferMLASparseTRTLLMMetadataBuilder"]:
+        return FlashInferMLASparseTRTLLMMetadataBuilder
+
     @classmethod
     def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
         return capability.major == 10
@@ -261,6 +266,18 @@ class FlashInferMLASparseMetadataBuilder(
         )
 
 
+class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilder):
+    """Metadata builder for the SM100 TRT-LLM sparse MLA kernel."""
+
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
+
+    def _build_req_id_per_token(
+        self,
+        common_attn_metadata: "CommonAttentionMetadata",
+    ) -> torch.Tensor:
+        return common_attn_metadata.token_to_req_indices(self.req_id_per_token_buffer)
+
+
 # Global workspace buffer (lazily initialized)
 _fi_sparse_workspace: torch.Tensor | None = None
 

vllm/v1/worker/gpu/attn_utils.py

@@ -130,8 +130,6 @@ def init_attn_backend(
 
     # Phase 3: create metadata builders and determine cudagraph support.
     attn_backend_workspace: torch.Tensor | None = None
-    min_cg_support = AttentionCGSupport.ALWAYS
-    min_cg_attn_backend = None
     for kv_cache_group_id, groups in enumerate(attn_groups):
         kernel_block_size = None
         if kv_cache_group_id < len(kernel_block_sizes):
@@ -150,7 +148,25 @@ def init_attn_backend(
             else:
                 if hasattr(builder, "set_workspace_buffer"):
                     builder.set_workspace_buffer(attn_backend_workspace)
-            # Check cudagraph support for the attention backend
+    attn_cg_support_info = get_attn_cg_support(attn_groups, vllm_config)
+    return attn_groups, attn_cg_support_info, kernel_block_sizes
+
+
+def get_attn_cg_support(
+    attn_groups: list[list[AttentionGroup]],
+    vllm_config: VllmConfig,
+    checked_layer_names: set[str] | None = None,
+) -> AttentionCGSupportInfo:
+    """Return the weakest CUDA graph support among the checked layers."""
+    min_cg_support = AttentionCGSupport.ALWAYS
+    min_cg_attn_backend = None
+    for groups in attn_groups:
+        for group in groups:
+            if checked_layer_names is not None and checked_layer_names.isdisjoint(
+                group.layer_names
+            ):
+                continue
+            builder = group.get_metadata_builder(0)
             cg_support = builder.get_cudagraph_support(
                 vllm_config,
                 group.kv_cache_spec,
@@ -158,15 +174,15 @@ def init_attn_backend(
             if cg_support.value < min_cg_support.value:
                 min_cg_support = cg_support
                 min_cg_attn_backend = group.backend.__name__
-
-    attn_cg_support_info = AttentionCGSupportInfo(
-        min_cg_support=min_cg_support, min_cg_attn_backend=min_cg_attn_backend
+    return AttentionCGSupportInfo(
+        min_cg_support=min_cg_support,
+        min_cg_attn_backend=min_cg_attn_backend,
     )
-    return attn_groups, attn_cg_support_info, kernel_block_sizes
 
 
 def get_query_lens_mismatch_unsupported_backend(
     attn_groups: list[list[AttentionGroup]],
+    checked_layer_names: set[str] | None = None,
 ) -> str | None:
     """Name the first backend needing the CPU query lengths to be exact, if any.
 
@@ -176,6 +192,10 @@ def get_query_lens_mismatch_unsupported_backend(
     """
     for groups in attn_groups:
         for group in groups:
+            if checked_layer_names is not None and checked_layer_names.isdisjoint(
+                group.layer_names
+            ):
+                continue
             if not group.backend.supports_device_cpu_query_lens_mismatch():
                 return group.backend.__name__
     return None

vllm/v1/worker/gpu/model_runner.py

@@ -559,12 +559,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 max_num_blocks = get_block_table_width(max_num_blocks, spec.block_size)
             max_num_blocks_per_group.append(max_num_blocks)
 
+        target_attn_layer_names = None
+        if isinstance(self.speculator, DraftModelSpeculator):
+            # Adaptive verification validates target attention separately.
+            target_attn_layer_names = {
+                layer_name
+                for group in self.kv_cache_config.kv_cache_groups
+                for layer_name in group.layer_names
+            } - self.speculator.draft_attn_layer_names
         self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend(
-            self.kv_cache_config, self.vllm_config, self.device
-        )
-        attn_cg_support = attn_cg_support.narrow(
-            *self.model_state.get_additional_cg_support()
+            self.kv_cache_config,
+            self.vllm_config,
+            self.device,
         )
+        additional_attn_cg_support = self.model_state.get_additional_cg_support()
+        attn_cg_support = attn_cg_support.narrow(*additional_attn_cg_support)
         # The speculator clears the flag at load time when the checkpoint has
         # no confidence head, so it holds the effective value.
         self.adaptive_verification = maybe_create_adaptive_verification_manager(
@@ -577,6 +586,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             query_start_loc=self.input_buffers.query_start_loc,
             num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step,
             max_total_logits=get_max_chunk_logits(self.vocab_size),
+            vllm_config=self.vllm_config,
+            target_layer_names=target_attn_layer_names,
+            additional_attn_cg_support=additional_attn_cg_support,
         )
 
         self.block_tables = BlockTables(

vllm/v1/worker/gpu/spec_decode/adaptive_verification.py

@@ -17,6 +17,7 @@ from vllm.v1.attention.backend import AttentionCGSupport
 from vllm.v1.utils import CpuGpuBuffer
 from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream
 from vllm.v1.worker.gpu.attn_utils import (
+    get_attn_cg_support,
     get_query_lens_mismatch_unsupported_backend,
 )
 from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
@@ -25,6 +26,7 @@ logger = init_logger(__name__)
 _PROFILE_REPLAYS = 5
 
 if TYPE_CHECKING:
+    from vllm.config import VllmConfig
     from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo
     from vllm.v1.worker.gpu.input_batch import InputBatch
     from vllm.v1.worker.gpu.states import RequestState
@@ -444,13 +446,19 @@ def maybe_create_adaptive_verification_manager(
     query_start_loc: torch.Tensor,
     num_bonus_tokens: int,
     max_total_logits: int,
+    vllm_config: "VllmConfig",
+    target_layer_names: set[str] | None = None,
+    additional_attn_cg_support: tuple[AttentionCGSupport, str | None] | None = None,
 ) -> AdaptiveVerificationManager | None:
     if not enable_adaptive_verification:
         return None
 
     # The selector rejects unsupported backends, but models that
     # hard-wire theirs (e.g. DeepSeek-V4) never go through it.
-    backend = get_query_lens_mismatch_unsupported_backend(attn_groups)
+    backend = get_query_lens_mismatch_unsupported_backend(
+        attn_groups,
+        checked_layer_names=target_layer_names,
+    )
     if backend is not None:
         raise ValueError(
             "Adaptive verification trims verification requests on device, which"
@@ -459,12 +467,23 @@ def maybe_create_adaptive_verification_manager(
             "use a backend that does."
         )
 
-    if attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS:
+    target_attn_cg_support = attn_cg_support
+    if target_layer_names is not None:
+        target_attn_cg_support = get_attn_cg_support(
+            attn_groups,
+            vllm_config,
+            checked_layer_names=target_layer_names,
+        )
+        if additional_attn_cg_support is not None:
+            target_attn_cg_support = target_attn_cg_support.narrow(
+                *additional_attn_cg_support
+            )
+    if target_attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS:
         raise ValueError(
             "Adaptive verification captures varlen decode cudagraphs, so every"
-            " attention builder must report AttentionCGSupport.ALWAYS, but "
-            f"{attn_cg_support.min_cg_attn_backend} reports "
-            f"{attn_cg_support.min_cg_support}. Pass "
+            " target attention builder must report AttentionCGSupport.ALWAYS, but "
+            f"{target_attn_cg_support.min_cg_attn_backend} reports "
+            f"{target_attn_cg_support.min_cg_support}. Pass "
             "enable_adaptive_verification=false in the speculative config, or "
             "use a backend that does."
         )