From 1a2754e28de38bdd3e9fae8299c07a2f18c6a541 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 13:17:50 +0000 Subject: [PATCH 1/3] fix(spec decode): harden DSpark and DFlash edge paths --- .../test_deepseek_v4_dspark_metadata.py | 46 ++++ .../test_acceptance_length_controller.py | 3 + .../test_dflash_cudagraph_lifetime.py | 65 ++++++ .../test_dflash_prefix_cache_masking.py | 170 ++++++++++++++ tests/v1/spec_decode/test_dynamic_sd_cug.py | 53 ++++- .../worker/test_gpu_sampling_states_seed.py | 35 +++ .../generate_attention_backend_docs.py | 4 +- vllm/platforms/cuda.py | 7 +- vllm/v1/attention/backends/flashinfer.py | 3 +- vllm/v1/attention/backends/mla/sparse_swa.py | 6 +- vllm/v1/core/sched/async_scheduler.py | 6 +- vllm/v1/core/sched/output.py | 6 +- vllm/v1/core/sched/scheduler.py | 4 +- vllm/v1/worker/gpu/cudagraph_utils.py | 12 +- vllm/v1/worker/gpu/model_runner.py | 19 +- vllm/v1/worker/gpu/sample/sampler.py | 3 +- vllm/v1/worker/gpu/sample/states.py | 12 +- .../gpu/spec_decode/dflash/speculator.py | 216 +++++++++++++++++- vllm/v1/worker/gpu/spec_decode/speculator.py | 10 +- vllm/v1/worker/gpu/states.py | 13 ++ 20 files changed, 659 insertions(+), 34 deletions(-) create mode 100644 tests/v1/attention/test_deepseek_v4_dspark_metadata.py create mode 100644 tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py create mode 100644 tests/v1/spec_decode/test_dflash_prefix_cache_masking.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/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/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..11fb1b66e85b --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_cudagraph_lifetime.py @@ -0,0 +1,65 @@ +# 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/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..1f8fbc547091 --- /dev/null +++ b/tests/v1/spec_decode/test_dflash_prefix_cache_masking.py @@ -0,0 +1,170 @@ +# 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). +""" + +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, +) + +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 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( + 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/tests/v1/spec_decode/test_dynamic_sd_cug.py b/tests/v1/spec_decode/test_dynamic_sd_cug.py index d75495ea606a..deb23908a9d1 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 @@ -326,3 +329,51 @@ 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/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/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: diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 12eab21e3e13..fd3aec903fec 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -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/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/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. 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() diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b893f640a8a1..c498c42c4553 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) @@ -524,15 +525,23 @@ 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) 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.req_states.num_cached_tokens_np, + ) + 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( @@ -824,7 +833,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..1f45a90b5f06 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -12,14 +12,16 @@ from collections.abc import Mapping from typing import Any +import numpy as np 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 @@ -81,6 +83,17 @@ 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 + ) + 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 self.sample_indices = torch.zeros( @@ -89,8 +102,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, :]. @@ -100,10 +116,37 @@ 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: + # 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 @@ -117,11 +160,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, @@ -147,6 +190,26 @@ def load_draft_model( ) return model + 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, model_state: ModelState, @@ -160,6 +223,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( @@ -236,6 +310,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]] @@ -304,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( @@ -370,6 +454,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, @@ -379,6 +464,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 + # 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( + 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. @@ -469,6 +575,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, @@ -575,8 +683,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 @@ -607,6 +727,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) @@ -630,6 +752,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, @@ -666,6 +790,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, @@ -677,3 +802,74 @@ 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. + # 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) + + +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/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, ) diff --git a/vllm/v1/worker/gpu/states.py b/vllm/v1/worker/gpu/states.py index 7f0ae33c8099..f82a73a94733 100644 --- a/vllm/v1/worker/gpu/states.py +++ b/vllm/v1/worker/gpu/states.py @@ -80,6 +80,16 @@ 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 + ) + self.num_cached_tokens_np = np.zeros(self.max_num_reqs, dtype=np.int32) + @property def num_reqs(self) -> int: return len(self.req_id_to_index) @@ -109,6 +119,8 @@ 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.num_cached_tokens_np[req_idx] = num_computed_tokens self.draft_tokens[req_idx].zero_() @@ -118,6 +130,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 56533f4a1abef65ecf79da313c85502853316e0c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 13:17:55 +0000 Subject: [PATCH 2/3] perf(dspark): add opt-in rowwise FP8 draft head --- .../spec_decode/test_dspark_fp8_draft_head.py | 144 ++++++++++++++++++ vllm/envs.py | 7 + vllm/model_executor/layers/fp8_draft_head.py | 89 +++++++++++ vllm/models/deepseek_v4/nvidia/dspark.py | 51 +++++++ .../v1/worker/gpu/spec_decode/dspark/utils.py | 7 + 5 files changed, 298 insertions(+) create mode 100644 tests/v1/spec_decode/test_dspark_fp8_draft_head.py create mode 100644 vllm/model_executor/layers/fp8_draft_head.py diff --git a/tests/v1/spec_decode/test_dspark_fp8_draft_head.py b/tests/v1/spec_decode/test_dspark_fp8_draft_head.py new file mode 100644 index 000000000000..8405957a8e9e --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_fp8_draft_head.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the rowwise-fp8 DSpark draft lm_head (VLLM_DSPARK_FP8_DRAFT_HEAD). + +CPU tests cover the quantize helper (roundtrip error bound) and draft argmax +agreement between the bf16 reference path and the fp8 path, using a float32 +emulation of torch._scaled_mm. The GPU test additionally checks the real +_scaled_mm kernel against the emulation and the bf16 reference. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fp8_draft_head import ( + _FP8_MAX, + Fp8DraftHead, + fp8_draft_head_logits, + fp8_draft_head_supported, + quantize_draft_head, +) + +VOCAB = 2048 +HIDDEN = 256 +NUM_TOKENS = 64 + +# float8_e4m3fn has a 3-bit mantissa: worst-case relative rounding error of a +# representable-range value is 2^-4 of its binade, i.e. <= (1/16) * value and +# <= (32/448) * rowmax absolute (half ulp of the top binade after rowwise +# scaling to [-448, 448]). +_ROUNDTRIP_REL_BOUND = 32.0 / 448.0 / 2.0 # half ulp at the top binade + + +def _dequant(head: Fp8DraftHead) -> torch.Tensor: + # row_scale is [1, num_rows] laid out for the GEMM epilogue; transpose to + # dequantize the [num_rows, hidden] weight. + return head.weight_fp8.float() * head.row_scale.float().t() + + +def _emulated_fp8_logits(x: torch.Tensor, head: Fp8DraftHead) -> torch.Tensor: + """Float32 emulation of fp8_draft_head_logits (no _scaled_mm needed).""" + act_max = x.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) + act_fp8 = (x * (_FP8_MAX / act_max)).to(torch.float8_e4m3fn) + out = act_fp8.float() @ head.weight_fp8.float().t() + out = out.to(x.dtype) + out = out * head.row_scale + out = out * (act_max / _FP8_MAX).to(x.dtype) + return out + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_rowwise_quant_roundtrip_error_bound(dtype: torch.dtype): + torch.manual_seed(0) + weight = torch.randn(VOCAB, HIDDEN, dtype=dtype) * 0.05 + # Give rows wildly different magnitudes to exercise the rowwise scales. + weight *= torch.logspace(-3, 1, VOCAB, dtype=dtype).unsqueeze(1) + + head = quantize_draft_head(weight) + assert head.weight_fp8.dtype == torch.float8_e4m3fn + assert head.weight_fp8.shape == weight.shape + assert head.row_scale.shape == (1, VOCAB) + + err = (weight.float() - _dequant(head)).abs() + row_max = weight.float().abs().amax(dim=1, keepdim=True) + # Small slack over the analytic half-ulp bound for the two roundings + # (scale multiply in fp32, then fp8 cast). + bound = row_max * (_ROUNDTRIP_REL_BOUND * 1.25) + 1e-8 + assert (err <= bound).all(), ( + f"rowwise fp8 roundtrip error {err.max().item()} exceeds bound" + ) + + +def test_draft_argmax_agreement_bf16_vs_fp8_emulated(): + """The fp8 draft head must (almost) always agree with bf16 on argmax. + + Uses the float32 emulation of _scaled_mm; disagreements are only + acceptable when the bf16 top-2 margin is within the fp8 error scale. + """ + torch.manual_seed(1234) + weight = torch.randn(VOCAB, HIDDEN, dtype=torch.bfloat16) * 0.02 + hidden = torch.randn(NUM_TOKENS, HIDDEN, dtype=torch.bfloat16) + + ref_logits = hidden.float() @ weight.float().t() + ref_argmax = ref_logits.argmax(dim=-1) + + head = quantize_draft_head(weight) + fp8_logits = _emulated_fp8_logits(hidden, head) + fp8_argmax = fp8_logits.argmax(dim=-1) + + # Primary correctness property: any disagreement must be a near-tie in + # the reference logits, i.e. the kind of flip that costs at most one + # rejected draft token during verification, never a wrong output. + top2 = ref_logits.topk(2, dim=-1).values + margin = (top2[:, 0] - top2[:, 1]).abs() + err_scale = ref_logits.abs().amax(dim=-1) * 2.0 * _ROUNDTRIP_REL_BOUND * 2.0 + disagree = ref_argmax != fp8_argmax + assert (margin[disagree] <= err_scale[disagree]).all(), ( + "fp8 draft head flipped an argmax with a large top-2 margin" + ) + + # I.i.d. Gaussian logits are the worst case for argmax stability (the + # top-2 gap of 2048 i.i.d. samples is tiny); real LM logits have much + # larger top-1 margins, where the fp8 head measured argmax-identical. + # Even in this adversarial regime agreement must stay high. + agree = (ref_argmax == fp8_argmax).float().mean().item() + assert agree >= 0.90, f"draft argmax agreement too low: {agree:.4f}" + + +def test_emulated_logits_close_to_reference(): + torch.manual_seed(7) + weight = torch.randn(VOCAB, HIDDEN, dtype=torch.bfloat16) * 0.02 + hidden = torch.randn(NUM_TOKENS, HIDDEN, dtype=torch.bfloat16) + + ref = (hidden.float() @ weight.float().t()).float() + head = quantize_draft_head(weight) + fp8 = _emulated_fp8_logits(hidden, head).float() + + scale = ref.abs().amax() + assert (fp8 - ref).abs().max() <= 0.05 * scale + + +@pytest.mark.skipif( + not fp8_draft_head_supported(), + reason="requires a CUDA device with fp8 support (SM89+)", +) +def test_fp8_draft_head_logits_cuda_matches_emulation(): + torch.manual_seed(42) + device = torch.device("cuda") + weight = torch.randn(VOCAB, HIDDEN, dtype=torch.bfloat16, device=device) + weight *= 0.02 + hidden = torch.randn(NUM_TOKENS, HIDDEN, dtype=torch.bfloat16, device=device) + + head = quantize_draft_head(weight) + real = fp8_draft_head_logits(hidden, head).float() + emulated = _emulated_fp8_logits(hidden, head).float() + + # _scaled_mm accumulates in fp32 like the emulation; only the final + # bf16 rounding differs. + scale = emulated.abs().amax() + assert (real - emulated).abs().max() <= 0.01 * scale + + # The CUDA path must make exactly the same draft choice as the emulation. + # Quantization-induced flips against BF16 are covered separately by the + # top-2 margin bound in test_draft_argmax_agreement_bf16_vs_fp8_emulated. + torch.testing.assert_close(real.argmax(dim=-1), emulated.argmax(dim=-1)) diff --git a/vllm/envs.py b/vllm/envs.py index 01e965dc58cc..b6addbca59d8 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -61,6 +61,7 @@ VLLM_USE_B12X_SPARSE_INDEXER: bool = False VLLM_USE_B12X_MHC: bool = False VLLM_USE_B12X_FP8_GEMM: bool = False + VLLM_DSPARK_FP8_DRAFT_HEAD: bool = False VLLM_USE_B12X_WO_PROJECTION: bool = False VLLM_USE_B12X_MOE: bool = False VLLM_USE_B12X_MINIMAX_M3_MSA: bool = False @@ -1062,6 +1063,12 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_USE_B12X_FP8_GEMM": lambda: bool( int(os.getenv("VLLM_USE_B12X_FP8_GEMM", "0")) ), + # Compute DSpark draft-proposal logits with a rowwise-fp8 copy of the + # shared target lm_head. Verification is unchanged, so accepted outputs + # retain target-model semantics. Requires fp8 tensor cores (SM89+). + "VLLM_DSPARK_FP8_DRAFT_HEAD": lambda: bool( + int(os.getenv("VLLM_DSPARK_FP8_DRAFT_HEAD", "0")) + ), # Use b12x for the DeepSeek V4 WO-A/WO-B fused projection. # This is separate from the generic FP8 linear switch for perf isolation. "VLLM_USE_B12X_WO_PROJECTION": lambda: bool( diff --git a/vllm/model_executor/layers/fp8_draft_head.py b/vllm/model_executor/layers/fp8_draft_head.py new file mode 100644 index 000000000000..8acd21fc066a --- /dev/null +++ b/vllm/model_executor/layers/fp8_draft_head.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Rowwise FP8 (e4m3) quantization of a speculative-decoding draft LM head. + +Drafters that share the target model's LM head (e.g. DSpark, +``has_own_lm_head=False``) pay a full ``[hidden, vocab]`` bf16 GEMM per draft +step just to propose tokens. On bandwidth-bound GPUs that weight read +dominates the draft loop. These helpers build a one-time rowwise-fp8 copy of +the (vocab-sharded) head and compute draft logits with a dynamic per-token +fp8 GEMM, halving the weight traffic. + +This is draft-time only by construction: the verify pass never sees the fp8 +weights, so accepted outputs are bitwise unchanged. The quantized logits are +used for the draft argmax/top-k proposal; a rare argmax flip only costs a +rejected draft token, never an incorrect sample. +""" + +from typing import NamedTuple + +import torch + +# Finite max of float8_e4m3fn. +_FP8_MAX = 448.0 + + +class Fp8DraftHead(NamedTuple): + """Rowwise-quantized copy of a (possibly vocab-sharded) LM head.""" + + # [num_local_vocab, hidden] fp8_e4m3, each row scaled to [-448, 448]. + weight_fp8: torch.Tensor + # [1, num_local_vocab] dequant scale per output row (rowmax / 448), + # kept in the activation dtype so the epilogue multiplies stay cheap. + row_scale: torch.Tensor + # fp32 scalar 1.0 for torch._scaled_mm (scaling is applied manually in + # the epilogue because both operands use dynamic per-row scales). + unit_scale: torch.Tensor + + +def fp8_draft_head_supported(device: torch.device | None = None) -> bool: + """torch._scaled_mm needs fp8 tensor cores (SM89+) on CUDA devices.""" + if not torch.cuda.is_available(): + return False + major, minor = torch.cuda.get_device_capability(device) + return (major, minor) >= (8, 9) + + +def quantize_draft_head(weight: torch.Tensor) -> Fp8DraftHead: + """Quantize a ``[num_local_vocab, hidden]`` head weight rowwise to fp8. + + Per-row (per vocab entry) symmetric scaling: ``w8 = w * (448 / rowmax)`` + stored as fp8_e4m3, ``row_scale = rowmax / 448`` for the epilogue. For a + vocab-parallel (sharded) head, pass the local shard; row scales are + per-local-row, so gather/argmax semantics downstream are unchanged. + """ + with torch.no_grad(): + w = weight.detach() + row_max = w.abs().amax(dim=1, keepdim=True).float().clamp(min=1e-6) + weight_fp8 = (w.float() * (_FP8_MAX / row_max)).to(torch.float8_e4m3fn) + row_scale = (row_max / _FP8_MAX).to(w.dtype).reshape(1, -1) + unit_scale = torch.ones(1, dtype=torch.float32, device=w.device) + return Fp8DraftHead(weight_fp8, row_scale, unit_scale) + + +def fp8_draft_head_logits( + hidden_states: torch.Tensor, + head: Fp8DraftHead, +) -> torch.Tensor: + """Local (shard) draft logits via dynamic-fp8 x rowwise-fp8 GEMM. + + Activations are quantized with a per-token amax scale, then + ``logits = _scaled_mm(a8, w8.T) * row_scale * (amax / 448)``. Output + dtype and shape match ``lm_head.quant_method.apply``: the caller is + responsible for the same TP gather / vocab-padding slice it would apply + to the unquantized local logits. Contains no data-dependent control + flow or allocations beyond the GEMM output, so it is safe to run inside + a captured CUDA graph as long as ``head`` was materialized beforehand. + """ + act_max = hidden_states.abs().amax(dim=-1, keepdim=True).clamp(min=1e-6) + act_fp8 = (hidden_states * (_FP8_MAX / act_max)).to(torch.float8_e4m3fn) + logits = torch._scaled_mm( + act_fp8, + head.weight_fp8.t(), + scale_a=head.unit_scale, + scale_b=head.unit_scale, + out_dtype=hidden_states.dtype, + ) + logits = logits * head.row_scale + logits = logits * (act_max / _FP8_MAX).to(hidden_states.dtype) + return logits diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index e4e258372a29..109cf2f7c6d6 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -15,6 +15,7 @@ import torch import torch.nn as nn +import vllm.envs as envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -25,6 +26,12 @@ hc_head_fused_kernel_tilelang, mhc_post_tilelang, ) +from vllm.model_executor.layers.fp8_draft_head import ( + Fp8DraftHead, + fp8_draft_head_logits, + fp8_draft_head_supported, + quantize_draft_head, +) from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -287,6 +294,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: prefix=maybe_prefix(prefix, "lm_head"), ) self.logits_processor = LogitsProcessor(self.config.vocab_size) + # Optional rowwise-fp8 copy of the (shared) lm_head, used ONLY for + # draft-proposal logits. Materialized eagerly at load time via + # maybe_init_fp8_draft_head(); None means the bf16 path is used. + self._fp8_draft_head: Fp8DraftHead | None = None # --- Hooks used by the speculator ------------------------------------- @@ -324,8 +335,48 @@ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: """Base logits U_k = lm_head(norm(head_hidden)).""" return self.logits_processor(self.lm_head, self.model.norm(hidden_states)) + def maybe_init_fp8_draft_head(self) -> None: + """Materialize the rowwise-fp8 draft lm_head copy (opt-in). + + Called by ``load_dspark_model`` right after the target's lm_head is + aliased onto this model. This must happen eagerly at load time, NOT + lazily on the first draft call: the whole DSpark draft step + (including ``compute_draft_logits``) runs inside a FULL captured + CUDA graph, so the quantization kernels must not fire during + capture/replay. + + TP: the lm_head is vocab-sharded, so each rank quantizes its local + shard with per-local-row scales; the draft path's gather and argmax + semantics are unchanged. + """ + if not envs.VLLM_DSPARK_FP8_DRAFT_HEAD: + return + if not fp8_draft_head_supported(self.lm_head.weight.device): + logger.warning( + "VLLM_DSPARK_FP8_DRAFT_HEAD is set but this device has no " + "fp8 support (SM89+ required); using the unquantized " + "draft lm_head." + ) + return + self._fp8_draft_head = quantize_draft_head(self.lm_head.weight) + logger.info_once( + "DSpark draft-proposal logits use a rowwise-fp8 copy of the " + "target lm_head (draft-time only; verify pass untouched)." + ) + def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: # Full-vocab draft: base logits, no d2t scatter. + if self._fp8_draft_head is not None: + # Draft-proposal-only fp8 path. Mirrors compute_logits -> + # LogitsProcessor._get_logits: local (shard) logits, then the + # same TP gather and vocab-padding slice. + local_logits = fp8_draft_head_logits( + self.model.norm(hidden_states), self._fp8_draft_head + ) + logits = self.logits_processor._gather_logits(local_logits) + if logits is not None: + logits = logits[..., : self.logits_processor.org_vocab_size] + return logits return self.compute_logits(hidden_states) def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor: diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index 5cb967a2fdfc..b06323f0c4cf 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -76,4 +76,11 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo del draft_model.lm_head draft_model.lm_head = target_lm_head + # Opt-in rowwise-fp8 draft head (VLLM_DSPARK_FP8_DRAFT_HEAD). Must run + # after the lm_head aliasing above and BEFORE CUDA graph capture: the + # draft step is captured whole, so the fp8 copy is materialized eagerly. + maybe_init_fp8_draft_head = getattr(draft_model, "maybe_init_fp8_draft_head", None) + if maybe_init_fp8_draft_head is not None: + maybe_init_fp8_draft_head() + return draft_model From d3d14156dcadbc6cd0144bad33566738026073b3 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 17 Jul 2026 13:19:10 +0000 Subject: [PATCH 3/3] perf(dspark): add load-aware compact verification capacity --- benchmarks/profile_dspark_sps_curve.py | 165 +++ tests/engine/test_arg_utils.py | 37 + .../DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml | 22 + tests/test_config.py | 69 ++ tests/v1/worker/test_gpu_block_table.py | 47 + ...test_gpu_model_runner_v2_draft_capacity.py | 1000 +++++++++++++++++ tests/v1/worker/test_mixed_warmup_gate.py | 24 +- vllm/config/speculative.py | 135 +++ vllm/engine/arg_utils.py | 12 + vllm/envs.py | 18 + .../layers/sparse_attn_indexer.py | 7 +- vllm/model_executor/models/qwen3_dspark.py | 68 +- .../deepseek_v4/common/ops/cache_utils.py | 40 +- vllm/models/deepseek_v4/nvidia/dspark.py | 29 +- .../deepseek_v4/nvidia/flashinfer_sparse.py | 8 +- vllm/models/deepseek_v4/sparse_mla.py | 2 +- vllm/utils/deep_gemm.py | 30 +- vllm/v1/attention/backend.py | 2 + vllm/v1/attention/backends/flashinfer.py | 52 +- vllm/v1/attention/backends/mla/indexer.py | 240 +++- vllm/v1/attention/backends/mla/sparse_swa.py | 16 +- vllm/v1/attention/backends/utils.py | 2 + vllm/v1/core/sched/scheduler.py | 15 +- vllm/v1/worker/gpu/attn_utils.py | 10 +- vllm/v1/worker/gpu/block_table.py | 11 + vllm/v1/worker/gpu/cudagraph_utils.py | 172 ++- vllm/v1/worker/gpu/dp_utils.py | 13 +- vllm/v1/worker/gpu/input_batch.py | 29 +- vllm/v1/worker/gpu/model_runner.py | 196 +++- vllm/v1/worker/gpu/model_states/default.py | 11 +- vllm/v1/worker/gpu/sample/bad_words.py | 6 +- vllm/v1/worker/gpu/sample/gumbel.py | 6 +- .../spec_decode/autoregressive/speculator.py | 2 + vllm/v1/worker/gpu/spec_decode/capacity.py | 980 ++++++++++++++++ .../spec_decode/causal_cascade/speculator.py | 1 + .../gpu/spec_decode/dflash/cudagraph.py | 1 + .../gpu/spec_decode/dflash/speculator.py | 186 ++- .../worker/gpu/spec_decode/dspark/capacity.py | 174 +++ .../gpu/spec_decode/dspark/online_sts.py | 192 ++++ .../gpu/spec_decode/dspark/speculator.py | 285 ++++- .../gpu/spec_decode/rejection_sampler.py | 3 + .../spec_decode/rejection_sampler_utils.py | 4 +- vllm/v1/worker/gpu/spec_decode/speculator.py | 20 +- vllm/v1/worker/gpu/warmup.py | 135 ++- vllm/v1/worker/ubatch_utils.py | 1 + 45 files changed, 4300 insertions(+), 178 deletions(-) create mode 100644 benchmarks/profile_dspark_sps_curve.py create mode 100644 tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml create mode 100644 tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py create mode 100644 vllm/v1/worker/gpu/spec_decode/capacity.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/capacity.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/online_sts.py diff --git a/benchmarks/profile_dspark_sps_curve.py b/benchmarks/profile_dspark_sps_curve.py new file mode 100644 index 000000000000..0c0b72782b5a --- /dev/null +++ b/benchmarks/profile_dspark_sps_curve.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Profile the engine step-rate curve for the DSpark prefix scheduler. + +Times the captured FULL cudagraph replays of the target verification step at +every captured batch token count and emits ``dspark_sps_curve`` breakpoints, +one per capture size. The scheduler linearly interpolates between +breakpoints, which amortizes cudagraph padding smoothly instead of +concentrating it into thresholds at capture-size boundaries. + +Example: + python benchmarks/profile_dspark_sps_curve.py \\ + --speculative-config '{"method": "dspark", "model": "...", ...}' \\ + --engine-args '{"tensor_parallel_size": 4, "max_num_seqs": 32}' + +Paste the printed ``dspark_sps_curve`` entry into --speculative-config. + +Caveats: replays run on whatever (dummy) buffer state capture left behind, so +data-dependent kernels (e.g. MoE routing) may be timed on unrepresentative +inputs, and per-step CPU/draft overhead is modeled only through the constant +``--overhead-ms``. Only the curve's shape matters to the scheduler. +""" + +import argparse +import json + + +def _time_fullgraph_replays(worker, iters: int, warmup: int) -> dict[int, float]: + """Worker-side: time FULL graph replay per batch token count (ms/step). + + Runs on every TP rank via collective_rpc so the collectives captured in + the graphs stay matched; every rank replays the same descs in the same + sorted order. Before timing each descriptor the input buffers are + refreshed into the same coherent dummy state capture used, so replays + never read stale metadata. + """ + import torch + + from vllm.v1.worker.gpu.cudagraph_utils import prepare_inputs_to_capture + + runner = worker.model_runner + mgr = runner.cudagraph_manager + assert mgr is not None and mgr.graphs, ( + "No FULL cudagraphs captured; run with a cudagraph_mode that captures " + "FULL decode graphs." + ) + # Prefer varlen spec-decode descs; fall back to all captured graphs. + descs = [d for d in mgr.graphs if d.max_req_tokens is not None] + if not descs: + descs = list(mgr.graphs.keys()) + # One desc per token count: the largest request count is the most + # representative shape under load. + by_tokens: dict[int, object] = {} + for d in descs: + cur = by_tokens.get(d.num_tokens) + if cur is None or (d.num_reqs or 0) > (cur.num_reqs or 0): + by_tokens[d.num_tokens] = d + + results: dict[int, float] = {} + for num_tokens in sorted(by_tokens): + desc = by_tokens[num_tokens] + num_reqs = desc.num_reqs or min(num_tokens, mgr.max_num_reqs) + prepare_inputs_to_capture( + num_reqs, + num_tokens, + runner.model_state, + runner.input_buffers, + runner.block_tables, + runner.attn_groups, + runner.kv_cache_config, + max_req_tokens=desc.max_req_tokens, + ) + graph = mgr.graphs[desc] + for _ in range(warmup): + graph.replay() + torch.accelerator.synchronize() + start = torch.Event(enable_timing=True) + end = torch.Event(enable_timing=True) + start.record() + for _ in range(iters): + graph.replay() + end.record() + torch.accelerator.synchronize() + results[num_tokens] = start.elapsed_time(end) / iters + return results + + +def curve_breakpoints( + ms_per_step: dict[int, float], overhead_ms: float +) -> list[list[float]]: + """Convert per-capture-size step times into ``dspark_sps_curve`` + breakpoints, one per capture size. The scheduler's table linearly + interpolates between them (and clamps at the ends).""" + return [ + [size, round(1000.0 / (ms_per_step[size] + overhead_ms), 3)] + for size in sorted(ms_per_step) + ] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("model", help="Target model (path or HF id)") + parser.add_argument( + "--speculative-config", + required=True, + help="JSON speculative config (same value you pass to vllm serve)", + ) + parser.add_argument( + "--engine-args", + default="{}", + help="JSON dict of extra vllm.LLM kwargs " + '(e.g. \'{"tensor_parallel_size": 4, "max_num_seqs": 32}\')', + ) + parser.add_argument("--iters", type=int, default=50) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument( + "--overhead-ms", + type=float, + default=0.0, + help="Constant per-step overhead (draft forward, sampling, CPU gap) " + "added to every measured step time before converting to a rate.", + ) + parser.add_argument("--output", help="Write the curve JSON to this file") + args = parser.parse_args() + if args.iters <= 0: + parser.error("--iters must be greater than zero") + if args.warmup < 0: + parser.error("--warmup must be non-negative") + if args.overhead_ms < 0: + parser.error("--overhead-ms must be non-negative") + + # The timing callable is shipped to the workers via collective_rpc, which + # requires the pickle fallback. Local profiling tool, trusted input. + import os + + os.environ.setdefault("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") + + from vllm import LLM + + llm = LLM( + model=args.model, + speculative_config=json.loads(args.speculative_config), + **json.loads(args.engine_args), + ) + per_rank = llm.collective_rpc( + _time_fullgraph_replays, kwargs={"iters": args.iters, "warmup": args.warmup} + ) + ms_per_step = per_rank[0] + + print("\nMeasured FULL-graph step times (rank 0):") + for size in sorted(ms_per_step): + print(f" B={size:5d} tokens: {ms_per_step[size]:8.3f} ms/step") + + curve = curve_breakpoints(ms_per_step, args.overhead_ms) + entry = {"dspark_sps_curve": curve} + print("\nAdd to --speculative-config:") + print(json.dumps(entry)) + if args.output: + with open(args.output, "w") as f: + json.dump(entry, f, indent=2) + print(f"\nWritten to {args.output}") + + +if __name__ == "__main__": + main() diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index b7895c179a3a..f882834599c1 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -225,6 +225,43 @@ def test_jit_monitor_mode_arg(mode): assert engine_args.create_observability_config().jit_monitor_mode == mode +def test_dspark_capacity_verification_mode_arg(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = parser.parse_args( + [ + "--spec-method", + "ngram", + "--spec-tokens", + "1", + "--dspark-capacity-verification-mode", + "mask", + ] + ) + + engine_args = EngineArgs.from_cli_args(args) + assert engine_args.dspark_capacity_verification_mode == "mask" + speculative_config = engine_args.create_speculative_config(None, None) + assert speculative_config is not None + assert speculative_config.dspark_capacity_verification_mode == "mask" + + +def test_dspark_capacity_verification_mode_conflicts_with_speculative_config(): + parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) + args = parser.parse_args( + [ + "--speculative-config", + '{"method":"ngram","num_speculative_tokens":1,' + '"dspark_capacity_verification_mode":"mask"}', + "--dspark-capacity-verification-mode", + "varlen", + ] + ) + + engine_args = EngineArgs.from_cli_args(args) + with pytest.raises(ValueError, match="dspark_capacity_verification_mode"): + engine_args.create_speculative_config(None, None) + + def test_hf_token_get_kwargs(): kwargs = get_kwargs(ModelConfig)["hf_token"] diff --git a/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml b/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml new file mode 100644 index 000000000000..c85c8f88d263 --- /dev/null +++ b/tests/evals/gsm8k/configs/DeepSeek-V4-Flash-DSpark-varlen-TP4.yaml @@ -0,0 +1,22 @@ +model_name: "deepseek-ai/DeepSeek-V4-Flash-DSpark" +accuracy_threshold: 0.92 +num_questions: 1319 +num_fewshot: 5 +startup_max_wait_seconds: 1800 +server_args: >- + --tokenizer-mode deepseek_v4 + --trust-remote-code + --dtype bfloat16 + --max-model-len 8192 + --tensor-parallel-size 4 + --enable-expert-parallel + --block-size 256 + --gpu-memory-utilization 0.5 + --kv-cache-dtype fp8 + --max-num-batched-tokens 16384 + --max-num-seqs 32 + --speculative-config '{"method":"dspark", + "model":"deepseek-ai/DeepSeek-V4-Flash-DSpark", + "attention_backend":"FLASH_ATTN","num_speculative_tokens":7, + "draft_sample_method":"probabilistic","dspark_confidence_threshold":0.0, + "dspark_budget_frac":0.5,"dspark_capacity_verification_mode":"varlen"}' diff --git a/tests/test_config.py b/tests/test_config.py index 30e2d5114874..3cf5ce17350f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1630,6 +1630,75 @@ def test_draft_sample_method_gumbel_is_rejected(): ) +def test_dspark_capacity_config_validation(): + speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + dspark_confidence_threshold=0.25, + dspark_budget_frac=0.5, + dspark_capacity_verification_mode="mask", + ) + assert speculative_config.dspark_confidence_threshold == 0.25 + assert speculative_config.dspark_budget_frac == 0.5 + assert speculative_config.dspark_capacity_verification_mode == "mask" + assert ( + SpeculativeConfig( + method="ngram", num_speculative_tokens=1 + ).dspark_capacity_verification_mode + == "varlen" + ) + assert ( + SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + dspark_capacity_verification_mode="compact", + ).dspark_capacity_verification_mode + == "varlen" + ) + assert ( + SpeculativeConfig( + method="ngram", num_speculative_tokens=1 + ).dspark_confidence_threshold + == 0.0 + ) + + for threshold in (-0.1, 1.1, float("nan"), float("inf")): + with pytest.raises(ValueError, match="dspark_confidence_threshold"): + SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + dspark_confidence_threshold=threshold, + ) + + for budget_frac in (0.0, -0.1, 1.1, float("nan"), float("inf")): + with pytest.raises(ValueError, match="dspark_budget_frac"): + SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + dspark_budget_frac=budget_frac, + ) + + for config_field, value in ( + ("dspark_confidence_temperature", float("nan")), + ("dspark_confidence_temperature", float("inf")), + ("dspark_sps_overhead_ms", float("nan")), + ("dspark_sps_overhead_ms", float("inf")), + ): + with pytest.raises(ValueError, match=config_field): + SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + **{config_field: value}, + ) + + with pytest.raises(ValueError, match="dspark_sps_curve"): + SpeculativeConfig( + method="ngram", + num_speculative_tokens=1, + dspark_sps_curve=[(1, float("nan"))], + ) + + def test_ir_op_priority_default(): """Test that IR op priority defaults are set correctly.""" from vllm.config.kernel import IrOpPriorityConfig diff --git a/tests/v1/worker/test_gpu_block_table.py b/tests/v1/worker/test_gpu_block_table.py index 31acd475adec..4c15915912ab 100644 --- a/tests/v1/worker/test_gpu_block_table.py +++ b/tests/v1/worker/test_gpu_block_table.py @@ -5,6 +5,7 @@ import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import PAD_SLOT_ID from vllm.v1.worker.gpu.block_table import BlockTables pytestmark = pytest.mark.skipif( @@ -130,3 +131,49 @@ def test_block_tables_apply_staged_writes_single_group(): block_tables.block_tables[0].gpu[0, :2], torch.tensor([1, 2], dtype=torch.int32, device=device), ) + + +def test_compute_slot_mappings_applies_padding_mask(): + device = torch.device("cuda") + block_tables = BlockTables( + block_sizes=[16], + max_num_reqs=2, + max_num_batched_tokens=8, + max_num_blocks_per_group=[4], + device=device, + kernel_block_sizes=[16], + ) + + block_tables.append_block_ids( + req_index=0, + new_block_ids=([2],), + overwrite=True, + ) + block_tables.append_block_ids( + req_index=1, + new_block_ids=([3],), + overwrite=True, + ) + block_tables.apply_staged_writes() + + idx_mapping = torch.tensor([0, 1], dtype=torch.int32, device=device) + query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device) + positions = torch.tensor([0, 1, 2, 0, 1], dtype=torch.int64, device=device) + is_padding = torch.tensor( + [False, True, False, False, True, False, False, False], + dtype=torch.bool, + device=device, + ) + + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + positions, + num_tokens_padded=8, + is_padding=is_padding, + ) + torch.accelerator.synchronize() + + assert slot_mappings.cpu().tolist() == [ + [32, PAD_SLOT_ID, 34, 48, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID, PAD_SLOT_ID] + ] diff --git a/tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py b/tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py new file mode 100644 index 000000000000..56976030baaa --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2_draft_capacity.py @@ -0,0 +1,1000 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from typing import Any + +import numpy as np +import pytest +import torch + +pytest.importorskip("triton") +if not torch.cuda.is_available(): + pytest.skip("CUDA required for draft capacity tests", allow_module_level=True) + +from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.models import qwen3_dspark +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.attention.backends.mla import indexer as indexer_module +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV32IndexerMetadataBuilder, + _needs_varlen_decode, + _uses_varlen_dspark_capacity, +) +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.spec_decode.capacity import ( + CapacityBasedVerificationManager, + DSparkDynamicDraftDepthController, + MaskedCapacityBasedVerificationManager, + VarlenCapacityBasedVerificationManager, +) +from vllm.v1.worker.gpu.spec_decode.dspark.capacity import ( + compute_draft_token_capacity_from_confidence, +) +from vllm.v1.worker.gpu.spec_decode.dspark.online_sts import DSparkOnlineSTS +from vllm.v1.worker.gpu.states import RequestState + + +def test_varlen_indexer_is_limited_to_active_dspark_capacity(): + def config(**overrides): + values = { + "use_dspark": lambda: True, + "dspark_capacity_verification_mode": "varlen", + "dspark_confidence_threshold": 0.0, + "dspark_budget_frac": 1.0, + "dspark_sps_curve": None, + } + values.update(overrides) + return SimpleNamespace( + speculative_config=SimpleNamespace(**values), + ) + + assert not _uses_varlen_dspark_capacity(config()) + assert _uses_varlen_dspark_capacity(config(dspark_confidence_threshold=0.5)) + assert _uses_varlen_dspark_capacity(config(dspark_budget_frac=0.5)) + assert _uses_varlen_dspark_capacity(config(dspark_sps_curve="auto")) + assert not _uses_varlen_dspark_capacity( + config(dspark_capacity_verification_mode="mask") + ) + assert not _uses_varlen_dspark_capacity(config(use_dspark=lambda: False)) + + +def test_varlen_indexer_cudagraph_support_requires_capacity_opt_in(monkeypatch): + monkeypatch.setattr( + indexer_module, "_supports_varlen_paged_mqa_logits", lambda: True + ) + + def config(*, threshold: float): + return SimpleNamespace( + speculative_config=SimpleNamespace( + use_dspark=lambda: True, + dspark_capacity_verification_mode="varlen", + dspark_confidence_threshold=threshold, + dspark_budget_frac=1.0, + dspark_sps_curve=None, + ) + ) + + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + config(threshold=0.0), None + ) + is AttentionCGSupport.UNIFORM_BATCH + ) + assert ( + DeepseekV32IndexerMetadataBuilder.get_cudagraph_support( + config(threshold=0.5), None + ) + is AttentionCGSupport.ALWAYS + ) + + +def test_varlen_indexer_skips_uniform_full_width_batches(): + assert not _needs_varlen_decode(True, True, 6, 6) + assert _needs_varlen_decode(True, False, 4, 4) + assert _needs_varlen_decode(True, False, 6, 6) + assert not _needs_varlen_decode(False, False, 6, 6) + + +def test_dynamic_draft_depth_tracks_capacity_and_probes_upward(): + controller = DSparkDynamicDraftDepthController(max_depth=5, observation_window=2) + attempted_5 = np.array([5, 5], dtype=np.int32) + capacities_2_3 = np.array([2, 3], dtype=np.int32) + + assert controller.observe(capacities_2_3, attempted_5) == 5 + assert controller.observe(capacities_2_3, attempted_5) == 3 + + attempted_3 = np.array([3, 3], dtype=np.int32) + capacities_3 = np.array([3, 3], dtype=np.int32) + for _ in range(controller._PROBE_AFTER_WINDOWS * 2 - 1): + assert controller.observe(capacities_3, attempted_3) == 3 + assert controller.observe(capacities_3, attempted_3) == 4 + + # A substantial load decrease resets the next proposal to the maximum. + assert ( + controller.observe( + np.array([3, 0], dtype=np.int32), + np.array([3, 0], dtype=np.int32), + ) + == 5 + ) + + +def test_dynamic_draft_depth_preserves_longest_useful_request(): + controller = DSparkDynamicDraftDepthController(max_depth=5, observation_window=2) + attempted = np.array([5, 5, 5], dtype=np.int32) + capacities = np.array([2, 3, 5], dtype=np.int32) + + assert controller.observe(capacities, attempted) == 5 + assert controller.observe(capacities, attempted) == 5 + + +def test_dynamic_draft_depth_applies_profiled_load_budget(): + controller = DSparkDynamicDraftDepthController(max_depth=5, observation_window=2) + controller.set_draft_token_budget(160) + attempted = np.full(64, 5, dtype=np.int32) + capacities = np.full(64, 5, dtype=np.int32) + + assert controller.observe(capacities, attempted) == 5 + assert controller.observe(capacities, attempted) == 3 + + # At half the load, the same profiled budget exposes the full K5 again. + attempted[32:] = 0 + capacities[32:] = 0 + assert controller.observe(capacities, attempted) == 5 + + +def test_capacity_kernel_supports_shorter_logical_width_than_storage_stride(): + device = torch.device("cuda") + confidence_probs = torch.full((2, 5), 0.01, device=device) + confidence_probs[:, :2] = torch.tensor([[0.9, 0.9], [0.8, 0.8]], device=device) + confidence_logits = torch.logit(confidence_probs) + capacities = torch.full((2,), -1, dtype=torch.int32, device=device) + survival = torch.empty_like(confidence_logits) + + compute_draft_token_capacity_from_confidence( + confidence_logits, + capacities, + min_survival_probability=0.7, + num_reqs=2, + num_speculative_steps=2, + runtime_num_reqs=torch.tensor([2], dtype=torch.int32, device=device), + survival_probs=survival, + ) + + torch.accelerator.synchronize() + assert capacities.cpu().tolist() == [2, 1] + + +def test_compute_draft_token_capacity_from_confidence_uses_global_prefix_order(): + device = torch.device("cuda") + confidence_probs = torch.tensor( + [ + [0.90, 0.90, 0.90], + [0.95, 0.10, 0.99], + [0.70, 0.70, 0.70], + ], + dtype=torch.float32, + device=device, + ) + confidence_logits = torch.logit(confidence_probs) + draft_token_capacity = torch.full((3,), -1, dtype=torch.int32, device=device) + survival_probs = torch.empty_like(confidence_logits) + + compute_draft_token_capacity_from_confidence( + confidence_logits, + draft_token_capacity, + min_survival_probability=0.75, + num_reqs=3, + num_speculative_steps=3, + runtime_num_reqs=torch.tensor([3], dtype=torch.int32, device=device), + survival_probs=survival_probs, + ) + + torch.accelerator.synchronize() + assert draft_token_capacity.cpu().tolist() == [2, 1, 0] + + +def test_compute_draft_token_capacity_uses_budgeted_global_prefix_order(): + device = torch.device("cuda") + confidence_probs = torch.tensor( + [ + [0.90, 0.80], + [0.80, 0.80], + ], + dtype=torch.float32, + device=device, + ) + confidence_logits = torch.logit(confidence_probs) + draft_token_capacity = torch.full((2,), -1, dtype=torch.int32, device=device) + survival_probs = torch.empty_like(confidence_logits) + + compute_draft_token_capacity_from_confidence( + confidence_logits, + draft_token_capacity, + min_survival_probability=0.0, + num_reqs=2, + num_speculative_steps=2, + runtime_num_reqs=torch.tensor([2], dtype=torch.int32, device=device), + survival_probs=survival_probs, + budget_frac=0.5, + ) + + torch.accelerator.synchronize() + # ceil(4 * 0.5) admits exactly two globally ranked prefixes. + assert draft_token_capacity.cpu().tolist() == [1, 1] + + +def test_compute_draft_token_capacity_keeps_threshold_ties(): + device = torch.device("cuda") + confidence_probs = torch.tensor( + [ + [0.90, 0.80], + [0.90, 0.80], + ], + dtype=torch.float32, + device=device, + ) + confidence_logits = torch.logit(confidence_probs) + draft_token_capacity = torch.full((2,), -1, dtype=torch.int32, device=device) + survival_probs = torch.empty_like(confidence_logits) + + compute_draft_token_capacity_from_confidence( + confidence_logits, + draft_token_capacity, + min_survival_probability=0.0, + num_reqs=2, + num_speculative_steps=2, + runtime_num_reqs=torch.tensor([2], dtype=torch.int32, device=device), + survival_probs=survival_probs, + budget_frac=0.25, + ) + + torch.accelerator.synchronize() + assert draft_token_capacity.cpu().tolist() == [1, 0] + + +def test_compute_draft_token_capacity_budget_is_hard_cap_under_ties(): + """Saturated (tied) survival scores must not escape the budget. + + With every confidence saturated to 1.0, a kth-score-threshold recount + would admit all tokens; the budget must remain a hard cap on the total. + """ + device = torch.device("cuda") + confidence_logits = torch.full((4, 7), 40.0, dtype=torch.float32, device=device) + draft_token_capacity = torch.full((4,), -1, dtype=torch.int32, device=device) + survival_probs = torch.empty_like(confidence_logits) + + compute_draft_token_capacity_from_confidence( + confidence_logits, + draft_token_capacity, + min_survival_probability=0.0, + num_reqs=4, + num_speculative_steps=7, + runtime_num_reqs=torch.tensor([4], dtype=torch.int32, device=device), + survival_probs=survival_probs, + budget_frac=0.5, + ) + + torch.accelerator.synchronize() + capacities = draft_token_capacity.cpu() + assert int(capacities.sum()) == int(4 * 7 * 0.5) + assert int(capacities.max()) <= 7 + + +def test_compute_draft_token_capacity_never_admits_zero_survival(): + """Zero-survival tokens are not candidates (DSpark Alg. 1: a_{r,j} > 0), + so leftover budget must not be spent past the first dead position.""" + device = torch.device("cuda") + # Positions 0-1 confident, position 2 dead (sigmoid underflows to an + # exact fp32 zero) -> survival is exactly 0 from there on. + confidence_logits = torch.full((2, 5), 40.0, dtype=torch.float32, device=device) + confidence_logits[:, 2] = -100.0 + draft_token_capacity = torch.full((2,), -1, dtype=torch.int32, device=device) + survival_probs = torch.empty_like(confidence_logits) + + compute_draft_token_capacity_from_confidence( + confidence_logits, + draft_token_capacity, + min_survival_probability=0.0, + num_reqs=2, + num_speculative_steps=5, + runtime_num_reqs=torch.tensor([2], dtype=torch.int32, device=device), + survival_probs=survival_probs, + budget_frac=0.9, + ) + + torch.accelerator.synchronize() + assert draft_token_capacity.cpu().tolist() == [2, 2] + + +def test_compute_draft_token_capacity_sps_curve_argmax(): + """With an SPS curve, verification lengths maximize tau * SPS(B) + (DSpark Alg. 1) instead of spending the whole budget.""" + device = torch.device("cuda") + confidence_probs = torch.tensor( + [ + [0.90, 0.80], + [0.60, 0.50], + ], + dtype=torch.float32, + device=device, + ) + confidence_logits = torch.logit(confidence_probs) + draft_token_capacity = torch.full((2,), -1, dtype=torch.int32, device=device) + survival_probs = torch.empty_like(confidence_logits) + # Survival: r0 [0.9, 0.72], r1 [0.6, 0.3]; admission order + # 0.9, 0.72, 0.6, 0.3 with B = 2 + k. + # SPS drops sharply after B=4 so theta peaks at k=2: + # k=0: 2.00*1.00, k=1: 2.90*0.95=2.755, k=2: 3.62*0.90=3.258, + # k=3: 4.22*0.20=0.844, k=4: 4.52*0.10=0.452. + sps_table = torch.tensor( + [1.0, 1.0, 1.0, 0.95, 0.90, 0.20, 0.10], dtype=torch.float32, device=device + ) + + from vllm.v1.worker.gpu.spec_decode.dspark.capacity import ( + compute_draft_token_capacity_from_confidence as compute, + ) + + compute( + confidence_logits, + draft_token_capacity, + min_survival_probability=0.0, + num_reqs=2, + num_speculative_steps=2, + runtime_num_reqs=torch.tensor([2], dtype=torch.int32, device=device), + survival_probs=survival_probs, + budget_frac=1.0, + sps_table=sps_table, + ) + + torch.accelerator.synchronize() + assert draft_token_capacity.cpu().tolist() == [2, 0] + + +def test_compute_draft_token_capacity_temperature_desaturates_zeros(): + """A confidence temperature > 1 keeps saturated-negative positions in the + candidate set (no exact-zero survival), so the budget is spent instead of + being truncated by miscalibrated zeros.""" + device = torch.device("cuda") + confidence_logits = torch.full((2, 5), 40.0, dtype=torch.float32, device=device) + confidence_logits[:, 2] = -100.0 + survival_probs = torch.empty_like(confidence_logits) + + kwargs = dict( + min_survival_probability=0.0, + num_reqs=2, + num_speculative_steps=5, + runtime_num_reqs=torch.tensor([2], dtype=torch.int32, device=device), + survival_probs=survival_probs, + budget_frac=0.9, + ) + capacity_t1 = torch.full((2,), -1, dtype=torch.int32, device=device) + compute_draft_token_capacity_from_confidence( + confidence_logits, capacity_t1, **kwargs + ) + capacity_t10 = torch.full((2,), -1, dtype=torch.int32, device=device) + compute_draft_token_capacity_from_confidence( + confidence_logits, capacity_t10, confidence_temperature=10.0, **kwargs + ) + + torch.accelerator.synchronize() + # T=1: exact-zero survival past position 2 truncates both requests. + assert capacity_t1.cpu().tolist() == [2, 2] + # T=10: sigmoid(-10) > 0, so the ceil(10*0.9) = 9 admission + # budget is fully spent. + assert capacity_t10.cpu().tolist() == [5, 4] + + +def test_online_sts_fits_order_preserving_temperatures(): + """Online STS fits per-position temperatures from rejection-sampler + outcomes: identity before data, softens over-confident positions, + sharpens under-confident ones, and never reorders candidates.""" + device = torch.device("cuda") + sts = DSparkOnlineSTS(max_num_reqs=4, num_steps=3, device=device) + + # Cold start: identity calibration. + probe = torch.tensor([[2.0, 1.0, -1.0]], dtype=torch.float32, device=device) + assert torch.equal(sts.calibrate(probe), probe) + + slots = torch.tensor([0, 1], dtype=torch.int32, device=device) + # Head claims p~0.88 everywhere (logit 2.0). + logits = torch.full((2, 3), 2.0, dtype=torch.float32, device=device) + # Alternate outcomes so pos0 accepts 50% (head over-confident there) + # while pos1/pos2 always accept once reached (head under-confident). + acc_hi = torch.tensor([3, 3], device=device) + acc_lo = torch.tensor([0, 0], device=device) + ver = torch.tensor([3, 3], device=device) + for _ in range(1000): + sts.stage_proposal(slots, logits) + sts.record(slots, acc_hi, ver) + sts.stage_proposal(slots, logits) + sts.record(slots, acc_lo, ver) + + torch.accelerator.synchronize() + temps = sts.temperatures.cpu() + calibrated = torch.sigmoid(sts.calibrate(logits)[0]).cpu() + # pos0 empirical 0.5 vs raw 0.88: temperature must soften (T >> 1, + # pushed toward the grid edge since sigmoid(2/T) -> 0.5+). + assert temps[0] > 2.0 + assert calibrated[0] < 0.65 + # pos1/pos2 empirical 1.0 (conditioned on the prefix surviving): + # temperature sharpens (T < 1). + assert temps[1] < 1.0 and temps[2] < 1.0 + assert calibrated[1] > 0.9 + + # Order preservation within every position, regardless of fit. + lo = torch.tensor([[0.5, 0.5, 0.5]], dtype=torch.float32, device=device) + hi = torch.tensor([[3.0, 3.0, 3.0]], dtype=torch.float32, device=device) + assert (sts.calibrate(hi) > sts.calibrate(lo)).all() + + +def test_online_sts_ignores_invalid_or_consumed_proposals(): + device = torch.device("cuda") + sts = DSparkOnlineSTS(max_num_reqs=2, num_steps=2, device=device) + slots = torch.tensor([0, 1], dtype=torch.int32, device=device) + logits = torch.ones((2, 2), dtype=torch.float32, device=device) + accepted = torch.tensor([2, 1], dtype=torch.int32, device=device) + verified = torch.tensor([2, 2], dtype=torch.int32, device=device) + + pristine_trials = sts.bin_trials.clone() + pristine_hits = sts.bin_hits.clone() + sts.stage_proposal(slots, logits) + sts.invalidate_all() + sts.record(slots, accepted, verified) + sts.stage_proposal(slots, logits, valid=False) + sts.record(slots, accepted, verified) + torch.accelerator.synchronize() + assert torch.equal(sts.bin_trials, pristine_trials) + assert torch.equal(sts.bin_hits, pristine_hits) + + sts.stage_proposal(slots, logits) + sts.record(slots, accepted, verified) + recorded_trials = sts.bin_trials.clone() + recorded_hits = sts.bin_hits.clone() + sts.record(slots, accepted, verified) + torch.accelerator.synchronize() + assert torch.equal(sts.bin_trials, recorded_trials) + assert torch.equal(sts.bin_hits, recorded_hits) + + +def test_qwen_dspark_confidence_head_honors_markov_mode(monkeypatch): + class CapturingLinear(torch.nn.Module): + def __init__(self, input_size: int, output_size: int, **_kwargs): + super().__init__() + self.input_size = input_size + self.output_size = output_size + self.last_input: torch.Tensor | None = None + + def forward(self, value: torch.Tensor) -> torch.Tensor: + assert value.shape[-1] == self.input_size + self.last_input = value + return value[..., : self.output_size] + + monkeypatch.setattr(qwen3_dspark, "ReplicatedLinear", CapturingLinear) + hidden = torch.randn(2, 4) + markov = torch.randn(2, 3) + + plain = qwen3_dspark.DSparkConfidenceHead(4, prefix="plain", include_markov=False) + plain(hidden, markov) + assert isinstance(plain.proj, CapturingLinear) + assert plain.proj.last_input is not None + assert plain.proj.last_input.shape == (2, 4) + assert torch.equal(plain.proj.last_input, hidden.float()) + + with_markov = qwen3_dspark.DSparkConfidenceHead( + 7, prefix="markov", include_markov=True + ) + with_markov(hidden, markov) + assert isinstance(with_markov.proj, CapturingLinear) + assert with_markov.proj.last_input is not None + assert with_markov.proj.last_input.shape == (2, 7) + assert torch.equal( + with_markov.proj.last_input, torch.cat([hidden, markov], dim=-1).float() + ) + + +def test_capacity_based_verification_manager_updates_cpu_capacities(): + device = torch.device("cuda") + req_states = RequestState( + max_num_reqs=4, + max_model_len=4, + max_num_batched_tokens=16, + num_speculative_steps=3, + vocab_size=32, + device=device, + ) + req_states.req_id_to_index = {"req0": 2, "req1": 0} + handler = VarlenCapacityBasedVerificationManager( + max_num_tokens=16, + req_states=req_states, + device=device, + ) + handler.add_request(2) + handler.add_request(0) + input_batch: Any = SimpleNamespace( + req_ids=["req0", "req1"], + idx_mapping_np=np.array([2, 0], dtype=np.int32), + num_tokens=0, + num_tokens_after_padding=0, + num_draft_tokens=0, + num_draft_tokens_per_req=None, + input_ids=torch.empty(0, dtype=torch.int32, device=device), + positions=torch.empty(0, dtype=torch.int64, device=device), + is_padding=torch.empty(0, dtype=torch.bool, device=device), + ) + draft_token_capacity = torch.tensor([1, 2], dtype=torch.int32, device=device) + + handler.trim_batch(input_batch) + handler.update_capacities(draft_token_capacity) + assert handler.copy_event_pending + + torch.accelerator.synchronize() + handler.trim_batch(input_batch) + assert handler.draft_token_capacity_np.tolist() == [2, 3, 1, 3] + + handler.update_capacities(draft_token_capacity) + torch.accelerator.synchronize() + del req_states.req_id_to_index["req0"] + handler.draft_token_capacity_np.fill(3) + handler.trim_batch(input_batch) + assert handler.draft_token_capacity_np.tolist() == [2, 3, 3, 3] + + +def test_capacity_update_canonicalizes_across_tp_before_staging(monkeypatch): + events: list[object] = [] + manager: Any = SimpleNamespace( + capacity_bypassed=False, + idx_mapping_np=np.array([0, 1], dtype=np.int32), + _flush_draft_token_capacity_copy=lambda: events.append("flush"), + _stage_draft_token_capacity_copy=lambda tensor: events.append( + ("stage", tensor.tolist()) + ), + ) + + class FakeTPGroup: + world_size = 2 + + @staticmethod + def broadcast(tensor: torch.Tensor, src: int = 0) -> None: + events.append(("broadcast", src)) + tensor.copy_(torch.tensor([3, 1], dtype=tensor.dtype)) + + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_tp_group", lambda: FakeTPGroup() + ) + + CapacityBasedVerificationManager.update_capacities( + manager, torch.tensor([1, 2], dtype=torch.int32) + ) + + assert events == ["flush", ("broadcast", 0), ("stage", [3, 1])] + + +def test_capacity_manager_bypasses_readback_below_profiled_knee(): + device = torch.device("cuda") + req_states = RequestState( + max_num_reqs=4, + max_model_len=4, + max_num_batched_tokens=16, + num_speculative_steps=3, + vocab_size=32, + device=device, + ) + handler = VarlenCapacityBasedVerificationManager( + max_num_tokens=16, + req_states=req_states, + device=device, + ) + handler.capacity_activation_batch_size = 2 + handler.draft_token_capacity_np.fill(0) + + assert not handler.should_apply_capacity(1) + assert handler.capacity_bypassed + assert handler.draft_token_capacity_np.tolist() == [3, 3, 3, 3] + + handler.idx_mapping_np = np.array([0], dtype=np.int32) + handler.update_capacities(torch.tensor([1], dtype=torch.int32, device=device)) + assert not handler.copy_event_pending + + assert handler.should_apply_capacity(2) + assert not handler.capacity_bypassed + + +def test_varlen_capacity_manager_warmup_compacts_inputs_in_place(): + device = torch.device("cuda") + req_states = RequestState( + max_num_reqs=4, + max_model_len=8, + max_num_batched_tokens=16, + num_speculative_steps=3, + vocab_size=32, + device=device, + ) + handler = VarlenCapacityBasedVerificationManager( + max_num_tokens=16, + req_states=req_states, + device=device, + ) + + handler.warmup(InputBuffers(4, 16, device)) + torch.accelerator.synchronize() + + +def test_varlen_capacity_manager_compacts_verifier_batch(): + device = torch.device("cuda") + req_states = RequestState( + max_num_reqs=4, + max_model_len=8, + max_num_batched_tokens=16, + num_speculative_steps=3, + vocab_size=32, + device=device, + ) + req_states.last_sampled_tokens[:2] = torch.tensor( + [[101], [201]], dtype=torch.int64, device=device + ) + req_states.draft_tokens[:2] = torch.tensor( + [[11, 12, 13], [21, 22, 23]], dtype=torch.int64, device=device + ) + handler = VarlenCapacityBasedVerificationManager( + max_num_tokens=16, + req_states=req_states, + device=device, + ) + handler.add_request(0) + handler.add_request(1) + handler.draft_token_capacity_np[:2] = np.array([1, 2], dtype=np.int32) + + input_ids = torch.tensor( + [101, 11, 12, 13, 201, 21, 22, 23, 0, 0], + dtype=torch.int32, + device=device, + ) + positions = torch.tensor( + [0, 1, 2, 3, 0, 1, 2, 3, 0, 0], + dtype=torch.int64, + device=device, + ) + is_padding = torch.zeros(10, dtype=torch.bool, device=device) + input_batch = InputBatch( + req_ids=["req0", "req1"], + num_reqs=2, + num_reqs_after_padding=2, + idx_mapping=torch.tensor([0, 1], dtype=torch.int32, device=device), + idx_mapping_np=np.array([0, 1], dtype=np.int32), + expanded_idx_mapping=torch.tensor( + [0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device + ), + expanded_local_pos=torch.tensor( + [0, 1, 2, 3, 0, 1, 2, 3], dtype=torch.int32, device=device + ), + num_scheduled_tokens=np.array([4, 4], dtype=np.int32), + max_query_len=4, + num_tokens=8, + num_tokens_after_padding=5, + num_draft_tokens=6, + num_draft_tokens_per_req=np.array([3, 3], dtype=np.int32), + query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + query_start_loc_np=np.array([0, 4, 8], dtype=np.int32), + seq_lens=torch.zeros(2, dtype=torch.int32, device=device), + seq_lens_cpu_upper_bound=torch.tensor([4, 4], dtype=torch.int32), + max_seq_len_upper_bound=4, + dcp_local_seq_lens=None, + num_computed_tokens_np=np.array([0, 0], dtype=np.int32), + prefill_len_np=np.array([0, 0], dtype=np.int32), + num_computed_prefill_tokens_np=np.array([0, 0], dtype=np.int32), + is_prefilling_np=np.array([False, False], dtype=np.bool_), + max_seq_len_np=None, + input_ids=input_ids[:5], + positions=positions[:5], + is_padding=is_padding[:5], + logits_indices=torch.arange(8, dtype=torch.int64, device=device), + cu_num_logits=torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + cu_num_logits_np=np.array([0, 4, 8], dtype=np.int32), + has_structured_output_reqs=False, + prompt_lens=None, + ) + + handler.trim_batch(input_batch) + + torch.accelerator.synchronize() + assert input_batch.num_scheduled_tokens.tolist() == [2, 3] + assert input_batch.num_draft_tokens_per_req.tolist() == [1, 2] + assert input_batch.num_tokens == 5 + assert input_batch.num_draft_tokens == 3 + assert input_batch.cu_num_logits_np.tolist() == [0, 2, 5] + assert input_batch.query_start_loc_np.tolist() == [0, 2, 5] + assert input_batch.input_ids.shape[0] == input_batch.num_tokens_after_padding + assert input_batch.input_ids[: input_batch.num_tokens].cpu().tolist() == [ + 101, + 11, + 201, + 21, + 22, + ] + assert input_batch.positions[: input_batch.num_tokens].cpu().tolist() == [ + 0, + 1, + 0, + 1, + 2, + ] + assert input_batch.seq_lens.cpu().tolist() == [2, 3] + assert input_batch.logits_indices.cpu().tolist() == [0, 1, 2, 3, 4] + assert ( + input_batch.is_padding[: input_batch.num_tokens].cpu().tolist() == [False] * 5 + ) + + +def test_masked_capacity_manager_marks_pruned_tokens_for_forward_and_sampler(): + device = torch.device("cuda") + req_states = RequestState( + max_num_reqs=4, + max_model_len=8, + max_num_batched_tokens=16, + num_speculative_steps=3, + vocab_size=32, + device=device, + ) + handler = MaskedCapacityBasedVerificationManager( + max_num_tokens=16, + req_states=req_states, + device=device, + ) + handler.add_request(0) + handler.add_request(1) + handler.draft_token_capacity_np[:2] = np.array([1, 2], dtype=np.int32) + + input_ids = torch.arange(16, dtype=torch.int32, device=device) + input_batch = InputBatch( + req_ids=["req0", "req1"], + num_reqs=2, + num_reqs_after_padding=2, + idx_mapping=torch.tensor([0, 1], dtype=torch.int32, device=device), + idx_mapping_np=np.array([0, 1], dtype=np.int32), + expanded_idx_mapping=torch.tensor( + [0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device + ), + expanded_local_pos=torch.tensor( + [0, 1, 2, 3, 0, 1, 2, 3], dtype=torch.int32, device=device + ), + num_scheduled_tokens=np.array([4, 4], dtype=np.int32), + max_query_len=4, + num_tokens=8, + num_tokens_after_padding=10, + num_draft_tokens=6, + num_draft_tokens_per_req=np.array([3, 3], dtype=np.int32), + query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + query_start_loc_np=np.array([0, 4, 8], dtype=np.int32), + seq_lens=torch.tensor([4, 4], dtype=torch.int32, device=device), + seq_lens_cpu_upper_bound=torch.tensor([4, 4], dtype=torch.int32), + max_seq_len_upper_bound=4, + dcp_local_seq_lens=None, + num_computed_tokens_np=np.array([0, 0], dtype=np.int32), + prefill_len_np=np.array([0, 0], dtype=np.int32), + num_computed_prefill_tokens_np=np.array([0, 0], dtype=np.int32), + is_prefilling_np=np.array([False, False], dtype=np.bool_), + max_seq_len_np=None, + input_ids=input_ids, + positions=torch.arange(16, dtype=torch.int64, device=device), + is_padding=torch.zeros(10, dtype=torch.bool, device=device), + logits_indices=torch.arange(8, dtype=torch.int64, device=device), + cu_num_logits=torch.tensor([0, 4, 8], dtype=torch.int32, device=device), + cu_num_logits_np=np.array([0, 4, 8], dtype=np.int32), + has_structured_output_reqs=False, + prompt_lens=None, + ) + + handler.trim_batch(input_batch) + slot_mappings = torch.arange(20, dtype=torch.int64, device=device).view(2, 10) + slot_mappings.masked_fill_( + input_batch.is_padding[: slot_mappings.shape[1]].unsqueeze(0), + PAD_SLOT_ID, + ) + draft_sampled = input_batch.input_ids[input_batch.logits_indices] + draft_sampled.masked_fill_(input_batch.is_padding[input_batch.logits_indices], -1) + + torch.accelerator.synchronize() + assert input_batch.num_scheduled_tokens.tolist() == [4, 4] + assert input_batch.num_draft_tokens_per_req.tolist() == [3, 3] + assert input_batch.is_padding.cpu().tolist() == [ + False, + False, + True, + True, + False, + False, + False, + True, + False, + False, + ] + assert slot_mappings.cpu().tolist() == [ + [0, 1, -1, -1, 4, 5, 6, -1, 8, 9], + [10, 11, -1, -1, 14, 15, 16, -1, 18, 19], + ] + assert draft_sampled.cpu().tolist() == [0, 1, -1, -1, 4, 5, 6, -1] + + +def test_capacity_cudagraph_dispatch_filters_by_max_query_len(): + manager = object.__new__(CudaGraphManager) + manager._graphs_captured = True + manager._resolve_effective_loras = lambda num_loras: num_loras + regular_desc = BatchExecutionDescriptor( + CUDAGraphMode.FULL, + num_tokens=12, + num_reqs=12, + uniform_token_count=6, + ) + full_capacity_desc = BatchExecutionDescriptor( + CUDAGraphMode.FULL, + num_tokens=15, + num_reqs=4, + max_req_tokens=6, + ) + piecewise_desc = BatchExecutionDescriptor( + CUDAGraphMode.PIECEWISE, + num_tokens=16, + num_reqs=None, + ) + manager._candidates = { + (11, 0): [ + regular_desc, + full_capacity_desc, + piecewise_desc, + ] + } + + desc = CudaGraphManager.dispatch( + manager, + num_reqs=4, + num_tokens=11, + uniform_token_count=None, + num_active_loras=0, + max_req_tokens=6, + ) + + assert desc is full_capacity_desc + + desc = CudaGraphManager.dispatch( + manager, + num_reqs=4, + num_tokens=11, + uniform_token_count=None, + num_active_loras=0, + max_req_tokens=7, + ) + + assert desc is piecewise_desc + + manager._candidates[(15, 0)] = [ + regular_desc, + full_capacity_desc, + piecewise_desc, + ] + desc = CudaGraphManager.dispatch( + manager, + num_reqs=4, + num_tokens=15, + uniform_token_count=None, + num_active_loras=0, + max_req_tokens=6, + ) + + assert desc is full_capacity_desc + + desc = CudaGraphManager.dispatch( + manager, + num_reqs=4, + num_tokens=11, + uniform_token_count=6, + num_active_loras=0, + ) + + assert desc is regular_desc + + +def test_varlen_cudagraph_capture_adds_full_desc(): + manager = object.__new__(CudaGraphManager) + manager.vllm_config = SimpleNamespace(speculative_config=None) + manager.compilation_config = SimpleNamespace( + cudagraph_capture_sizes=[5], + max_cudagraph_capture_size=16, + ) + manager.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + manager.decode_query_len = 4 + manager.varlen_spec_decode = True + manager.max_num_reqs = 16 + manager.lora_capture_cases = [0] + manager._candidates = {} + manager._capture_descs = {} + + manager._init_candidates() + + assert any( + desc.max_req_tokens == manager.decode_query_len + for desc in manager._capture_descs[CUDAGraphMode.FULL] + ) + + +def test_varlen_cudagraph_prefers_uniform_below_capacity_knee(monkeypatch): + monkeypatch.setenv("VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE", "4") + manager = object.__new__(CudaGraphManager) + manager.vllm_config = SimpleNamespace(speculative_config=None) + manager.compilation_config = SimpleNamespace( + cudagraph_capture_sizes=[12, 16], + max_cudagraph_capture_size=64, + ) + manager.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + manager.decode_query_len = 4 + manager.varlen_spec_decode = True + manager.max_num_reqs = 16 + manager.lora_capture_cases = [0] + manager._lora_dispatch_map = {} + manager._max_lora_case = 0 + manager._candidates = {} + manager._capture_descs = {} + + manager._init_candidates() + manager._graphs_captured = True + + low_load_desc = manager.dispatch( + num_reqs=3, + num_tokens=12, + uniform_token_count=4, + num_active_loras=0, + max_req_tokens=4, + ) + assert low_load_desc.uniform_token_count == 4 + assert low_load_desc.num_reqs == 3 + + capacity_desc = manager.dispatch( + num_reqs=4, + num_tokens=13, + uniform_token_count=None, + num_active_loras=0, + max_req_tokens=4, + ) + assert capacity_desc.uniform_token_count is None + assert capacity_desc.max_req_tokens == 4 + assert capacity_desc.num_tokens == 16 + + +def test_varlen_cudagraph_dispatch_skips_incompatible_uniform_grid(): + manager = object.__new__(CudaGraphManager) + manager.vllm_config = SimpleNamespace(speculative_config=None) + manager.compilation_config = SimpleNamespace( + cudagraph_capture_sizes=[96, 104], + max_cudagraph_capture_size=512, + ) + manager.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + manager.decode_query_len = 4 + manager.varlen_spec_decode = True + manager.max_num_reqs = 128 + manager.lora_capture_cases = [0] + manager._lora_dispatch_map = {} + manager._max_lora_case = 0 + manager._candidates = {} + manager._capture_descs = {} + + manager._init_candidates() + manager._graphs_captured = True + + desc = manager.dispatch( + num_reqs=32, + num_tokens=97, + uniform_token_count=None, + num_active_loras=0, + max_req_tokens=4, + ) + + assert desc.cg_mode == CUDAGraphMode.FULL + assert desc.num_tokens == 104 + assert desc.max_req_tokens == 4 diff --git a/tests/v1/worker/test_mixed_warmup_gate.py b/tests/v1/worker/test_mixed_warmup_gate.py index 6941df6773c1..2ef784f96c91 100644 --- a/tests/v1/worker/test_mixed_warmup_gate.py +++ b/tests/v1/worker/test_mixed_warmup_gate.py @@ -6,13 +6,35 @@ import pytest -from vllm.v1.worker.gpu.warmup import run_mixed_prefill_decode_warmup +from vllm.v1.worker.gpu.warmup import ( + _derive_dspark_draft_token_budget, + _stable_sps_step_ms, + run_mixed_prefill_decode_warmup, +) def _fail(*args, **kwargs): raise AssertionError("worker callback must not run when warmup is skipped") +def test_sps_profile_median_rejects_single_stall(): + assert _stable_sps_step_ms([10.0, 10.2, 75.0, 9.9, 10.1]) == pytest.approx(10.1) + + +def test_dspark_dynamic_budget_uses_upper_load_sps_knee(): + curve = [ + (6, 200.0), + (12, 80.0), + (24, 67.71), + (48, 58.52), + (96, 51.29), + (192, 40.33), + (384, 23.72), + ] + + assert _derive_dspark_draft_token_budget(curve, max_draft_depth=5) == 160 + + @pytest.mark.parametrize("max_num_reqs", [1, 0]) def test_mixed_warmup_skipped_for_single_seq(max_num_reqs): """A mixed prefill+decode step needs >=2 requests; with max_num_reqs < 2 diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 1b87da7f61b9..b23e041803f2 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -3,6 +3,8 @@ import copy import functools +import math +import os from collections.abc import Callable from typing import TYPE_CHECKING, Any, Literal, get_args @@ -81,6 +83,7 @@ ] RejectionSampleMethod = Literal["standard", "synthetic", "block"] DraftSampleMethod = Literal["greedy", "probabilistic"] +DSparkCapacityVerificationMode = Literal["varlen", "mask"] @config @@ -249,6 +252,56 @@ class SpeculativeConfig: synthetic_acceptance_rates. Only valid when rejection_sample_method is 'synthetic'. Mutually exclusive with synthetic_acceptance_rates.""" + dspark_confidence_threshold: float = 0.0 + """Minimum DSpark cumulative prefix-survival probability for keeping a + per-request draft prefix. Set to 0.0 to use budget-based global top-k + allocation.""" + + dspark_budget_frac: float = 1.0 + """Fraction of the full per-request draft-token budget available to the + DSpark global prefix allocator.""" + + dspark_capacity_verification_mode: DSparkCapacityVerificationMode = "varlen" + """How DSpark capacity-pruned target verification tokens are handled.""" + + dspark_confidence_temperature: float = 1.0 + """Temperature applied to the DSpark confidence-head logits before the + survival-probability computation. The released heads can emit saturated + logits whose sigmoids round to exact 0/1 in fp32; exact zeros remove + tokens from the capacity allocator's candidate set entirely. A + temperature > 1 desaturates the logits (order-preserving) so every token + stays a candidate with a usable ranking. A stopgap until checkpoints ship + Sequential-Temperature-Scaling-calibrated heads.""" + + dspark_online_sts: bool = True + """Calibrate the DSpark confidence head online with per-position + temperatures (the paper's Sequential Temperature Scaling, fitted at + serving time): per draft position, track binned empirical conditional + acceptance from the rejection sampler's outcomes (exponentially decayed) + and fit the temperature minimizing the calibration error. Order + preserving; identity until outcomes accumulate. Only active together + with a capacity verification mode.""" + + dspark_sps_curve: list[tuple[int, float]] | str | None = None + """Profiled engine step-rate curve for the DSpark hardware-aware prefix + scheduler, as ``(batch_num_tokens, steps_per_sec)`` breakpoints with + strictly increasing token counts (linearly interpolated, clamped at the + ends), or the string ``"auto"`` to profile the curve at engine init: + after CUDA graph capture, warmup decode steps are timed per power-of-two + request count (wall clock, so worker-side host prep and the draft step + are included) and rank 0's measurements are broadcast so all TP ranks + build the identical table. When set, verification lengths are chosen by + maximizing expected throughput ``tau * SPS(B)`` (DSpark Algorithm 1) + instead of spending the whole ``dspark_budget_frac`` budget; the budget + still acts as an upper bound on total admissions. Only the curve's shape + matters, so any consistent rate unit works.""" + + dspark_sps_overhead_ms: float = 0.0 + """Constant per-step overhead in milliseconds added to the step times + measured by ``dspark_sps_curve="auto"``, covering costs the init-time + profiling cannot see (scheduler/IPC above the worker). Flattens the + curve, making the theta-argmax verify more aggressively as it grows.""" + @staticmethod def _acceptance_length_to_rates(length: float, n: int) -> list[float]: """Mean acceptance length to unconditional per-position rates, using @@ -1220,6 +1273,13 @@ def _parse_attention_backend(cls, value: Any) -> Any: return AttentionBackendEnum[value.upper()] return value + @field_validator("dspark_capacity_verification_mode", mode="before") + @classmethod + def _parse_dspark_capacity_verification_mode(cls, value: Any) -> Any: + if value == "compact": + return "varlen" + return value + @model_validator(mode="after") def _verify_args(self) -> Self: if self.tensor_parallel_size is not None: @@ -1261,6 +1321,81 @@ def _verify_args(self) -> Self: "are only valid with rejection_sample_method='synthetic'." ) + if not math.isfinite(self.dspark_confidence_threshold) or not ( + 0.0 <= self.dspark_confidence_threshold <= 1.0 + ): + raise ValueError( + "dspark_confidence_threshold must be in [0, 1], got " + f"{self.dspark_confidence_threshold}." + ) + if not math.isfinite(self.dspark_budget_frac) or not ( + 0.0 < self.dspark_budget_frac <= 1.0 + ): + raise ValueError( + f"dspark_budget_frac must be in (0, 1], got {self.dspark_budget_frac}." + ) + if ( + not math.isfinite(self.dspark_confidence_temperature) + or self.dspark_confidence_temperature <= 0.0 + ): + raise ValueError( + "dspark_confidence_temperature must be > 0, got " + f"{self.dspark_confidence_temperature}." + ) + if ( + not math.isfinite(self.dspark_sps_overhead_ms) + or self.dspark_sps_overhead_ms < 0.0 + ): + raise ValueError( + "dspark_sps_overhead_ms must be >= 0, got " + f"{self.dspark_sps_overhead_ms}." + ) + if isinstance(self.dspark_sps_curve, str): + if self.dspark_sps_curve != "auto": + raise ValueError( + "dspark_sps_curve must be a list of (batch_num_tokens, " + f'steps_per_sec) pairs or "auto", got ' + f"{self.dspark_sps_curve!r}." + ) + elif self.dspark_sps_curve is not None: + self.dspark_sps_curve = [ + (int(b), float(s)) for b, s in self.dspark_sps_curve + ] + batch_sizes = [b for b, _ in self.dspark_sps_curve] + if not self.dspark_sps_curve or any( + b <= 0 or not math.isfinite(s) or s <= 0.0 + for b, s in self.dspark_sps_curve + ): + raise ValueError( + "dspark_sps_curve entries must have positive batch token " + f"counts and rates, got {self.dspark_sps_curve}." + ) + if batch_sizes != sorted(set(batch_sizes)): + raise ValueError( + "dspark_sps_curve batch token counts must be strictly " + f"increasing, got {batch_sizes}." + ) + + if ( + self.method == "dspark" + and self.dspark_capacity_verification_mode == "mask" + and ( + self.dspark_confidence_threshold > 0.0 + or self.dspark_budget_frac < 1.0 + or self.dspark_sps_curve is not None + ) + and "VLLM_MOE_SKIP_PADDING" not in os.environ + ): + # Mask mode keeps pruned verify rows in the batch as padding; the + # pruning only saves work if MoE kernels skip those rows. Set + # here (frontend) so spawned workers inherit it before their env + # caches freeze. Set VLLM_MOE_SKIP_PADDING=0 to override. + logger.info( + "DSpark mask capacity mode: defaulting VLLM_MOE_SKIP_PADDING=1 " + "so MoE kernels skip pruned verify rows." + ) + os.environ["VLLM_MOE_SKIP_PADDING"] = "1" + if self.draft_model_config: self._maybe_apply_virtual_tp_to_draft() self.draft_model_config.verify_with_parallel_config( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 70cc2b0acc5a..ccfd1986c7fe 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -630,6 +630,7 @@ class EngineArgs: spec_method: str | None = None spec_model: str | None = None spec_tokens: int | None = None + dspark_capacity_verification_mode: Literal["varlen", "mask"] | None = None diffusion_config: dict[str, Any] | None = None show_hidden_metrics_for_version: str | None = ( @@ -1530,6 +1531,12 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: vllm_group.add_argument( "--spec-tokens", **speculative_kwargs["num_speculative_tokens"] ) + vllm_group.add_argument( + "--dspark-capacity-verification-mode", + choices=["varlen", "mask"], + default=None, + help=speculative_kwargs["dspark_capacity_verification_mode"]["help"], + ) vllm_kwargs["diffusion_config"]["type"] = optional_type(json.loads) vllm_group.add_argument( "--diffusion-config", "-dc", **vllm_kwargs["diffusion_config"] @@ -1736,6 +1743,11 @@ def create_speculative_config( ("--spec-method", "method", self.spec_method), ("--spec-model", "model", self.spec_model), ("--spec-tokens", "num_speculative_tokens", self.spec_tokens), + ( + "--dspark-capacity-verification-mode", + "dspark_capacity_verification_mode", + self.dspark_capacity_verification_mode, + ), ): if value is None: continue diff --git a/vllm/envs.py b/vllm/envs.py index b6addbca59d8..6529ac851388 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -174,6 +174,9 @@ VLLM_SERVER_DEV_MODE: bool = False VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128 VLLM_MLA_DISABLE: bool = False + VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH: bool = False + VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH_WINDOW: int = 8 + VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE: int = 0 VLLM_RAY_PER_WORKER_GPUS: float = 1.0 VLLM_RAY_BUNDLE_INDICES: str = "" VLLM_CUDART_SO_PATH: str | None = None @@ -1459,6 +1462,21 @@ def _resolve_rust_frontend_path() -> str | None: ), # If set, vLLM will disable the MLA attention optimizations. "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))), + # Physically shorten DSpark's next draft block from the historical + # confidence-scheduled verification capacity. This captures/replays one + # draft graph per K instead of computing Kmax and slicing afterwards. + "VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH": lambda: bool( + int(os.getenv("VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH", "0")) + ), + "VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH_WINDOW": lambda: int( + os.getenv("VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH_WINDOW", "8") + ), + # Capture low-load DSpark draft graphs without confidence/capacity kernels. + # 0 keeps capacity active at every batch size; a positive value must match + # the profiled saturation knee used by the verification manager. + "VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE": lambda: int( + os.getenv("VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE", "0") + ), # If set, vLLM will pick up the provided Flash Attention MLA # Number of GPUs per worker in Ray, if it is set to be a fraction, # it allows ray to schedule multiple actors on a single GPU, diff --git a/vllm/model_executor/layers/sparse_attn_indexer.py b/vllm/model_executor/layers/sparse_attn_indexer.py index 1bf922bda9cb..70af878ba658 100644 --- a/vllm/model_executor/layers/sparse_attn_indexer.py +++ b/vllm/model_executor/layers/sparse_attn_indexer.py @@ -2019,6 +2019,7 @@ def sparse_attn_indexer( batch_size = padded_q_quant_decode_tokens.shape[0] next_n = padded_q_quant_decode_tokens.shape[1] num_padded_tokens = batch_size * next_n + paged_weights = weights[:num_padded_tokens] seq_lens = decode_metadata.seq_lens[:batch_size] # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; @@ -2037,7 +2038,7 @@ def sparse_attn_indexer( logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits( padded_q_quant_cast, kv_cache, - weights[:num_padded_tokens], + paged_weights, seq_lens_xpu, decode_metadata.block_table, schedule_metadata, @@ -2049,12 +2050,14 @@ def sparse_attn_indexer( logits = fp8_fp4_paged_mqa_logits( (padded_q_quant_cast, padded_q_scale), kv_cache, - weights[:num_padded_tokens], + paged_weights, seq_lens, decode_metadata.block_table, schedule_metadata, max_model_len=max_model_len, clean_logits=False, + # SM100 varlen path when set; None keeps the non-varlen path. + indices=decode_metadata.indices, ) num_rows = logits.shape[0] topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 219819759ac0..c2b75a72c7cc 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -10,8 +10,9 @@ DFlash Qwen3 draft (see qwen3_dflash.py). DSpark adds: * ``markov_head``: low-rank V x r / r x V transition bias added to the base logits, sampled left-to-right by the speculator (the sequential stage). + * ``confidence_head``: per-position acceptance-probability estimate. -DSparkMarkovHead is shared with the DSV4-style DSpark model. +DSparkMarkovHead and DSparkConfidenceHead are shared with the DSV4-style DSpark model. """ from collections.abc import Iterable @@ -21,6 +22,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ( ParallelLMHead, @@ -67,8 +69,42 @@ def bias(self, markov_embed: torch.Tensor, logits_processor) -> torch.Tensor: return logits_processor(self.markov_w2, markov_embed) +class DSparkConfidenceHead(nn.Module): + """Per-position acceptance-probability head: w^T [h_k; W1[x_{k-1}]] (+ bias). + + Returns the pre-sigmoid score; the scheduler applies the sigmoid + + temperature calibration. fp32 for a stable confidence estimate. ``bias`` + differs across checkpoints (DeepSeek-V4 DSpark: no bias; Qwen3 DSpark: + bias). + """ + + def __init__( + self, + input_dim: int, + prefix: str, + bias: bool = False, + include_markov: bool = True, + ) -> None: + super().__init__() + self.include_markov = include_markov + self.proj = ReplicatedLinear( + input_dim, + 1, + bias=bias, + return_bias=False, + params_dtype=torch.float32, + prefix=maybe_prefix(prefix, "proj"), + ) + + def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor) -> torch.Tensor: + x = ( + torch.cat([hidden, markov_embed], dim=-1) if self.include_markov else hidden + ).float() + return self.proj(x).squeeze(-1) + + class Qwen3DSparkModel(DFlashQwen3Model): - """DFlash Qwen3 backbone + DSpark Markov head.""" + """DFlash Qwen3 backbone + DSpark Markov / confidence heads.""" def __init__( self, @@ -90,6 +126,18 @@ def __init__( config.markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) + self.confidence_head: DSparkConfidenceHead | None = None + if getattr(config, "enable_confidence_head", False): + include_markov = getattr(config, "confidence_head_with_markov", False) + input_dim = config.hidden_size + if include_markov: + input_dim += config.markov_rank + self.confidence_head = DSparkConfidenceHead( + input_dim, + prefix=maybe_prefix(prefix, "confidence_head"), + bias=True, + include_markov=include_markov, + ) class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): @@ -146,11 +194,19 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed, self.logits_processor) + def compute_confidence( + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor + ) -> torch.Tensor | None: + if self.model.confidence_head is None: + return None + return self.model.confidence_head(head_hidden, markov_embed) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights = {} includes_embed_tokens = False includes_lm_head = False includes_draft_id_mapping = False + includes_confidence_head = False for name, loaded_weight in weights: # t2d is training-only; the draft remaps via d2t at sampling time. if "t2d" in name: @@ -164,22 +220,26 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): includes_embed_tokens = True if "lm_head" in name: includes_lm_head = True + if "confidence_head" in name: + includes_confidence_head = True model_weights[name] = loaded_weight # Sets has_own_embed_tokens / has_own_lm_head so load_dspark_model # knows whether to keep these or alias the target's. process_eagle_weight(self, name) # mask_embedding is an unused placeholder param; DSpark masks via the vocab row. - # confidence_head is not wired into inference yet; skip its weights. # embed_tokens / lm_head are optional; when omitted they are shared from # the target by load_dspark_model, so skip the unloaded params here. - skip_substrs = ["mask_embedding", "confidence_head"] + skip_substrs = ["mask_embedding"] if not includes_embed_tokens: skip_substrs.append("embed_tokens") if not includes_lm_head: skip_substrs.append("lm_head") if not includes_draft_id_mapping: skip_substrs.append("draft_id_to_target_id") + if self.model.confidence_head is None or not includes_confidence_head: + self.model.confidence_head = None + skip_substrs.append("confidence_head") loader = AutoWeightsLoader(self, skip_substrs=skip_substrs) loader.load_weights(model_weights.items()) self.model._build_fused_kv_buffers() diff --git a/vllm/models/deepseek_v4/common/ops/cache_utils.py b/vllm/models/deepseek_v4/common/ops/cache_utils.py index a92a0a5bb1ed..0f64b9638b1f 100644 --- a/vllm/models/deepseek_v4/common/ops/cache_utils.py +++ b/vllm/models/deepseek_v4/common/ops/cache_utils.py @@ -590,8 +590,7 @@ def _compute_dcp_global_topk_indices_and_lens_kernel( virtual_block_offsets // CP_KV_CACHE_INTERLEAVE_SIZE ) % DCP_WORLD_SIZE == DCP_RANK local_block_offsets = ( - virtual_block_offsets - // (DCP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) + virtual_block_offsets // (DCP_WORLD_SIZE * CP_KV_CACHE_INTERLEAVE_SIZE) ) * CP_KV_CACHE_INTERLEAVE_SIZE + ( virtual_block_offsets % CP_KV_CACHE_INTERLEAVE_SIZE ) @@ -758,16 +757,17 @@ def build_flashinfer_mixed_sparse_indices( ) -> tuple[torch.Tensor, torch.Tensor]: """Build the FlashInfer DSV4 sparse-index matrix for decode-first batches. - Produces ``sparse_indices`` of shape ``[num_tokens, window_size + - padded_topk]`` (the first ``window_size`` columns are SWA slot ids, the rest - are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length per - token). Decode tokens read precomputed SWA/compressed indices; prefill tokens - derive their SWA window from the position and translate local compressed - indices to global slots via the block tables. + Produces ``sparse_indices`` of shape ``[num_tokens, swa_index_width + + padded_topk]`` (the first ``swa_index_width`` columns are SWA slot ids, the + rest are compressed/top-k slot ids) and ``sparse_topk_lens`` (active length + per token). Decode tokens read precomputed SWA/compressed indices; prefill + tokens derive their SWA window from the position and translate local + compressed indices to global slots via the block tables. """ assert decode_swa_indices.dtype == torch.int32 assert decode_swa_indices.dim() == 2 - assert decode_swa_indices.shape[-1] == window_size + swa_index_width = decode_swa_indices.shape[-1] + assert swa_index_width >= window_size if decode_compressed_topk_lens is not None: assert decode_compressed_topk_lens.dtype == torch.int32 assert prefill_topk_indices.dtype == torch.int32 @@ -816,7 +816,7 @@ def build_flashinfer_mixed_sparse_indices( padded_topk = max(topk, decode_compressed_topk) padded_topk = (padded_topk + 3) // 4 * 4 sparse_indices = torch.empty( - (num_tokens, window_size + padded_topk), + (num_tokens, swa_index_width + padded_topk), dtype=torch.int32, device=decode_swa_indices.device, ) @@ -826,7 +826,7 @@ def build_flashinfer_mixed_sparse_indices( if num_tokens == 0: return sparse_indices, sparse_topk_lens - window_block_size = triton.next_power_of_2(max(window_size, 1)) + window_block_size = triton.next_power_of_2(max(swa_index_width, 1)) topk_block_size = triton.next_power_of_2(max(padded_topk, 1)) max_block_size = max(window_block_size, topk_block_size) num_warps = 4 if max_block_size >= 256 else 1 @@ -863,6 +863,7 @@ def build_flashinfer_mixed_sparse_indices( compressed_span, NUM_DECODE_TOKENS=num_decode_tokens, WINDOW_SIZE=window_size, + SWA_INDEX_WIDTH=swa_index_width, COMPRESS_RATIO=compress_ratio, TOP_K=topk, PADDED_TOP_K=padded_topk, @@ -930,6 +931,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( compressed_block_span, NUM_DECODE_TOKENS, WINDOW_SIZE: tl.constexpr, + SWA_INDEX_WIDTH: tl.constexpr, COMPRESS_RATIO: tl.constexpr, TOP_K: tl.constexpr, PADDED_TOP_K: tl.constexpr, @@ -943,9 +945,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( token_idx = tl.program_id(0) if token_idx < NUM_DECODE_TOKENS: - for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < WINDOW_SIZE + mask = offset < SWA_INDEX_WIDTH values = tl.load( decode_swa_indices_ptr + token_idx * decode_swa_stride + offset, mask=mask, @@ -991,7 +993,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + WINDOW_SIZE + + SWA_INDEX_WIDTH + offset, values, mask=mask, @@ -1005,7 +1007,7 @@ def _build_flashinfer_mixed_sparse_indices_kernel( else: compressed_len = tl.full((), DECODE_COMPRESSED_TOPK, dtype=tl.int32) - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + compressed_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + compressed_len) return prefill_idx = token_idx - NUM_DECODE_TOKENS @@ -1021,9 +1023,9 @@ def _build_flashinfer_mixed_sparse_indices_kernel( swa_start_pos = pos - swa_len + 1 topk_len = tl.minimum((pos + 1) // COMPRESS_RATIO, TOP_K) - for i in range(0, WINDOW_SIZE, WINDOW_BLOCK_SIZE): + for i in range(0, SWA_INDEX_WIDTH, WINDOW_BLOCK_SIZE): offset = i + tl.arange(0, WINDOW_BLOCK_SIZE) - mask = offset < WINDOW_SIZE + mask = offset < SWA_INDEX_WIDTH pos_offset = swa_start_pos + offset block_indices = pos_offset // swa_block_size block_numbers = tl.load( @@ -1067,10 +1069,10 @@ def _build_flashinfer_mixed_sparse_indices_kernel( tl.store( sparse_indices_ptr + token_idx * sparse_indices_stride - + WINDOW_SIZE + + SWA_INDEX_WIDTH + offset, slot_ids, mask=mask, ) - tl.store(sparse_topk_lens_ptr + token_idx, WINDOW_SIZE + topk_len) + tl.store(sparse_topk_lens_ptr + token_idx, SWA_INDEX_WIDTH + topk_len) diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 109cf2f7c6d6..02c5fbf6832e 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -44,6 +44,7 @@ ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.qwen3_dspark import ( + DSparkConfidenceHead, DSparkMarkovHead, ) from vllm.model_executor.models.utils import maybe_prefix @@ -103,7 +104,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ] ) - # Heads: final norm + hc_head, and the Markov head + # Heads: final norm + hc_head, and the Markov + confidence heads # Loaded from the "final" MTP layer weights (mtp.*) in the target checkpoint self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) hc_dim = self.hc_mult * config.hidden_size @@ -126,6 +127,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: config.dspark_markov_rank, prefix=maybe_prefix(prefix, "markov_head"), ) + self.confidence_head: DSparkConfidenceHead | None = None + if getattr(config, "enable_confidence_head", True): + self.confidence_head = DSparkConfidenceHead( + config.hidden_size + config.dspark_markov_rank, + prefix=maybe_prefix(prefix, "confidence_head"), + ) def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -388,6 +395,13 @@ def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor: def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed, self.logits_processor) + def compute_confidence( + self, head_hidden: torch.Tensor, markov_embed: torch.Tensor + ) -> torch.Tensor | None: + if self.model.confidence_head is None: + return None + return self.model.confidence_head(head_hidden, markov_embed) + # --- Weight loading ---------------------------------------------------- def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: @@ -426,6 +440,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() + loaded_confidence_head = False tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() @@ -438,6 +453,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: if mapped is None: continue name = mapped + if "confidence_head." in name: + loaded_confidence_head = True # ``.scale`` -> per-method scale suffix. if name.endswith(".scale"): @@ -474,8 +491,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: continue # Stacked rules only apply to decoder-layer weights. Head-stack params - # (main_proj/norm/hc_head/markov_head) load directly — otherwise e.g. - # "markov_w1" would collide with the "w1" shard rule. + # (main_proj/norm/hc_head/markov_head/confidence_head) load directly — + # otherwise e.g. "markov_w1" would collide with the "w1" shard rule. is_layer_param = name.startswith("model.layers.") for param_name, weight_name, stacked_shard_id in stacked_params_mapping: if not is_layer_param or weight_name not in name: @@ -505,6 +522,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loaded_params.add(name) self._finalize_moe() + if self.model.confidence_head is not None and not loaded_confidence_head: + self.model.confidence_head = None logger.info_once("DSpark draft model loaded: %d params", len(loaded_params)) return loaded_params @@ -522,8 +541,7 @@ def _remap_dspark_name(self, name: str) -> str | None: return None stage = int(m.group(1)) rest = m.group(2) - # The confidence head is not wired into inference yet; drop its weights. - if rest.startswith("confidence_head."): + if rest.startswith("confidence_head.") and self.model.confidence_head is None: return None # Head-stack params live at model level (mtp.last), context combiner at # model level (mtp.0); everything else is a per-layer decoder block. @@ -533,6 +551,7 @@ def _remap_dspark_name(self, name: str) -> str | None: "hc_head_base", "hc_head_scale", "markov_head.", + "confidence_head.", ) if rest.startswith(("main_proj.", "main_norm.")) or rest.startswith( head_prefixes diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 1848c1930db0..7fefa7f2fb84 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py @@ -301,20 +301,18 @@ def _build_sparse_index_metadata( assert swa_metadata.decode_swa_indices is not None assert swa_metadata.block_table is not None + decode_swa_width = swa_metadata.decode_swa_indices.shape[-1] decode_swa_indices = swa_metadata.decode_swa_indices.reshape( - num_decode_tokens, self.window_size + num_decode_tokens, decode_swa_width ) decode_compressed_topk_lens = None decode_compressed_indices_are_local = False decode_is_valid_token = None if swa_only: - assert self.topk_indices_buffer is not None compressed_kv_cache = swa_k_cache decode_compressed_indices = None - prefill_topk_indices = self.topk_indices_buffer[ - num_decode_tokens:num_tokens, :0 - ] + prefill_topk_indices = decode_swa_indices.new_empty(num_prefill_tokens, 0) compressed_block_table = None compressed_block_size = swa_metadata.block_size top_k = 0 diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index b67e23e523de..1cb3ccf552ab 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/vllm/models/deepseek_v4/sparse_mla.py @@ -134,7 +134,7 @@ class DeepseekV4FlashMLAMetadata(AttentionMetadata): class DeepseekV4FlashMLAMetadataBuilder( AttentionMetadataBuilder[DeepseekV4FlashMLAMetadata] ): - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS def __init__( self, diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 0a1644bcc5c6..995b02e895ef 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -542,15 +542,22 @@ def fp8_fp4_mqa_logits( def get_paged_mqa_logits_metadata( - context_lens: torch.Tensor, block_size: int, num_sms: int + context_lens: torch.Tensor, + block_size: int, + num_sms: int, + indices: torch.Tensor | None = None, ) -> torch.Tensor: """Build scheduling metadata for paged MQA logits. Args: context_lens: Tensor of shape [B], dtype int32; effective context length - per batch element. + per batch element. For the varlen path this is [total_tokens, 1]. block_size: KV-cache block size in tokens (e.g., 64). num_sms: Number of SMs available. 132 for Hopper + indices: Optional per-row request id, shape [total_tokens], int32. When + provided, selects the SM100 varlen path (next_n == 1 rows, + block_size in {32, 64}); adjacent equal values form one run. Must be + passed to `fp8_fp4_paged_mqa_logits` unchanged. Returns: Backend-specific tensor consumed by `fp8_fp4_paged_mqa_logits` to @@ -559,7 +566,11 @@ def get_paged_mqa_logits_metadata( _lazy_init() if _get_paged_mqa_logits_metadata_impl is None: return _missing() - return _get_paged_mqa_logits_metadata_impl(context_lens, block_size, num_sms) + # Only forward `indices` when set so non-varlen callers stay byte-identical. + kwargs = {} if indices is None else {"indices": indices} + return _get_paged_mqa_logits_metadata_impl( + context_lens, block_size, num_sms, **kwargs + ) def fp8_fp4_paged_mqa_logits( @@ -571,6 +582,7 @@ def fp8_fp4_paged_mqa_logits( schedule_metadata: torch.Tensor, max_model_len: int, clean_logits: bool, + indices: torch.Tensor | None = None, ) -> torch.Tensor: """Compute MQA logits using a paged KV-cache. @@ -582,19 +594,24 @@ def fp8_fp4_paged_mqa_logits( q: Tuple ``(q_values, q_scale)``. FP8 path: q_values is [B, next_n, H, D] float8_e4m3fn and q_scale is None. FP4 path: q_values is packed uint8 and q_scale is the companion - block-scale tensor. + block-scale tensor. For the varlen path B is total_tokens and + next_n == 1. kv_cache: Paged KV-cache. FP8 layout is [num_blocks, block_size, 1, D+4], dtype `torch.uint8`, with the last 4 bytes per (block, pos) storing the float dequant scale. weights: Tensor of shape [B * next_n, H], dtype `torch.float32`. context_lens: Tensor of shape [B], dtype int32; effective context length - for each batch element. + for each batch element. For the varlen path this is + [total_tokens, 1]. block_tables: Tensor of shape [B, max_blocks], dtype int32; maps logical block indices to physical blocks in the paged cache. schedule_metadata: Returned by `get_paged_mqa_logits_metadata`; used to distribute work across SMs. max_model_len: Maximum sequence length used to size the logits output. clean_logits: Whether to clean the unfilled logits into `-inf`. + indices: Optional per-row request id, shape [total_tokens], int32. When + provided, selects the SM100 varlen path (next_n == 1); must match + the tensor passed to `get_paged_mqa_logits_metadata`. Returns: Logits tensor of shape [B * next_n, max_model_len], dtype @@ -603,6 +620,8 @@ def fp8_fp4_paged_mqa_logits( _lazy_init() if _fp8_fp4_paged_mqa_logits_impl is None: return _missing() + # Only forward `indices` when set so non-varlen callers stay byte-identical. + kwargs = {} if indices is None else {"indices": indices} return _fp8_fp4_paged_mqa_logits_impl( q, kv_cache, @@ -612,6 +631,7 @@ def fp8_fp4_paged_mqa_logits( schedule_metadata, max_model_len, clean_logits=clean_logits, + **kwargs, ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index ef3ccb6f55b3..9b9d3b6794f1 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -531,6 +531,7 @@ class CommonAttentionMetadata: slot_mapping: torch.Tensor causal: bool | torch.Tensor = True + max_req_tokens: int = 0 # Needed by FastPrefillAttentionBuilder logits_indices_padded: torch.Tensor | None = None @@ -714,6 +715,7 @@ def unpadded( causal=self.causal[:num_actual_reqs] if isinstance(self.causal, torch.Tensor) else self.causal, + max_req_tokens=self.max_req_tokens, logits_indices_padded=self.logits_indices_padded, num_logits_indices=self.num_logits_indices, encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens), diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index fd3aec903fec..02ca01c3b79c 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/vllm/v1/attention/backends/flashinfer.py @@ -1064,14 +1064,39 @@ def build( num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens causal = common_attn_metadata.causal + uses_spec_reorder = self.reorder_batch_threshold > 1 + force_prefill = False if causal: - num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( - split_decodes_and_prefills( + decode_threshold = self.reorder_batch_threshold + if decode_threshold > 1: + query_lens = ( + common_attn_metadata.query_start_loc_cpu[1:] + - common_attn_metadata.query_start_loc_cpu[:-1] + ) + has_compact_spec_decode = torch.any( + (query_lens > 0) + & (query_lens != common_attn_metadata.max_query_len) + ) + if has_compact_spec_decode: + decode_threshold = 1 + uses_spec_reorder = False + force_prefill = True + if force_prefill: + num_decodes = 0 + num_prefills = num_reqs + num_decode_tokens = 0 + num_prefill_tokens = num_actual_tokens + else: + ( + num_decodes, + num_prefills, + num_decode_tokens, + num_prefill_tokens, + ) = split_decodes_and_prefills( common_attn_metadata, - decode_threshold=self.reorder_batch_threshold, + decode_threshold=decode_threshold, require_uniform=True, ) - ) else: # FlashInfer decode/TRTLLM paths cannot express non-causal # query-query attention, so DFlash runs as native prefill. @@ -1092,7 +1117,6 @@ def build( # - Prefill (FI native or TRTLLM) # - Decode (FI native, XQA, or trtllm-gen) use_cascade = common_prefix_len > 0 - uses_spec_reorder = self.reorder_batch_threshold > 1 # Page sizes >= 128 must use trtllm-gen; force it for prefill too. prefill_force_trtllm = ( True if page_size >= 128 else self.attention_config.use_trtllm_attention @@ -1290,9 +1314,12 @@ def build( if prefill_use_trtllm: # Create GPU versions - qo_indptr_prefill_gpu = ( - qo_indptr[prefill_start:] - qo_indptr[prefill_start] - ) + if prefill_start == 0: + qo_indptr_prefill_gpu = qo_indptr[: num_prefills + 1] + else: + qo_indptr_prefill_gpu = ( + qo_indptr[prefill_start:] - qo_indptr[prefill_start] + ) # Compute cum_seq_lens_kv on GPU to avoid CPU sync. # This is the cumulative sum of the number of KV cache # blocks per prefill request. @@ -1845,6 +1872,15 @@ def forward( workspace_buffer = _get_trtllm_workspace_buffer() block_tables_prefill = attn_metadata.prefill.block_tables seq_lens_prefill = attn_metadata.prefill.seq_lens + cum_seq_lens_kv = attn_metadata.prefill.cum_seq_lens_kv + cum_seq_lens_kv[:1] = 0 + page_size = kv_cache_permute.shape[-2] + num_blocks_per_req = (seq_lens_prefill + page_size - 1) // page_size + torch.cumsum( + num_blocks_per_req, + dim=0, + out=cum_seq_lens_kv[1:], + ) # This path needs to be enabled with VLLM_KV_CACHE_LAYOUT = HND assert get_kv_cache_layout() == "HND" diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 88f86d06ba73..5c06493e2cbf 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -231,10 +231,13 @@ class DeepSeekV32IndexerDecodeMetadata: max_seq_len: int | None = None global_seq_lens: torch.Tensor | None = None # Live scorer window (max compressed context across the batch) in cache - # tokens, computed host-side in build() — a metadata tensor read by the + # tokens, computed host-side in build(); a metadata tensor read by the # captured indexer kernel, never an in-kernel reduction. None => b12x uses # the capacity cap. active_width: torch.Tensor | None = None + # Per-flattened-row request id for the SM100 varlen paged kernel; None + # selects the non-varlen paged path. + indices: torch.Tensor | None = None @dataclass @@ -269,6 +272,60 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig): return max_model_len * 40 +def _supports_varlen_paged_mqa_logits() -> bool: + if ( + envs.VLLM_USE_B12X_SPARSE_INDEXER + and current_platform.is_cuda() + and current_platform.is_device_capability_family(120) + ): + # B12X consumes the already-flattened rank-1 seq_lens and repeated + # block-table rows directly, so it does not need DeepGEMM's indices. + return True + return ( + current_platform.is_cuda() + and current_platform.is_device_capability_family(100) + and has_deep_gemm() + ) + + +def _uses_varlen_dspark_capacity(vllm_config: VllmConfig) -> bool: + spec_config = vllm_config.speculative_config + return bool( + spec_config is not None + and spec_config.use_dspark() + and spec_config.dspark_capacity_verification_mode == "varlen" + and ( + spec_config.dspark_confidence_threshold > 0.0 + or spec_config.dspark_budget_frac < 1.0 + or spec_config.dspark_sps_curve is not None + ) + ) + + +def _needs_varlen_decode( + use_varlen_decode: bool, + all_uniform_width: bool, + max_decode_len: int, + max_query_len: int, +) -> bool: + """Use the compact scorer only when verification is actually ragged. + + Args: + use_varlen_decode: Whether the active backend supports varlen decode. + all_uniform_width: Whether every request has the same query width. + max_decode_len: Largest decode width in the batch. + max_query_len: Largest query width in the batch. + + Returns: + Whether this batch requires the varlen decode path. + """ + return ( + use_varlen_decode + and not all_uniform_width + and (max_decode_len > 1 or max_query_len > 1) + ) + + class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): reorder_batch_threshold: int = 1 @@ -278,6 +335,10 @@ def get_cudagraph_support( vllm_config: VllmConfig, kv_cache_spec: AttentionSpec, ) -> AttentionCGSupport: + if _supports_varlen_paged_mqa_logits() and _uses_varlen_dspark_capacity( + vllm_config + ): + return AttentionCGSupport.ALWAYS return AttentionCGSupport.UNIFORM_BATCH def __init__(self, *args, **kwargs): @@ -326,10 +387,18 @@ def __init__(self, *args, **kwargs): self.use_flattening = not current_platform.is_device_capability_family( 100 ) and next_n not in (1, 2) + # SM100 supports the varlen paged MQA logits kernel (indices-selected, + # next_n == 1 rows). Only compact spec-decode verification batches opt + # into it; uniform DFlash draft proposal should keep the native path. + self.use_varlen = ( + _supports_varlen_paged_mqa_logits() + and _uses_varlen_dspark_capacity(self.vllm_config) + ) logger.info_once( - "DSA indexer decode path: use_flattening=%s " + "DSA indexer decode path: use_flattening=%s use_varlen=%s " "(next_n=%d, use_fp4_indexer_cache=%s)", self.use_flattening, + self.use_varlen, next_n, self.use_fp4_indexer_cache, ) @@ -358,6 +427,12 @@ def __init__(self, *args, **kwargs): dtype=torch.int32, device=self.device, ) + # Per-row request ids for the SM100 varlen paged kernel. + self.decode_indices_buffer = torch.zeros( + (scheduler_config.max_num_batched_tokens,), + dtype=torch.int32, + device=self.device, + ) self.arange_buffer = torch.arange( max( scheduler_config.max_num_seqs * next_n, @@ -448,26 +523,51 @@ def _prepare_decode_tensors( use_native: bool, next_n: int, max_decode_len: int, + force_flatten: bool = False, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]: """Expand seq_lens/block_table/decode_lens for the decode kernels. - Flatten path (not use_native, max_decode_len > 1): + The flatten path (not use_native, max_decode_len > 1 or force_flatten) + expands each multi-token decode request into individual single-token + entries so the kernel always sees next_n=1. ``force_flatten`` keeps the + varlen path on the flatten buffers even for all-single-token batches so + captured CUDA graphs always read the same tensors at replay. + + The native path (use_native or max_decode_len == 1) preserves plain + decode or spec-decode with 2D per-token context lengths. + + Args: + seq_lens: Per-request context lengths. + block_table: Per-request KV block table. + decode_lens: Device tensor with each request's decode width. + decode_lens_cpu: CPU mirror of ``decode_lens``. + query_start_loc: Cumulative query offsets. + num_decodes: Number of decode requests. + num_decode_tokens: Total active decode tokens. + use_native: Whether the backend accepts the native layout. + next_n: Native query width expected by the backend. + max_decode_len: Largest decode width in the batch. + force_flatten: Whether to preserve the flatten-buffer address even + for a uniform single-token batch. + + Returns: + A tuple of prepared sequence lengths, block table, decode lengths, + effective batch size, and whether kernel-side padding is required. + + Layout details: Each multi-token decode request is expanded into individual single-token entries so the kernel always sees next_n=1. - - Native path (use_native or max_decode_len == 1): - Plain decode or spec-decode with 2D per-token context lengths. - - Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding). - seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, max_decode_len) - for native MTP. + ``seq_lens`` is 1D ``(batch_size,)`` for flatten/plain and 2D + ``(B, max_decode_len)`` for native MTP. """ min_decode_len = int(decode_lens_cpu.min().item()) - if not use_native and max_decode_len > 1: + if not use_native and (max_decode_len > 1 or force_flatten): assert self.decode_seq_lens_buffer.dim() == 1 - if min_decode_len == max_decode_len: - # Uniform decode lengths. - num_decode_tokens = num_decodes * max_decode_len + if ( + min_decode_len == max_decode_len + and num_decodes * max_decode_len == num_decode_tokens + ): + # Uniform decode lengths with no cudagraph token padding. _prepare_uniform_decode_kernel[(num_decode_tokens,)]( seq_lens, self.decode_seq_lens_buffer, @@ -511,7 +611,19 @@ def _prepare_decode_tensors( self.decode_seq_lens_buffer[:actual_expanded] = ( expanded_offsets + self.arange_buffer[:actual_expanded] + 1 ) - self.decode_seq_lens_buffer[actual_expanded:] = 0 + # FULL graphs may pad the compact varlen token batch past the + # final real row. B12X paged scoring cannot launch a zero-K + # row, so point graph-only rows at one safe compressed token; + # their slot mappings stay padded and their outputs are ignored. + padding_seq_len = ( + self.compress_ratio + if force_flatten and envs.VLLM_USE_B12X_SPARSE_INDEXER + else 0 + ) + self.decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = ( + padding_seq_len + ) + self.decode_seq_lens_buffer[num_decode_tokens:] = 0 seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] # Give each of the flattened entries the same block table row as the @@ -666,6 +778,68 @@ def _maybe_build_b12x_schedule_metadata( out=self.scheduler_metadata_buffer, ) + def _build_varlen_decode_indices( + self, + decode_lens: torch.Tensor, + decode_lens_cpu: torch.Tensor, + num_decodes: int, + num_decode_tokens: int, + max_decode_len: int, + ) -> torch.Tensor: + """Per-flattened-row request id for the SM100 varlen paged kernel. + + Rows are in request-then-token order (matching the per-token expansion + in ``_prepare_decode_tensors``); adjacent equal ids form one run. The + result always has ``num_decode_tokens`` rows so it matches the + (possibly cudagraph-padded) context_lens rows. + ``decode_lens`` must be the original per-request counts, read before the + expansion overwrites the buffer. + + Args: + decode_lens: Device tensor with each request's decode width. + decode_lens_cpu: CPU mirror of ``decode_lens``. + num_decodes: Number of decode requests. + num_decode_tokens: Number of flattened decode rows. + max_decode_len: Largest decode width in the batch. + + Returns: + A persistent tensor mapping each flattened row to its request id. + """ + indices = self.decode_indices_buffer[:num_decode_tokens] + if max_decode_len <= 1: + # One query token per request: row r is request r, and any + # qsl-padded rows past the last real request naturally form + # singleton runs. Copy into the persistent buffer: captured CUDA + # graphs bake this buffer's address, so returning arange_buffer + # directly would leave the graph reading stale ids. + indices.copy_(self.arange_buffer[:num_decode_tokens]) + return indices + + min_decode_len = int(decode_lens_cpu.min().item()) + if ( + min_decode_len == max_decode_len + and num_decodes * max_decode_len == num_decode_tokens + ): + # Uniform with no token padding: row r belongs to request + # r // max_decode_len. Static closed form, no device sync. + indices.copy_(self.arange_buffer[:num_decode_tokens] // max_decode_len) + else: + # Variable (eager only): repeat each request id by its decode_len. + # Pad the tail with non-merging trailing ids so masked pad rows form + # singleton runs instead of extending the last real request's run. + actual_expanded = int(decode_lens_cpu.sum().item()) + indices[:actual_expanded] = torch.repeat_interleave( + self.arange_buffer[:num_decodes], + decode_lens, + output_size=actual_expanded, + ) + if actual_expanded < num_decode_tokens: + pad = num_decode_tokens - actual_expanded + indices[actual_expanded:num_decode_tokens] = ( + num_decodes + self.arange_buffer[:pad] + ) + return indices + def build( self, common_prefix_len: int, @@ -681,12 +855,17 @@ def build( block_table = common_attn_metadata.block_table_tensor dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens use_dcp_local_kv = self.dcp_world_size > 1 and dcp_local_seq_lens is not None + use_varlen_decode = self.use_varlen and common_attn_metadata.max_req_tokens > 0 + # Short extends ride the decode path (default): their per-token causal + # context is the same shape as spec-verify rows, and the boundary must + # agree with the other DSv4 builders (sparse_swa/sparse_mla) so that + # varlen FULL cudagraphs, which are captured all-decode, stay valid. num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.reorder_batch_threshold, - require_uniform=not self.use_flattening, + require_uniform=not (self.use_flattening or use_varlen_decode), ) ) @@ -809,7 +988,17 @@ def build( max_decode_len = int(decode_lens_cpu.max().item()) next_n = 1 + self.num_speculative_tokens - use_native = not self.use_flattening and max_decode_len <= next_n + use_varlen = _needs_varlen_decode( + use_varlen_decode=use_varlen_decode, + all_uniform_width=bool( + (decode_lens_cpu == max_decode_len).all().item() + ), + max_decode_len=max_decode_len, + max_query_len=common_attn_metadata.max_query_len, + ) + use_native = ( + not (self.use_flattening or use_varlen) and max_decode_len <= next_n + ) global_seq_lens_for_decode = self._prepare_global_decode_seq_lens( global_seq_lens=global_seq_lens_for_decode, @@ -821,6 +1010,18 @@ def build( max_decode_len=max_decode_len, ) + # Build the varlen per-row request ids from the original per-request + # decode_lens, before _prepare_decode_tensors overwrites the buffer. + decode_indices = None + if use_varlen: + decode_indices = self._build_varlen_decode_indices( + decode_lens=decode_lens, + decode_lens_cpu=decode_lens_cpu, + num_decodes=num_decodes, + num_decode_tokens=num_decode_tokens, + max_decode_len=max_decode_len, + ) + seq_lens, block_table, decode_lens, batch_size, requires_padding = ( self._prepare_decode_tensors( seq_lens=seq_lens, @@ -833,11 +1034,12 @@ def build( use_native=use_native, next_n=next_n, max_decode_len=max_decode_len, + force_flatten=use_varlen, ) ) seq_lens_is_buffer_view = (use_native and next_n > 1) or ( - not use_native and max_decode_len > 1 + not use_native and (max_decode_len > 1 or use_varlen) ) # Uncompressed DCP localizes after per-token expansion. Compressed @@ -940,6 +1142,7 @@ def build( seq_lens, self.kv_cache_spec.storage_block_size, self.num_sms, + indices=decode_indices, ) schedule_metadata = self.scheduler_metadata_buffer @@ -950,6 +1153,7 @@ def build( requires_padding=requires_padding, schedule_metadata=schedule_metadata, max_seq_len=decode_topk_max_seq_len, + indices=decode_indices, global_seq_lens=global_seq_lens_for_decode, active_width=active_width, ) diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index a890d62cb528..56d7f3bc10ac 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -291,7 +291,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): # Base threshold: query_len <= 1 is decode reorder_batch_threshold: int = 1 - _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS supports_exact_metadata_reuse: bool = True def __init__(self, *args, **kwargs): @@ -814,6 +814,13 @@ def _compute_swa_indices_and_lens_kernel( is_valid = tl.load(is_valid_token_ptr + token_idx) if not is_valid: tl.store(swa_lens_ptr + pid, 0) + for i in range(0, window_size, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + -1, + mask=offset < window_size, + ) return req_idx = tl.load(token_to_req_indices_ptr + token_idx) @@ -880,6 +887,13 @@ def _compute_dspark_noncausal_swa_indices_kernel( is_valid = tl.load(is_valid_token_ptr + token_idx) if not is_valid: tl.store(swa_lens_ptr + pid, 0) + for i in range(0, index_width, TRITON_BLOCK_SIZE): + offset = i + tl.arange(0, TRITON_BLOCK_SIZE) + tl.store( + swa_indices_ptr + pid * swa_indices_stride + offset, + -1, + mask=offset < index_width, + ) return req_idx = tl.load(token_to_req_indices_ptr + token_idx) diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 1e12f43caacb..baa276146a2b 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -414,6 +414,7 @@ def make_local_attention_virtual_batches( block_table_tensor=block_table_local, slot_mapping=common_attn_metadata.slot_mapping, causal=True, + max_req_tokens=common_attn_metadata.max_req_tokens, seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=torch.from_numpy(num_computed_tokens_local), @@ -482,6 +483,7 @@ def make_kv_sharing_fast_prefill_common_attn_metadata( block_table_tensor=common_attn_metadata.block_table_tensor, slot_mapping=common_attn_metadata.slot_mapping, causal=True, + max_req_tokens=common_attn_metadata.max_req_tokens, seq_lens_cpu_upper_bound=common_attn_metadata.seq_lens_cpu_upper_bound, _seq_lens_cpu=common_attn_metadata._seq_lens_cpu, _num_computed_tokens_cpu=common_attn_metadata._num_computed_tokens_cpu, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 84d63d6693f6..0aed1d5809d1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1584,6 +1584,9 @@ def update_from_output( adaptive_num_drafts = 0 adaptive_num_draft_tokens = 0 adaptive_num_accepted_tokens = 0 + acceptance_length_controller = getattr( + self, "acceptance_length_controller", None + ) for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 request = self.requests.get(req_id) @@ -1625,7 +1628,7 @@ def update_from_output( num_sampled = self.num_sampled_tokens_per_step num_accepted = max(len(generated_token_ids) - num_sampled, 0) num_rejected = num_draft_tokens - num_accepted - if self.acceptance_length_controller is not None: + if acceptance_length_controller is not None: adaptive_num_drafts += 1 adaptive_num_draft_tokens += num_draft_tokens adaptive_num_accepted_tokens += num_accepted @@ -1791,8 +1794,8 @@ def update_from_output( # Invariant: EngineCore returns no partial prefill outputs. assert not prompt_logprobs_tensors - if self.acceptance_length_controller is not None: - update = self.acceptance_length_controller.observe_batch( + if acceptance_length_controller is not None: + update = acceptance_length_controller.observe_batch( num_drafts=adaptive_num_drafts, num_draft_tokens=adaptive_num_draft_tokens, num_accepted_tokens=adaptive_num_accepted_tokens, @@ -1809,13 +1812,13 @@ def update_from_output( update.num_spec_tokens, update.mean_num_accepted_tokens, update.mean_num_draft_tokens, - self.acceptance_length_controller.observation_window, + acceptance_length_controller.observation_window, ) if spec_decoding_stats is not None: spec_decoding_stats.current_num_spec_tokens = ( - self.acceptance_length_controller.num_spec_tokens - if self.acceptance_length_controller is not None + acceptance_length_controller.num_spec_tokens + if acceptance_length_controller is not None else scheduler_output.resolve_num_spec_tokens_to_schedule( self.num_spec_tokens ) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 6b96ef1f9f0c..2823bec02bfa 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -325,7 +325,8 @@ def _reshape_kv_cache( # quantized cache dtype's (possibly packed) layout. layer_cache_dtype = ( "auto" - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE + if cache_dtype != "fp8_ds_mla" + and kv_cache_spec.kv_quant_mode == KVQuantMode.NONE and not isinstance(kv_cache_spec, TQFullAttentionSpec) else cache_dtype ) @@ -494,7 +495,8 @@ def _update_hybrid_attention_layout( # but it keeps both call sites consistent for skip layers. layer_cache_dtype = ( "auto" - if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE + if cache_dtype != "fp8_ds_mla" + and kv_cache_spec.kv_quant_mode == KVQuantMode.NONE and not isinstance(kv_cache_spec, TQFullAttentionSpec) else cache_dtype ) @@ -593,6 +595,8 @@ def build_attn_metadata( for_cudagraph_capture: bool = False, causal: bool | torch.Tensor | Mapping[int, bool] = True, rswa_prefix_lens: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, + max_req_tokens: int = 0, ) -> dict[str, Any]: seq_lens = seq_lens[:num_reqs] if dcp_local_seq_lens is not None: @@ -645,6 +649,8 @@ def build_attn_metadata( mm_req_doc_ranges=mm_req_doc_ranges, rswa_prefix_lens=rswa_prefix_lens, batch_topology=batch_topology, + max_req_tokens=max_req_tokens, + is_prefilling=is_prefilling, **common_attn_metadata_extra_kwargs, ) diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 8c2b2630366f..85ce1a8f9ae8 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -172,14 +172,19 @@ def compute_slot_mappings( query_start_loc: torch.Tensor, positions: torch.Tensor, num_tokens_padded: int, + is_padding: torch.Tensor | None = None, ) -> torch.Tensor: num_reqs = idx_mapping.shape[0] num_groups = self.num_kv_cache_groups + apply_padding_mask = is_padding is not None + if is_padding is None: + is_padding = self.slot_mappings _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)]( self.max_num_batched_tokens, idx_mapping, query_start_loc, positions, + is_padding, self.block_table_ptrs, self.block_table_strides, self.num_blocks.gpu, @@ -192,6 +197,7 @@ def compute_slot_mappings( CP_SIZE=self.cp_size, CP_INTERLEAVE=self.cp_interleave, PAD_ID=PAD_SLOT_ID, + APPLY_PADDING_MASK=apply_padding_mask, TRITON_BLOCK_SIZE=1024, # type: ignore ) return self.slot_mappings[:, :num_tokens_padded] @@ -261,6 +267,7 @@ def _compute_slot_mappings_kernel( idx_mapping, # [num_reqs] query_start_loc, # [num_reqs + 1] pos, # [num_tokens] + is_padding, # [num_tokens] block_table_ptrs, # [num_kv_cache_groups] block_table_strides, # [num_kv_cache_groups] num_blocks_ptr, # [num_kv_cache_groups, max_num_reqs] @@ -273,6 +280,7 @@ def _compute_slot_mappings_kernel( CP_SIZE: tl.constexpr, CP_INTERLEAVE: tl.constexpr, PAD_ID: tl.constexpr, + APPLY_PADDING_MASK: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, ): # kv cache group id @@ -330,6 +338,9 @@ def _compute_slot_mappings_kernel( local_offsets = rounds * CP_INTERLEAVE + remainder slot_ids = block_numbers * block_size + local_offsets slot_ids = tl.where(is_local, slot_ids, PAD_ID) + if APPLY_PADDING_MASK: + padding = tl.load(is_padding + offset, mask=offset < end_idx, other=True) + slot_ids = tl.where(padding, PAD_ID, slot_ids) slot_ids = tl.where(valid_block, slot_ids, PAD_ID) tl.store(slot_mapping_ptr + offset, slot_ids, mask=offset < end_idx) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index e92f56e159f8..c50cf6679ba5 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -11,6 +11,7 @@ import torch.nn as nn from tqdm import tqdm +import vllm.envs as envs from vllm.compilation.b12x_capture import ( b12x_cuda_graph_prewarm_enabled, guard_b12x_kernel_resolution, @@ -63,6 +64,7 @@ class BatchExecutionDescriptor: num_tokens: int num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs) uniform_token_count: int | None = None + max_req_tokens: int | None = None num_active_loras: int = 0 @@ -84,6 +86,7 @@ def _is_compatible( num_tokens: int, uniform_token_count: int | None, num_active_loras: int, + max_req_tokens: int, ) -> bool: # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count # desc.num_reqs=None means no request padding needed (PIECEWISE) @@ -92,6 +95,7 @@ def _is_compatible( desc.uniform_token_count is None or desc.uniform_token_count == uniform_token_count ) + and (desc.max_req_tokens is None or desc.max_req_tokens >= max_req_tokens) and (desc.num_reqs is None or desc.num_reqs >= num_reqs) and desc.num_tokens >= num_tokens and desc.num_active_loras == num_active_loras @@ -122,6 +126,7 @@ def __init__( cudagraph_mode: CUDAGraphMode, decode_query_len: int, lora_capture_cases: list[int] | None = None, + varlen_spec_decode: bool = False, ): self.vllm_config = vllm_config self.device = device @@ -130,6 +135,7 @@ def __init__( assert self.compilation_config is not None self.cudagraph_mode = cudagraph_mode self.decode_query_len = decode_query_len + self.varlen_spec_decode = varlen_spec_decode self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -146,6 +152,9 @@ def __init__( self._graphs_captured = False self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {} + self._exact_uniform_candidates: dict[ + tuple[int, int], list[BatchExecutionDescriptor] + ] = {} self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {} self._init_candidates() @@ -202,6 +211,9 @@ def _init_candidates(self) -> None: descs_by_mode: defaultdict[CUDAGraphMode, list[BatchExecutionDescriptor]] = ( defaultdict(list) ) + exact_uniform_descs: defaultdict[ + tuple[int, int], list[BatchExecutionDescriptor] + ] = defaultdict(list) # When using Dynamic SD, num_speculative_tokens is the max number of # draft tokens. The scheduler might use a smaller number so we need @@ -250,40 +262,83 @@ def _init_candidates(self) -> None: n + num_new_sampled_tokens_per_step for n in range(1, self.vllm_config.num_speculative_tokens + 1) ] + elif ( + speculative_config + and speculative_config.use_dspark() + and envs.VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH + ): + # The confidence-capacity controller selects the next physical + # DSpark width on the worker. Capture every K so dispatch can + # avoid computing and then slicing the unused draft suffix. + num_new_sampled_tokens_per_step = ( + self.decode_query_len - self.vllm_config.num_speculative_tokens + ) + decode_query_lens = [ + n + num_new_sampled_tokens_per_step + for n in range(1, self.vllm_config.num_speculative_tokens + 1) + ] else: decode_query_lens = [self.decode_query_len] - for num_tokens, num_active_loras in product( - capture_sizes, self.lora_capture_cases + def decode_descs( + num_tokens: int, + num_active_loras: int, ): - # Capture uniform decode specfifc graphs if required - # (i.e. separate decode routine) - if separate_decode_routine and decode_mode: - for decode_query_len in decode_query_lens: - rounded_num_tokens = round_up(num_tokens, decode_query_len) - rounded_num_reqs = rounded_num_tokens // decode_query_len - - if ( - rounded_num_tokens > max_decode_tokens - or rounded_num_tokens > max_cg_capture_size - or rounded_num_reqs > self.max_num_reqs - ): + if self.varlen_spec_decode: + if num_tokens > max_decode_tokens or num_tokens > max_cg_capture_size: + return + max_requests = min(num_tokens, self.max_num_reqs) + min_requests = (num_tokens + self.decode_query_len - 1) // ( + self.decode_query_len + ) + request_counts = { + min_requests, + (max_requests + 1) // 2, + (3 * max_requests + 3) // 4, + max_requests, + } + for num_reqs in sorted(request_counts): + if num_reqs * self.decode_query_len < num_tokens: continue - - desc = BatchExecutionDescriptor( + yield BatchExecutionDescriptor( cg_mode=decode_mode, - num_tokens=rounded_num_tokens, - num_reqs=rounded_num_reqs, - uniform_token_count=decode_query_len, + num_tokens=num_tokens, + num_reqs=num_reqs, + max_req_tokens=self.decode_query_len, num_active_loras=num_active_loras, ) + return - # avoid duplicate graphs - if desc not in descs_by_mode[decode_mode]: - descs_by_mode[decode_mode].append(desc) - descs_by_token_lora[ - (rounded_num_tokens, num_active_loras) - ].append(desc) + for decode_query_len in decode_query_lens: + rounded_num_tokens = round_up(num_tokens, decode_query_len) + rounded_num_reqs = rounded_num_tokens // decode_query_len + + if ( + rounded_num_tokens > max_decode_tokens + or rounded_num_tokens > max_cg_capture_size + or rounded_num_reqs > self.max_num_reqs + ): + continue + + yield BatchExecutionDescriptor( + cg_mode=decode_mode, + num_tokens=rounded_num_tokens, + num_reqs=rounded_num_reqs, + uniform_token_count=decode_query_len, + num_active_loras=num_active_loras, + ) + + for num_tokens, num_active_loras in product( + capture_sizes, self.lora_capture_cases + ): + if separate_decode_routine and decode_mode: + for desc in decode_descs(num_tokens, num_active_loras): + if desc in descs_by_mode[desc.cg_mode]: + continue + descs_by_mode[desc.cg_mode].append(desc) + descs_by_token_lora[(desc.num_tokens, num_active_loras)].append( + desc + ) if mixed_mode: # for PIECEWISE graphs there is no limit on requests when replaying @@ -303,11 +358,46 @@ def _init_candidates(self) -> None: descs_by_mode[mixed_mode].append(desc) descs_by_token_lora[(num_tokens, num_active_loras)].append(desc) + # Capacity-based DSpark uses generic varlen target graphs once the + # batch is large enough to benefit from compaction. Below that knee, + # the capacity manager deliberately bypasses compaction and verifies + # the full draft width. Capture exact uniform graphs for that range so + # the bypass does not replay a more expensive varlen graph. Keep them + # out of the generic token buckets: adding their intermediate token + # counts there would split varlen padding ranges and force eager runs. + # Dispatch checks this exact-match map before the generic candidates. + capacity_activation_batch_size = envs.VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE + if ( + separate_decode_routine + and decode_mode + and self.varlen_spec_decode + and capacity_activation_batch_size > 1 + ): + max_uniform_reqs = min( + self.max_num_reqs, + capacity_activation_batch_size - 1, + ) + for num_reqs in range(1, max_uniform_reqs + 1): + num_tokens = self.decode_query_len * num_reqs + if num_tokens > max_decode_tokens or num_tokens > max_cg_capture_size: + continue + for num_active_loras in self.lora_capture_cases: + desc = BatchExecutionDescriptor( + cg_mode=decode_mode, + num_tokens=num_tokens, + num_reqs=num_reqs, + uniform_token_count=self.decode_query_len, + num_active_loras=num_active_loras, + ) + if desc not in descs_by_mode[decode_mode]: + descs_by_mode[decode_mode].append(desc) + exact_uniform_descs[(num_tokens, num_active_loras)].append(desc) + # Guarantee the small-request grid for every selectable decode query # length, independent of the configured capture-size list: a missing # (depth, num_reqs) point would otherwise pad requests or fall off # the FULL-graph path. - if separate_decode_routine and decode_mode: + if separate_decode_routine and decode_mode and not self.varlen_spec_decode: for decode_query_len, num_reqs in product( decode_query_lens, range(1, min(self.max_num_reqs, 32) + 1) ): @@ -329,6 +419,8 @@ def _init_candidates(self) -> None: if not descs_by_token_lora: return + self._exact_uniform_candidates = dict(exact_uniform_descs) + all_token_counts = sorted({k[0] for k in descs_by_token_lora}) current_range_start = 0 for token_cg_size in all_token_counts: @@ -437,19 +529,32 @@ def dispatch( num_tokens: int, uniform_token_count: int | None, num_active_loras: int, + max_req_tokens: int = 0, ) -> BatchExecutionDescriptor: """Find matching cudagraph descriptor from priority-ordered candidates.""" effective_loras = self._resolve_effective_loras(num_active_loras) key = (num_tokens, effective_loras) - if self._graphs_captured and num_tokens > 0 and key in self._candidates: - for desc in self._candidates[key]: + if self._graphs_captured and num_tokens > 0: + if uniform_token_count is not None: + for desc in getattr(self, "_exact_uniform_candidates", {}).get(key, ()): + if _is_compatible( + desc, + num_reqs, + num_tokens, + uniform_token_count, + effective_loras, + max_req_tokens, + ): + return desc + for desc in self._candidates.get(key, ()): if _is_compatible( desc, num_reqs, num_tokens, uniform_token_count, effective_loras, + max_req_tokens, ): return desc return BatchExecutionDescriptor( @@ -500,6 +605,7 @@ def __init__( cudagraph_mode: CUDAGraphMode, decode_query_len: int, lora_capture_cases: list[int] | None = None, + varlen_spec_decode: bool = False, ): super().__init__( vllm_config, @@ -507,6 +613,7 @@ def __init__( cudagraph_mode, decode_query_len, lora_capture_cases=lora_capture_cases, + varlen_spec_decode=varlen_spec_decode, ) self.hidden_states: torch.Tensor | None = None self.aux_hidden_states: list[torch.Tensor] = [] @@ -573,6 +680,7 @@ def create_forward_fn( attn_groups, kv_cache_config, skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE), + max_req_tokens=desc.max_req_tokens, ) # Capture with dummy rows marked as padding. @@ -666,8 +774,14 @@ def prepare_inputs_to_capture( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, skip_attn: bool = False, + max_req_tokens: int | None = None, ) -> AttentionState: - input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers) + input_batch = InputBatch.make_dummy( + num_reqs, + num_tokens, + input_buffers, + max_req_tokens=max_req_tokens, + ) 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/dp_utils.py b/vllm/v1/worker/gpu/dp_utils.py index ee9b924ba13a..888bdcc3a764 100644 --- a/vllm/v1/worker/gpu/dp_utils.py +++ b/vllm/v1/worker/gpu/dp_utils.py @@ -19,6 +19,7 @@ def sync_cudagraph_and_dp_padding( num_tokens: int, num_reqs: int, uniform_token_count: int | None, + max_req_tokens: int | None, dp_size: int, dp_rank: int, num_active_loras: int = 0, @@ -30,15 +31,17 @@ def sync_cudagraph_and_dp_padding( """ assert dp_size > 1, "DP size must be greater than 1" group = get_dp_group().cpu_group - tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu") + tensor = torch.zeros(4, dp_size, dtype=torch.int32, device="cpu") tensor[0][dp_rank] = num_tokens tensor[1][dp_rank] = desired_batch_desc.cg_mode.value tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None) + tensor[3][dp_rank] = max_req_tokens or 0 # (0 means None) dist.all_reduce(tensor, group=group) num_tokens_across_dp = tensor[0] cg_mode_across_dp = tensor[1] uniform_token_counts_across_dp = tensor[2] + max_req_tokens_across_dp = tensor[3] if torch.all(num_tokens_across_dp == 0).item(): synced_desc = BatchExecutionDescriptor( @@ -68,6 +71,9 @@ def sync_cudagraph_and_dp_padding( uniform_token_counts_across_dp == synced_uniform_token_count ): synced_uniform_token_count = None + synced_max_req_tokens: int | None = int(max_req_tokens_across_dp.max()) + if synced_max_req_tokens == 0: + synced_max_req_tokens = None # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs # so we don't perform request padding for PIECEWISE graphs. @@ -77,6 +83,7 @@ def sync_cudagraph_and_dp_padding( synced_num_tokens, synced_uniform_token_count, num_active_loras=num_active_loras, + max_req_tokens=synced_max_req_tokens or 0, ) # Update num_tokens_across_dp to reflect padded size. @@ -90,6 +97,7 @@ def dispatch_cg_and_sync_dp( num_reqs: int, num_tokens: int, uniform_token_count: int | None, + max_req_tokens: int | None, dp_size: int, dp_rank: int, need_eager: bool = False, @@ -100,6 +108,7 @@ def dispatch_cg_and_sync_dp( cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, num_reqs=num_reqs, + max_req_tokens=max_req_tokens or None, num_active_loras=num_active_loras, ) else: @@ -112,6 +121,7 @@ def dispatch_cg_and_sync_dp( num_tokens, uniform_token_count, num_active_loras=num_active_loras, + max_req_tokens=max_req_tokens or 0, ) if dp_size == 1: @@ -123,6 +133,7 @@ def dispatch_cg_and_sync_dp( num_tokens, num_reqs, uniform_token_count, + max_req_tokens, dp_size, dp_rank, num_active_loras=num_active_loras, diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 015f54c48e55..f533e9ca5a50 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -101,12 +101,16 @@ class InputBatch: # [num_reqs] per-request prompt length, only populated for R-SWA. prompt_lens: torch.Tensor | None + max_req_tokens: int | None = None + valid_num_draft_tokens_per_req: np.ndarray | None = None + @classmethod def make_dummy( cls, num_reqs: int, num_tokens: int, input_buffers: InputBuffers, + max_req_tokens: int | None = None, ) -> "InputBatch": assert 0 < num_reqs <= num_tokens device = input_buffers.device @@ -117,13 +121,27 @@ def make_dummy( expanded_idx_mapping = idx_mapping expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device) - num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32) - num_scheduled_tokens[-1] += num_tokens % num_reqs + if max_req_tokens is None: + num_scheduled_tokens = np.full( + num_reqs, num_tokens // num_reqs, dtype=np.int32 + ) + num_scheduled_tokens[-1] += num_tokens % num_reqs + else: + assert num_tokens <= num_reqs * max_req_tokens + num_scheduled_tokens = np.ones(num_reqs, dtype=np.int32) + remaining = num_tokens - num_reqs + for i in range(num_reqs - 1, -1, -1): + num_tokens_for_req = min(remaining, max_req_tokens - 1) + num_scheduled_tokens[i] += num_tokens_for_req + remaining -= num_tokens_for_req + if remaining == 0: + break assert int(num_scheduled_tokens.sum()) == num_tokens # seq_len equals to query_len - input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs - input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs + input_buffers.seq_lens[:num_reqs].copy_( + torch.from_numpy(num_scheduled_tokens).to(device=device) + ) # Pad for full CUDA graph mode. input_buffers.seq_lens[num_reqs:] = 0 seq_lens = input_buffers.seq_lens[:num_reqs] @@ -145,7 +163,7 @@ def make_dummy( input_buffers.is_padding[:num_tokens].fill_(True) is_padding = input_buffers.is_padding[:num_tokens] - logits_indices = query_start_loc[1:] - 1 + logits_indices = torch.clamp(query_start_loc[1:] - 1, min=0) cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32) cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32) # Dummy: seq_len == query_len (fresh-prefill shape). @@ -183,6 +201,7 @@ def make_dummy( cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=False, prompt_lens=None, + max_req_tokens=max_req_tokens, ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index c498c42c4553..02a0e276942a 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -105,6 +105,12 @@ from vllm.v1.worker.gpu.sample.sampler import Sampler from vllm.v1.worker.gpu.shutdown import free_before_shutdown from vllm.v1.worker.gpu.spec_decode import init_speculator +from vllm.v1.worker.gpu.spec_decode.capacity import ( + CapacityBasedVerificationManager, + check_dspark_tp_consistency, + count_valid_draft_tokens, + make_capacity_based_verification_manager, +) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( set_eagle3_aux_hidden_state_layers, ) @@ -256,6 +262,11 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): vocab_size=self.vocab_size, device=self.device, ) + # Constructed in init_attn_backend, once the final verification mode + # is known (varlen requires full CUDA graph support). + self.verification_capacity_manager: CapacityBasedVerificationManager | None = ( + None + ) self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -494,6 +505,17 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: self.attn_groups, attn_cg_support, self.kernel_block_sizes = init_attn_backend( self.kv_cache_config, self.vllm_config, self.device ) + if self.speculator is not None and self.speculator.use_draft_token_capacity: + assert self.speculative_config is not None + self.verification_capacity_manager = ( + make_capacity_based_verification_manager( + self.speculative_config.dspark_capacity_verification_mode, + attn_cg_support, + self.max_num_tokens, + self.req_states, + self.device, + ) + ) self.block_tables = BlockTables( block_sizes=block_sizes, max_num_reqs=self.max_num_reqs, @@ -524,6 +546,10 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: cudagraph_mode, decode_query_len=self.decode_query_len, lora_capture_cases=self.lora_capture_cases, + varlen_spec_decode=( + self.verification_capacity_manager is not None + and self.verification_capacity_manager.varlen_spec_decode + ), ) check_attention_cp_compatibility(self.vllm_config) if isinstance(self.speculator, DraftModelSpeculator): @@ -872,6 +898,8 @@ def add_requests(self, scheduler_output: SchedulerOutput) -> None: max_tokens=sampling_params.max_tokens if sampling_params else 1, # type: ignore[arg-type] ) req_index = self.req_states.req_id_to_index[req_id] + if self.verification_capacity_manager is not None: + self.verification_capacity_manager.add_request(req_index) if self.encoder_cache is not None: self.encoder_cache.add_request(req_id, new_req_data.mm_features) @@ -933,14 +961,6 @@ def prepare_inputs( ) -> InputBatch: num_tokens = scheduler_output.total_num_scheduled_tokens num_tokens_after_padding = batch_desc.num_tokens - assert num_tokens > 0 - if envs.VLLM_MOE_SKIP_PADDING: - # Mark trailing cudagraph-padding rows so kernels can skip work for - # them when supported. - self.input_buffers.is_padding[:num_tokens].fill_(False) - self.input_buffers.is_padding[num_tokens:num_tokens_after_padding].fill_( - True - ) num_tokens_per_req = scheduler_output.num_scheduled_tokens num_reqs = len(num_tokens_per_req) @@ -956,6 +976,7 @@ def prepare_inputs( # Get the number of draft tokens for each request. draft_tokens = scheduler_output.scheduled_spec_decode_tokens num_draft_tokens_per_req = None + valid_num_draft_tokens_per_req = None if not draft_tokens: # No draft token scheduled (common case). total_num_draft_tokens = 0 @@ -974,6 +995,18 @@ def prepare_inputs( dtype=np.int32, count=num_reqs, ) + if scheduler_output.has_structured_output_requests: + valid_num_draft_tokens_per_req = count_valid_draft_tokens( + [draft_tokens.get(req_id, ()) for req_id in req_ids], + num_reqs, + ) + elif self.verification_capacity_manager is not None: + # Without structured outputs the scheduler only sees -1 + # placeholders (real draft ids stay on the GPU), so every + # scheduled draft slot counts. The capacity manager needs this + # bound so trim_batch prunes exactly what get_num_tokens + # predicted at graph dispatch. + valid_num_draft_tokens_per_req = num_draft_tokens_per_req num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens @@ -988,6 +1021,8 @@ def prepare_inputs( idx_mapping, total_num_logits, cu_num_logits, max_expand_len ) + assert num_tokens > 0 + # Get query_start_loc. # num_reqs_padded is None for PIECEWISE graphs (no request padding needed) num_reqs_padded = batch_desc.num_reqs or num_reqs @@ -1028,17 +1063,6 @@ def prepare_inputs( seq_lens = self.input_buffers.seq_lens[:num_reqs_padded] dcp_local_seq_lens = None - if self.use_dcp: - # Prepare dcp local seq_lens. - prepare_dcp_local_seq_lens( - self.input_buffers.dcp_local_seq_lens, - self.input_buffers.seq_lens, - num_reqs, - self.dcp_size, - self.dcp_rank, - self.cp_interleave, - ) - dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[:num_reqs_padded] # Some input token ids are directly read from the last sampled tokens # and draft tokens. Also, get the logits indices to sample tokens from. @@ -1076,7 +1100,18 @@ def prepare_inputs( # prompt_lens is only used in R-SWA case. prompt_lens = self.req_states.prompt_len.gpu[idx_mapping] - return InputBatch( + max_req_tokens = batch_desc.max_req_tokens + if ( + max_req_tokens is None + and draft_tokens + and self.verification_capacity_manager is not None + and self.verification_capacity_manager.varlen_spec_decode + ): + # Keep the compact varlen attention path for PIECEWISE/eager + # verify steps, where the descriptor carries no request bound. + max_req_tokens = int(num_scheduled_tokens.max()) + + input_batch = InputBatch( req_ids=req_ids, num_reqs=num_reqs, num_reqs_after_padding=num_reqs_padded, @@ -1109,7 +1144,41 @@ def prepare_inputs( cu_num_logits_np=cu_num_logits_np, has_structured_output_reqs=scheduler_output.has_structured_output_requests, prompt_lens=prompt_lens, + max_req_tokens=max_req_tokens, + valid_num_draft_tokens_per_req=valid_num_draft_tokens_per_req, ) + # InputBuffers are reused across real, dummy, and captured batches. + # Clear stale padding before a capacity manager optionally marks a + # subset of active rows as intentionally skipped. + self.input_buffers.is_padding[:num_tokens].fill_(False) + if self.verification_capacity_manager is not None: + input_batch = self.verification_capacity_manager.trim_batch(input_batch) + if self.use_dcp: + # Prepare dcp local seq_lens. + prepare_dcp_local_seq_lens( + self.input_buffers.dcp_local_seq_lens, + self.input_buffers.seq_lens, + input_batch.num_reqs, + self.dcp_size, + self.dcp_rank, + self.cp_interleave, + ) + input_batch.dcp_local_seq_lens = self.input_buffers.dcp_local_seq_lens[ + : input_batch.num_reqs_after_padding + ] + num_tokens = input_batch.num_tokens + num_tokens_after_padding = input_batch.num_tokens_after_padding + assert 0 < num_tokens <= num_tokens_after_padding, ( + f"Batch has {num_tokens} tokens after trimming but was dispatched " + f"for {num_tokens_after_padding}" + ) + if envs.VLLM_MOE_SKIP_PADDING: + # Mark trailing cudagraph-padding rows so kernels can skip work for + # them when supported. + self.input_buffers.is_padding[num_tokens:num_tokens_after_padding].fill_( + True + ) + return input_batch def prepare_attn( self, input_batch: InputBatch @@ -1126,6 +1195,7 @@ def prepare_attn( input_batch.query_start_loc, input_batch.positions, num_tokens_padded=input_batch.num_tokens_after_padding, + is_padding=input_batch.is_padding, ) return block_tables, slot_mappings @@ -1177,6 +1247,23 @@ def sample( self.speculator.draft_logits, ) + online_sts = self.speculator.online_sts if self.speculator else None + num_sampled = sampler_output.num_sampled + if ( + online_sts is not None + and num_sampled is not None + and self.verification_capacity_manager is not None + and not self.verification_capacity_manager.capacity_bypassed + and input_batch.num_draft_tokens_per_req is not None + ): + num_bonus = self.model_state.num_new_sampled_tokens_per_step + num_logits = input_batch.cu_num_logits[1:] - input_batch.cu_num_logits[:-1] + online_sts.record( + input_batch.idx_mapping, + num_sampled[: input_batch.num_reqs] - num_bonus, + num_logits - num_bonus, + ) + return sampler_output, sampler_output.num_sampled, sampler_output.num_rejected def postprocess_sampled( @@ -1238,6 +1325,12 @@ def execute_model( num_toks = scheduler_output.total_num_scheduled_tokens max_query_len = max(scheduler_output.num_scheduled_tokens.values()) uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len) + # Per-request token bound for graph dispatch: varlen spec-decode graphs + # are captured for at most `max_req_tokens` tokens per request, so a + # batch may only replay one if its longest request fits. + max_req_tokens = max_query_len + skip_compiled = False + verification_capacity_manager = self.verification_capacity_manager num_active_loras = 0 if self.lora_config: @@ -1246,24 +1339,69 @@ def execute_model( self.lora_config, self.lora_state, req_ids, dummy_run ) - skip_compiled = False if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: # Encoder-decoder models such as Whisper should run eager/non-compiled # when encoder inputs are scheduled, because this step updates # cross-attention cache with dynamic encoder outputs. skip_compiled = True + apply_verification_capacity = True + if ( + verification_capacity_manager is not None + and verification_capacity_manager.varlen_spec_decode + and not dummy_run + ): + capacity_was_bypassed = verification_capacity_manager.capacity_bypassed + apply_verification_capacity = ( + verification_capacity_manager.should_apply_capacity( + num_reqs, + scheduler_output.has_structured_output_requests, + ) + ) + if not apply_verification_capacity and not capacity_was_bypassed: + online_sts = self.speculator.online_sts if self.speculator else None + if online_sts is not None: + # The next high-load verification must not join against a + # proposal whose low-load graph bypassed confidence logits. + online_sts.invalidate_all() + use_varlen_capacity = ( + verification_capacity_manager is not None + and verification_capacity_manager.varlen_spec_decode + and bool(scheduler_output.scheduled_spec_decode_tokens) + and not dummy_run + and apply_verification_capacity + ) + if use_varlen_capacity: + assert verification_capacity_manager is not None + # Dispatch using the compacted verifier shape. The batch is + # trimmed later, but graph selection happens here. + uniform_tok_count = None + num_toks = verification_capacity_manager.get_num_tokens( + scheduler_output.num_scheduled_tokens, + scheduler_output.scheduled_spec_decode_tokens, + scheduler_output.has_structured_output_requests, + ) + if verification_capacity_manager.tp_check_level: + assert self.speculator is not None + check_dspark_tp_consistency( + num_toks, verification_capacity_manager, self.speculator + ) + with record_function_or_nullcontext("vllm:v2/target/dispatch"): batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.cudagraph_manager, num_reqs, num_toks, uniform_tok_count, + max_req_tokens, self.dp_size, self.dp_rank, need_eager=is_profile or skip_compiled, num_active_loras=num_active_loras, ) + if use_varlen_capacity: + assert verification_capacity_manager is not None + verification_capacity_manager.maybe_log_dispatch(num_toks, batch_desc) if batch_desc.num_tokens == 0: # All DP ranks have zero tokens to run. @@ -1308,6 +1446,7 @@ def execute_model( batch_desc.num_reqs or num_reqs, batch_desc.num_tokens, self.input_buffers, + max_req_tokens=batch_desc.max_req_tokens, ) phase = _profile_batch_phase(input_batch, dummy_run=True) if not skip_attn_for_dummy_run: @@ -1528,6 +1667,12 @@ def sample_tokens( return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output) num_spec_tokens_to_schedule = execute_model_state.num_spec_tokens_to_schedule + if self.verification_capacity_manager is not None: + num_spec_tokens_to_schedule = ( + self.verification_capacity_manager.recommended_draft_depth( + num_spec_tokens_to_schedule + ) + ) # Last rank: sample tokens phase = _profile_batch_phase(input_batch) @@ -1654,6 +1799,15 @@ def sample_tokens( # persistent fixed-width buffer untouched, but report an # empty draft list to the scheduler for the next iteration. draft_tokens_for_next_step = draft_tokens + if ( + self.verification_capacity_manager is not None + and not self.verification_capacity_manager.capacity_bypassed + ): + draft_token_capacity = self.speculator.compute_capacities(input_batch) + assert draft_token_capacity is not None + self.verification_capacity_manager.update_capacities( + draft_token_capacity + ) if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does diff --git a/vllm/v1/worker/gpu/model_states/default.py b/vllm/v1/worker/gpu/model_states/default.py index 867ff26fe028..7a06aaebc9e6 100644 --- a/vllm/v1/worker/gpu/model_states/default.py +++ b/vllm/v1/worker/gpu/model_states/default.py @@ -148,13 +148,20 @@ def prepare_attn( num_reqs = input_batch.num_reqs num_tokens = input_batch.num_tokens query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np) - max_query_len = input_batch.max_query_len + max_query_len = ( + input_batch.max_req_tokens or 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 else: max_seq_len = seq_lens_cpu_upper_bound[:num_reqs].max().item() + is_prefilling = torch.from_numpy(input_batch.is_prefilling_np) + if num_reqs != input_batch.num_reqs: + padded_is_prefilling = torch.zeros(num_reqs, dtype=torch.bool) + padded_is_prefilling[: input_batch.num_reqs] = is_prefilling + is_prefilling = padded_is_prefilling req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None if ( self.supports_mm_inputs @@ -185,5 +192,7 @@ def prepare_attn( mm_req_doc_ranges=req_doc_ranges, for_cudagraph_capture=for_capture, rswa_prefix_lens=input_batch.prompt_lens, + is_prefilling=is_prefilling, + max_req_tokens=input_batch.max_req_tokens or 0, ) return attn_metadata diff --git a/vllm/v1/worker/gpu/sample/bad_words.py b/vllm/v1/worker/gpu/sample/bad_words.py index b5517dee1b18..59cc1b8a05bd 100644 --- a/vllm/v1/worker/gpu/sample/bad_words.py +++ b/vllm/v1/worker/gpu/sample/bad_words.py @@ -152,9 +152,11 @@ def _bad_words_kernel( from_spec_input = actual_pos >= output_len if from_spec_input: spec_offset = actual_pos - output_len - actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset) + actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset).to( + tl.int64 + ) else: - actual = tl.load(output_base + actual_pos) + actual = tl.load(output_base + actual_pos).to(tl.int64) match = match & (expected == actual) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index f30a9b887b16..590cf47ba62e 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/vllm/v1/worker/gpu/sample/gumbel.py @@ -11,9 +11,7 @@ # in `tl.constexpr(...)`. We can only do that when Triton is actually # available — on the CPU worker path `tl` is a placeholder whose `constexpr` # attribute is `None`, and `tl.constexpr(...)` would crash at import time. -_TL_RAND_MIN = ( - tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 -) +_TL_RAND_MIN = tl.constexpr(4.6566127342e-10) if HAS_TRITON else 4.6566127342e-10 _FP64_ONE_MINUS_EPS = ( tl.constexpr(0.9999999999999999) if HAS_TRITON else 0.9999999999999999 ) @@ -227,7 +225,7 @@ def _gumbel_sample_kernel( USE_FP64=USE_FP64, PER_TOKEN_COL=PER_TOKEN_COL, ) - token_id = block_idx * BLOCK_SIZE + idx + token_id = tl.minimum(block_idx * BLOCK_SIZE + idx, vocab_size - 1) tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id) tl.store(local_max_ptr + token_idx * local_max_stride + block_idx, value) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 91f57f474b1b..874cebab907c 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -240,6 +240,7 @@ def propose( num_reqs, num_tokens, uniform_token_count, + max_req_tokens=None, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, @@ -329,6 +330,7 @@ def rebuild_draft_prefill_attn_state() -> tuple[ num_reqs, num_reqs, uniform_token_count=1, + max_req_tokens=None, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, diff --git a/vllm/v1/worker/gpu/spec_decode/capacity.py b/vllm/v1/worker/gpu/spec_decode/capacity.py new file mode 100644 index 000000000000..336eee512997 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/capacity.py @@ -0,0 +1,980 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os +import tempfile +from collections.abc import Sequence +from typing import TYPE_CHECKING + +import numpy as np +import torch + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.triton_utils import tl, triton +from vllm.v1.attention.backend import AttentionCGSupport +from vllm.v1.worker.gpu.async_utils import async_copy_to_np +from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu +from vllm.v1.worker.gpu.input_batch import ( + combine_sampled_and_draft_tokens, + expand_idx_mapping, + prepare_pos_seq_lens, + prepare_prefill_inputs, +) + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from vllm.v1.worker.gpu.attn_utils import AttentionCGSupportInfo + from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers + from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + from vllm.v1.worker.gpu.states import RequestState + + +@triton.jit +def _compact_token_inputs_kernel( + source_input_ids_ptr, + source_positions_ptr, + output_input_ids_ptr, + output_positions_ptr, + old_query_start_loc_ptr, + new_query_start_loc_ptr, + BLOCK_SIZE: tl.constexpr, +): + req_idx = tl.program_id(0) + old_start = tl.load(old_query_start_loc_ptr + req_idx) + new_start = tl.load(new_query_start_loc_ptr + req_idx) + new_end = tl.load(new_query_start_loc_ptr + req_idx + 1) + offsets = tl.arange(0, BLOCK_SIZE) + mask = offsets < (new_end - new_start) + input_ids = tl.load(source_input_ids_ptr + old_start + offsets, mask=mask) + positions = tl.load(source_positions_ptr + old_start + offsets, mask=mask) + tl.store(output_input_ids_ptr + new_start + offsets, input_ids, mask=mask) + tl.store(output_positions_ptr + new_start + offsets, positions, mask=mask) + + +def get_draft_token_capacities( + idx_mapping_np: np.ndarray, + draft_token_capacity_np: np.ndarray, + valid_draft_tokens_per_req: np.ndarray | None = None, +) -> np.ndarray: + capacities = draft_token_capacity_np[idx_mapping_np] + if valid_draft_tokens_per_req is not None: + capacities = np.minimum(capacities, valid_draft_tokens_per_req) + return capacities + + +def count_valid_draft_tokens( + draft_token_lists: Sequence[Sequence[int]], + num_reqs: int, +) -> np.ndarray: + """Grammar-validated draft ids round-trip through the scheduler; negative + ids mark invalidated drafts.""" + return np.fromiter( + (sum(token_id >= 0 for token_id in tokens) for tokens in draft_token_lists), + dtype=np.int32, + count=num_reqs, + ) + + +class DSparkDynamicDraftDepthController: + """Select the next physical draft width from historical capacities. + + Capacity is produced by the confidence scheduler for a proposal and is + available on the host when that proposal is verified one engine step + later. Reusing it for the next proposal is non-anticipating: the selected + width cannot depend on tokens generated by that proposal. + + A single physical width is shared by the whole batch, so it must cover the + longest useful proposal in that batch. The controller therefore follows + the maximum useful capacity over a short observation window, rather than + the mean. Reduced depths are periodically probed one token upward because + observed capacities are capped by the currently exposed draft width. + """ + + _PROBE_AFTER_WINDOWS = 8 + + def __init__(self, max_depth: int, observation_window: int) -> None: + if max_depth < 1: + raise ValueError("max_depth must be at least one") + if observation_window < 1: + raise ValueError("observation_window must be at least one") + self.max_depth = max_depth + self.observation_window = observation_window + self.depth = max_depth + self._steps = 0 + self._max_capacity = 0 + self._capacity_sum = 0 + self._attempted_sum = 0 + self._num_drafts = 0 + self._batch_size_high_watermark: int | None = None + self._saturated_windows = 0 + self._draft_token_budget: int | None = None + + def set_draft_token_budget(self, draft_token_budget: int) -> None: + if draft_token_budget < 1: + raise ValueError("draft_token_budget must be at least one") + self._draft_token_budget = draft_token_budget + if self._should_log(): + logger.info( + "DSpark dynamic physical-depth budget set to %d draft tokens", + draft_token_budget, + ) + + def _load_limited_depth(self, batch_size: int) -> int: + if self._draft_token_budget is None: + return self.max_depth + return max( + 1, + min( + self.max_depth, + (self._draft_token_budget + batch_size - 1) // batch_size, + ), + ) + + def _reset_observation(self) -> None: + self._steps = 0 + self._max_capacity = 0 + self._capacity_sum = 0 + self._attempted_sum = 0 + self._num_drafts = 0 + + @staticmethod + def _should_log() -> bool: + return ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ) + + def _reset_to_max_depth(self, reason: str) -> None: + previous_depth = self.depth + self.depth = self.max_depth + self._saturated_windows = 0 + self._reset_observation() + if self.depth != previous_depth and self._should_log(): + logger.info( + "DSpark physical draft depth changed %d -> %d (%s)", + previous_depth, + self.depth, + reason, + ) + + def observe(self, capacities: np.ndarray, attempted: np.ndarray) -> int: + valid = attempted > 0 + batch_size = int(valid.sum()) + if batch_size == 0: + self._reset_to_max_depth("idle boundary") + self._batch_size_high_watermark = None + return self.depth + + # A lower-concurrency request group should not inherit a short width + # selected under heavier load. Ignore one-request scheduling jitter, + # but reset after a material (25%) drop from the recent high water mark. + high_watermark = self._batch_size_high_watermark + if high_watermark is not None and batch_size * 4 <= high_watermark * 3: + self._reset_to_max_depth( + f"active requests decreased {high_watermark} -> {batch_size}" + ) + self._batch_size_high_watermark = batch_size + else: + self._batch_size_high_watermark = max(high_watermark or 0, batch_size) + + attempted_valid = attempted[valid].astype(np.int64, copy=False) + capacities_valid = np.minimum(capacities[valid], attempted_valid) + self._max_capacity = max( + self._max_capacity, + int(capacities_valid.max(initial=0)), + ) + self._capacity_sum += int(capacities_valid.sum()) + self._attempted_sum += int(attempted_valid.sum()) + self._num_drafts += batch_size + self._steps += 1 + if self._steps < self.observation_window: + return self.depth + + mean_capacity = self._capacity_sum / self._num_drafts + mean_attempted = self._attempted_sum / self._num_drafts + load_limited_depth = self._load_limited_depth(batch_size) + target = max(1, min(load_limited_depth, self._max_capacity)) + + previous_depth = self.depth + if target < self.depth: + self.depth = target + self._saturated_windows = 0 + elif target > self.depth: + self.depth += 1 + self._saturated_windows = 0 + elif self.depth < load_limited_depth: + self._saturated_windows += 1 + if self._saturated_windows >= self._PROBE_AFTER_WINDOWS: + self.depth += 1 + self._saturated_windows = 0 + + if self.depth != previous_depth and self._should_log(): + logger.info( + "DSpark physical draft depth changed %d -> %d " + "(max capacity %d, mean capacity %.3f, " + "mean attempted %.3f, load limit %d, window %d)", + previous_depth, + self.depth, + self._max_capacity, + mean_capacity, + mean_attempted, + load_limited_depth, + self.observation_window, + ) + self._reset_observation() + return self.depth + + +class CapacityBasedVerificationManager: + def __init__( + self, + max_num_tokens: int, + req_states: "RequestState", + device: torch.device, + ): + self.max_num_tokens = max_num_tokens + self.device = device + self.req_states = req_states + # Debug (VLLM_DSPARK_TP_CHECK={1,2}): cross-check capacity-derived + # batch shapes (=2 also GPU-side capacity/STS state) across TP ranks + # each step via check_dspark_tp_consistency; divergence otherwise + # surfaces as a collective-size-mismatch hang far downstream. + self.tp_check_level = int(os.environ.get("VLLM_DSPARK_TP_CHECK", "0") or "0") + self.draft_token_capacity_np = np.full( + req_states.max_num_reqs, + req_states.num_speculative_steps, + dtype=np.int32, + ) + self.copy_stream = torch.cuda.Stream(device) + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) + + self.req_ids: list[str] = [] + self.idx_mapping_np: np.ndarray | None = None + self.copied_draft_token_capacity_np: np.ndarray | None = None + self.copied_req_ids: list[str] = [] + self.copied_idx_mapping_np: np.ndarray | None = None + self.num_draft_tokens: int = 0 + self.copy_event_pending = False + self.varlen_spec_decode = False + self.capacity_log_interval = int( + os.environ.get("VLLM_DSPARK_CAPACITY_LOG_INTERVAL", "0") or "0" + ) + self._capacity_log_step = 0 + self._capacity_log_snapshot: np.ndarray | None = None + self.dynamic_draft_depth_controller = ( + DSparkDynamicDraftDepthController( + req_states.num_speculative_steps, + envs.VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH_WINDOW, + ) + if envs.VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH + else None + ) + configured_activation_batch_size = ( + envs.VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE + ) + if configured_activation_batch_size < 0: + raise ValueError( + "VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE must be >= 0, got " + f"{configured_activation_batch_size}." + ) + self.capacity_activation_batch_size = max(1, configured_activation_batch_size) + self._capacity_activation_is_configured = configured_activation_batch_size > 0 + self.capacity_bypassed = False + + def add_request(self, req_idx: int) -> None: + self.draft_token_capacity_np[req_idx] = self.req_states.num_speculative_steps + + def _stage_draft_token_capacity_copy( + self, + draft_token_capacity: torch.Tensor, + ) -> None: + self.num_draft_tokens = self.req_states.num_speculative_steps + self.copied_draft_token_capacity_np = None + self.copied_req_ids = self.req_ids + assert self.idx_mapping_np is not None + self.copied_idx_mapping_np = self.idx_mapping_np + self.copy_event_pending = False + + current_stream = torch.cuda.current_stream(self.device) + self.copy_stream.wait_stream(current_stream) + with torch.cuda.stream(self.copy_stream): + self.copied_draft_token_capacity_np = async_copy_to_np(draft_token_capacity) + draft_token_capacity.record_stream(self.copy_stream) + self.copy_event.record() + self.copy_event_pending = True + + def update_capacities(self, draft_token_capacity: torch.Tensor) -> None: + if self.capacity_bypassed: + return + self._flush_draft_token_capacity_copy() + if torch.distributed.is_initialized(): + from vllm.distributed.parallel_state import get_tp_group + + tp_group = get_tp_group() + if tp_group.world_size > 1: + # Varlen dispatch must select the same physical graph on every + # TP rank. Canonicalize the tiny capacity vector before its + # asynchronous host readback so rank-local confidence jitter + # cannot turn into a collective shape mismatch. + tp_group.broadcast(draft_token_capacity, src=0) + assert self.idx_mapping_np is not None + self._stage_draft_token_capacity_copy(draft_token_capacity) + + def get_num_tokens( + self, + num_tokens_per_req: dict[str, int], + draft_tokens: dict[str, list[int]], + has_structured_output_requests: bool = False, + ) -> int: + raise NotImplementedError + + def _flush_draft_token_capacity_copy(self) -> None: + if self.copied_draft_token_capacity_np is None: + return + if self.copy_event_pending: + # Block until the staged copy lands: batch shapes derived from the + # capacities must be identical on every TP rank, so an + # opportunistic query() (timing-dependent per rank) could diverge + # graph dispatch. + self.copy_event.synchronize() + self.copy_event_pending = False + capacities = np.clip( + self.copied_draft_token_capacity_np, 0, self.num_draft_tokens + ) + num_copied = capacities.shape[0] + req_ids = self.copied_req_ids[:num_copied] + assert self.copied_idx_mapping_np is not None + idx_mapping_np = self.copied_idx_mapping_np[:num_copied] + req_id_to_index = self.req_states.req_id_to_index + active = np.fromiter( + (req_id in req_id_to_index for req_id in req_ids), + dtype=np.bool_, + count=len(req_ids), + ) + self.draft_token_capacity_np[idx_mapping_np[active]] = capacities[active] + self.copied_draft_token_capacity_np = None + + def warmup(self, input_buffers: "InputBuffers") -> None: + pass + + def _remember_batch(self, input_batch: "InputBatch") -> None: + self.req_ids = list(input_batch.req_ids) + self.idx_mapping_np = input_batch.idx_mapping_np + + def _remember_capacity_log_snapshot(self, capacities: np.ndarray) -> None: + if self.capacity_log_interval > 0: + self._capacity_log_snapshot = capacities.copy() + + def _observe_dynamic_draft_depth( + self, + input_batch: "InputBatch", + capacities: np.ndarray, + ) -> None: + controller = self.dynamic_draft_depth_controller + if controller is None: + return + attempted = input_batch.valid_num_draft_tokens_per_req + if attempted is None: + attempted = input_batch.num_draft_tokens_per_req + if attempted is not None: + controller.observe(capacities, attempted) + + def recommended_draft_depth(self, default: int) -> int: + controller = self.dynamic_draft_depth_controller + return default if controller is None else min(default, controller.depth) + + def set_dynamic_draft_token_budget(self, draft_token_budget: int) -> None: + controller = self.dynamic_draft_depth_controller + if controller is not None: + controller.set_draft_token_budget(draft_token_budget) + max_depth = self.req_states.num_speculative_steps + profiled_activation_batch_size = max( + 1, (draft_token_budget + max_depth - 1) // max_depth + ) + if not self._capacity_activation_is_configured: + self.capacity_activation_batch_size = profiled_activation_batch_size + if controller._should_log(): + logger.info( + "DSpark capacity readback activates at %d requests " + "(profiled threshold %d, draft-token budget %d, K=%d)", + self.capacity_activation_batch_size, + profiled_activation_batch_size, + draft_token_budget, + max_depth, + ) + if ( + self._capacity_activation_is_configured + and self.capacity_activation_batch_size + != profiled_activation_batch_size + ): + logger.warning( + "Configured DSpark capacity activation batch size %d " + "does not match profiled threshold %d.", + self.capacity_activation_batch_size, + profiled_activation_batch_size, + ) + + def should_apply_capacity( + self, + num_reqs: int, + has_structured_output_requests: bool = False, + ) -> bool: + apply_capacity = ( + has_structured_output_requests + or num_reqs >= self.capacity_activation_batch_size + ) + bypass = not apply_capacity + if bypass == self.capacity_bypassed: + return apply_capacity + + if bypass: + # Drain a copy staged by the previous high-load step once. Low-load + # steps then stay entirely GPU/graph driven with no host readback. + self._flush_draft_token_capacity_copy() + self.draft_token_capacity_np.fill(self.req_states.num_speculative_steps) + self.capacity_bypassed = bypass + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + logger.info( + "DSpark capacity readback %s at %d active requests " + "(activation threshold %d)", + "enabled" if apply_capacity else "bypassed", + num_reqs, + self.capacity_activation_batch_size, + ) + return apply_capacity + + def maybe_log_dispatch(self, num_tokens: int, batch_desc: object) -> None: + """Log the effective capacity and graph shape for opt-in diagnostics.""" + if self.capacity_log_interval <= 0: + return + self._capacity_log_step += 1 + if self._capacity_log_step % self.capacity_log_interval: + return + if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: + return + capacities = self._capacity_log_snapshot + if capacities is None or capacities.size == 0: + return + histogram = np.bincount( + capacities, + minlength=self.req_states.num_speculative_steps + 1, + ) + logger.info( + "DSpark capacity dispatch: reqs=%d kept=%d/%d mean=%.3f hist=%s " + "tokens=%d padded=%s cg=%s max_req=%s", + capacities.size, + int(capacities.sum()), + capacities.size * self.req_states.num_speculative_steps, + float(capacities.mean()), + histogram.tolist(), + num_tokens, + getattr(batch_desc, "num_tokens", None), + getattr(batch_desc, "cg_mode", None), + getattr(batch_desc, "max_req_tokens", None), + ) + + @staticmethod + def _get_num_bonus_tokens(input_batch: "InputBatch") -> int: + num_logits = np.diff(input_batch.cu_num_logits_np) + num_bonus_tokens_per_req = num_logits - input_batch.num_draft_tokens_per_req + num_bonus_tokens = int(num_bonus_tokens_per_req[0]) + assert np.all(num_bonus_tokens_per_req == num_bonus_tokens) + return num_bonus_tokens + + def _set_token_views( + self, + input_batch: "InputBatch", + num_tokens: int | None = None, + ) -> "InputBatch": + n = input_batch.num_tokens_after_padding if num_tokens is None else num_tokens + input_batch.input_ids = input_batch.input_ids.as_strided((n,), (1,)) + input_batch.positions = input_batch.positions.as_strided((n,), (1,)) + input_batch.is_padding = input_batch.is_padding.as_strided((n,), (1,)) + return input_batch + + def trim_batch( + self, + input_batch: "InputBatch", + ) -> "InputBatch": + raise NotImplementedError + + +class VarlenCapacityBasedVerificationManager(CapacityBasedVerificationManager): + def __init__( + self, + max_num_tokens: int, + req_states: "RequestState", + device: torch.device, + ): + super().__init__(max_num_tokens, req_states, device) + self.varlen_spec_decode = True + + def get_num_tokens( + self, + num_tokens_per_req: dict[str, int], + draft_tokens: dict[str, list[int]], + has_structured_output_requests: bool = False, + ) -> int: + self._flush_draft_token_capacity_copy() + num_reqs = len(num_tokens_per_req) + req_ids = sorted( + num_tokens_per_req, + key=num_tokens_per_req.get, # type: ignore[arg-type] + ) + num_scheduled_tokens = np.fromiter( + (num_tokens_per_req[req_id] for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + draft_token_lists = [draft_tokens.get(req_id, ()) for req_id in req_ids] + num_draft_tokens_per_req = np.fromiter( + (len(tokens) for tokens in draft_token_lists), + dtype=np.int32, + count=num_reqs, + ) + if has_structured_output_requests: + valid_num_draft_tokens_per_req = count_valid_draft_tokens( + draft_token_lists, num_reqs + ) + else: + # Otherwise the scheduler only sees -1 placeholders (real draft ids + # stay on the GPU), so every scheduled slot counts. + valid_num_draft_tokens_per_req = num_draft_tokens_per_req + idx_mapping_np = np.fromiter( + (self.req_states.req_id_to_index[req_id] for req_id in req_ids), + dtype=np.int32, + count=num_reqs, + ) + capacities = get_draft_token_capacities( + idx_mapping_np, + self.draft_token_capacity_np, + valid_num_draft_tokens_per_req, + ) + self._remember_capacity_log_snapshot(capacities) + total_num_draft_tokens = int(capacities.sum()) + return int( + num_scheduled_tokens.sum() + - num_draft_tokens_per_req.sum() + + total_num_draft_tokens + ) + + def warmup(self, input_buffers: "InputBuffers") -> None: + max_query_len = self.req_states.num_speculative_steps + 1 + num_reqs = max( + 1, + min( + self.req_states.max_num_reqs, + self.max_num_tokens // max_query_len, + ), + ) + lengths = set() + max_warmup_query_len = min(self.max_num_tokens, 2048) + block_size = 1 + while block_size <= max_warmup_query_len: + lengths.add(block_size) + block_size *= 2 + lengths.add(max_query_len) + + idx_mapping_np = np.arange(num_reqs, dtype=np.int32) + idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + for query_len in sorted(lengths): + reqs_for_len = max(1, min(num_reqs, self.max_num_tokens // query_len)) + num_tokens = reqs_for_len * query_len + query_start_loc_np = np.arange( + 0, + num_tokens + 1, + query_len, + dtype=np.int32, + ) + query_start_loc = input_buffers.query_start_loc[: reqs_for_len + 1] + async_copy_to_gpu(query_start_loc_np, out=query_start_loc) + input_ids = input_buffers.input_ids[:num_tokens] + positions = input_buffers.positions[:num_tokens] + seq_lens = input_buffers.seq_lens + + _compact_token_inputs_kernel[(reqs_for_len,)]( + input_ids, + positions, + input_ids, + positions, + query_start_loc, + query_start_loc, + BLOCK_SIZE=triton.next_power_of_2(query_len), + ) + prepare_pos_seq_lens( + idx_mapping[:reqs_for_len], + query_start_loc, + self.req_states.num_computed_tokens.gpu, + positions, + seq_lens, + ) + expand_idx_mapping( + idx_mapping[:reqs_for_len], + num_tokens, + query_start_loc, + query_len, + ) + combine_sampled_and_draft_tokens( + input_ids, + idx_mapping[:reqs_for_len], + self.req_states.last_sampled_tokens, + query_start_loc, + seq_lens, + self.req_states.prefill_len.gpu, + self.req_states.draft_tokens, + query_start_loc, + num_tokens, + ) + + def _rewrite_compact_batch( + self, + input_batch: "InputBatch", + num_scheduled_tokens: np.ndarray, + num_draft_tokens_per_req: np.ndarray, + num_bonus_tokens: int, + ) -> None: + num_tokens = int(num_scheduled_tokens.sum()) + old_num_tokens = input_batch.num_tokens + old_query_start_loc = async_copy_to_gpu( + input_batch.query_start_loc_np[: input_batch.num_reqs + 1], + device=self.device, + ) + # The active views may already be trimmed below the old packed size, + # while their backing allocations still contain every source row. + source_input_ids = input_batch.input_ids.as_strided( + (old_num_tokens,), (1,) + ).clone() + source_positions = input_batch.positions.as_strided( + (old_num_tokens,), (1,) + ).clone() + + input_batch.num_scheduled_tokens = num_scheduled_tokens + input_batch.num_tokens = num_tokens + input_batch.num_draft_tokens_per_req = num_draft_tokens_per_req + input_batch.num_draft_tokens = int(num_draft_tokens_per_req.sum()) + + num_logits = num_draft_tokens_per_req + num_bonus_tokens + input_batch.cu_num_logits_np = np.empty( + input_batch.num_reqs + 1, dtype=np.int32 + ) + input_batch.cu_num_logits_np[0] = 0 + np.cumsum(num_logits, out=input_batch.cu_num_logits_np[1:]) + input_batch.cu_num_logits = async_copy_to_gpu( + input_batch.cu_num_logits_np, + device=self.device, + ) + ( + input_batch.expanded_idx_mapping, + input_batch.expanded_local_pos, + ) = expand_idx_mapping( + input_batch.idx_mapping, + int(input_batch.cu_num_logits_np[-1]), + input_batch.cu_num_logits, + max(1, int(num_logits.max())), + ) + + query_start_loc_np = np.empty(self.req_states.max_num_reqs + 1, dtype=np.int32) + query_start_loc_np[0] = 0 + np.cumsum( + num_scheduled_tokens, + out=query_start_loc_np[1 : input_batch.num_reqs + 1], + ) + query_start_loc_np[input_batch.num_reqs + 1 :] = input_batch.num_tokens + input_batch.query_start_loc_np = query_start_loc_np[ + : input_batch.num_reqs_after_padding + 1 + ] + async_copy_to_gpu( + input_batch.query_start_loc_np, + out=input_batch.query_start_loc, + ) + _compact_token_inputs_kernel[(input_batch.num_reqs,)]( + source_input_ids, + source_positions, + input_batch.input_ids, + input_batch.positions, + old_query_start_loc, + input_batch.query_start_loc, + BLOCK_SIZE=triton.next_power_of_2(max(1, int(num_scheduled_tokens.max()))), + ) + old_query_start_loc.record_stream(torch.cuda.current_stream(self.device)) + source_input_ids.record_stream(torch.cuda.current_stream(self.device)) + source_positions.record_stream(torch.cuda.current_stream(self.device)) + + if np.any(input_batch.is_prefilling_np): + prepare_prefill_inputs( + input_batch.input_ids, + self.req_states.next_prefill_tokens, + input_batch.idx_mapping, + input_batch.query_start_loc, + self.req_states.all_token_ids.gpu, + self.req_states.prefill_len.gpu, + self.req_states.num_computed_tokens.gpu, + ) + prepare_pos_seq_lens( + input_batch.idx_mapping, + input_batch.query_start_loc, + self.req_states.num_computed_tokens.gpu, + input_batch.positions, + input_batch.seq_lens, + ) + input_batch.logits_indices = combine_sampled_and_draft_tokens( + input_batch.input_ids, + input_batch.idx_mapping, + self.req_states.last_sampled_tokens, + input_batch.query_start_loc, + input_batch.seq_lens, + self.req_states.prefill_len.gpu, + self.req_states.draft_tokens, + input_batch.cu_num_logits, + int(input_batch.cu_num_logits_np[-1]), + num_bonus_tokens, + ) + + seq_lens_cpu_upper_bound_np = np.zeros( + input_batch.num_reqs_after_padding, + dtype=np.int32, + ) + np.add( + input_batch.num_computed_tokens_np, + num_scheduled_tokens, + out=seq_lens_cpu_upper_bound_np[: input_batch.num_reqs], + ) + input_batch.seq_lens_cpu_upper_bound = torch.from_numpy( + seq_lens_cpu_upper_bound_np + ) + input_batch.is_padding[: input_batch.num_tokens].fill_(False) + + def trim_batch( + self, + input_batch: "InputBatch", + ) -> "InputBatch": + self._remember_batch(input_batch) + self._set_token_views( + input_batch, + max(input_batch.num_tokens, input_batch.num_tokens_after_padding), + ) + input_batch.is_padding[: input_batch.num_tokens].fill_(False) + if not self.capacity_bypassed: + self._flush_draft_token_capacity_copy() + if ( + input_batch.num_draft_tokens == 0 + or input_batch.num_draft_tokens_per_req is None + ): + return self._set_token_views(input_batch) + + if self.capacity_bypassed: + attempted = input_batch.valid_num_draft_tokens_per_req + if attempted is None: + attempted = input_batch.num_draft_tokens_per_req + self._remember_capacity_log_snapshot(attempted) + self._observe_dynamic_draft_depth(input_batch, attempted) + return self._set_token_views(input_batch) + + num_bonus_tokens = self._get_num_bonus_tokens(input_batch) + num_draft_tokens_per_req = get_draft_token_capacities( + input_batch.idx_mapping_np, + self.draft_token_capacity_np, + input_batch.valid_num_draft_tokens_per_req, + ) + self._observe_dynamic_draft_depth(input_batch, num_draft_tokens_per_req) + if int(num_draft_tokens_per_req.sum()) != input_batch.num_draft_tokens: + num_scheduled_tokens = ( + input_batch.num_scheduled_tokens + - input_batch.num_draft_tokens_per_req + + num_draft_tokens_per_req + ) + self._rewrite_compact_batch( + input_batch, + num_scheduled_tokens, + num_draft_tokens_per_req, + num_bonus_tokens, + ) + self._remember_batch(input_batch) + return self._set_token_views(input_batch) + + +class MaskedCapacityBasedVerificationManager(CapacityBasedVerificationManager): + def __init__( + self, + max_num_tokens: int, + req_states: "RequestState", + device: torch.device, + ): + super().__init__(max_num_tokens, req_states, device) + self.forward_skip_mask = torch.empty( + max_num_tokens, dtype=torch.bool, device=device + ) + self.forward_skip_mask_np = np.zeros(max_num_tokens, dtype=np.bool_) + self.forward_skip_mask_len = 0 + self.has_forward_skip_mask = False + + def _prepare_forward_skip_mask( + self, + input_batch: "InputBatch", + num_bonus_tokens: int, + ) -> None: + assert input_batch.num_draft_tokens_per_req is not None + self.forward_skip_mask_np[: input_batch.num_tokens_after_padding] = False + capacities = get_draft_token_capacities( + input_batch.idx_mapping_np, + self.draft_token_capacity_np, + input_batch.valid_num_draft_tokens_per_req, + ) + self._observe_dynamic_draft_depth(input_batch, capacities) + for req_idx, (num_draft_tokens, capacity) in enumerate( + zip(input_batch.num_draft_tokens_per_req, capacities) + ): + num_kept = int(capacity) + if num_kept == num_draft_tokens: + continue + start = int(input_batch.query_start_loc_np[req_idx]) + prune_start = start + num_bonus_tokens + num_kept + prune_end = start + num_bonus_tokens + int(num_draft_tokens) + self.forward_skip_mask_np[prune_start:prune_end] = True + self.forward_skip_mask_len = input_batch.num_tokens_after_padding + self.has_forward_skip_mask = bool( + np.any(self.forward_skip_mask_np[: input_batch.num_tokens]) + ) + if self.has_forward_skip_mask: + async_copy_to_gpu( + self.forward_skip_mask_np[: self.forward_skip_mask_len], + out=self.forward_skip_mask[: self.forward_skip_mask_len], + ) + + def trim_batch( + self, + input_batch: "InputBatch", + ) -> "InputBatch": + self._flush_draft_token_capacity_copy() + self._remember_batch(input_batch) + input_batch.is_padding[: input_batch.num_tokens].fill_(False) + self.forward_skip_mask_len = 0 + self.has_forward_skip_mask = False + if ( + input_batch.num_draft_tokens == 0 + or input_batch.num_draft_tokens_per_req is None + ): + return self._set_token_views(input_batch) + + num_bonus_tokens = self._get_num_bonus_tokens(input_batch) + self._prepare_forward_skip_mask(input_batch, num_bonus_tokens) + if self.has_forward_skip_mask: + input_batch.is_padding[: input_batch.num_tokens].logical_or_( + self.forward_skip_mask[: input_batch.num_tokens] + ) + input_batch.input_ids[: input_batch.num_tokens].masked_fill_( + input_batch.is_padding[: input_batch.num_tokens], + 0, + ) + return self._set_token_views(input_batch) + + +def check_dspark_tp_consistency( + num_toks: int, + manager: CapacityBasedVerificationManager, + speculator: "DSparkSpeculator", +) -> None: + """Debug-only (VLLM_DSPARK_TP_CHECK={1,2}): fail fast, with per-rank + state dumps, if the capacity-derived dispatch shape diverges across TP. + + Request->slot binding may legitimately differ across ranks, so slot-keyed + state is re-keyed by req_id before hashing. + """ + import hashlib + + from vllm.distributed.parallel_state import get_tp_group + + tp_group = get_tp_group() + if tp_group.world_size <= 1: + return + req_states = manager.req_states + capacities = manager.draft_token_capacity_np + req_ids = sorted(req_states.req_id_to_index) + slots_np = np.fromiter( + (req_states.req_id_to_index[req_id] for req_id in req_ids), + dtype=np.int64, + count=len(req_ids), + ) + slots = torch.from_numpy(slots_np).to(manager.device) + + def gpu_md5(x: torch.Tensor) -> str: + return hashlib.md5(x.cpu().numpy().tobytes()).hexdigest() + + sts = speculator.online_sts + payload: list[object] = [ + num_toks, + hashlib.md5(capacities[slots_np].tobytes()).hexdigest(), + ] + if manager.tp_check_level >= 2: + # GPU-side state hashes; cheap here since the capacity flush just + # drained the previous step. + if speculator.draft_logits is not None: + payload.append(gpu_md5(speculator.draft_logits[slots])) + payload.append(gpu_md5(req_states.draft_tokens[slots])) + if sts is not None: + payload.append(gpu_md5(sts.logits_by_state[slots])) + payload.append(gpu_md5(sts.bin_trials)) + payload.append(gpu_md5(sts.temperatures)) + payload_tuple = tuple(payload) + all_payloads: list[tuple | None] = [None] * tp_group.world_size + torch.distributed.all_gather_object( + all_payloads, payload_tuple, group=tp_group.cpu_group + ) + if all(p == all_payloads[0] for p in all_payloads): + return + dump = { + "rank": tp_group.rank_in_group, + "payloads": all_payloads, + "req_id_to_index": dict(req_states.req_id_to_index), + "draft_token_capacity_np": capacities.copy(), + "draft_tokens_tail": req_states.draft_tokens[-8:].cpu(), + "seeds_tail": speculator.seeds[-8:].cpu(), + "confidence_logits": speculator.draft_token_confidence_logits.cpu(), + "sts_bin_trials": sts.bin_trials.cpu() if sts else None, + "sts_bin_hits": sts.bin_hits.cpu() if sts else None, + "sts_temperatures": sts.temperatures.cpu() if sts else None, + "sts_logits_by_state": sts.logits_by_state.cpu() if sts else None, + } + with tempfile.NamedTemporaryFile( + prefix=f"dspark_tp_divergence_rank{tp_group.rank_in_group}_", + suffix=".pt", + delete=False, + ) as dump_file: + path = dump_file.name + torch.save(dump, path) + raise RuntimeError( + f"DSpark capacity TP divergence: rank {tp_group.rank_in_group} " + f"payloads {all_payloads}; state dumped to {path}" + ) + + +def make_capacity_based_verification_manager( + mode: str, + attn_cg_support: "AttentionCGSupportInfo", + max_num_tokens: int, + req_states: "RequestState", + device: torch.device, +) -> CapacityBasedVerificationManager: + if mode == "varlen" and attn_cg_support.min_cg_support != AttentionCGSupport.ALWAYS: + logger.info_once( + "Falling back to masked DSpark capacity verification because " + "%s reports CUDA graph support %s.", + attn_cg_support.min_cg_attn_backend, + attn_cg_support.min_cg_support.name, + ) + mode = "mask" + if mode == "varlen": + logger.info_once("Using compact varlen DSpark capacity verification.") + return VarlenCapacityBasedVerificationManager( + max_num_tokens, + req_states, + device, + ) + if mode == "mask": + return MaskedCapacityBasedVerificationManager( + max_num_tokens, + req_states, + device, + ) + raise ValueError(f"Unknown DSpark capacity verification mode: {mode}") diff --git a/vllm/v1/worker/gpu/spec_decode/causal_cascade/speculator.py b/vllm/v1/worker/gpu/spec_decode/causal_cascade/speculator.py index 122b3fa78f9c..b6c3ab47eeea 100644 --- a/vllm/v1/worker/gpu/spec_decode/causal_cascade/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/causal_cascade/speculator.py @@ -606,6 +606,7 @@ def _run_embedded_mtp_prefill( num_reqs, num_tokens, uniform_token_count, + max_req_tokens=None, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py index a4b4033cafa9..e39baec83691 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/cudagraph.py @@ -105,6 +105,7 @@ def create_forward_fn( slot_mappings, num_tokens_across_dp, cg_mode, + num_query_per_req=desc.uniform_token_count, ) return fwd, attn_state diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 1f45a90b5f06..b9b065d09c30 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -10,12 +10,13 @@ """ from collections.abc import Mapping -from typing import Any +from typing import Any, NamedTuple import numpy as np import torch import torch.nn as nn +import vllm.envs as envs from vllm.config import VllmConfig, replace from vllm.config.compilation import CUDAGraphMode from vllm.forward_context import BatchDescriptor, set_forward_context @@ -41,6 +42,14 @@ logger = init_logger(__name__) +class _DFlashInputBatch(NamedTuple): + num_reqs: int + num_scheduled_tokens: np.ndarray + positions: torch.Tensor + query_start_loc: torch.Tensor + idx_mapping: torch.Tensor + + class DFlashSpeculator(DraftModelSpeculator): _speculator_name = "DFlash" # For logging, so we can share methods with subclasses @@ -120,6 +129,19 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): # 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] = [] + self.dynamic_physical_depth = ( + envs.VLLM_DSPARK_DYNAMIC_DRAFT_DEPTH and self._speculator_name == "DSpark" + ) + + def _query_len_for_speculative_steps(self, num_speculative_steps: int) -> int: + return ( + num_speculative_steps + if self.sample_from_anchor + else 1 + num_speculative_steps + ) + + def _speculative_steps_for_query_len(self, query_len: int) -> int: + return query_len if self.sample_from_anchor else query_len - 1 @property def attn_vllm_config(self) -> VllmConfig: @@ -177,6 +199,98 @@ def capture(self, attn_states: dict | None = None) -> None: progress_bar_desc=f"Capturing {self._speculator_name.lower()} CUDA graphs", ) + def _warmup_prepare_inputs_kernel(self) -> None: + if self.draft_kv_cache_group_id < 0: + return + + target_query_lens = { + self.num_query_per_req + 1, + 32, + 128, + 256, + 1024, + } + draft_query_lens = ( + range(1, self.num_query_per_req + 1) + if self.dynamic_physical_depth + else (self.num_query_per_req,) + ) + for target_query_len in sorted(target_query_lens): + num_reqs = max( + 1, + min( + self.max_num_reqs, + self.max_num_tokens // target_query_len, + ), + ) + num_tokens = num_reqs * target_query_len + positions = torch.arange( + num_tokens, + dtype=torch.int64, + device=self.device, + ) + query_start_loc = torch.arange( + 0, + num_tokens + 1, + target_query_len, + dtype=torch.int32, + device=self.device, + ) + idx_mapping = torch.arange( + num_reqs, + dtype=torch.int32, + device=self.device, + ) + input_batch = _DFlashInputBatch( + num_reqs, + np.full( + num_reqs, + target_query_len, + dtype=np.int32, + ), + positions, + query_start_loc, + idx_mapping, + ) + num_sampled = torch.zeros(num_reqs, dtype=torch.int32, device=self.device) + num_rejected = torch.zeros_like(num_sampled) + last_sampled = torch.zeros( + self.max_num_reqs, + dtype=torch.int32, + device=self.device, + ) + next_prefill_tokens = torch.zeros_like(last_sampled) + + for draft_query_len in draft_query_lens: + num_speculative_steps = self._speculative_steps_for_query_len( + draft_query_len + ) + for i, gid in enumerate(self.draft_kv_cache_group_ids): + prepare_dflash_inputs( + self.input_buffers, + self.block_tables.slot_mappings[gid], + self.context_positions, + self._context_slot_mappings[i], + self.sample_indices, + self.sample_pos, + self.sample_idx_mapping, + input_batch, + num_sampled, + num_rejected, + last_sampled, + 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, + draft_query_len, + num_speculative_steps, + self.max_num_reqs, + self.max_num_tokens, + self.max_model_len, + self.sample_from_anchor, + ) + def load_draft_model( self, target_model: nn.Module, @@ -302,7 +416,15 @@ def _generate_draft( slot_mappings: dict[str, torch.Tensor] | None, num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + is_profile: bool = False, + num_query_per_req: int | None = None, ) -> None: + if num_query_per_req is None: + num_speculative_steps = self.num_speculative_steps + else: + num_speculative_steps = self._speculative_steps_for_query_len( + num_query_per_req + ) last_hidden_states = self._run_model( num_tokens_padded, attn_metadata, @@ -313,7 +435,7 @@ def _generate_draft( if torch.cuda.is_current_stream_capturing(): self._captured_backbone_outputs.append(last_hidden_states) - num_sample = num_reqs * self.num_speculative_steps + num_sample = num_reqs * num_speculative_steps sample_hidden_states = last_hidden_states[self.sample_indices[:num_sample]] # sample_pos is the predicted token's position Q; verification keys # Gumbel by the predecessor (Q-1). sample_draft adds +1, so pass Q-2. @@ -326,8 +448,8 @@ def _generate_draft( self.sample_col[:num_sample], self.draft_logits, ) - self.draft_tokens[:num_reqs] = draft_tokens.view( - num_reqs, self.num_speculative_steps + self.draft_tokens[:num_reqs, :num_speculative_steps] = draft_tokens.view( + num_reqs, num_speculative_steps ) def _build_draft_attn_metadata( @@ -340,12 +462,13 @@ def _build_draft_attn_metadata( ) -> dict[str, Any] | None: if not self.draft_attn_layer_names: return None - assert num_query_per_req is None # Omitted for DFlash, read from self instead + if num_query_per_req is None: + num_query_per_req = self.num_query_per_req return super()._build_draft_attn_metadata( num_reqs, num_reqs_padded, num_tokens_padded, - num_query_per_req=self.num_query_per_req, + num_query_per_req=num_query_per_req, causal=causal, ) @@ -388,11 +511,20 @@ def propose( ) 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( - max_seq_len + self.num_query_per_req, self.max_model_len + active_num_speculative_steps = self.num_speculative_steps + if self.dynamic_physical_depth and num_speculative_tokens is not None: + active_num_speculative_steps = num_speculative_tokens + if not 1 <= active_num_speculative_steps <= self.num_speculative_steps: + raise ValueError( + "DSpark physical draft depth must be between 1 and " + f"{self.num_speculative_steps}, got {active_num_speculative_steps}." + ) + active_query_len = self._query_len_for_speculative_steps( + active_num_speculative_steps ) + num_query_tokens = num_reqs * active_query_len + max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() + self.draft_max_seq_len = min(max_seq_len + active_query_len, self.max_model_len) # NOTE: To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of input_ids and @@ -431,8 +563,10 @@ def propose( slot_mappings=None, num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=CUDAGraphMode.NONE, + is_profile=is_profile, + num_query_per_req=active_query_len, ) - return self.draft_tokens[:num_reqs] + return self.draft_tokens[:num_reqs, :active_num_speculative_steps] # The query slot mapping is written into the shared BlockTables slot_mappings. # That buffer's address is what the captured CUDA graph reads from at replay. @@ -456,8 +590,8 @@ def propose( 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, + active_query_len, + active_num_speculative_steps, self.max_num_reqs, self.max_num_tokens, self.max_model_len, @@ -509,7 +643,8 @@ def propose( self.query_cudagraph_manager, num_reqs, num_query_tokens, - uniform_token_count=self.num_query_per_req, + uniform_token_count=active_query_len, + max_req_tokens=None, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, @@ -524,6 +659,7 @@ def propose( num_reqs=num_reqs, num_reqs_padded=num_reqs_padded, num_tokens_padded=num_tokens_padded, + num_query_per_req=active_query_len, causal=self._group_causal, ) draft_slot_mappings_by_layer = build_slot_mappings_by_layer( @@ -546,9 +682,11 @@ def propose( draft_slot_mappings_by_layer, num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=batch_desc.cg_mode, + is_profile=is_profile, + num_query_per_req=active_query_len, ) - return self.draft_tokens[:num_reqs] + return self.draft_tokens[:num_reqs, :active_num_speculative_steps] @triton.jit @@ -600,6 +738,7 @@ def _prepare_dflash_inputs_kernel( num_rejected = tl.load(num_rejected_ptr + req_idx) valid_ctx_end = ctx_end - num_rejected + num_valid_ctx = valid_ctx_end - ctx_start num_sampled = tl.load(num_sampled_ptr + req_idx) if num_sampled > 0: @@ -613,17 +752,18 @@ def _prepare_dflash_inputs_kernel( j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) is_ctx = j < num_ctx - is_query = (j >= num_ctx) & (j < num_ctx + num_query_per_req) - query_off = j - num_ctx + is_valid_ctx = j < num_valid_ctx + is_query = (j >= num_valid_ctx) & (j < num_valid_ctx + num_query_per_req) + query_off = j - num_valid_ctx # --- Context positions / slots --- ctx_pos_idx = ctx_start + tl.where(is_ctx, j, 0) - ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0) + ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_valid_ctx, other=0) ctx_block_num = ctx_pos // block_size ctx_block_num = tl.minimum(ctx_block_num, block_table_stride - 1) ctx_block_id = tl.load( block_table_ptr + req_idx * block_table_stride + ctx_block_num, - mask=is_ctx, + mask=is_valid_ctx, other=0, ).to(tl.int64) # Sliding-window draft KV: old context positions can be evicted and point @@ -636,7 +776,11 @@ def _prepare_dflash_inputs_kernel( PAD_SLOT_ID, ) tl.store(out_context_positions_ptr + ctx_start + j, ctx_pos, mask=is_ctx) - tl.store(out_context_slot_mapping_ptr + ctx_start + j, ctx_slot, mask=is_ctx) + tl.store( + out_context_slot_mapping_ptr + ctx_start + j, + tl.where(is_valid_ctx, ctx_slot, PAD_SLOT_ID), + mask=is_ctx, + ) # --- Query positions / input_ids / slots --- query_pos = last_valid_pos + 1 + query_off @@ -740,7 +884,7 @@ def prepare_dflash_inputs( sample_indices: torch.Tensor, sample_pos: torch.Tensor, sample_idx_mapping: torch.Tensor, - input_batch: InputBatch, + input_batch: InputBatch | _DFlashInputBatch, # [num_reqs] num_sampled: torch.Tensor, # [num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/capacity.py b/vllm/v1/worker/gpu/spec_decode/dspark/capacity.py new file mode 100644 index 000000000000..0a3b600ae77f --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/capacity.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np +import torch + +from vllm.triton_utils import tl, triton + + +@triton.jit +def _compute_prefix_survival_probabilities_kernel( + confidence_logits_ptr, + survival_probs_ptr, + inv_temperature, + CONFIDENCE_STRIDE: tl.constexpr, + SURVIVAL_STRIDE: tl.constexpr, + NUM_SPECULATIVE_STEPS: tl.constexpr, +): + req_idx = tl.program_id(0) + survival_prob = tl.full((), 1.0, tl.float32) + for step in tl.static_range(0, NUM_SPECULATIVE_STEPS): + confidence_logit = tl.load( + confidence_logits_ptr + req_idx * CONFIDENCE_STRIDE + step + ).to(tl.float32) + confidence_prob = 1.0 / (1.0 + tl.exp(-confidence_logit * inv_temperature)) + survival_prob *= confidence_prob + tl.store( + survival_probs_ptr + req_idx * SURVIVAL_STRIDE + step, + survival_prob, + ) + + +@triton.jit +def _allocate_draft_token_capacity_kernel( + survival_probs_ptr, + capacity_ptr, + runtime_num_reqs_ptr, + sps_table_ptr, + min_survival_probability, + REQ_BLOCK: tl.constexpr, + NUM_SPECULATIVE_STEPS: tl.constexpr, + MAX_ADMISSIONS: tl.constexpr, + USE_BUDGET: tl.constexpr, + BUDGET_FRAC: tl.constexpr, + USE_SPS: tl.constexpr, + SURVIVAL_STRIDE: tl.constexpr, +): + offsets = tl.arange(0, REQ_BLOCK) + runtime_num_reqs = tl.load(runtime_num_reqs_ptr).to(tl.int32) + active = offsets < runtime_num_reqs + + if USE_BUDGET: + total_admissions = runtime_num_reqs * NUM_SPECULATIVE_STEPS + max_admissions = tl.minimum( + tl.ceil(total_admissions.to(tl.float32) * BUDGET_FRAC).to(tl.int32), + total_admissions, + ) + # DSpark Algorithm 1: greedy global admission over the candidate set + # {(r, j) | survival(r, j) > 0}, sorted by prefix-survival score. The + # admission counts ARE the capacities, so the spent budget is exactly + # sum(capacities). (Re-deriving capacities from the kth-score + # threshold would blow past the budget whenever scores tie, e.g. + # saturated sigmoids, and zero-survival tokens are never candidates.) + # With an SPS curve, stop at the admission count k* maximizing + # expected throughput theta = tau * SPS(B), where after k admissions + # tau = R + sum of admitted survival scores and B = R + k + # verification tokens (one bonus token per request). + lengths = tl.full((REQ_BLOCK,), 0, tl.int32) + if USE_SPS: + tau = runtime_num_reqs * 1.0 + best_theta = tau * tl.load(sps_table_ptr + runtime_num_reqs) + best_lengths = lengths + for admission_idx in tl.range(0, MAX_ADMISSIONS): + has_next = ( + active + & (admission_idx < max_admissions) + & (lengths < NUM_SPECULATIVE_STEPS) + ) + next_scores = tl.load( + survival_probs_ptr + offsets * SURVIVAL_STRIDE + lengths, + mask=has_next, + other=-1.0, + ) + best_score, best_idx = tl.max(next_scores, axis=0, return_indices=True) + admit = best_score > 0.0 + lengths += tl.where(admit & (offsets == best_idx), 1, 0) + if USE_SPS: + tau += tl.where(admit, best_score, 0.0) + sps = tl.load(sps_table_ptr + runtime_num_reqs + admission_idx + 1) + theta = tau * sps + better = admit & (theta > best_theta) + best_theta = tl.where(better, theta, best_theta) + best_lengths = tl.where(better, lengths, best_lengths) + capacities = best_lengths if USE_SPS else lengths + else: + capacities = tl.full((REQ_BLOCK,), 0, tl.int32) + for step in tl.static_range(0, NUM_SPECULATIVE_STEPS): + scores = tl.load( + survival_probs_ptr + offsets * SURVIVAL_STRIDE + step, + mask=active, + other=-1.0, + ) + capacities += tl.where(scores >= min_survival_probability, 1, 0) + + tl.store(capacity_ptr + offsets, capacities, mask=active) + + +def build_sps_table( + sps_curve: list[tuple[int, float]], + max_batch_tokens: int, + device: torch.device, +) -> torch.Tensor: + """Densify (batch_num_tokens, steps_per_sec) breakpoints into a lookup + table indexed by verification batch token count, linearly interpolated + and clamped at the ends.""" + xs = np.array([b for b, _ in sps_curve], dtype=np.float64) + ys = np.array([s for _, s in sps_curve], dtype=np.float64) + table = np.interp(np.arange(max_batch_tokens + 1), xs, ys) + return torch.tensor(table, dtype=torch.float32, device=device) + + +def compute_draft_token_capacity_from_confidence( + confidence_logits: torch.Tensor, + draft_token_capacity: torch.Tensor, + min_survival_probability: float, + num_reqs: int, + num_speculative_steps: int, + runtime_num_reqs: torch.Tensor, + survival_probs: torch.Tensor | None = None, + budget_frac: float = 1.0, + sps_table: torch.Tensor | None = None, + confidence_temperature: float = 1.0, +) -> None: + if num_reqs == 0 or num_speculative_steps == 0: + return + if survival_probs is None: + survival_probs = torch.empty_like(confidence_logits) + _compute_prefix_survival_probabilities_kernel[(num_reqs,)]( + confidence_logits, + survival_probs, + 1.0 / confidence_temperature, + CONFIDENCE_STRIDE=confidence_logits.stride(0), + SURVIVAL_STRIDE=survival_probs.stride(0), + NUM_SPECULATIVE_STEPS=num_speculative_steps, + ) + # Even when the budget covers every token, zero-survival tokens are + # not admission candidates (DSpark Alg. 1), so always run the kernel. + use_budget = min_survival_probability <= 0.0 + use_sps = use_budget and sps_table is not None + # Pow2-padded so one compiled variant serves all runtime_num_reqs values + # under CUDA graph capture. + kernel_num_reqs = triton.next_power_of_2(max(num_reqs, 1)) + if use_sps: + assert sps_table is not None + # The theta scan reads SPS(B) for B up to kernel_num_reqs * (1 + n_spec). + assert sps_table.shape[0] > kernel_num_reqs * (1 + num_speculative_steps), ( + f"SPS table has {sps_table.shape[0]} entries but batch token " + f"counts can reach {kernel_num_reqs * (1 + num_speculative_steps)}" + ) + if sps_table is None: + sps_table = draft_token_capacity + _allocate_draft_token_capacity_kernel[(1,)]( + survival_probs, + draft_token_capacity, + runtime_num_reqs, + sps_table, + min_survival_probability, + REQ_BLOCK=kernel_num_reqs, + NUM_SPECULATIVE_STEPS=num_speculative_steps, + MAX_ADMISSIONS=kernel_num_reqs * num_speculative_steps, + USE_BUDGET=use_budget, + BUDGET_FRAC=budget_frac, + USE_SPS=use_sps, + SURVIVAL_STRIDE=survival_probs.stride(0), + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/online_sts.py b/vllm/v1/worker/gpu/spec_decode/dspark/online_sts.py new file mode 100644 index 000000000000..b14bc9aff80d --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/online_sts.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + +import torch + +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +class DSparkOnlineSTS: + """Online Sequential Temperature Scaling for the DSpark capacity scheduler. + + The paper (Section 3.2.1) calibrates each position's conditional survival + probability with a per-position temperature chosen to minimize the + Expected Calibration Error of the cumulative product on a validation set + — an order-preserving transform. This is the serving-time analogue: per + position it maintains binned empirical conditional acceptance + P(accept_k | prefix accepted, position verified) from the rejection + sampler's own outcomes (exponential decay), and each step fits the + temperature that minimizes the trial-weighted ECE of + sigmoid(logit / T_k) against those bins. Fitting each conditional + directly is the chain-rule equivalent of the paper's sequential + cumulative-product fit, with censored online observations instead of a + held-out set. + + Observations are censored by capacity (unverified positions yield no + trials); under light load the theta-argmax verifies every candidate, + which is where the observation mass comes from. With few observations + the fitted temperature is blended toward 1.0 (identity), so cold-start + behaves like the raw confidence head. + + All buffers are persistent: ``temperatures`` is read inside the captured + draft graph, while ``record()`` runs eagerly each step. All reductions + are one-hot sums (no atomics) so the state stays bitwise identical + across TP ranks. + """ + + DECAY = 0.999 + PRIOR_WEIGHT = 64.0 + NUM_BINS = 16 + LOGIT_RANGE = 8.0 + # Log-spaced temperature grid; 1.0 is on the grid so a well-calibrated + # head fits the identity exactly. + TEMP_GRID_MIN = 0.125 + TEMP_GRID_MAX = 8.0 + TEMP_GRID_SIZE = 49 + + def __init__(self, max_num_reqs: int, num_steps: int, device: torch.device): + self.num_steps = num_steps + # Raw head logits of each request's latest proposal, by req-state + # slot, so verification outcomes (one step later, possibly reordered) + # can be joined back to the confidences that produced them. + self.logits_by_state = torch.zeros( + max_num_reqs, num_steps, dtype=torch.float32, device=device + ) + self.proposal_valid_by_state = torch.zeros( + max_num_reqs, dtype=torch.bool, device=device + ) + # EMA counters per (position, logit bin). + self.bin_trials = torch.zeros( + num_steps, self.NUM_BINS, dtype=torch.float32, device=device + ) + self.bin_hits = torch.zeros_like(self.bin_trials) + # Per-position temperatures; persistent so captured graphs see + # updates. Identity until observations accumulate. + self.temperatures = torch.ones(num_steps, dtype=torch.float32, device=device) + + self._steps = torch.arange(num_steps, device=device) + self._bins = torch.arange(self.NUM_BINS, device=device) + bin_width = 2 * self.LOGIT_RANGE / self.NUM_BINS + self._bin_mids = ( + -self.LOGIT_RANGE + (self._bins.to(torch.float32) + 0.5) * bin_width + ) + self._temp_grid = torch.logspace( + torch.log10(torch.tensor(self.TEMP_GRID_MIN)), + torch.log10(torch.tensor(self.TEMP_GRID_MAX)), + self.TEMP_GRID_SIZE, + device=device, + ) + # sigmoid(mid_b / T) for every (T, bin) pair, fixed for the run. + self._grid_probs = torch.sigmoid( + self._bin_mids.unsqueeze(0) / self._temp_grid.unsqueeze(1) + ) + self._log_temp_grid = self._temp_grid.log() + self._log_interval = int( + os.environ.get("VLLM_DSPARK_STS_LOG_INTERVAL", "0") or "0" + ) + self._record_count = 0 + + def stage_proposal( + self, + req_state_indices: torch.Tensor, + logits: torch.Tensor, + *, + valid: bool = True, + ) -> None: + """Remember the raw head logits of the current proposal.""" + if valid: + rows = self.logits_by_state[req_state_indices] + rows.zero_() + rows[:, : logits.shape[1]] = logits + self.logits_by_state[req_state_indices] = rows + self.proposal_valid_by_state[req_state_indices] = valid + + def invalidate_all(self) -> None: + """Discard proposals whose confidence logits were bypassed or profiled.""" + self.proposal_valid_by_state.zero_() + + def calibrate( + self, logits: torch.Tensor, out: torch.Tensor | None = None + ) -> torch.Tensor: + """Apply the per-position temperatures (order-preserving).""" + return torch.div(logits, self.temperatures[: logits.shape[-1]], out=out) + + def record( + self, + req_state_indices: torch.Tensor, + num_accepted: torch.Tensor, + num_verified: torch.Tensor, + ) -> None: + """Fold one verification step's outcomes into the calibration. + + Args: + req_state_indices: [num_reqs] req-state slot of each request. + num_accepted: [num_reqs] accepted draft tokens this step. + num_verified: [num_reqs] draft tokens that were verified + (post-capacity), zero for rows without drafts. + """ + logits = self.logits_by_state[req_state_indices] + proposal_valid = self.proposal_valid_by_state[req_state_indices] + bin_width = 2 * self.LOGIT_RANGE / self.NUM_BINS + bin_idx = ((logits + self.LOGIT_RANGE) / bin_width).long() + bin_idx.clamp_(0, self.NUM_BINS - 1) + + k = self._steps.unsqueeze(0) + # Position k (0-based) is evaluated iff the k-token prefix before it + # was accepted and it was inside the verified capacity. + trial = k < torch.minimum(num_accepted + 1, num_verified).unsqueeze(1) + trial.logical_and_(proposal_valid.unsqueeze(1)) + hit = (k < num_accepted.unsqueeze(1)) & trial + + # One-hot reduction (deterministic; index_add_ atomics are not). + onehot = bin_idx.unsqueeze(-1) == self._bins # [reqs, steps, bins] + # Do not age the calibration state when every row was invalid. Keeping + # this as a device scalar avoids a host synchronization on the hot path. + decay = 1.0 - trial.any().to(torch.float32) * (1.0 - self.DECAY) + self.bin_trials.mul_(decay).add_( + (trial.unsqueeze(-1) & onehot).sum(0).to(torch.float32) + ) + self.bin_hits.mul_(decay).add_( + (hit.unsqueeze(-1) & onehot).sum(0).to(torch.float32) + ) + self.proposal_valid_by_state[req_state_indices] = False + + # Per-position 1D grid search: T_k minimizing trial-weighted ECE of + # sigmoid(mid_b / T) against the empirical bin acceptance. + emp = self.bin_hits / self.bin_trials.clamp(min=1e-6) + err = (self._grid_probs.unsqueeze(1) - emp).abs() # [T, steps, bins] + ece = (err * self.bin_trials).sum(-1) # [T, steps] + log_t = self._log_temp_grid[ece.argmin(0)] + # Blend toward the identity until enough outcomes accumulate. + total = self.bin_trials.sum(-1) + torch.exp(log_t * (total / (total + self.PRIOR_WEIGHT)), out=self.temperatures) + self._maybe_log_diagnostics(total) + + def _maybe_log_diagnostics(self, total_trials: torch.Tensor) -> None: + if self._log_interval <= 0: + return + self._record_count += 1 + if self._record_count % self._log_interval: + return + if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: + return + + denom = total_trials.clamp(min=1e-6) + empirical_mean = self.bin_hits.sum(-1) / denom + calibrated_bin_probs = torch.sigmoid( + self._bin_mids.unsqueeze(0) / self.temperatures.unsqueeze(1) + ) + predicted_mean = (calibrated_bin_probs * self.bin_trials).sum(-1) / denom + raw_logit_mean = (self._bin_mids.unsqueeze(0) * self.bin_trials).sum(-1) / denom + logger.info( + "DSpark online STS: trials=%s empirical_cond=%s " + "predicted_cond=%s temperatures=%s raw_logit_mean=%s", + [round(x, 1) for x in total_trials.tolist()], + [round(x, 4) for x in empirical_mean.tolist()], + [round(x, 4) for x in predicted_mean.tolist()], + [round(x, 4) for x in self.temperatures.tolist()], + [round(x, 4) for x in raw_logit_mean.tolist()], + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 0236017cf2fc..0f0c58a2e7eb 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -27,10 +27,18 @@ import torch +import vllm.envs as envs from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.triton_utils import triton +from vllm.v1.worker.gpu.input_batch import InputBatch 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.dspark.capacity import ( + build_sps_table, + compute_draft_token_capacity_from_confidence, +) +from vllm.v1.worker.gpu.spec_decode.dspark.online_sts import DSparkOnlineSTS from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model @@ -65,26 +73,113 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): self.num_speculative_steps, dtype=torch.int32, device=device ) - self._anchor_idx = ( - torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) - * self.num_query_per_req - ) - # Reduced-vocab probabilistic drafting only; set in load_draft_model. self._d2t_scatter_index: torch.Tensor | None = None self._draft_scatter_buf: torch.Tensor | None = None + self.draft_token_confidence_logits = torch.empty( + self.max_num_reqs, + self.num_speculative_steps, + dtype=torch.float32, + device=device, + ) + self.draft_token_survival_probs = torch.empty_like( + self.draft_token_confidence_logits + ) + self.draft_token_capacity = torch.full( + (self.max_num_reqs,), + self.num_speculative_steps, + dtype=torch.int32, + device=device, + ) + self.capacity_activation_batch_size = ( + envs.VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE + ) + if self.capacity_activation_batch_size < 0: + raise ValueError( + "VLLM_DSPARK_CAPACITY_ACTIVATION_BATCH_SIZE must be >= 0, got " + f"{self.capacity_activation_batch_size}." + ) + self._runtime_num_reqs_for_capacity = torch.zeros( + (1,), + dtype=torch.int32, + device=device, + ) + self.draft_token_valid_lengths = torch.empty( + (self.max_num_reqs,), + dtype=torch.int32, + device=device, + ) + self._last_num_speculative_steps = self.num_speculative_steps + self._last_proposal_confidence_valid = False + self.min_survival_probability = ( + self.speculative_config.dspark_confidence_threshold + ) + self.capacity_budget_frac = self.speculative_config.dspark_budget_frac + self.confidence_temperature = ( + self.speculative_config.dspark_confidence_temperature + ) + sps_curve = self.speculative_config.dspark_sps_curve + self.sps_table: torch.Tensor | None = None + self.wants_auto_sps_curve = sps_curve == "auto" + if sps_curve is not None: + # Sized for the pow2-padded request count the allocator kernel + # can index under CUDA graph capture. + padded_reqs = triton.next_power_of_2(max(self.max_num_reqs, 1)) + max_batch_tokens = padded_reqs * (1 + self.num_speculative_steps) + if self.wants_auto_sps_curve: + # Flat placeholder (theta argmax verifies everything) until + # the post-capture profiling refreshes the contents in place; + # the captured allocator kernel bakes this buffer's address. + self.sps_table = torch.ones( + max_batch_tokens + 1, dtype=torch.float32, device=device + ) + else: + assert isinstance(sps_curve, list) + self.sps_table = build_sps_table( + sps_curve, + max_batch_tokens, + device, + ) + self.use_draft_token_capacity = ( + self.min_survival_probability > 0.0 + or self.capacity_budget_frac < 1.0 + or self.sps_table is not None + ) + self.online_sts: DSparkOnlineSTS | None = None + if self.use_draft_token_capacity and self.speculative_config.dspark_online_sts: + self.online_sts = DSparkOnlineSTS( + self.max_num_reqs, self.num_speculative_steps, device + ) + # Calibrated survival buffer consumed by the capacity kernels + # inside the captured draft graph. + self.calibrated_confidence_logits = torch.zeros_like( + self.draft_token_confidence_logits + ) + def load_draft_model( self, target_model: torch.nn.Module, target_attn_layer_names: set[str], ) -> torch.nn.Module: model = load_dspark_model(target_model, self.vllm_config) + confidence_head = getattr( + getattr(model, "model", None), "confidence_head", None + ) + if self.use_draft_token_capacity and ( + getattr(model, "compute_confidence", None) is None + or confidence_head is None + ): + raise ValueError( + "DSpark draft-token capacity requires a draft model with a " + f"confidence head; {type(model).__name__} does not implement " + "compute_confidence." + ) # Reduced draft vocab: probabilistic rejection sampling indexes draft # logits by target id, so precompute the draft->target column map and a # scratch buffer to scatter logits into target vocab before sampling. - if self.draft_logits is not None and model.draft_id_to_target_id is not None: - d2t = model.draft_id_to_target_id + d2t = getattr(model, "draft_id_to_target_id", None) + if self.draft_logits is not None and d2t is not None: self._d2t_scatter_index = ( torch.arange(d2t.shape[0], device=d2t.device) + d2t ) @@ -98,27 +193,56 @@ def load_draft_model( ) return model - def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: + def _sample_sequential( + self, + num_reqs: int, + head_hidden: torch.Tensor, + num_speculative_steps: int, + num_query_per_req: int, + is_profile: bool = False, + use_capacity: bool = True, + ) -> None: # Sequential Markov sampling over the backbone's output hidden states. - n_spec = self.num_speculative_steps + n_spec = num_speculative_steps num_sample = num_reqs * n_spec # Per-(req, position) head hidden, ordered (req, step). sample_hidden = head_hidden[self.sample_indices[:num_sample]] + sample_hidden = sample_hidden.view(num_reqs, n_spec, -1) # Draft-vocab logits; sampled ids are remapped to target vocab below. - base_logits = self.model.compute_draft_logits(sample_hidden) + base_logits = self.model.compute_draft_logits( + sample_hidden.reshape(num_sample, -1) + ) vocab_size = base_logits.shape[-1] base_logits = base_logits.view(num_reqs, n_spec, vocab_size) idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) + confidence_logits = self.draft_token_confidence_logits[:num_reqs, :n_spec] + min_survival_probability = self.min_survival_probability + use_confidence_capacity = self.use_draft_token_capacity and use_capacity # Anchor (bonus) token per request = the input id at query offset 0, - # read via the precomputed persistent index (fixed buffer for capture). - prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] + # laid out as one row per request in the draft query block. + prev = self.input_buffers.input_ids[ + : num_reqs * num_query_per_req : num_query_per_req + ] + valid_prefix = torch.ones(num_reqs, dtype=torch.bool, device=self.device) + valid_lengths = self.draft_token_valid_lengths[:num_reqs] + valid_lengths.zero_() for i in range(n_spec): # Sequential stage: Markov bias from the previously sampled token. markov_embed = self.model.markov_embed(prev) + if use_confidence_capacity: + confidence_i = self.model.compute_confidence( + sample_hidden[:, i], markov_embed + ) + if confidence_i is None: + raise RuntimeError( + "DSpark draft-token capacity requires loaded " + "confidence-head weights." + ) + confidence_logits[:, i] = confidence_i bias = self.model.markov_bias(markov_embed) logits_i = base_logits[:, i] + bias if self.draft_logits is not None: @@ -146,9 +270,129 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: draft_sampled_i = self.model.map_draft_to_target( logits_i.argmax(dim=-1) ) + valid_prefix.logical_and_( + (draft_sampled_i >= 0) & (draft_sampled_i < self.vocab_size) + ) + draft_sampled_i = torch.where( + valid_prefix, draft_sampled_i, torch.zeros_like(draft_sampled_i) + ) + valid_lengths.add_(valid_prefix.to(torch.int32)) self.draft_tokens[:num_reqs, i] = draft_sampled_i prev = draft_sampled_i + if use_confidence_capacity and not is_profile: + capacity_confidence = self.draft_token_confidence_logits + capacity_temperature = self.confidence_temperature + if self.online_sts is not None: + self.online_sts.calibrate( + confidence_logits, + out=self.calibrated_confidence_logits[:num_reqs, :n_spec], + ) + capacity_confidence = self.calibrated_confidence_logits + capacity_temperature = 1.0 + compute_draft_token_capacity_from_confidence( + capacity_confidence, + self.draft_token_capacity, + min_survival_probability, + num_reqs, + n_spec, + self._runtime_num_reqs_for_capacity, + self.draft_token_survival_probs, + self.capacity_budget_frac, + sps_table=self.sps_table, + confidence_temperature=capacity_temperature, + ) + else: + self.draft_token_capacity[:num_reqs].fill_(n_spec) + torch.minimum( + self.draft_token_capacity[:num_reqs], + valid_lengths, + out=self.draft_token_capacity[:num_reqs], + ) + + def set_sps_curve(self, sps_curve: list[tuple[int, float]]) -> None: + """Refresh the SPS lookup table in place (its address is baked into + the captured allocator kernel).""" + assert self.sps_table is not None + dense = build_sps_table( + sps_curve, self.sps_table.shape[0] - 1, self.sps_table.device + ) + self.sps_table.copy_(dense) + + def compute_capacities(self, input_batch: InputBatch) -> torch.Tensor | None: + if not self.use_draft_token_capacity: + return None + num_reqs = input_batch.num_reqs + if self.online_sts is not None: + # Join key for verification outcomes arriving next step. Staged + # eagerly (not in the captured graph): a padded replay would + # index_put through stale padding-row ids, and -1 sentinels wrap + # to the last row, so neither is safe for a scatter by slot. + n_spec = self._last_num_speculative_steps + self.online_sts.stage_proposal( + self.sample_idx_mapping[: num_reqs * n_spec : n_spec], + self.draft_token_confidence_logits[:num_reqs, :n_spec], + valid=self._last_proposal_confidence_valid, + ) + return self.draft_token_capacity[:num_reqs] + + def warmup_capacity_kernels(self) -> None: + self._warmup_prepare_inputs_kernel() + if not self.use_draft_token_capacity: + return + + self.draft_token_confidence_logits.zero_() + sizes = {self.max_num_reqs} + num_reqs = 1 + while num_reqs < self.max_num_reqs: + sizes.add(num_reqs) + num_reqs *= 2 + for num_reqs in sorted(sizes): + self._runtime_num_reqs_for_capacity.fill_(num_reqs) + compute_draft_token_capacity_from_confidence( + self.draft_token_confidence_logits, + self.draft_token_capacity, + self.min_survival_probability, + num_reqs, + self.num_speculative_steps, + self._runtime_num_reqs_for_capacity, + self.draft_token_survival_probs, + self.capacity_budget_frac, + sps_table=self.sps_table, + confidence_temperature=self.confidence_temperature, + ) + + def propose( + self, + input_batch: InputBatch, + *args, + num_speculative_tokens: int | None = None, + **kwargs, + ) -> torch.Tensor: + if self.use_draft_token_capacity: + self._runtime_num_reqs_for_capacity.fill_(input_batch.num_reqs) + self._last_proposal_confidence_valid = bool( + self.use_draft_token_capacity + and not kwargs.get("is_profile", False) + and not kwargs.get("dummy_run", False) + and not self._has_unaligned_cached_prefix(input_batch) + and ( + self.capacity_activation_batch_size <= 0 + or input_batch.num_reqs >= self.capacity_activation_batch_size + ) + ) + self._last_num_speculative_steps = ( + num_speculative_tokens + if self.dynamic_physical_depth and num_speculative_tokens is not None + else self.num_speculative_steps + ) + return super().propose( + input_batch, + *args, + num_speculative_tokens=num_speculative_tokens, + **kwargs, + ) + def _generate_draft( self, num_reqs: int, @@ -157,7 +401,12 @@ def _generate_draft( slot_mappings: dict[str, torch.Tensor] | None, num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + is_profile: bool = False, + num_query_per_req: int | None = None, ) -> None: + if num_query_per_req is None: + num_query_per_req = self.num_query_per_req + num_speculative_steps = self._speculative_steps_for_query_len(num_query_per_req) # Full draft step (captured under CUDA graph): parallel backbone forward # then sequential Markov sampling over its hidden state outputs. head_hidden = self._run_model( @@ -167,4 +416,14 @@ def _generate_draft( num_tokens_across_dp, cudagraph_runtime_mode, ) - self._sample_sequential(num_reqs, head_hidden) + self._sample_sequential( + num_reqs, + head_hidden, + num_speculative_steps, + num_query_per_req, + is_profile=is_profile, + use_capacity=( + self.capacity_activation_batch_size <= 0 + or num_reqs >= self.capacity_activation_batch_size + ), + ) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 8b1147f5c5e2..4ad919589100 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -124,6 +124,9 @@ def __call__( num_nans = get_num_nans(logits) if self.sampler.compute_nans else None draft_sampled = input_batch.input_ids[input_batch.logits_indices] + draft_sampled.masked_fill_( + input_batch.is_padding[input_batch.logits_indices], -1 + ) pos = input_batch.positions[input_batch.logits_indices] processed_logits = self.sampler.apply_sampling_params( logits, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index a845a556ca5e..2fc00ea2c703 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -240,7 +240,7 @@ def _compute_local_logits_stats_kernel( other=float("-inf"), ).to(tl.float32) value, idx = tl.max(target_logits, axis=0, return_indices=True) - token_id = block_idx * BLOCK_SIZE + idx + token_id = tl.minimum(block_idx * BLOCK_SIZE + idx, vocab_size - 1) tl.store( target_local_argmax_ptr + logit_idx * target_local_argmax_stride @@ -791,7 +791,7 @@ def _resample_kernel( HAS_ACTIVE_ROW_LIMIT=False, USE_FP64=USE_FP64, ) - token_id = block_idx * BLOCK_SIZE + idx + token_id = tl.minimum(block_idx * BLOCK_SIZE + idx, vocab_size - 1) tl.store( resampled_local_argmax_ptr + req_idx * resampled_local_argmax_stride diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 6bcdea5ee484..50d85f43a350 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod from collections.abc import Mapping -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn @@ -28,10 +28,25 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +if TYPE_CHECKING: + from vllm.v1.worker.gpu.spec_decode.dspark.online_sts import DSparkOnlineSTS + logger = init_logger(__name__) class BaseSpeculator(ABC): + # Draft-token capacity surface, implemented by speculators with a + # confidence head (see DSparkSpeculator). + use_draft_token_capacity: bool = False + online_sts: "DSparkOnlineSTS | None" = None + wants_auto_sps_curve: bool = False + + def warmup_capacity_kernels(self) -> None: # noqa: B027 + pass + + def set_sps_curve(self, sps_curve: list[tuple[int, float]]) -> None: + raise NotImplementedError + @abstractmethod def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: pass @@ -74,6 +89,9 @@ def propose( ) -> torch.Tensor: pass + def compute_capacities(self, input_batch: InputBatch) -> torch.Tensor | None: + return None + class DraftModelSpeculator(BaseSpeculator): def __init__(self, vllm_config: VllmConfig, device: torch.device): diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index a02c34a6972e..6908e833cbc0 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -200,7 +200,8 @@ def warmup_kernels( # DCP rank receives zero local KV tokens. min_dcp_prompt_len = max( 1, - getattr(model_runner, "dcp_size", 1) * getattr(model_runner, "cp_interleave", 1), + getattr(model_runner, "dcp_size", 1) + * getattr(model_runner, "cp_interleave", 1), ) prompt_len = max(decode_query_len + 1, min_dcp_prompt_len) prompt_token_ids = list(range(prompt_len)) @@ -342,9 +343,141 @@ def _alloc_blocks(num_blocks: int) -> list[int]: worker_execute_model(decode_output) worker_sample_tokens(None) + if model_runner.verification_capacity_manager is not None: + model_runner.verification_capacity_manager.warmup( + model_runner.input_buffers + ) + assert model_runner.speculator is not None + model_runner.speculator.warmup_capacity_kernels() + if model_runner.speculator.wants_auto_sps_curve: + _profile_sps_curve(model_runner) + model_runner.kv_connector.set_disabled(True) + # Clean up - process finish_req_ids. cleanup_output = SchedulerOutput.make_empty() cleanup_output.finished_req_ids = set(req_ids) worker_execute_model(cleanup_output) model_runner.kv_connector.set_disabled(False) torch.accelerator.synchronize() + + +def _stable_sps_step_ms(samples: list[float]) -> float: + if not samples or any(ms <= 0.0 for ms in samples): + raise ValueError(f"SPS timing samples must be positive, got {samples}.") + return float(np.median(samples)) + + +def _derive_dspark_draft_token_budget( + sps_curve: list[tuple[int, float]], + max_draft_depth: int, +) -> int: + """Convert the upper-load SPS knee into a physical draft-token budget.""" + if len(sps_curve) < 2: + raise ValueError("DSpark SPS curve needs at least two points") + if max_draft_depth < 1: + raise ValueError("max_draft_depth must be at least one") + + # Startup profiles powers-of-two request counts at the maximum depth. Use + # the largest relative SPS drop in the upper half of that load range as the + # saturation knee. Remove the target bonus-token share from the knee to get + # the corresponding draft-only token budget. + first_drop = max(1, len(sps_curve) // 2) + knee_end = max( + range(first_drop, len(sps_curve)), + key=lambda i: (sps_curve[i - 1][1] - sps_curve[i][1]) / sps_curve[i - 1][1], + ) + knee_tokens = sps_curve[knee_end - 1][0] + return cdiv(knee_tokens * max_draft_depth, max_draft_depth + 1) + + +def _profile_sps_curve( + model_runner: GPUModelRunner, + warmup_iters: int = 5, + timed_iters: int = 10, + timed_rounds: int = 5, +) -> None: + """Profile the engine step-rate curve for ``dspark_sps_curve="auto"``. + + Times uniform-decode dummy runs per power-of-two request count — the same + self-contained path DP idle steps use (``execute_dummy_batch``), which + replays the captured verify graph AND the full DSpark draft step + (``propose(dummy_run=True)``) with no request bookkeeping. Runs after + graph capture with the placeholder flat table active, so the theta-argmax + verifies every candidate and B is exactly reqs * decode_query_len. + Real-step host prep, sampling, and true attention lengths are not + visible here; account for them via ``dspark_sps_overhead_ms``. Rank 0's + measurements are broadcast so every TP rank builds the identical table + (capacities feed batch-shape decisions, which must agree across ranks). + """ + import time + + from vllm.distributed.parallel_state import get_tp_group + + decode_query_len = model_runner.decode_query_len + max_reqs = min( + model_runner.max_num_reqs, + model_runner.max_num_tokens // decode_query_len, + ) + req_counts = [] + count = 1 + while count < max_reqs: + req_counts.append(count) + count *= 2 + req_counts.append(max_reqs) + + step_ms = [] + step_ms_samples: list[list[float]] = [] + for num_reqs in req_counts: + num_tokens = num_reqs * decode_query_len + for _ in range(warmup_iters): + model_runner._dummy_run(num_tokens, uniform_decode=True) + torch.accelerator.synchronize() + samples = [] + for _ in range(timed_rounds): + start = time.perf_counter() + for _ in range(timed_iters): + model_runner._dummy_run(num_tokens, uniform_decode=True) + torch.accelerator.synchronize() + samples.append((time.perf_counter() - start) * 1000.0 / timed_iters) + # A lazy kernel initialization or a transient host stall must not become + # the scheduler's permanent cost model. Multiple independent windows + # make those events visible, and the median rejects an isolated one. + step_ms_samples.append(samples) + step_ms.append(_stable_sps_step_ms(samples)) + + timings = torch.tensor(step_ms, dtype=torch.float64, device=model_runner.device) + tp_group = get_tp_group() + if tp_group.world_size > 1: + tp_group.broadcast(timings, src=0) + step_ms = timings.cpu().tolist() + + assert model_runner.speculative_config is not None + overhead_ms = model_runner.speculative_config.dspark_sps_overhead_ms + sps_curve = [ + (num_reqs * decode_query_len, 1000.0 / (ms + overhead_ms)) + for num_reqs, ms in zip(req_counts, step_ms) + ] + assert model_runner.speculator is not None + model_runner.speculator.set_sps_curve(sps_curve) + capacity_manager = model_runner.verification_capacity_manager + if capacity_manager is not None: + draft_token_budget = _derive_dspark_draft_token_budget( + sps_curve, + model_runner.num_speculative_steps, + ) + capacity_manager.set_dynamic_draft_token_budget(draft_token_budget) + logger.info( + "DSpark auto-profiled dynamic draft-token budget: %d", + draft_token_budget, + ) + logger.info( + "DSpark SPS profile windows (tokens, ms/step samples): %s", + [ + (num_reqs * decode_query_len, [round(ms, 3) for ms in samples]) + for num_reqs, samples in zip(req_counts, step_ms_samples) + ], + ) + logger.info( + "DSpark auto-profiled SPS curve (tokens, steps/s): %s", + [(b, round(s, 2)) for b, s in sps_curve], + ) diff --git a/vllm/v1/worker/ubatch_utils.py b/vllm/v1/worker/ubatch_utils.py index f4a76529023c..b7758b3f8725 100644 --- a/vllm/v1/worker/ubatch_utils.py +++ b/vllm/v1/worker/ubatch_utils.py @@ -242,6 +242,7 @@ def _make_metadata_with_slice( max_seq_len=max_seq_len, block_table_tensor=block_table_tensor, slot_mapping=slot_mapping, + max_req_tokens=attn_metadata.max_req_tokens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, _seq_lens_cpu=seq_lens_cpu, _num_computed_tokens_cpu=num_computed_tokens_cpu,