From 7de96050c530009eb8e81065aea76570fc6ed15b Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Tue, 25 Aug 2026 13:31:15 -0400 Subject: [PATCH] [Spec Decode] Enable adaptive DSpark on SM100 sparse MLA (#52783) Signed-off-by: mgoin --- .../test_dspark_noncausal_sparse_mla.py | 62 ++++++++++++- .../spec_decode/test_adaptive_verification.py | 58 ++++++++++++ tests/v1/worker/test_attn_utils.py | 91 +++++++++++++++++++ .../backends/mla/flashinfer_mla_sparse.py | 17 ++++ vllm/v1/worker/gpu/attn_utils.py | 34 +++++-- vllm/v1/worker/gpu/model_runner.py | 20 +++- .../gpu/spec_decode/adaptive_verification.py | 29 +++++- 7 files changed, 293 insertions(+), 18 deletions(-) diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py index 71579c025b..8fd74314d1 100644 --- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py +++ b/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], diff --git a/tests/v1/spec_decode/test_adaptive_verification.py b/tests/v1/spec_decode/test_adaptive_verification.py index 99708f50f8..d92f78a715 100644 --- a/tests/v1/spec_decode/test_adaptive_verification.py +++ b/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. diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 01ad759bcd..ba0f13d119 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/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( diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 3cf3d9541d..512627559a 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/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 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index f603487ace..d4ab1e5066 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3bdccebfe2..32fbd63fe8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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( diff --git a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py b/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py index 12dc8c7c65..d3b3743d67 100644 --- a/vllm/v1/worker/gpu/spec_decode/adaptive_verification.py +++ b/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." )