From 0071867bc76af93adadec708a4dfbf5ef1cd0d2e Mon Sep 17 00:00:00 2001 From: EanWang211123 Date: Mon, 6 Jul 2026 20:32:18 +0800 Subject: [PATCH 01/11] [fix] fix step0 dsd support for dspkv4-dspark Signed-off-by: EanWang211123 --- tests/v1/spec_decode/test_dynamic_sd_cug.py | 49 +++++++++++++++++++++ vllm/v1/worker/gpu/cudagraph_utils.py | 12 +++-- 2 files changed, 58 insertions(+), 3 deletions(-) diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py index d75495ea606a..718630c0522a 100644 --- a/tests/v1/spec_decode/test_dynamic_sd_cug.py +++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py @@ -326,3 +326,52 @@ def test_dynamic_sd_only_captures_scheduled_query_lengths(monkeypatch): assert desc.num_tokens == num_tokens assert desc.num_reqs is None assert desc.num_active_loras == 0 + + +def test_dynamic_sd_skips_zero_draft_tokens_in_cudagraph_schedule(monkeypatch): + """K=0 in the DSD schedule must not produce decode_query_len=0. + + DSpark (anchor-as-first) passes ``num_query_per_req == num_speculative_tokens`` + to the draft CudaGraphManager, so ``num_new_sampled_tokens_per_step`` recovers + as 0. A schedule entry with K=0 would otherwise crash during candidate init. + """ + + max_num_seqs = 128 + max_spec_tokens = 5 + + monkeypatch.setattr( + gpu_cudagraph_utils, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + + vllm_config = _create_vllm_config_for_dsd( + max_num_seqs=max_num_seqs, + max_spec_tokens=max_spec_tokens, + cudagraph_mode="FULL_AND_PIECEWISE", + use_dynamic_sd=True, + num_spec_per_batch_size=[ + (1, 32, 5), + (33, 64, 3), + (65, 96, 1), + (97, 128, 0), + ], + ) + draft_decode_query_len = max_spec_tokens + + manager = gpu_cudagraph_utils.CudaGraphManager( + vllm_config=vllm_config, + device=torch.device("cpu"), + cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE, + decode_query_len=draft_decode_query_len, + ) + + scheduled_query_lens = {5, 3, 1} + captured_query_lens = { + desc.uniform_token_count + for descs in manager._candidates.values() + for desc in descs + if desc.cg_mode == CUDAGraphMode.FULL + and desc.uniform_token_count is not None + } + assert captured_query_lens == scheduled_query_lens diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index f54bfcc797b9..e92f56e159f8 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -226,9 +226,15 @@ def _init_candidates(self) -> None: self.decode_query_len - self.vllm_config.num_speculative_tokens ) # Each entry is (range_start, range_end, num_speculative_tokens). - decode_query_lens = [ - x[2] + num_new_sampled_tokens_per_step for x in num_spec_per_batch_size - ] + # K=0 disables drafting at that concurrency; no draft graph is + # needed, and a zero query length would break capture bucketing. + decode_query_lens = sorted( + { + x[2] + num_new_sampled_tokens_per_step + for x in num_spec_per_batch_size + if x[2] + num_new_sampled_tokens_per_step > 0 + } + ) elif ( speculative_config and speculative_config.uses_acceptance_length_adaptation() From 4b48898a1923f5f28fcb70a93c57c20a93a84045 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 10 Jul 2026 02:20:05 +0000 Subject: [PATCH 02/11] [Spec Decode] Harden DSpark metadata and TP sampling state --- .../test_deepseek_v4_dspark_metadata.py | 46 +++++++++++++++++++ .../worker/test_gpu_sampling_states_seed.py | 35 ++++++++++++++ vllm/v1/attention/backends/mla/sparse_swa.py | 6 ++- vllm/v1/worker/gpu/model_runner.py | 5 +- vllm/v1/worker/gpu/sample/sampler.py | 3 +- vllm/v1/worker/gpu/sample/states.py | 12 ++++- .../gpu/spec_decode/dflash/speculator.py | 15 ++++-- 7 files changed, 111 insertions(+), 11 deletions(-) create mode 100644 tests/v1/attention/test_deepseek_v4_dspark_metadata.py create mode 100644 tests/v1/worker/test_gpu_sampling_states_seed.py diff --git a/tests/v1/attention/test_deepseek_v4_dspark_metadata.py b/tests/v1/attention/test_deepseek_v4_dspark_metadata.py new file mode 100644 index 000000000000..c4b9437b0d2f --- /dev/null +++ b/tests/v1/attention/test_deepseek_v4_dspark_metadata.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import torch + +from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadataBuilder, +) +from vllm.v1.kv_cache_interface import MLAAttentionSpec + + +def test_dspark_swa_decode_threshold_matches_target_verification() -> None: + """DSpark verifies 1 + K target tokens, not the generic 1 + 2K.""" + speculative_config = SimpleNamespace( + num_speculative_tokens=5, + parallel_drafting=True, + use_dspark=lambda: True, + ) + hf_config = SimpleNamespace(sliding_window=128, compress_ratios=[1, 4, 128]) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(max_model_len=4096, hf_config=hf_config), + scheduler_config=SimpleNamespace(max_num_batched_tokens=16), + speculative_config=speculative_config, + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + cp_kv_cache_interleave_size=1, + ), + ) + kv_cache_spec = MLAAttentionSpec( + block_size=256, + num_kv_heads=1, + head_size=512, + dtype=torch.bfloat16, + ) + + builder = DeepseekSparseSWAMetadataBuilder( + kv_cache_spec, + ["placeholder"], + vllm_config, + torch.device("cpu"), + ) + + assert builder.decode_threshold == 6 diff --git a/tests/v1/worker/test_gpu_sampling_states_seed.py b/tests/v1/worker/test_gpu_sampling_states_seed.py new file mode 100644 index 000000000000..b87638312488 --- /dev/null +++ b/tests/v1/worker/test_gpu_sampling_states_seed.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import torch + +from vllm.sampling_params import SamplingParams +from vllm.v1.worker.gpu.sample import states + + +class _HostBackedTensor: + def __init__(self, size: int, dtype: torch.dtype): + self.cpu = torch.zeros(size, dtype=dtype) + self.np = self.cpu.numpy() + self.gpu = self.cpu + + def copy_to_uva(self, n: int | None = None) -> torch.Tensor: + return self.gpu[:n] if n is not None else self.gpu + + +def test_fallback_seeds_do_not_depend_on_global_numpy_rng(monkeypatch) -> None: + monkeypatch.setattr(states, "UvaBackedTensor", _HostBackedTensor) + rank0 = states.SamplingStates(4, 128, seed=17) + + np.random.seed(1234) + np.random.random(1000) + rank1 = states.SamplingStates(4, 128, seed=17) + + params = SamplingParams(seed=None) + for req_idx in range(4): + rank0.add_request(req_idx, params) + np.random.random(req_idx + 1) + rank1.add_request(req_idx, params) + + np.testing.assert_array_equal(rank0.seeds.np, rank1.seeds.np) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 5c0116caadbf..a890d62cb528 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -306,8 +306,10 @@ def __init__(self, *args, **kwargs): self.vllm_config.scheduler_config.max_num_batched_tokens ) - # Keep the split aligned with target verification. DSpark verifies the - # sampled token plus K draft tokens, even though it drafts in parallel. + # Keep the decode/prefill split identical to the DeepSeek V4 C128A + # metadata and indexer. Target verification contains the bonus token + # plus N speculative tokens even for parallel drafters such as DSpark; + # the generic parallel-drafting threshold (1 + 2N) is not applicable. spec_config = self.vllm_config.speculative_config self.num_speculative_tokens = ( spec_config.num_speculative_tokens if spec_config else 0 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b893f640a8a1..7e6e8df69355 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -381,6 +381,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: logprobs_mode=self.model_config.logprobs_mode, num_speculative_tokens=self.decode_query_len, use_fp64_gumbel=self.model_config.use_fp64_gumbel, + seed=self.model_config.seed, ) custom = self.model_state.custom_sampler(self.sampler) @@ -824,7 +825,9 @@ def finish_requests(self, scheduler_output: SchedulerOutput) -> None: preempted_req_ids = scheduler_output.preempted_req_ids if preempted_req_ids: finished_req_ids = finished_req_ids.union(preempted_req_ids) - for req_id in finished_req_ids: + # A set's order can differ across TP processes. Recycle slots in a + # deterministic order so request-to-slot state stays rank-aligned. + for req_id in sorted(finished_req_ids): self._remove_request(req_id) def free_states(self, scheduler_output: SchedulerOutput) -> None: diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index b269de9eaed0..09fd92e3264a 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -37,6 +37,7 @@ def __init__( logprobs_mode: LogprobsMode = "raw_logprobs", num_speculative_tokens: int = 1, use_fp64_gumbel: bool = False, + seed: int | None = None, ): if logprobs_mode not in ("processed_logprobs", "raw_logprobs"): raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}") @@ -45,7 +46,7 @@ def __init__( self.use_fp64_gumbel = use_fp64_gumbel self.req_states = req_states - self.sampling_states = SamplingStates(max_num_reqs, vocab_size) + self.sampling_states = SamplingStates(max_num_reqs, vocab_size, seed) self.penalties_state = PenaltiesState(req_states) self.logit_bias_state = LogitBiasState(max_num_reqs, device) self.bad_words_state = BadWordsState(req_states) diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index fe4dee6a6b10..9a42a1191833 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -15,10 +15,14 @@ class SamplingStates: - def __init__(self, max_num_reqs: int, vocab_size: int): + def __init__(self, max_num_reqs: int, vocab_size: int, seed: int | None = None): self.max_num_reqs = max_num_reqs self.vocab_size = vocab_size + # Every TP rank must derive the same fallback request seeds. A private + # stream avoids rank-local consumers perturbing NumPy's global RNG. + self._fallback_seed_rng = np.random.default_rng(seed if seed is not None else 0) + self.temperature = UvaBackedTensor(max_num_reqs, dtype=torch.float32) self.top_k = UvaBackedTensor(max_num_reqs, dtype=torch.int32) self.top_p = UvaBackedTensor(max_num_reqs, dtype=torch.float32) @@ -50,7 +54,11 @@ def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None: seed = sampling_params.seed self.seeds_set[req_idx] = seed is not None if seed is None: - seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX) + seed = int( + self._fallback_seed_rng.integers( + _NP_INT64_MIN, _NP_INT64_MAX, dtype=np.int64 + ) + ) self.seeds.np[req_idx] = seed num_logprobs = sampling_params.logprobs diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index c4ac1c3ede70..93ee8ce0aaef 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -89,8 +89,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.sample_pos = torch.zeros( max_num_sampled_tokens, dtype=torch.int64, device=device ) - self.sample_idx_mapping = torch.zeros( - max_num_sampled_tokens, dtype=torch.int32, device=device + # -1 marks an inert sampling row. CUDA graph capture can execute the + # full buffer before a real batch has populated it, so zero would make + # every padding row race while scattering into request slot 0. + self.sample_idx_mapping = torch.full( + (max_num_sampled_tokens,), -1, dtype=torch.int32, device=device ) # [0, 1, ..., N-1, 0, 1, ..., N-1, ...] -> the per-token column index into # draft_logits[req, step, :]. @@ -117,11 +120,11 @@ def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: def capture(self, attn_states: dict | None = None) -> None: logger.info("Capturing model for %s speculator...", self._speculator_name) - # Reset sampling indices to zero to prevent stale values from prior - # dummy runs from being baked into the captured graph. + # Reset sampling indices to prevent stale values from prior dummy runs + # from being baked into the captured graph. Mapping rows stay inert. self.sample_indices.zero_() self.sample_pos.zero_() - self.sample_idx_mapping.zero_() + self.sample_idx_mapping.fill_(-1) assert self.query_cudagraph_manager is not None self.query_cudagraph_manager.capture( self._generate_draft, @@ -607,6 +610,8 @@ def _prepare_dflash_inputs_kernel( for i in range(q_pad_start, max_num_tokens, BLOCK_SIZE): block = i + tl.arange(0, BLOCK_SIZE) mask = block < max_num_tokens + tl.store(out_input_ids_ptr + block, 0, mask=mask) + tl.store(out_query_positions_ptr + block, 0, mask=mask) tl.store(out_query_slot_mapping_ptr + block, PAD_SLOT_ID, mask=mask) From 357fa29f004cec425bcb28bc69d524bfb66473d6 Mon Sep 17 00:00:00 2001 From: giorgiopiatti-dfinity Date: Tue, 7 Jul 2026 21:04:35 +0200 Subject: [PATCH 03/11] [Bugfix][Spec Decode] Mask cache-restored tokens out of DFlash draft context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DFlash/DSpark build the draft's context KV from target aux hidden states, which only exist for tokens that flow through a target forward pass. Tokens restored from the prefix cache (or a KV connector) at request (re)admission never do, so their draft KV slots are never written — yet the draft attends over the full sequence. With automatic prefix caching and a long shared prefix, the draft reads thousands of uninitialized slots and acceptance collapses to ~0.3% (position-0 only); the same workload with unique prompts reaches ~20%. MTP is unaffected (no context KV), which hid the interaction. Fix: track per request-slot how many tokens were restored at the last (re)admission (RequestState.num_cached_tokens) and hide the restored whole blocks from the draft's attention — the prep kernel shortens the draft seq_lens and a new kernel left-shifts the draft block-table rows in place (safe: input_block_tables are regathered every step, and the shift runs after slot mappings are computed from the unshifted table). Draft KV stores post-RoPE keys at absolute positions, so no position rewriting is needed. Requests without cache hits and dense DFlash/DSpark setups are unaffected (shift 0). Up to block_size - 1 restored slots stay visible when the restored count is not block-aligned (e.g. full-prompt hits). The draft loses the cached prefix from its context (bounded by its training window anyway) in exchange for prefix caching and speculative decoding composing at all. A durable alternative — letting the draft KV cache group participate in prefix-cache block reuse — is left for a follow-up RFC. Co-Authored-By: Claude Fable 5 Signed-off-by: giorgiopiatti-dfinity --- .../test_dflash_prefix_cache_masking.py | 155 ++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 6 + .../gpu/spec_decode/dflash/speculator.py | 131 ++++++++++++++- vllm/v1/worker/gpu/states.py | 11 ++ 4 files changed, 301 insertions(+), 2 deletions(-) create mode 100644 tests/v1/spec_decode/test_dflash_prefix_cache_masking.py diff --git a/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py new file mode 100644 index 000000000000..5fddac1c802c --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DFlash/DSpark draft context masking under prefix caching. + +Cache-restored tokens never flow through the target forward, so the draft's +context KV is never written for them. shift_draft_block_tables hides those +slots from the draft's attention by left-shifting each request's block-table +row by the restored whole blocks (seq_lens is shortened to match by +_prepare_dflash_inputs_kernel). +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( + shift_draft_block_tables, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="Requires CUDA" +) + +DEVICE = "cuda" +BLOCK_SIZE = 16 +MAX_BLOCKS = 64 +MAX_NUM_REQS = 8 + + +def _make_block_table(num_reqs: int) -> torch.Tensor: + # Distinct block ids per (request, slot) so shifts are detectable. + table = torch.arange( + MAX_NUM_REQS * MAX_BLOCKS, dtype=torch.int32, device=DEVICE + ).view(MAX_NUM_REQS, MAX_BLOCKS) + return table[:num_reqs].contiguous() + + +@pytest.mark.parametrize( + "num_cached,expected_shift", + [ + (0, 0), # no cache hit: no-op + (BLOCK_SIZE * 3, 3), # block-aligned hit (the common APC case) + (BLOCK_SIZE * 3 + 5, 3), # unaligned: floor to whole blocks + (BLOCK_SIZE - 1, 0), # less than one block: no-op + ], +) +def test_shift_single_request(num_cached: int, expected_shift: int): + block_table = _make_block_table(1) + original = block_table.clone() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE) + num_cached_tokens = torch.full( + (MAX_NUM_REQS,), num_cached, dtype=torch.int32, device=DEVICE + ) + + seq_lens = torch.full( + (idx_mapping.shape[0],), + MAX_BLOCKS * BLOCK_SIZE, + dtype=torch.int32, + device=DEVICE, + ) + shift_draft_block_tables( + block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE + ) + + kept = MAX_BLOCKS - expected_shift + torch.testing.assert_close( + block_table[0, :kept], original[0, expected_shift:] + ) + + +def test_shift_per_request_and_idx_mapping(): + # Requests in batch order 0..3 map to request-state slots 3..0, with a + # different cached count per slot. Each row must shift by its own count. + num_reqs = 4 + block_table = _make_block_table(num_reqs) + original = block_table.clone() + idx_mapping = torch.tensor([3, 2, 1, 0], dtype=torch.int32, device=DEVICE) + # Slot i has i whole cached blocks. + num_cached_tokens = torch.zeros( + MAX_NUM_REQS, dtype=torch.int32, device=DEVICE + ) + num_cached_tokens[:4] = ( + torch.arange(4, dtype=torch.int32, device=DEVICE) * BLOCK_SIZE + ) + + seq_lens = torch.full( + (idx_mapping.shape[0],), + MAX_BLOCKS * BLOCK_SIZE, + dtype=torch.int32, + device=DEVICE, + ) + shift_draft_block_tables( + block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE + ) + + for batch_idx in range(num_reqs): + shift = int(idx_mapping[batch_idx]) # slot id == cached blocks + kept = MAX_BLOCKS - shift + torch.testing.assert_close( + block_table[batch_idx, :kept], + original[batch_idx, shift:], + msg=f"batch row {batch_idx} (slot {shift})", + ) + + +def test_shift_large_row_in_place_overlap(): + # Shift smaller than the copy chunk (1024) exercises the overlapping + # in-place load-before-store path on a long row. + max_blocks = 4096 + block_table = ( + torch.arange(max_blocks, dtype=torch.int32, device=DEVICE) + .unsqueeze(0) + .contiguous() + ) + original = block_table.clone() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE) + num_cached_tokens = torch.full( + (1,), 7 * BLOCK_SIZE, dtype=torch.int32, device=DEVICE + ) + + seq_lens = torch.full( + (idx_mapping.shape[0],), + max_blocks * BLOCK_SIZE, + dtype=torch.int32, + device=DEVICE, + ) + shift_draft_block_tables( + block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE + ) + + torch.testing.assert_close( + block_table[0, : max_blocks - 7], original[0, 7:] + ) + + +def test_shift_copy_bounded_by_seq_len(): + # Only the blocks referenced by the shifted sequence move; the tail of the + # row must stay untouched (perf guard for long-context block tables). + block_table = _make_block_table(1) + original = block_table.clone() + idx_mapping = torch.zeros(1, dtype=torch.int32, device=DEVICE) + num_cached_tokens = torch.full( + (MAX_NUM_REQS,), 4 * BLOCK_SIZE, dtype=torch.int32, device=DEVICE + ) + # Shifted draft length of 3.5 blocks -> exactly 4 blocks copied. + seq_lens = torch.full( + (1,), 3 * BLOCK_SIZE + BLOCK_SIZE // 2, dtype=torch.int32, device=DEVICE + ) + + shift_draft_block_tables( + block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE + ) + + torch.testing.assert_close(block_table[0, :4], original[0, 4:8]) + torch.testing.assert_close(block_table[0, 4:], original[0, 4:]) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7e6e8df69355..bb1c1069dec7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -534,6 +534,12 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.speculator.set_attn( self.model_state, self.kv_cache_config, self.block_tables ) + if hasattr(self.speculator, "set_num_cached_tokens"): + # DFlash/DSpark mask cache-restored tokens out of the draft's + # context (their draft context KV was never computed). + self.speculator.set_num_cached_tokens( + self.req_states.num_cached_tokens.gpu + ) self.kv_caches: list[torch.Tensor] = [] kv_caches_dict = init_kv_cache( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 93ee8ce0aaef..c7e9466dffe5 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -81,6 +81,16 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.max_num_tokens, dtype=torch.int64, device=device ) + # Per-request-slot count of tokens whose KV was restored (e.g. from the + # prefix cache) at the request's last (re)admission, indexed by + # req_state_idx. The target never ran a forward pass over them, so + # their draft context KV was never computed; the prep kernel and the + # block-table shift in propose() hide them from the draft's attention. + # The runner replaces this zeros fallback via set_num_cached_tokens. + self.num_cached_tokens = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=device + ) + # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens). max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps self.sample_indices = torch.zeros( @@ -150,6 +160,13 @@ def load_draft_model( ) return model + def set_num_cached_tokens(self, num_cached_tokens: torch.Tensor) -> None: + """Register the runner's per-request-slot cache-restored token counts. + + Indexed by req_state_idx; see the buffer comment in __init__. + """ + self.num_cached_tokens = num_cached_tokens + def set_attn( self, model_state: ModelState, @@ -163,6 +180,17 @@ def set_attn( ] assert self.draft_kv_cache_group_ids, "No draft attention groups found." self.draft_kv_cache_group_id = self.draft_kv_cache_group_ids[0] + # The shared seq_lens buffer carries the cache-shifted draft sequence + # lengths (see _prepare_dflash_inputs_kernel), which only works if all + # draft groups shift by the same number of slots per cached block. + draft_block_sizes = { + self.block_tables.kernel_block_sizes[gid] + for gid in self.draft_kv_cache_group_ids + } + assert len(draft_block_sizes) == 1, ( + "DFlash requires a uniform block size across draft KV cache " + f"groups, got {draft_block_sizes}." + ) # Per-group context slot buffers for the precompute (one row per group). self._context_slot_mappings = torch.zeros( @@ -373,6 +401,7 @@ def propose( next_prefill_tokens, self.block_tables.input_block_tables[gid], self.block_tables.kernel_block_sizes[gid], + self.num_cached_tokens, self.parallel_drafting_token_id, self.num_query_per_req, self.num_speculative_steps, @@ -382,6 +411,27 @@ def propose( self.sample_from_anchor, ) + # Cache-restored tokens (e.g. prefix-cache hits) never flowed through + # the target forward, so their draft context KV was never written and + # their cache slots hold garbage. Hide them from the draft's + # attention: shift each draft block-table row left by the restored + # whole blocks (the prep kernel shortened seq_lens to match). Runs + # after prepare_dflash_inputs because the slot mappings index the + # unshifted table; in-place is safe because input_block_tables are + # regathered from the persistent block tables every step. Up to + # block_size - 1 restored slots may remain visible when the restored + # count is not block-aligned (e.g. a full-prompt cache hit). Skipped + # for dummy runs, whose idx_mapping does not reference live requests. + if not dummy_run: + for gid in self.draft_kv_cache_group_ids: + shift_draft_block_tables( + self.block_tables.input_block_tables[gid], + input_batch.idx_mapping, + self.num_cached_tokens, + self.input_buffers.seq_lens, + self.block_tables.kernel_block_sizes[gid], + ) + # Pre-insert context K/V into the cache. Runs eagerly outside the captured graph # because the context shape varies per step. During dummy runs the block tables # are placeholders, so we skip the cache write to avoid clobbering real entries. @@ -472,6 +522,8 @@ def _prepare_dflash_inputs_kernel( # Block table for slot mapping lookup. block_table_ptr, block_table_stride, + # [max_num_reqs] cache-restored token counts, indexed by req_state_idx. + num_cached_tokens_ptr, # Scalars parallel_drafting_token_id, block_size, @@ -578,8 +630,20 @@ def _prepare_dflash_inputs_kernel( tl.store(out_query_start_loc_ptr + req_idx, query_base) # seq_lens is the absolute sequence length the draft attention # reads up to (context + query), not just the count of accepted - # tokens this step. - tl.store(out_seq_lens_ptr + req_idx, last_valid_pos + 1 + num_query_per_req) + # tokens this step — minus the cache-restored whole blocks, which + # hold no draft KV and are shifted out of the block table (see + # shift_draft_block_tables). + num_cached = tl.load(num_cached_tokens_ptr + req_state_idx) + num_shifted_slots = (num_cached // block_size) * block_size + # The clamp guards dummy runs, where req_state_idx may point at a + # stale slot whose cached count exceeds the dummy sequence length. + tl.store( + out_seq_lens_ptr + req_idx, + tl.maximum( + last_valid_pos + 1 + num_query_per_req - num_shifted_slots, + num_query_per_req, + ), + ) if req_idx == num_reqs - 1: # Pad per-request buffers to max_num_reqs for CUDA graph safety. last_query_end = num_reqs * num_query_per_req @@ -635,6 +699,8 @@ def prepare_dflash_inputs( # [max_num_reqs, max_num_blocks] block_table: torch.Tensor, block_size: int, + # [max_num_reqs] + num_cached_tokens: torch.Tensor, parallel_drafting_token_id: int, num_query_per_req: int, num_speculative_steps: int, @@ -671,6 +737,7 @@ def prepare_dflash_inputs( num_rejected, block_table, block_table.stride(0), + num_cached_tokens, parallel_drafting_token_id, block_size, num_query_per_req, @@ -682,3 +749,63 @@ def prepare_dflash_inputs( PAD_SLOT_ID=PAD_SLOT_ID, BLOCK_SIZE=BLOCK_SIZE, ) + + +@triton.jit +def _shift_draft_block_tables_kernel( + block_table_ptr, + block_table_stride, + idx_mapping_ptr, + num_cached_tokens_ptr, + seq_lens_ptr, + block_size, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) + num_cached = tl.load(num_cached_tokens_ptr + req_state_idx) + shift = num_cached // block_size + if shift == 0: + return + row_ptr = block_table_ptr + req_idx.to(tl.int64) * block_table_stride + # Only the blocks the shifted sequence still references need to move; + # seq_lens holds the cache-shifted draft length (written by + # _prepare_dflash_inputs_kernel, which must run first). + seq_len = tl.load(seq_lens_ptr + req_idx) + num_needed = (seq_len + block_size - 1) // block_size + num_remaining = tl.minimum(block_table_stride - shift, num_needed) + # In-place left shift is safe: iterations run in ascending order and each + # loads its chunk (from offset + shift) before storing (at offset), so no + # store ever precedes a load of the same element. + for i in tl.range(0, num_remaining, BLOCK_SIZE): + offset = i + tl.arange(0, BLOCK_SIZE) + mask = offset < num_remaining + block_ids = tl.load(row_ptr + offset + shift, mask=mask, other=0) + tl.store(row_ptr + offset, block_ids, mask=mask) + + +def shift_draft_block_tables( + # [max_num_reqs, max_num_blocks] + block_table: torch.Tensor, + # [num_reqs] + idx_mapping: torch.Tensor, + # [max_num_reqs] + num_cached_tokens: torch.Tensor, + # [num_reqs] cache-shifted draft sequence lengths + seq_lens: torch.Tensor, + block_size: int, +) -> None: + """Shift each request's draft block-table row left by its cache-restored + whole blocks, hiding slots that hold no draft context KV from the draft's + attention. Must run after prepare_dflash_inputs (slot mappings index the + unshifted table, and seq_lens must already hold the shifted lengths).""" + num_reqs = idx_mapping.shape[0] + _shift_draft_block_tables_kernel[(num_reqs,)]( + block_table, + block_table.stride(0), + idx_mapping, + num_cached_tokens, + seq_lens, + block_size, + BLOCK_SIZE=1024, # type: ignore + ) diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 7f0ae33c8099..1bc10444654d 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -80,6 +80,15 @@ def __init__( self.max_num_reqs, dtype=torch.int32, device=device ) + # Tokens whose KV was restored (e.g. from the prefix cache) rather than + # computed at the request's most recent (re)admission. The target never + # runs a forward pass over them, so speculators that derive per-token + # state from target hidden states (DFlash/DSpark context KV) have + # nothing for these positions. + self.num_cached_tokens = StagedWriteTensor( + self.max_num_reqs, dtype=torch.int32, device=device + ) + @property def num_reqs(self) -> int: return len(self.req_id_to_index) @@ -109,6 +118,7 @@ def add_request( self.num_computed_prefill_tokens[req_idx] = num_computed_tokens self.num_computed_tokens_np[req_idx] = num_computed_tokens self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens) + self.num_cached_tokens.stage_write_elem(req_idx, num_computed_tokens) self.draft_tokens[req_idx].zero_() @@ -118,6 +128,7 @@ def apply_staged_writes(self) -> None: self.total_len.apply_write() self.all_token_ids.apply_write() self.num_computed_tokens.apply_write() + self.num_cached_tokens.apply_write() def remove_request(self, req_id: str) -> int | None: """Return the freed slot index, or None if the request was not found.""" From ed5f7c84f820ed22da0cdad4414bfa982219d08a Mon Sep 17 00:00:00 2001 From: mgoin Date: Thu, 9 Jul 2026 20:23:24 +0000 Subject: [PATCH 04/11] Prefer FlashAttn over FlashInfer for SM100f non-causal attention Signed-off-by: mgoin --- tools/pre_commit/generate_attention_backend_docs.py | 4 +++- vllm/platforms/cuda.py | 7 ++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 6f6e1341ac14..4d72c086a0be 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -1341,7 +1341,9 @@ def _get_backends_from_return(stmts: list) -> list[str]: def _is_sm100_check(test: ast.expr) -> bool: - """Check if test is `something.major == 10`.""" + """Check if test is `something.major == 10`, possibly inside an `and`.""" + if isinstance(test, ast.BoolOp) and isinstance(test.op, ast.And): + return any(_is_sm100_check(value) for value in test.values) return ( isinstance(test, ast.Compare) and isinstance(test.left, ast.Attribute) diff --git a/vllm/platforms/cuda.py b/vllm/platforms/cuda.py index 9eac95e03249..4059a7e4f50d 100644 --- a/vllm/platforms/cuda.py +++ b/vllm/platforms/cuda.py @@ -85,6 +85,7 @@ def _get_backend_priorities( device_capability: DeviceCapability, num_heads: int | None = None, kv_cache_dtype: CacheDType | None = None, + use_non_causal: bool = False, ) -> list[AttentionBackendEnum]: """Get backend priorities with lazy import to avoid circular dependency.""" from vllm.utils.torch_utils import is_quantized_kv_cache @@ -141,7 +142,10 @@ def _get_backend_priorities( AttentionBackendEnum.FLASHMLA_SPARSE, ] else: - if device_capability.major == 10: + # SM100f defaults to FlashInfer for TRTLLM causal attention, but its non-causal + # cutlass path (used for dflash attention) is known to have problems. + # So prefer FlashAttention when non-causal on SM100f. + if device_capability.major == 10 and not use_non_causal: return [ AttentionBackendEnum.FLASHINFER, AttentionBackendEnum.FLASH_ATTN, @@ -368,6 +372,7 @@ def get_valid_backends( device_capability, num_heads, attn_selector_config.kv_cache_dtype, + attn_selector_config.use_non_causal, ) for priority, backend in enumerate(backend_priorities): try: From 4b248b030380590c814800516b4b347c82ec7886 Mon Sep 17 00:00:00 2001 From: mgoin Date: Thu, 9 Jul 2026 22:59:34 +0000 Subject: [PATCH 05/11] [Spec Decode] Never full-graph-capture non-causal FlashInfer draft attention Non-causal draft attention (DFlash/DSpark) skips trtllm-gen and runs the FlashInfer prefill wrapper, whose run() is not replay-safe once plan() changes; replaying a full CUDA graph then returns wrong output or an illegal memory access. Only claim UNIFORM_BATCH cudagraph support for causal attention, build draft attention metadata under the draft's attention config, and fall back to eager draft attention when full graphs are unsupported. Signed-off-by: mgoin Co-Authored-By: Claude Fable 5 --- vllm/v1/attention/backends/flashinfer.py | 13 ++++---- vllm/v1/worker/gpu/model_runner.py | 7 +++-- .../gpu/spec_decode/dflash/speculator.py | 30 +++++++++++++++++-- vllm/v1/worker/gpu/spec_decode/speculator.py | 10 +++++-- 4 files changed, 46 insertions(+), 14 deletions(-) diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 12eab21e3e13..36b1d724eaf2 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -9,17 +9,17 @@ import numpy as np import torch +from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache +from flashinfer.prefill import trtllm_batch_context_with_kv_cache +from flashinfer.utils import FP4Tensor +from typing_extensions import override + from flashinfer import ( BatchDecodeWithPagedKVCacheWrapper, BatchPrefillWithPagedKVCacheWrapper, BatchPrefillWithRaggedKVCacheWrapper, MultiLevelCascadeAttentionWrapper, ) -from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache -from flashinfer.prefill import trtllm_batch_context_with_kv_cache -from flashinfer.utils import FP4Tensor -from typing_extensions import override - from vllm import _custom_ops as custom_ops from vllm import envs from vllm.config import ( @@ -885,7 +885,8 @@ def get_cudagraph_support( has_trtllm_support = False break - if has_trtllm_support: + # trtllm-gen only supports causal attention. + if has_trtllm_support and not vllm_config.attention_config.use_non_causal: return AttentionCGSupport.UNIFORM_BATCH else: return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index bb1c1069dec7..798cd037f357 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -525,9 +525,6 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, ) - if self.speculator is not None: - self.speculator.init_cudagraph_manager(cudagraph_mode) - check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): # HACK(woosuk) @@ -540,6 +537,10 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.speculator.set_num_cached_tokens( self.req_states.num_cached_tokens.gpu ) + if self.speculator is not None: + # After set_attn, so the speculator can size its cudagraph mode + # to its own attention support. + self.speculator.init_cudagraph_manager(cudagraph_mode) self.kv_caches: list[torch.Tensor] = [] kv_caches_dict = init_kv_cache( diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index c7e9466dffe5..66cf4663d6af 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -15,11 +15,12 @@ import torch import torch.nn as nn -from vllm.config import VllmConfig +from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger from vllm.triton_utils import tl, triton +from vllm.v1.attention.backend import AttentionCGSupport from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer @@ -114,9 +115,32 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.query_cudagraph_manager: DFlashCudaGraphManager | None = None self.draft_kv_cache_group_id: int = -1 + @property + def attn_vllm_config(self) -> VllmConfig: + # The draft's attention differs from the target's in causality. + return replace( + self.vllm_config, + attention_config=replace( + self.vllm_config.attention_config, + use_non_causal=not self.dflash_causal, + ), + ) + def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - # PIECEWISE cudagraphs are not supported for dflash - if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL: + wants_full = cudagraph_mode.decode_mode() == CUDAGraphMode.FULL + supports_full = ( + self.attn_cg_support.min_cg_support.value + >= AttentionCGSupport.UNIFORM_BATCH.value + ) + if wants_full and not supports_full: + logger.warning( + "%s draft attention (%s) does not support full CUDA graphs; " + "running the draft eagerly.", + self._speculator_name, + self.attn_cg_support.min_cg_attn_backend, + ) + # PIECEWISE cudagraphs are not supported for dflash. + if wants_full and supports_full: cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY else: cudagraph_mode = CUDAGraphMode.NONE diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 320e05fd4792..6bcdea5ee484 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -185,6 +185,12 @@ def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: num_unpadded_tokens, ) + @property + def attn_vllm_config(self) -> VllmConfig: + """Config for the draft's attention metadata builders. Overridden by + speculators whose attention mode differs from the target's.""" + return self.vllm_config + def set_attn( self, model_state: ModelState, @@ -193,9 +199,9 @@ def set_attn( ) -> None: self.model_state = model_state self.kv_cache_config = kv_cache_config - self.attn_groups, _, _ = init_attn_backend( + self.attn_groups, self.attn_cg_support, _ = init_attn_backend( kv_cache_config, - self.vllm_config, + self.attn_vllm_config, self.device, active_layer_names=self.draft_attn_layer_names, ) From 1d13bf67648a1deaa53a38d910812ad5e357c842 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 09:36:38 +0000 Subject: [PATCH 06/11] fix(spec decode): preserve explicit zero adaptive depth --- tests/v1/spec_decode/test_acceptance_length_controller.py | 3 +++ vllm/v1/core/sched/async_scheduler.py | 6 +++--- vllm/v1/core/sched/output.py | 6 ++++-- vllm/v1/core/sched/scheduler.py | 4 +++- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/tests/v1/spec_decode/test_acceptance_length_controller.py b/tests/v1/spec_decode/test_acceptance_length_controller.py index 42b020db62a6..e53baa3d0d5b 100644 --- a/tests/v1/spec_decode/test_acceptance_length_controller.py +++ b/tests/v1/spec_decode/test_acceptance_length_controller.py @@ -268,6 +268,9 @@ def test_synthetic_scheduler_output_uses_default_speculative_depth(): output.num_spec_tokens_to_schedule = 2 assert output.resolve_num_spec_tokens_to_schedule(default=5) == 2 + output.num_spec_tokens_to_schedule = 0 + assert output.resolve_num_spec_tokens_to_schedule(default=5) == 0 + def test_runner_v2_autoregressive_drafter_stops_at_adaptive_depth(monkeypatch): monkeypatch.setattr(AutoRegressiveSpeculator, "__abstractmethods__", frozenset()) diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index d1c652c46efa..41a231eaa448 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -20,9 +20,9 @@ def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: super()._update_after_schedule(scheduler_output) spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens # Use the latest num of scheduled draft tokens in next step as placeholder. - self._spec_token_placeholders = [ - -1 - ] * scheduler_output.num_spec_tokens_to_schedule + self._spec_token_placeholders = [-1] * ( + scheduler_output.resolve_num_spec_tokens_to_schedule(self.num_spec_tokens) + ) for req_id in scheduler_output.num_scheduled_tokens: request = self.requests[req_id] if request.is_prefill_chunk: diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index ad6da702e26b..7cc60558f0dc 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -242,7 +242,7 @@ class SchedulerOutput: # Dynamic speculative decoding: optimal K chosen by scheduler. # Number of spec tokens to schedule for the next step. - num_spec_tokens_to_schedule: int = 0 + num_spec_tokens_to_schedule: int | None = None @classmethod def make_empty(cls) -> "SchedulerOutput": @@ -260,7 +260,9 @@ def make_empty(cls) -> "SchedulerOutput": def resolve_num_spec_tokens_to_schedule(self, default: int) -> int: """Resolve the speculative depth for real and synthetic outputs.""" - return self.num_spec_tokens_to_schedule or default + if self.num_spec_tokens_to_schedule is None: + return default + return self.num_spec_tokens_to_schedule @dataclass diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index a553b7946a51..84d63d6693f6 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1816,7 +1816,9 @@ def update_from_output( spec_decoding_stats.current_num_spec_tokens = ( self.acceptance_length_controller.num_spec_tokens if self.acceptance_length_controller is not None - else scheduler_output.num_spec_tokens_to_schedule + else scheduler_output.resolve_num_spec_tokens_to_schedule( + self.num_spec_tokens + ) ) # Remove the stopped requests from the running and waiting queues. From 7916aaf88ceef29ec1dbd7b31ccd4b83a632fe45 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 10 Jul 2026 09:03:46 +0000 Subject: [PATCH 07/11] spec_decode: retain DFlash CUDA graph backbone outputs --- .../test_dflash_cudagraph_lifetime.py | 66 +++++++++++++++++++ .../gpu/spec_decode/dflash/speculator.py | 6 ++ 2 files changed, 72 insertions(+) create mode 100644 tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py diff --git a/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py b/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py new file mode 100644 index 000000000000..1d28950a762a --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator + + +def _make_speculator() -> SimpleNamespace: + hidden_states = torch.randn(2, 8) + return SimpleNamespace( + _run_model=Mock(return_value=hidden_states), + _captured_backbone_outputs=[], + num_speculative_steps=2, + sample_indices=torch.tensor([0, 1]), + sample_pos=torch.tensor([1, 2]), + sample_idx_mapping=torch.tensor([0, 0]), + temperature=torch.ones(1), + seeds=torch.zeros(1, dtype=torch.int64), + sample_col=torch.tensor([0, 1]), + draft_logits=None, + sample_draft=Mock(return_value=torch.tensor([11, 12])), + draft_tokens=torch.zeros(1, 2, dtype=torch.int64), + ) + + +def test_dflash_retains_backbone_output_during_cudagraph_capture(monkeypatch): + speculator = _make_speculator() + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + DFlashSpeculator._generate_draft( + speculator, + num_reqs=1, + num_tokens_padded=2, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + + assert len(speculator._captured_backbone_outputs) == 1 + assert ( + speculator._captured_backbone_outputs[0] + is speculator._run_model.return_value + ) + + +def test_dflash_does_not_retain_eager_backbone_output(monkeypatch): + speculator = _make_speculator() + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + DFlashSpeculator._generate_draft( + speculator, + num_reqs=1, + num_tokens_padded=2, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + ) + + assert speculator._captured_backbone_outputs == [] diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 66cf4663d6af..64db4234aa0b 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -114,6 +114,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.query_cudagraph_manager: DFlashCudaGraphManager | None = None self.draft_kv_cache_group_id: int = -1 + # Manual CUDA graphs keep raw addresses for model intermediates. Retain + # each captured backbone output so its storage cannot be recycled while + # a graph still reads it during sampling. + self._captured_backbone_outputs: list[torch.Tensor] = [] @property def attn_vllm_config(self) -> VllmConfig: @@ -291,6 +295,8 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) + if torch.cuda.is_current_stream_capturing(): + self._captured_backbone_outputs.append(last_hidden_states) num_sample = num_reqs * self.num_speculative_steps sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] From 7a8830ecfd156cda59f12d17dd6d8ecfe30c1cf8 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 09:38:48 +0000 Subject: [PATCH 08/11] test(spec decode): use current dynamic-depth predicate --- tests/v1/spec_decode/test_dynamic_sd_cug.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py index 718630c0522a..5f6b22e6cac9 100644 --- a/tests/v1/spec_decode/test_dynamic_sd_cug.py +++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py @@ -62,7 +62,10 @@ def _create_vllm_config_for_dsd( vllm_config.num_speculative_tokens = max_spec_tokens speculative_config = MagicMock() - speculative_config.uses_dynamic_speculative_decoding.return_value = use_dynamic_sd + speculative_config.uses_batch_size_dynamic_speculative_decoding.return_value = ( + use_dynamic_sd + ) + speculative_config.uses_acceptance_length_adaptation.return_value = False if use_dynamic_sd: # DSD reads the per-batch-size schedule; a schedule entry with K # speculative tokens maps to decode query length K + 1. By default From 3571604ddd0a6447df1d1b05dd5ddfbd2c5fb392 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 09:42:48 +0000 Subject: [PATCH 09/11] fix(dflash): serialize overlapping block-table shift loads --- vllm/v1/worker/gpu/spec_decode/dflash/speculator.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 64db4234aa0b..bec2310f4c9a 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -807,10 +807,21 @@ def _shift_draft_block_tables_kernel( # In-place left shift is safe: iterations run in ascending order and each # loads its chunk (from offset + shift) before storing (at offset), so no # store ever precedes a load of the same element. - for i in tl.range(0, num_remaining, BLOCK_SIZE): + # Keep iterations strictly ordered. Compiler software pipelining may start + # a store before a later overlapping source load has completed. + for i in tl.range( + 0, + num_remaining, + BLOCK_SIZE, + num_stages=1, + loop_unroll_factor=1, + ): offset = i + tl.arange(0, BLOCK_SIZE) mask = offset < num_remaining block_ids = tl.load(row_ptr + offset + shift, mask=mask, other=0) + # Source and destination overlap for shifts smaller than BLOCK_SIZE. + # Ensure every lane has consumed its source before any lane stores. + tl.debug_barrier() tl.store(row_ptr + offset, block_ids, mask=mask) From 5d06d3d48a2e1ccc620d0192e80f4fb5dd6bfd3f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 09:44:44 +0000 Subject: [PATCH 10/11] style: normalize DSpark correctness tests --- .../test_dflash_cudagraph_lifetime.py | 3 +-- .../test_dflash_prefix_cache_masking.py | 16 ++++------------ tests/v1/spec_decode/test_dynamic_sd_cug.py | 3 +-- vllm/v1/attention/backends/flashinfer.py | 10 +++++----- 4 files changed, 11 insertions(+), 21 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py b/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py index 1d28950a762a..11fb1b66e85b 100644 --- a/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py +++ b/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py @@ -44,8 +44,7 @@ def test_dflash_retains_backbone_output_during_cudagraph_capture(monkeypatch): assert len(speculator._captured_backbone_outputs) == 1 assert ( - speculator._captured_backbone_outputs[0] - is speculator._run_model.return_value + speculator._captured_backbone_outputs[0] is speculator._run_model.return_value ) diff --git a/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py index 5fddac1c802c..a93ab4bbe0b0 100644 --- a/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py +++ b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py @@ -17,9 +17,7 @@ shift_draft_block_tables, ) -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda(), reason="Requires CUDA" -) +pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") DEVICE = "cuda" BLOCK_SIZE = 16 @@ -63,9 +61,7 @@ def test_shift_single_request(num_cached: int, expected_shift: int): ) kept = MAX_BLOCKS - expected_shift - torch.testing.assert_close( - block_table[0, :kept], original[0, expected_shift:] - ) + torch.testing.assert_close(block_table[0, :kept], original[0, expected_shift:]) def test_shift_per_request_and_idx_mapping(): @@ -76,9 +72,7 @@ def test_shift_per_request_and_idx_mapping(): original = block_table.clone() idx_mapping = torch.tensor([3, 2, 1, 0], dtype=torch.int32, device=DEVICE) # Slot i has i whole cached blocks. - num_cached_tokens = torch.zeros( - MAX_NUM_REQS, dtype=torch.int32, device=DEVICE - ) + num_cached_tokens = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE) num_cached_tokens[:4] = ( torch.arange(4, dtype=torch.int32, device=DEVICE) * BLOCK_SIZE ) @@ -128,9 +122,7 @@ def test_shift_large_row_in_place_overlap(): block_table, idx_mapping, num_cached_tokens, seq_lens, BLOCK_SIZE ) - torch.testing.assert_close( - block_table[0, : max_blocks - 7], original[0, 7:] - ) + torch.testing.assert_close(block_table[0, : max_blocks - 7], original[0, 7:]) def test_shift_copy_bounded_by_seq_len(): diff --git a/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py index 5f6b22e6cac9..deb23908a9d1 100644 --- a/tests/v1/spec_decode/test_dynamic_sd_cug.py +++ b/tests/v1/spec_decode/test_dynamic_sd_cug.py @@ -374,7 +374,6 @@ def test_dynamic_sd_skips_zero_draft_tokens_in_cudagraph_schedule(monkeypatch): desc.uniform_token_count for descs in manager._candidates.values() for desc in descs - if desc.cg_mode == CUDAGraphMode.FULL - and desc.uniform_token_count is not None + if desc.cg_mode == CUDAGraphMode.FULL and desc.uniform_token_count is not None } assert captured_query_lens == scheduled_query_lens diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 36b1d724eaf2..fd3aec903fec 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -9,17 +9,17 @@ import numpy as np import torch -from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache -from flashinfer.prefill import trtllm_batch_context_with_kv_cache -from flashinfer.utils import FP4Tensor -from typing_extensions import override - from flashinfer import ( BatchDecodeWithPagedKVCacheWrapper, BatchPrefillWithPagedKVCacheWrapper, BatchPrefillWithRaggedKVCacheWrapper, MultiLevelCascadeAttentionWrapper, ) +from flashinfer.decode import fast_decode_plan, trtllm_batch_decode_with_kv_cache +from flashinfer.prefill import trtllm_batch_context_with_kv_cache +from flashinfer.utils import FP4Tensor +from typing_extensions import override + from vllm import _custom_ops as custom_ops from vllm import envs from vllm.config import ( From caa795ec33d4f04b5266ad9f9719e1192c362968 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 10:03:21 +0000 Subject: [PATCH 11/11] fix(dflash): fail closed on partial restored KV blocks --- .../test_dflash_prefix_cache_masking.py | 23 ++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 3 +- .../gpu/spec_decode/dflash/speculator.py | 31 ++++++++++++++++--- vllm/v1/worker/gpu/states.py | 2 ++ 4 files changed, 54 insertions(+), 5 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py index a93ab4bbe0b0..1f8fbc547091 100644 --- a/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py +++ b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py @@ -9,11 +9,15 @@ _prepare_dflash_inputs_kernel). """ +from types import SimpleNamespace + +import numpy as np import pytest import torch from vllm.platforms import current_platform from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( + DFlashSpeculator, shift_draft_block_tables, ) @@ -25,6 +29,25 @@ MAX_NUM_REQS = 8 +def test_unaligned_cached_prefix_detection(): + speculator = SimpleNamespace( + num_cached_tokens_np=np.array([32, 35, 64], dtype=np.int32), + block_tables=SimpleNamespace(kernel_block_sizes=[16]), + ) + + aligned = SimpleNamespace( + idx_mapping_np=np.array([0, 2], dtype=np.int32), + num_reqs=2, + ) + unaligned = SimpleNamespace( + idx_mapping_np=np.array([0, 1], dtype=np.int32), + num_reqs=2, + ) + + assert not DFlashSpeculator._has_unaligned_cached_prefix(speculator, aligned) + assert DFlashSpeculator._has_unaligned_cached_prefix(speculator, unaligned) + + def _make_block_table(num_reqs: int) -> torch.Tensor: # Distinct block ids per (request, slot) so shifts are detectable. table = torch.arange( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 798cd037f357..c498c42c4553 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -535,7 +535,8 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: # DFlash/DSpark mask cache-restored tokens out of the draft's # context (their draft context KV was never computed). self.speculator.set_num_cached_tokens( - self.req_states.num_cached_tokens.gpu + self.req_states.num_cached_tokens.gpu, + self.req_states.num_cached_tokens_np, ) if self.speculator is not None: # After set_attn, so the speculator can size its cudagraph mode diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index bec2310f4c9a..1f45a90b5f06 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -12,6 +12,7 @@ from collections.abc import Mapping from typing import Any +import numpy as np import torch import torch.nn as nn @@ -91,6 +92,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.num_cached_tokens = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=device ) + self.num_cached_tokens_np = np.zeros(self.max_num_reqs, dtype=np.int32) # Per-mask-token sampling buffers. Flattened from (num_reqs, num_spec_tokens). max_num_sampled_tokens = self.max_num_reqs * self.num_speculative_steps @@ -188,12 +190,25 @@ def load_draft_model( ) return model - def set_num_cached_tokens(self, num_cached_tokens: torch.Tensor) -> None: + def set_num_cached_tokens( + self, + num_cached_tokens: torch.Tensor, + num_cached_tokens_np: np.ndarray, + ) -> None: """Register the runner's per-request-slot cache-restored token counts. Indexed by req_state_idx; see the buffer comment in __init__. """ self.num_cached_tokens = num_cached_tokens + self.num_cached_tokens_np = num_cached_tokens_np + + def _has_unaligned_cached_prefix(self, input_batch: InputBatch) -> bool: + req_state_indices = input_batch.idx_mapping_np[: input_batch.num_reqs] + cached = self.num_cached_tokens_np[req_state_indices] + return any( + np.any(cached % block_size != 0) + for block_size in self.block_tables.kernel_block_sizes + ) def set_attn( self, @@ -365,6 +380,14 @@ def propose( ) -> torch.Tensor: num_reqs = input_batch.num_reqs num_target_tokens = input_batch.num_tokens + if not dummy_run and self._has_unaligned_cached_prefix(input_batch): + logger.warning_once( + "DFlash/DSpark drafting is disabled for a batch containing a " + "block-unaligned cache-restored prefix because draft KV is " + "not available for the restored partial block." + ) + self.draft_tokens[:num_reqs].fill_(-1) + return self.draft_tokens[:num_reqs] num_query_tokens = num_reqs * self.num_query_per_req max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() self.draft_max_seq_len = min( @@ -449,9 +472,9 @@ def propose( # after prepare_dflash_inputs because the slot mappings index the # unshifted table; in-place is safe because input_block_tables are # regathered from the persistent block tables every step. Up to - # block_size - 1 restored slots may remain visible when the restored - # count is not block-aligned (e.g. a full-prompt cache hit). Skipped - # for dummy runs, whose idx_mapping does not reference live requests. + # Non-aligned restored prefixes fail closed before this path because a + # block-table shift cannot hide the residual partial block. Skipped for + # dummy runs, whose idx_mapping does not reference live requests. if not dummy_run: for gid in self.draft_kv_cache_group_ids: shift_draft_block_tables( diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 1bc10444654d..f82a73a94733 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -88,6 +88,7 @@ def __init__( self.num_cached_tokens = StagedWriteTensor( self.max_num_reqs, dtype=torch.int32, device=device ) + self.num_cached_tokens_np = np.zeros(self.max_num_reqs, dtype=np.int32) @property def num_reqs(self) -> int: @@ -119,6 +120,7 @@ def add_request( self.num_computed_tokens_np[req_idx] = num_computed_tokens self.num_computed_tokens.stage_write_elem(req_idx, num_computed_tokens) self.num_cached_tokens.stage_write_elem(req_idx, num_computed_tokens) + self.num_cached_tokens_np[req_idx] = num_computed_tokens self.draft_tokens[req_idx].zero_()