From e3614d9d4a306b0daef0fb648ce09b53f989b99c Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:02:56 +0200 Subject: [PATCH 1/3] [Spec Decode][MRV2] DSpark speculator for Model Runner V2 Port of vllm-project/vllm's MRV2 DSparkSpeculator (vllm#46995 by @benchislett, speculator as of 2026-09-30) onto this fork's V2 DFlash speculator: anchor-as-first query block (SAMPLE_FROM_ANCHOR), sequential Markov sampling inside the drafter's FULL CUDA graph, DeepSeek-V4 context slots per cache layer, drafter-wide staging buffers, and the V1 proposer's weight contract (the drafter keeps its own embedding copy under PP). Admits method 'dspark' on MRV2, wires its aux target layers and noise token, and feeds drafters the target hidden state unless the target provides an override (DeepSeek-V4 MTP only). The layering baseline records the new speculator's references to the DFlash base it extends, as it does for the V1 proposer (vllm/v1/spec_decode/dspark.py). Co-authored-by: Benjamin Chislett Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJxVXkjSHhfq2D71DoYSK7 --- tools/pre_commit/layering_baseline.json | 3 + vllm/config/vllm.py | 4 +- vllm/v1/worker/gpu/model_runner.py | 36 ++-- vllm/v1/worker/gpu/spec_decode/__init__.py | 4 + .../v1/worker/gpu/spec_decode/dflash/utils.py | 17 +- .../worker/gpu/spec_decode/dspark/__init__.py | 2 + .../gpu/spec_decode/dspark/speculator.py | 189 ++++++++++++++++++ .../gpu/spec_decode/eagle/eagle3_utils.py | 5 + vllm/v1/worker/gpu/spec_decode/utils.py | 4 +- 9 files changed, 239 insertions(+), 25 deletions(-) create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/__init__.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/speculator.py diff --git a/tools/pre_commit/layering_baseline.json b/tools/pre_commit/layering_baseline.json index 4fcc6a1db0..a799fb757b 100644 --- a/tools/pre_commit/layering_baseline.json +++ b/tools/pre_commit/layering_baseline.json @@ -539,6 +539,9 @@ "vllm/v1/worker/gpu/spec_decode/__init__.py": { "model": 16 }, + "vllm/v1/worker/gpu/spec_decode/dspark/speculator.py": { + "model": 6 + }, "vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py": { "model": 6 }, diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 23ccd0b38c..388596859e 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2440,13 +2440,15 @@ def _get_v2_model_runner_unsupported_features(self) -> list[str]: "eagle3", "mtp", "dflash", + "dspark", ): unsupported.append(f"speculative method '{speculative_config.method}'") # V2 EagleSpeculator does not support parallel_drafting (required by PEagle) + parallel_drafting_methods = ("dflash", "dspark") if ( speculative_config.parallel_drafting - and speculative_config.method != "dflash" + and speculative_config.method not in parallel_drafting_methods ): unsupported.append("parallel drafting for speculative decoding") diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 2dd1256d73..c1db487c45 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -236,8 +236,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): if self.is_last_pp_rank: self.speculator = init_speculator(self.vllm_config, self.device) - if self.speculative_config.method in ("eagle3", "dflash"): - # EAGLE3 and DFlash may require auxiliary target hidden states. + if self.speculative_config.method in ("eagle3", "dflash", "dspark"): + # EAGLE3, DFlash and DSpark may require auxiliary target hidden states. self.use_aux_hidden_state_outputs = True if self.use_pp and self.speculative_config.method == "eagle3": raise ValueError( @@ -697,6 +697,20 @@ def _warmup_sm70_aux_kernels(self) -> None: if warmed: logger.info_once("SM70 V2 auxiliary kernel warmup finished: %s", warmed) + def _speculator_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Hidden state fed to the drafter. + + The target may override it (DeepSeek V4 MTP needs the pre-hc_head + residual) with a persistent buffer sized at max_num_batched_tokens, + sliced to the active token count that propose() expects. Targets + return no buffer for other draft methods (e.g. DSpark). + """ + get_override = getattr(self.model, "get_mtp_target_hidden_states", None) + override = get_override() if get_override is not None else None + if override is None: + return hidden_states + return override[: hidden_states.shape[0]] + @torch.inference_mode() @step_eplb_after(is_dummy=True) def _dummy_run( @@ -780,14 +794,7 @@ def _dummy_run( ), ) - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + spec_hidden_states = self._speculator_hidden_states(hidden_states) self.speculator.propose( input_batch=input_batch, attn_metadata=attn_metadata, @@ -1883,14 +1890,7 @@ def sample_tokens( requires_host_token_state = bool( getattr(self.speculator, "requires_host_token_state", False) ) - # Let the target override the hidden state fed to the drafter - # (e.g. DeepSeek V4 MTP needs the pre-hc_head residual). The - # target returns a persistent buffer sized at max_num_batched_tokens; - # slice to the active token count that propose() expects. - spec_hidden_states = hidden_states - if hasattr(self.model, "get_mtp_target_hidden_states"): - pre_hc_hidden_states = self.model.get_mtp_target_hidden_states() - spec_hidden_states = pre_hc_hidden_states[: hidden_states.shape[0]] # type: ignore[union-attr] + spec_hidden_states = self._speculator_hidden_states(hidden_states) mtp_draft_start = self._step_profiler.start_context(mtp_profile_ctx) if self._ple_draft_ring is not None: # The committed state now ends at each request's bonus token. diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index ee25697099..ee1d7e5cef 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -46,6 +46,10 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator return DFlashSpeculator(vllm_config, device) + if speculative_config.method == "dspark": + from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + return DSparkSpeculator(vllm_config, device) if speculative_config.use_eagle(): from vllm.v1.worker.gpu.spec_decode.eagle.speculator import EagleSpeculator diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index c9d025eeee..be35e404d5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -35,7 +35,12 @@ def _validate_dflash_shared_weights( ) -def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: +def get_parallel_draft_model(vllm_config: VllmConfig, model_tag: str) -> nn.Module: + """Build a DFlash-family (parallel query block) drafter from its own config. + + Shared by DFlash and DSpark; each caller then shares target weights under + its own contract. + """ from vllm.compilation.backends import set_model_tag from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal @@ -73,10 +78,12 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo else vllm_config.cache_config ), ) - with set_model_tag("dflash_head"): - dflash_model = get_model( - vllm_config=draft_vllm_config, model_config=draft_model_config - ) + with set_model_tag(model_tag): + return get_model(vllm_config=draft_vllm_config, model_config=draft_model_config) + + +def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: + dflash_model = get_parallel_draft_model(vllm_config, "dflash_head") target_language_model = ( target_model.get_language_model() diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py b/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py new file mode 100644 index 0000000000..0ef5a83ef1 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DSpark speculator for Model Runner V2: semi-autoregressive parallel drafting. + +DSpark drafts a block of ``num_speculative_tokens`` tokens in one parallel pass +(the DFlash machinery: context-KV precompute plus a query-block forward), then +adds intra-block dependency with a light sequential Markov head. + +Differences from DFlash: + * Anchor-as-first-prediction: each request emits exactly N query tokens + (anchor + N-1 noise tokens), every query position predicts the next token + (the ``SAMPLE_FROM_ANCHOR`` path of the shared prepare-inputs kernel). + * Sequential Markov sampling: tokens are sampled left to right, adding a bias + derived from the previously sampled token at each step. + +The draft step, backbone forward and Markov sampling, runs inside the drafter's +FULL CUDA graph. Follows vllm-project/vllm's MRV2 DSparkSpeculator. +""" + +from typing import Any + +import torch +import torch.nn as nn + +from vllm.config import VllmConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.parallel_state import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.models.shared_weights import ( + LEGACY_DRAFT_WEIGHTS, + share_embeddings, + share_lm_head, +) +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.dflash.utils import get_parallel_draft_model + +logger = init_logger(__name__) + + +class DSparkSpeculator(DFlashSpeculator): + _speculator_name = "DSpark" + + def __init__(self, vllm_config: VllmConfig, device: torch.device): + super().__init__(vllm_config, device) + speculative_config = self.speculative_config + if ( + speculative_config.dspark_confidence_threshold > 0.0 + or speculative_config.dspark_max_verification_tokens is not None + ): + raise ValueError( + "DSpark confidence-scheduled verification is not supported by " + "Model Runner V2; set VLLM_USE_V2_MODEL_RUNNER=0 to use it." + ) + # DSpark consumes the target aux states combined to hidden_size by + # main_proj, not DeepSeek-V4 MTP's hc_mult-wide residual for which the + # base class widens hidden_size: size the staging buffers to the drafter. + self.hidden_size = self.draft_model_config.get_hidden_size() + self.hidden_states = torch.zeros( + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device + ) + self.inputs_embeds = torch.zeros( + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device + ) + # The anchor is the first prediction: N query positions per request. + self.sample_from_anchor = True + self.num_query_per_req = self.draft_block + # Anchor (bonus) token of each request = input id at query offset 0. + # A persistent index keeps the read graph-stable. + self._anchor_idx = ( + torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) + * self.num_query_per_req + ) + # Row s holds column s for every request: the draft_logits column each + # sequential step writes, contiguous for the sampling kernel. + self._step_cols = ( + torch.arange(self.draft_block, dtype=torch.int32, device=device) + .repeat_interleave(self.max_num_reqs) + .view(self.draft_block, self.max_num_reqs) + ) + + def load_draft_model( + self, + target_model: nn.Module, + target_attn_layer_names: set[str], + ) -> nn.Module: + model = get_parallel_draft_model(self.vllm_config, "dspark_head") + target_language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + # The DeepSeek-V4 drafter loads its own copy of the target embedding + # from the checkpoint. With pipeline parallelism the target embedding + # lives on the first stage, so the drafter keeps that copy: the same + # contract as the V1 proposer. + share_embeddings( + model, + target_language_model, + LEGACY_DRAFT_WEIGHTS, + pp_size=get_pp_group().world_size, + log=logger, + ) + share_lm_head( + model, + target_language_model, + target_language_model, + LEGACY_DRAFT_WEIGHTS, + log=logger, + ) + return model + + def _precompute_context_kv( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + slots: torch.Tensor | list[torch.Tensor | None] | None, + ) -> None: + # The DeepSeek-V4 drafter takes its context slots per cache layer name. + slots_by_layer: dict[str, torch.Tensor] | None = None + if slots is not None: + layer_names = self.model.get_draft_kv_cache_layer_names() + per_layer = slots if isinstance(slots, list) else [slots] * len(layer_names) + slots_by_layer = { + name: layer_slots + for name, layer_slots in zip(layer_names, per_layer) + if layer_slots is not None + } + self.model.precompute_and_store_context_kv( + hidden_states, positions, slots_by_layer + ) + + def _sample_step( + self, + logits: torch.Tensor, + idx_mapping: torch.Tensor, + sample_pos: torch.Tensor, + step: int, + ) -> torch.Tensor: + if self.draft_logits is None: + return self.model.map_draft_to_target(logits.argmax(dim=-1)) + # sample_pos is the predicted token's position P. Sampling keys a draw + # by the position before the sampled token, P-1. + sampled = gumbel_sample( + logits, + idx_mapping, + self.temperature, + self.seeds, + sample_pos - 1, + apply_temperature=True, + is_drafting=True, + output_processed_logits=self.draft_logits, + output_processed_logits_col=self._step_cols[step, : logits.shape[0]], + use_fp64=self.use_fp64_gumbel, + ) + return self.model.map_draft_to_target(sampled) + + def _generate_draft( + self, + num_reqs: int, + num_tokens_padded: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> None: + head_hidden = self._run_model( + num_tokens_padded, + attn_metadata, + slot_mappings, + num_tokens_across_dp, + cudagraph_runtime_mode, + ) + num_steps = self.draft_block + num_sample = num_reqs * num_steps + sample_hidden = head_hidden[self.sample_indices[:num_sample]] + base_logits = self.model.compute_logits(sample_hidden) + base_logits = base_logits.view(num_reqs, num_steps, base_logits.shape[-1]) + idx_mapping = self.sample_idx_mapping[:num_sample].view(num_reqs, num_steps) + sample_pos = self.sample_pos[:num_sample].view(num_reqs, num_steps) + + prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]].to(torch.long) + for step in range(num_steps): + markov_embed = self.model.markov_embed(prev) + logits = base_logits[:, step] + self.model.markov_bias(markov_embed) + prev = self._sample_step( + logits, idx_mapping[:, step], sample_pos[:, step], step + ) + self.draft_tokens[:num_reqs, step] = prev diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py index 8c8d56ba31..717364dd0c 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/eagle3_utils.py @@ -49,6 +49,11 @@ def get_eagle3_aux_layers_from_config( target_layer_ids = getattr(hf_config, "target_layer_ids", None) if target_layer_ids: layer_ids = [i + 1 for i in target_layer_ids] + if not layer_ids: + # DSpark config uses zero-based decoder-layer IDs. + dspark_layer_ids = getattr(hf_config, "dspark_target_layer_ids", None) + if dspark_layer_ids: + layer_ids = [i + 1 for i in dspark_layer_ids] if layer_ids and isinstance(layer_ids, (list, tuple)): return tuple(layer_ids) return None diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index be6fb3d8aa..4ef2311983 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -66,7 +66,9 @@ def get_parallel_drafting_token_id(hf_config) -> int: return int(hf_config.pard_token) if hasattr(hf_config, "ptd_token_id"): return int(hf_config.ptd_token_id) + if hasattr(hf_config, "dspark_noise_token_id"): + return int(hf_config.dspark_noise_token_id) raise ValueError( "Model config must specify dflash_config.mask_token_id, pard_token, " - "or ptd_token_id for parallel drafting." + "ptd_token_id, or dspark_noise_token_id for parallel drafting." ) From 966659632960622339a4218ef236457c87eb4300 Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:02:56 +0200 Subject: [PATCH 2/3] [Perf][MRV2] Compressed-index decode graphs at context buckets Capture DeepSeek-V4 decode graphs at the SM70 compressed-index context buckets as the V1 dispatcher does, with the capture metadata bounded by the bucket, and replay the smallest bucket that covers the live context. Without them the graphs were captured at the model's maximum length and the 35k-token decode step took 91-92 ms instead of 70 ms. The buckets come from the same SM70 graph plan the V1 dispatcher reads; the layering baseline accepts that one import. Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJxVXkjSHhfq2D71DoYSK7 --- tools/pre_commit/layering_baseline.json | 2 +- vllm/v1/worker/gpu/cudagraph_utils.py | 58 +++++++++++++++++++++- vllm/v1/worker/gpu/input_batch.py | 21 +++++++- vllm/v1/worker/gpu/model_states/default.py | 5 +- 4 files changed, 81 insertions(+), 5 deletions(-) diff --git a/tools/pre_commit/layering_baseline.json b/tools/pre_commit/layering_baseline.json index a799fb757b..653d2c16ce 100644 --- a/tools/pre_commit/layering_baseline.json +++ b/tools/pre_commit/layering_baseline.json @@ -514,7 +514,7 @@ }, "vllm/v1/worker/gpu/cudagraph_utils.py": { "model": 11, - "platform": 33 + "platform": 34 }, "vllm/v1/worker/gpu/mamba_align.py": { "model": 2 diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index bc95173f6e..858513b992 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -623,9 +623,57 @@ def __init__( scope="process", ) + # Compressed-index (DeepSeek-V4) decode graphs captured at each context + # bucket, as the V1 dispatcher does: the indexer and sparse attention + # are bounded by the bucket instead of the model's maximum length. + self._compressed_context_graphs: dict[ + BatchExecutionDescriptor, tuple[BatchExecutionDescriptor, ...] + ] = {} + from vllm.platforms.sm70.graph_policy import resolve_graph_plan + + compressed_buckets = resolve_graph_plan(vllm_config).compressed_buckets + if compressed_buckets: + descs = self._capture_descs.get(CUDAGraphMode.FULL, []) + for desc in list(descs): + # One request with one decode (or verification) block. + if ( + desc.attention_context_bucket is None + and desc.num_reqs == 1 + and desc.num_tokens == self.decode_query_len + ): + variants = tuple( + replace(desc, attention_context_bucket=bucket) + for bucket in compressed_buckets + ) + self._compressed_context_graphs[desc] = variants + descs.extend(variants) + logger.info_once( + "Compressed-index decode graph variants captured at context " + "buckets %s for %d decode shapes.", + compressed_buckets, + len(self._compressed_context_graphs), + scope="process", + ) + self._compressed_bucket_descs = { + variant + for variants in self._compressed_context_graphs.values() + for variant in variants + } + def select_attention_graph( self, desc: BatchExecutionDescriptor, cpu_upper_bounds: torch.Tensor ) -> BatchExecutionDescriptor: + variants = self._compressed_context_graphs.get(desc) + if variants: + # The smallest captured bucket that covers the live context. + if cpu_upper_bounds.device.type == "cpu" and cpu_upper_bounds.numel(): + upper = int(cpu_upper_bounds.max()) + for bucket_desc in variants: + bucket = bucket_desc.attention_context_bucket + assert bucket is not None + if 0 < upper <= bucket and bucket_desc in self.graphs: + return bucket_desc + return desc short = self._qsa_short_graphs.get(desc) if ( short is not None @@ -724,6 +772,11 @@ def create_forward_fn( attn_groups, kv_cache_config, skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), + context_len=( + desc.attention_context_bucket + if desc in self._compressed_bucket_descs + else None + ), ) def forward_fn(cg_mode: CUDAGraphMode) -> None: @@ -813,8 +866,11 @@ def prepare_inputs_to_capture( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, skip_attn: bool = False, + context_len: int | None = None, ) -> CapturedAttentionState: - input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) + input_batch = InputBatch.make_dummy( + num_reqs, num_tokens, input_buffers, context_len=context_len + ) input_block_tables = block_tables.get_dummy_block_tables(num_reqs) slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens) slot_mappings_by_layer = build_slot_mappings_by_layer( diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index b4310bd220..5b0ecbc040 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -99,6 +99,9 @@ class InputBatch: is_incomplete_prefilling_np: np.ndarray | None = None # Dummy profiling and graph-capture inputs must never write KV slots. is_dummy_batch: bool = False + # Context bound a context-bucket graph is captured at; the attention + # metadata is built for this length instead of the model's maximum. + capture_context_len: int | None = None @classmethod def make_dummy( @@ -106,7 +109,13 @@ def make_dummy( num_reqs: int, num_tokens: int, input_buffers: InputBuffers, + context_len: int | None = None, ) -> "InputBatch": + """Dummy batch for profiling and graph capture. + + ``context_len`` (a context-bucket graph's bound) gives every request + that sequence length; otherwise seq_len equals query_len. + """ assert 0 < num_reqs <= num_tokens device = input_buffers.device @@ -146,6 +155,15 @@ def make_dummy( cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) # Dummy: seq_len == query_len (fresh-prefill shape). seq_lens_cpu_upper_bound = torch.from_numpy(num_scheduled_tokens.copy()) + num_computed_tokens_np = np.zeros(num_reqs, dtype=np.int32) + if context_len is not None: + assert context_len >= int(num_scheduled_tokens.max()) + # query_start_loc above already holds the query offsets. + input_buffers.seq_lens[:num_reqs] = context_len + seq_lens_cpu_upper_bound = torch.full( + (num_reqs,), context_len, dtype=torch.int32 + ) + num_computed_tokens_np = context_len - num_scheduled_tokens return cls( req_ids=req_ids, num_reqs=num_reqs, @@ -164,7 +182,7 @@ def make_dummy( seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=None, - num_computed_tokens_np=np.zeros(num_reqs, dtype=np.int32), + num_computed_tokens_np=num_computed_tokens_np, prefill_len_np=np.zeros(num_reqs, dtype=np.int32), num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32), is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_), @@ -176,6 +194,7 @@ def make_dummy( cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=False, is_dummy_batch=True, + capture_context_len=context_len, ) diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 86a1886083..42b449f6d0 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -179,8 +179,9 @@ def prepare_attn( max_query_len = input_batch.num_scheduled_tokens.max().item() seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound if for_capture: - # Capture with worst-case max_seq_len so the graph is valid at any replay. - max_seq_len = self.max_model_len + # Capture with worst-case max_seq_len so the graph is valid at any + # replay; a context-bucket graph replays only within its bound. + max_seq_len = input_batch.capture_context_len or self.max_model_len else: max_seq_len = int(seq_lens_cpu_upper_bound[:num_reqs].max().item()) attn_metadata = build_attn_metadata( From 450e65c5262eae2853d97144e5f6fb12cde854e9 Mon Sep 17 00:00:00 2001 From: Peuqui <43776522+Peuqui@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:02:56 +0200 Subject: [PATCH 3/3] [Test][MRV2] Cover the DSpark speculator and compressed-index context buckets Context slots per cache layer, greedy Markov step mapping, dummy batches at a context bound, and the smallest-covering bucket selection; the long-attention graph fakes gain the new manager attribute. Signed-off-by: Peuqui <43776522+Peuqui@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01GJxVXkjSHhfq2D71DoYSK7 --- .../test_mrv2_dspark_context_buckets.py | 111 ++++++++++++++++++ .../worker/test_sm70_long_attention_graphs.py | 3 + 2 files changed, 114 insertions(+) create mode 100644 tests/v1/worker/test_mrv2_dspark_context_buckets.py diff --git a/tests/v1/worker/test_mrv2_dspark_context_buckets.py b/tests/v1/worker/test_mrv2_dspark_context_buckets.py new file mode 100644 index 0000000000..bd0990a8fc --- /dev/null +++ b/tests/v1/worker/test_mrv2_dspark_context_buckets.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import numpy as np +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + ModelCudaGraphManager, +) +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + +class _RecordingDraft: + def __init__(self, layer_names): + self._layer_names = layer_names + self.slots = "unset" + + def get_draft_kv_cache_layer_names(self): + return self._layer_names + + def precompute_and_store_context_kv(self, hidden, positions, slots): + self.slots = slots + + def map_draft_to_target(self, ids): + return ids + 100 + + +def _speculator(model) -> DSparkSpeculator: + speculator = object.__new__(DSparkSpeculator) + speculator.model = model + speculator.draft_logits = None + return speculator + + +def test_context_slots_are_passed_per_cache_layer(): + model = _RecordingDraft(["layer.a", "layer.b"]) + speculator = _speculator(model) + hidden, positions = torch.zeros(3, 4), torch.zeros(3, dtype=torch.int64) + first, shared = torch.tensor([1, 2, 3]), torch.tensor([7, 8, 9]) + + speculator._precompute_context_kv(hidden, positions, [first, None]) + assert model.slots == {"layer.a": first} + + speculator._precompute_context_kv(hidden, positions, shared) + assert model.slots == {"layer.a": shared, "layer.b": shared} + + speculator._precompute_context_kv(hidden, positions, None) + assert model.slots is None + + +def test_greedy_step_maps_the_draft_argmax_to_target_ids(): + speculator = _speculator(_RecordingDraft([])) + logits = torch.tensor([[0.1, 0.9, 0.0], [0.7, 0.2, 0.1]]) + tokens = speculator._sample_step( + logits, torch.zeros(2, dtype=torch.int32), torch.ones(2, dtype=torch.int64), 0 + ) + assert tokens.tolist() == [101, 100] + + +def test_dummy_batch_at_a_context_bound_keeps_query_offsets(): + buffers = InputBuffers( + max_num_reqs=4, max_num_tokens=16, device=torch.device("cpu") + ) + batch = InputBatch.make_dummy(2, 12, buffers, context_len=4096) + assert batch.query_start_loc.tolist() == [0, 6, 12] + assert batch.seq_lens.tolist() == [4096, 4096] + assert batch.seq_lens_cpu_upper_bound.tolist() == [4096, 4096] + assert batch.num_computed_tokens_np.tolist() == [4090, 4090] + assert batch.capture_context_len == 4096 + + plain = InputBatch.make_dummy(2, 12, buffers) + assert plain.seq_lens.tolist() == [6, 6] + assert plain.capture_context_len is None + assert np.array_equal(plain.num_computed_tokens_np, np.zeros(2, dtype=np.int32)) + + +def _manager_with_buckets(buckets): + desc = BatchExecutionDescriptor(CUDAGraphMode.FULL, 6, 1, uniform_token_count=6) + variants = tuple( + BatchExecutionDescriptor( + CUDAGraphMode.FULL, 6, 1, uniform_token_count=6, attention_context_bucket=b + ) + for b in buckets + ) + manager = object.__new__(ModelCudaGraphManager) + manager._compressed_context_graphs = {desc: variants} + manager._qsa_short_graphs = {} + manager._long_attention_graphs = {} + manager.graphs = {variant: SimpleNamespace() for variant in variants} + return manager, desc, variants + + +def test_replay_selects_the_smallest_covering_context_bucket(): + manager, desc, (b2k, b4k, b16k) = _manager_with_buckets((2048, 4096, 16384)) + select = manager.select_attention_graph + assert select(desc, torch.tensor([300], dtype=torch.int32)) == b2k + assert select(desc, torch.tensor([2049], dtype=torch.int32)) == b4k + assert select(desc, torch.tensor([16384], dtype=torch.int32)) == b16k + # Beyond the largest bucket, or without a host bound: the full-length graph. + assert select(desc, torch.tensor([16385], dtype=torch.int32)) == desc + assert select(desc, torch.tensor([], dtype=torch.int32)) == desc + + +def test_replay_skips_buckets_that_were_not_captured(): + manager, desc, (b2k, b4k) = _manager_with_buckets((2048, 4096)) + del manager.graphs[b2k] + assert manager.select_attention_graph(desc, torch.tensor([100])) == b4k diff --git a/tests/v1/worker/test_sm70_long_attention_graphs.py b/tests/v1/worker/test_sm70_long_attention_graphs.py index 8164f5b15e..162a5fa378 100644 --- a/tests/v1/worker/test_sm70_long_attention_graphs.py +++ b/tests/v1/worker/test_sm70_long_attention_graphs.py @@ -26,6 +26,7 @@ def graph_pair(): ordinary = BatchExecutionDescriptor(CUDAGraphMode.FULL, 8, 1, 8) bounded = replace(ordinary, attention_context_bucket=BUILTIN_MAX_CONTEXT) manager._qsa_short_graphs = {} + manager._compressed_context_graphs = {} manager._long_attention_graphs = {ordinary: bounded} manager.graphs = {ordinary: object(), bounded: object()} return manager, ordinary, bounded @@ -51,6 +52,7 @@ def test_request_major_q8_context_boundary(num_reqs): ordinary = BatchExecutionDescriptor(CUDAGraphMode.FULL, num_reqs * 8, num_reqs, 8) bounded = replace(ordinary, attention_context_bucket=BUILTIN_MAX_CONTEXT) manager._qsa_short_graphs = {} + manager._compressed_context_graphs = {} manager._long_attention_graphs = {ordinary: bounded} manager.graphs = {ordinary: object(), bounded: object()} @@ -89,6 +91,7 @@ def test_scalar_graph_retains_short_and_oversized_fallback(): ordinary = BatchExecutionDescriptor(CUDAGraphMode.FULL, 1, 1, 1) bounded = replace(ordinary, attention_context_bucket=262144) manager._qsa_short_graphs = {} + manager._compressed_context_graphs = {} manager._long_attention_graphs = {ordinary: bounded} manager.graphs = {ordinary: object(), bounded: object()} for upper in (1, 1024, 131071, 262145):