diff --git a/docs/features/speculative_decoding/README.md b/docs/features/speculative_decoding/README.md index 58d1df9dce..7213ef41ec 100644 --- a/docs/features/speculative_decoding/README.md +++ b/docs/features/speculative_decoding/README.md @@ -17,6 +17,7 @@ vLLM supports a variety of methods of speculative decoding. Model-based methods - [Suffix Decoding](suffix.md) - [Hidden State Extraction](extract_hidden_states.md) - [Custom Proposer Backend (Experimental)](#custom-proposer-backend-experimental) +- [Dynamic Speculative Decoding](dynamic_speculative_decoding.md) ## Method Selection at a Glance @@ -33,6 +34,7 @@ depend on your model family, traffic pattern, hardware, and sampling settings. | N-gram | Low to medium gain | Medium gain | Lightweight and easy to enable. | | Suffix decoding | Low to medium gain | Medium gain | No extra draft model; dynamic speculation depth. | | Custom Proposer | Varies | Varies | Bring your own proposer class (experimental). | +| Dynamic Speculative Decoding | High gain | Higher than base SD method | Useful for RL or workload with fluctuating QPS | For reproducible measurements in your environment, use [`examples/features/speculative_decoding/spec_decode_offline.py`](../../../examples/features/speculative_decoding/spec_decode_offline.py) diff --git a/docs/features/speculative_decoding/dynamic_speculative_decoding.md b/docs/features/speculative_decoding/dynamic_speculative_decoding.md new file mode 100644 index 0000000000..eecf789d6d --- /dev/null +++ b/docs/features/speculative_decoding/dynamic_speculative_decoding.md @@ -0,0 +1,78 @@ +# Dynamic Speculative Decoding + +## Why is Dynamic SD needed? + +SD methods need to verify K tokens for each sequence during decoding. As BS increases, the effective BS becomes BS\*K which increases the compute requirement during verification. When this BS\*K goes beyond a critical BS then SD negatively impacts the decode speed (TPOT). DSD helps by tuning the K to an optimal value such that we continue to reap the benefits from SD. + +## Use cases + +* Variable concurrency workload using same deployment. K would decrease as concurrency increases. +* During RL rollout where we start off with high BS but then end up with small BS due to very few long tail request which end up generating a lot of tokens stalling the progress of the current rollout. Here K would go up during the end of rollout. + +## `--speculative-config` schema + +To use Dynamic SD, add `num_speculative_tokens_per_batch_size` to the config of an SD method which is a list of list. Here, an entry is `[start_bs, end_bs, optimal_K]` which means when the concurrency is within range `[start_bs, end_bs]` then `optimal_K` number of draft tokens are used. For e.g., + +```bash +--speculative-config '{ + "method": "eagle", + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "num_speculative_tokens_per_batch_size": [ + [1, 64, 3], + [65, 128, 1], + [129, 512, 0] + ] + }' +``` + +implies that: + +* K=3 will be used when the concurrency is in range [1, 64] +* K=1 will be used when the concurrency is in range [65, 128] +* K=0 will be used when the concurrency is in range [129, 512], i.e., no draft tokens will be produced. + +## Online Examples + +### Dynamic SD Eagle Drafter + +```bash +VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --speculative-config '{ + "method": "eagle", + "model": "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "num_speculative_tokens_per_batch_size": [ + [1, 64, 3], + [65, 128, 1], + [129, 512, 0] + ] + }' +``` + +### Dynamic SD Eagle3 Drafter + +```bash +VLLM_USE_V2_MODEL_RUNNER=0 vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --speculative-config '{ + "method": "eagle3", + "model": "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", + "num_speculative_tokens": 3, + "num_speculative_tokens_per_batch_size": [ + [1, 16, 5], + [17, 32, 4], + [33, 64, 3], + [65, 128, 1], + [129, 512, 0] + ] + }' + +``` + +## Limitations + +* only tested with Eagle and Eagle-3. Other SD methods may or may not work out of the box +* only usable with Model Runner V1 +* not compatible with full cuda graph so we force piece-wise cuda graph with this feature + +We are working on enabling it on MRv2 with full cuda graph support. diff --git a/tests/v1/spec_decode/test_dynamic_sd.py b/tests/v1/spec_decode/test_dynamic_sd.py new file mode 100644 index 0000000000..fe9f30ba25 --- /dev/null +++ b/tests/v1/spec_decode/test_dynamic_sd.py @@ -0,0 +1,218 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for the Dynamic SD batch-size schedule helpers.""" + +import pytest + +from tests.v1.core.utils import create_requests, create_scheduler +from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup +from vllm.v1.structured_output import StructuredOutputManager + + +def _make_lookup( + num_speculative_tokens_per_batch_size: list[tuple[int, int, int]], + *, + max_batch_size: int = 256, + runtime_num_speculative_tokens: int = 3, +) -> list[int]: + return build_dynamic_sd_schedule_lookup( + num_speculative_tokens_per_batch_size=num_speculative_tokens_per_batch_size, + vllm_max_batch_size=max_batch_size, + vllm_num_speculative_tokens=runtime_num_speculative_tokens, + ) + + +def _make_scheduler_with_dynamic_sd( + schedule: list[tuple[int, int, int]], + *, + max_num_seqs: int = 16, + max_num_batched_tokens: int = 8192, + runtime_num_speculative_tokens: int = 3, +) -> Scheduler: + base_scheduler = create_scheduler( + max_num_seqs=max_num_seqs, + max_num_batched_tokens=max_num_batched_tokens, + num_speculative_tokens=runtime_num_speculative_tokens, + ) + + speculative_config = base_scheduler.vllm_config.speculative_config + assert speculative_config is not None + speculative_config.num_speculative_tokens_per_batch_size = schedule + + return Scheduler( + vllm_config=base_scheduler.vllm_config, + kv_cache_config=base_scheduler.kv_cache_config, + block_size=base_scheduler.block_size, + log_stats=True, + structured_output_manager=StructuredOutputManager(base_scheduler.vllm_config), + ) + + +def _add_requests_and_schedule( + scheduler: Scheduler, num_requests: int, *, num_tokens: int = 10 +): + requests = create_requests(num_requests=num_requests, num_tokens=num_tokens) + for request in requests: + scheduler.add_request(request) + return scheduler.schedule() + + +def test_dynamic_sd_uses_batch_size_schedule(): + dynamic_sd_lookup = _make_lookup( + [ + (1, 16, 3), + (32, 128, 2), + (256, 2048, 0), + ] + ) + + assert dynamic_sd_lookup[1] == 3 + assert dynamic_sd_lookup[16] == 3 + assert dynamic_sd_lookup[17] == 3 + assert dynamic_sd_lookup[31] == 3 + assert dynamic_sd_lookup[32] == 2 + assert dynamic_sd_lookup[128] == 2 + assert dynamic_sd_lookup[129] == 2 + assert dynamic_sd_lookup[255] == 2 + assert dynamic_sd_lookup[256] == 0 + + +def test_dynamic_sd_requires_schedule_starting_at_batch_size_one(): + with pytest.raises(ValueError, match="must start at 1"): + _make_lookup([(2, 16, 3)]) + + +def test_dynamic_sd_clamps_k_to_runtime_max(): + dynamic_sd_lookup = _make_lookup( + [(1, 256, 4)], + runtime_num_speculative_tokens=3, + ) + + assert dynamic_sd_lookup[1] == 3 + assert dynamic_sd_lookup[256] == 3 + + +def test_dynamic_sd_rejects_invalid_schedule_entry(): + with pytest.raises(ValueError, match="3-item sequence"): + _make_lookup([(1, 16, 3), (32, 64)]) # type: ignore[list-item] + + +def test_dynamic_sd_rejects_overlapping_ranges(): + with pytest.raises(ValueError, match="non-overlapping and sorted"): + _make_lookup([(1, 16, 3), (16, 32, 2)]) + + +def test_dynamic_sd_rejects_negative_k(): + with pytest.raises(ValueError, match="values must be >= 0"): + _make_lookup([(1, 16, -1)]) + + +def test_dynamic_sd_rejects_empty_schedule(): + with pytest.raises(ValueError, match="must not be empty"): + _make_lookup([]) + + +def test_dynamic_sd_requires_schedule_config(): + with pytest.raises( + ValueError, match="num_speculative_tokens_per_batch_size is required" + ): + build_dynamic_sd_schedule_lookup( + None, + vllm_max_batch_size=256, + vllm_num_speculative_tokens=3, + ) + + +def test_dynamic_sd_lookup_rejects_invalid_batch_size_queries(): + dynamic_sd_lookup = _make_lookup([(1, 256, 3)]) + + assert dynamic_sd_lookup[0] == 0 + with pytest.raises(IndexError): + _ = dynamic_sd_lookup[257] + + +def test_scheduler_initializes_dynamic_sd_lookup_from_speculative_config(): + scheduler = _make_scheduler_with_dynamic_sd( + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], + runtime_num_speculative_tokens=3, + ) + + assert scheduler.dynamic_sd_lookup is not None + assert scheduler.num_spec_tokens == 3 + + +def test_scheduler_uses_dsd_k_based_on_number_of_scheduled_requests(): + test_cases = [ + (4, 3), + (64, 2), + (256, 0), + ] + + for num_requests, expected_k in test_cases: + scheduler = _make_scheduler_with_dynamic_sd( + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], + max_num_seqs=num_requests, + max_num_batched_tokens=num_requests * 10, + runtime_num_speculative_tokens=3, + ) + output = _add_requests_and_schedule(scheduler, num_requests) + + assert len(output.num_scheduled_tokens) == num_requests + assert output.num_spec_tokens_to_schedule == expected_k + + +def test_scheduler_clamps_dsd_k_to_runtime_num_speculative_tokens(): + scheduler = _make_scheduler_with_dynamic_sd( + [(1, 256, 5)], + max_num_seqs=16, + max_num_batched_tokens=160, + runtime_num_speculative_tokens=3, + ) + output = _add_requests_and_schedule(scheduler, 16) + + assert len(output.num_scheduled_tokens) == 16 + assert output.num_spec_tokens_to_schedule == 3 + + +def test_scheduler_falls_back_to_static_k_when_dsd_not_configured(): + scheduler = create_scheduler( + max_num_seqs=4, + max_num_batched_tokens=40, + num_speculative_tokens=3, + ) + output = _add_requests_and_schedule(scheduler, 4) + + assert scheduler.dynamic_sd_lookup is None + assert output.num_spec_tokens_to_schedule == 3 + + +def test_scheduler_uses_static_k_when_no_requests_are_scheduled(): + scheduler = _make_scheduler_with_dynamic_sd( + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], + runtime_num_speculative_tokens=3, + ) + output = scheduler.schedule() + + assert len(output.num_scheduled_tokens) == 0 + assert output.num_spec_tokens_to_schedule == 3 + + +def test_scheduler_rejects_bad_dsd_config_at_construction(): + with pytest.raises(ValueError, match="must start at 1"): + _make_scheduler_with_dynamic_sd([(2, 16, 3)]) + + +def test_scheduler_passes_max_num_seqs_as_dsd_runtime_batch_limit(): + scheduler = _make_scheduler_with_dynamic_sd( + [(1, 16, 3), (64, 128, 2), (256, 4096, 0)], + max_num_seqs=16, + max_num_batched_tokens=160, + runtime_num_speculative_tokens=3, + ) + output = _add_requests_and_schedule(scheduler, 16) + + assert scheduler.dynamic_sd_lookup is not None + assert len(scheduler.dynamic_sd_lookup) == 17 + assert len(output.num_scheduled_tokens) == 16 + assert output.num_spec_tokens_to_schedule == 3 diff --git a/tests/v1/spec_decode/test_eagle.py b/tests/v1/spec_decode/test_eagle.py index 32f9dcc86a..848130725a 100644 --- a/tests/v1/spec_decode/test_eagle.py +++ b/tests/v1/spec_decode/test_eagle.py @@ -969,6 +969,7 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch): proposer.draft_attn_groups = [mock_attn_group] result = proposer.propose( + num_speculative_tokens=num_speculative_tokens, target_token_ids=target_token_ids, target_positions=target_positions, target_hidden_states=target_hidden_states, @@ -1071,6 +1072,7 @@ def test_propose_stores_probabilistic_draft_probs(monkeypatch): sampling_metadata.all_greedy = False result = proposer.propose( + num_speculative_tokens=num_speculative_tokens, target_token_ids=torch.randint(0, vocab_size, (total_tokens,), device=device), target_positions=torch.cat( [ diff --git a/tests/v1/spec_decode/test_extract_hidden_states.py b/tests/v1/spec_decode/test_extract_hidden_states.py index 2a67257b09..6b4e53ced6 100644 --- a/tests/v1/spec_decode/test_extract_hidden_states.py +++ b/tests/v1/spec_decode/test_extract_hidden_states.py @@ -255,6 +255,7 @@ def test_propose(): # Call propose draft_tokens = proposer.propose( + num_speculative_tokens=1, sampled_token_ids=sampled_token_ids, target_hidden_states=target_hidden_states, common_attn_metadata=common_attn_metadata, @@ -321,6 +322,7 @@ def test_propose_different_layer_counts(num_hidden_layers): ).unsqueeze(-1) draft_tokens = proposer.propose( + num_speculative_tokens=1, sampled_token_ids=sampled_token_ids, target_hidden_states=target_hidden_states, common_attn_metadata=common_attn_metadata, diff --git a/tests/v1/spec_decode/test_mtp.py b/tests/v1/spec_decode/test_mtp.py index 7c478f81d8..e334371f6d 100644 --- a/tests/v1/spec_decode/test_mtp.py +++ b/tests/v1/spec_decode/test_mtp.py @@ -205,6 +205,7 @@ def test_mtp_propose(num_speculative_tokens, monkeypatch): # Run propose result = proposer.propose( + num_speculative_tokens=num_speculative_tokens, target_token_ids=target_token_ids, target_positions=target_positions, target_hidden_states=target_hidden_states, diff --git a/tests/v1/spec_decode/test_ngram.py b/tests/v1/spec_decode/test_ngram.py index 7d2a07ddce..459edddd1c 100644 --- a/tests/v1/spec_decode/test_ngram.py +++ b/tests/v1/spec_decode/test_ngram.py @@ -81,6 +81,7 @@ def test_ngram_proposer(): # No match. token_ids_cpu = np.array([[1, 2, 3, 4, 5]]) result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -90,6 +91,7 @@ def test_ngram_proposer(): # No match for 4-gram. token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]]) result = get_ngram_proposer(min_n=4, max_n=4, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -99,6 +101,7 @@ def test_ngram_proposer(): # No match for 4-gram but match for 3-gram. token_ids_cpu = np.array([[1, 2, 3, 4, 1, 2, 3]]) result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -109,6 +112,7 @@ def test_ngram_proposer(): # In this case, the proposer should return the 4-gram match. token_ids_cpu = np.array([[2, 3, 4, 5, 1, 2, 3, 4, 1, 2, 3, 4]]) result = get_ngram_proposer(min_n=3, max_n=4, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -118,6 +122,7 @@ def test_ngram_proposer(): # Match for 2-gram and 3-gram, but not 4-gram. token_ids_cpu = np.array([[3, 4, 5, 2, 3, 4, 1, 2, 3, 4]]) result = get_ngram_proposer(min_n=2, max_n=4, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -127,6 +132,7 @@ def test_ngram_proposer(): # Multiple 3-gram matched, but always pick the first one. token_ids_cpu = np.array([[1, 2, 3, 100, 1, 2, 3, 200, 1, 2, 3, 300, 1, 2, 3]]) result = get_ngram_proposer(min_n=3, max_n=3, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -136,6 +142,7 @@ def test_ngram_proposer(): # check empty input token_ids_cpu = np.array([[]]) result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0]], num_tokens_no_spec=np.array([len(c) for c in token_ids_cpu]), token_ids_cpu=token_ids_cpu, @@ -147,6 +154,7 @@ def test_ngram_proposer(): # second request has 3 tokens and no match. Padded with -1 for max len 5 token_ids_cpu = np.array([[1, 2, 3, 1, 2], [4, 5, 6, -1, -1]]) result = get_ngram_proposer(min_n=2, max_n=2, k=2).propose( + num_speculative_tokens=2, sampled_token_ids=[[0], [1]], num_tokens_no_spec=np.array([5, 3]), token_ids_cpu=token_ids_cpu, @@ -166,6 +174,7 @@ def test_ngram_proposer(): num_tokens_no_spec = np.array([5, 3, 5], dtype=np.int32) sampled_token_ids = [[2], [], [8]] # Empty list for request 1 simulates prefill result = proposer.propose( + num_speculative_tokens=2, sampled_token_ids=sampled_token_ids, num_tokens_no_spec=num_tokens_no_spec, token_ids_cpu=token_ids_cpu, @@ -195,6 +204,7 @@ def test_ngram_proposer(): input_2[:3] = [4, 5, 6] token_ids_cpu = np.array([input_1, input_2]) result = ngram_proposer.propose( + num_speculative_tokens=2, sampled_token_ids=[[0], [1]], num_tokens_no_spec=np.array([len(input_1), 3]), token_ids_cpu=token_ids_cpu, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index a4d5b1302e..eba8653d63 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -157,6 +157,14 @@ class SpeculativeConfig: target_parallel_config: SkipValidation[ParallelConfig] = None # type: ignore """The parallel configuration for the target model.""" + # dynamic speculative decoding control + num_speculative_tokens_per_batch_size: list[tuple[int, int, int]] | None = None + """Batch-size schedule used to dynamically choose speculative-token count. + + Each entry is ``(range_start, range_end, num_speculative_tokens)`` with an + inclusive batch-size range. + """ + # params generated in the post-init stage draft_model_config: SkipValidation[ModelConfig] = None # type: ignore """The configuration of the draft model initialized internal.""" @@ -1073,6 +1081,9 @@ class SpeculativeConfig: def use_dflash(self) -> bool: return self.method == "dflash" + def uses_dynamic_speculative_decoding(self) -> bool: + return self.num_speculative_tokens_per_batch_size is not None + def uses_draft_model(self) -> bool: return self.method == "draft_model" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 6122476abb..308e1626ba 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -764,6 +764,23 @@ class VllmConfig: apply_recursive(self, defaults) + def _maybe_override_dynamic_sd_cudagraph_mode(self) -> None: + speculative_config = self.speculative_config + if ( + speculative_config is None + or not speculative_config.uses_dynamic_speculative_decoding() + or not self.compilation_config.cudagraph_mode.has_full_cudagraphs() + ): + return + + logger.warning_once( + "Dynamic speculative decoding changes the target verification " + "length at runtime. Overriding cudagraph_mode from %s to " + "PIECEWISE for reliability.", + self.compilation_config.cudagraph_mode.name, + ) + self.compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE + def _post_init_kv_transfer_config(self) -> None: """Update KVTransferConfig based on top-level configs in VllmConfig. @@ -1153,6 +1170,8 @@ class VllmConfig: "optimization level defaults." ) + self._maybe_override_dynamic_sd_cudagraph_mode() + if ( self.compilation_config.cudagraph_mode.requires_piecewise_compilation() and self.compilation_config.mode != CompilationMode.VLLM_COMPILE @@ -2005,6 +2024,9 @@ class VllmConfig: elif speculative_config.method not in ("eagle", "eagle3", "mtp", "dflash"): unsupported.append(f"speculative method '{speculative_config.method}'") + if speculative_config.uses_dynamic_speculative_decoding(): + unsupported.append("dynamic speculative decoding") + # V2 EagleSpeculator does not support parallel_drafting (for P-Eagle) # DFlash uses parallel drafting natively in V2 via DFlashSpeculator. if ( diff --git a/vllm/v1/core/sched/async_scheduler.py b/vllm/v1/core/sched/async_scheduler.py index a79e84289a..d1c652c46e 100644 --- a/vllm/v1/core/sched/async_scheduler.py +++ b/vllm/v1/core/sched/async_scheduler.py @@ -19,6 +19,10 @@ class AsyncScheduler(Scheduler): 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 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 b2e9dd8b17..0c1b9d34c5 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -240,6 +240,10 @@ class SchedulerOutput: # preventing stale NaN/data from corrupting attention or SSM computation. new_block_ids_to_zero: list[int] | None = None + # 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 + @classmethod def make_empty(cls) -> "SchedulerOutput": return cls( diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 6bae149a83..3b63ba3210 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -56,6 +56,7 @@ from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate +from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm.v1.structured_output import StructuredOutputManager from vllm.v1.utils import record_function_or_nullcontext @@ -218,7 +219,14 @@ class Scheduler(SchedulerInterface): self.use_eagle = False self.num_spec_tokens = vllm_config.num_speculative_tokens self.num_lookahead_tokens = 0 + self.dynamic_sd_lookup: list[int] | None = None if speculative_config is not None: + if speculative_config.num_speculative_tokens_per_batch_size: + self.dynamic_sd_lookup = build_dynamic_sd_schedule_lookup( + speculative_config.num_speculative_tokens_per_batch_size, + vllm_max_batch_size=self.scheduler_config.max_num_seqs, + vllm_num_speculative_tokens=self.num_spec_tokens, + ) if speculative_config.use_eagle(): self.use_eagle = True self.num_lookahead_tokens = self.num_spec_tokens @@ -995,6 +1003,13 @@ class Scheduler(SchedulerInterface): else None ) + # Dynamic speculative decoding: compute optimal K + num_spec_tokens_to_schedule = self.num_spec_tokens + if self.dynamic_sd_lookup is not None and len(num_scheduled_tokens) > 0: + num_spec_tokens_to_schedule = self.dynamic_sd_lookup[ + len(num_scheduled_tokens) + ] + scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, @@ -1011,6 +1026,7 @@ class Scheduler(SchedulerInterface): finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), new_block_ids_to_zero=new_block_ids_to_zero, + num_spec_tokens_to_schedule=num_spec_tokens_to_schedule, ) # NOTE(Kuntai): this function is designed for multiple purposes: diff --git a/vllm/v1/spec_decode/dynamic/__init__.py b/vllm/v1/spec_decode/dynamic/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/vllm/v1/spec_decode/dynamic/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/v1/spec_decode/dynamic/utils.py b/vllm/v1/spec_decode/dynamic/utils.py new file mode 100644 index 0000000000..de869b19a7 --- /dev/null +++ b/vllm/v1/spec_decode/dynamic/utils.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +DynamicSDSchedule = list[tuple[int, int, int]] + + +def validate_and_normalize_dynamic_sd_schedule( + num_speculative_tokens_per_batch_size: object, +) -> DynamicSDSchedule: + """Validate and normalize a Dynamic SD batch-size schedule. + + The schedule is expressed as a list of inclusive ranges: + + ``[(range_start, range_end, num_speculative_tokens), ...]`` + """ + if num_speculative_tokens_per_batch_size is None: + raise ValueError( + "num_speculative_tokens_per_batch_size is required for " + "dynamic speculative decoding." + ) + if not isinstance(num_speculative_tokens_per_batch_size, list): + raise ValueError( + "num_speculative_tokens_per_batch_size must be a non-empty list of " + "(range_start, range_end, num_speculative_tokens) entries." + ) + if not num_speculative_tokens_per_batch_size: + raise ValueError("num_speculative_tokens_per_batch_size must not be empty.") + + parsed_schedule: DynamicSDSchedule = [] + for entry in num_speculative_tokens_per_batch_size: + if not isinstance(entry, list | tuple) or len(entry) != 3: + raise ValueError( + "Each num_speculative_tokens_per_batch_size entry must be a " + "3-item sequence: (range_start, range_end, num_speculative_tokens)." + ) + + range_start, range_end, num_speculative_tokens = ( + int(entry[0]), + int(entry[1]), + int(entry[2]), + ) + + if range_start <= 0 or range_end <= 0: + raise ValueError( + f"Batch-size range ({range_start}, {range_end}) must be positive." + ) + if range_start > range_end: + raise ValueError( + "Batch-size range start must be <= end for " + f"({range_start}, {range_end}, {num_speculative_tokens})." + ) + if num_speculative_tokens < 0: + raise ValueError( + "num_speculative_tokens_per_batch_size values must be >= 0." + ) + + parsed_schedule.append((range_start, range_end, num_speculative_tokens)) + + parsed_schedule.sort(key=lambda entry: entry[0]) + + previous_end = 0 + for range_start, range_end, _ in parsed_schedule: + if range_start <= previous_end: + raise ValueError("Batch-size ranges must be non-overlapping and sorted.") + previous_end = range_end + + first_range_start = parsed_schedule[0][0] + if first_range_start != 1: + raise ValueError( + "The first batch-size range must start at 1 so every runtime " + "batch size has a defined schedule." + ) + + return parsed_schedule + + +def build_dynamic_sd_schedule_lookup( + num_speculative_tokens_per_batch_size: object, + vllm_max_batch_size: int, + vllm_num_speculative_tokens: int, +) -> list[int]: + """Expand the configured schedule into a dense batch_size -> K lookup. + + "dense_schedule" means a 1-indexed lookup table where index ``batch_size`` + stores the exact K to use for that runtime batch size. This lets the + scheduler do a simple array lookup instead of searching the configured + ranges on every scheduling step. + """ + if vllm_max_batch_size <= 0: + raise ValueError("vllm_max_batch_size must be > 0.") + if vllm_num_speculative_tokens <= 0: + raise ValueError("vllm_num_speculative_tokens must be > 0.") + + parsed_schedule = validate_and_normalize_dynamic_sd_schedule( + num_speculative_tokens_per_batch_size + ) + + # Index 0 is intentionally unused so that valid runtime batch sizes can be + # looked up directly as dense_schedule[batch_size]. + dense_schedule = [0] * (vllm_max_batch_size + 1) + next_batch_size = 1 + last_num_speculative_tokens: int | None = None + + for range_start, range_end, num_speculative_tokens in parsed_schedule: + if range_start > next_batch_size and last_num_speculative_tokens is not None: + # Fill any gap before the next configured range by carrying forward + # the previous K. For example, [(1, 16, 3), (32, 128, 2)] should map + # batch sizes 17-31 to K=3. + for batch_size in range( + next_batch_size, + min(range_start, vllm_max_batch_size + 1), + ): + dense_schedule[batch_size] = min( + vllm_num_speculative_tokens, + last_num_speculative_tokens, + ) + + # Fill the current configured inclusive range with its K value. + for batch_size in range( + max(range_start, next_batch_size), + min(range_end, vllm_max_batch_size) + 1, + ): + dense_schedule[batch_size] = min( + vllm_num_speculative_tokens, + num_speculative_tokens, + ) + + next_batch_size = max(next_batch_size, range_end + 1) + last_num_speculative_tokens = num_speculative_tokens + + if next_batch_size > vllm_max_batch_size: + break + + if last_num_speculative_tokens is None: + raise ValueError( + "num_speculative_tokens_per_batch_size must contain at least " + "one valid batch-size range." + ) + + # Fill the tail after the final configured range by carrying forward the + # last K through vllm_max_batch_size. + for batch_size in range(next_batch_size, vllm_max_batch_size + 1): + dense_schedule[batch_size] = min( + vllm_num_speculative_tokens, + last_num_speculative_tokens, + ) + + return dense_schedule diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index c3cb3c8aae..a0a1f03c71 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -29,7 +29,10 @@ class ExtractHiddenStatesProposer: def __init__(self, vllm_config: VllmConfig, device): assert vllm_config.speculative_config is not None - assert vllm_config.speculative_config.num_speculative_tokens == 1 + self.num_speculative_tokens = ( + vllm_config.speculative_config.num_speculative_tokens + ) + assert self.num_speculative_tokens == 1 if vllm_config.speculative_config.disable_padded_drafter_batch: raise ValueError( "disable_padded_drafter_batch is not supported with " @@ -82,6 +85,7 @@ class ExtractHiddenStatesProposer: def propose( self, + num_speculative_tokens: int, sampled_token_ids: torch.Tensor, target_hidden_states: list[torch.Tensor], common_attn_metadata: CommonAttentionMetadata, @@ -112,6 +116,7 @@ class ExtractHiddenStatesProposer: - Draft tokens matching sampled tokens, shape [batch_size, 1] - KV connector output (if KV transfer is active), else None """ + assert num_speculative_tokens == self.num_speculative_tokens assert self.model is not None and isinstance(target_hidden_states, list) # target_hidden_states is a list of tensors (one per layer) diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index 88e3030d2e..e11798ce6b 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -434,6 +434,7 @@ class SpecDecodeBaseProposer: def propose( self, + num_speculative_tokens, # [num_tokens] target_token_ids: torch.Tensor, # [num_tokens] or [3, num_tokens] when M-RoPE is enabled @@ -451,6 +452,7 @@ class SpecDecodeBaseProposer: | list[dict[str, torch.Tensor]] | None = None, ) -> torch.Tensor: + self.num_speculative_tokens = num_speculative_tokens self._last_draft_probs = None batch_size = common_attn_metadata.batch_size() @@ -521,6 +523,17 @@ class SpecDecodeBaseProposer: sample_hidden_states = last_hidden_states[token_indices_to_sample] + # No draft tokens requested (e.g. Dynamic SD decided K=0). + # The prefill forward pass above already ran to keep the drafter + # KV cache in sync, so just return an empty tensor. + if self.num_speculative_tokens == 0: + return torch.empty( + batch_size, + 0, + device=sample_hidden_states.device, + dtype=torch.int64, + ) + # Early exit if there is only one draft token to be generated. if self.num_speculative_tokens == 1 or self.parallel_drafting: draft_token_ids, draft_probs = self._sample_draft_tokens( diff --git a/vllm/v1/spec_decode/medusa.py b/vllm/v1/spec_decode/medusa.py index 80b0f0a987..7adf7cff5f 100644 --- a/vllm/v1/spec_decode/medusa.py +++ b/vllm/v1/spec_decode/medusa.py @@ -35,15 +35,18 @@ class MedusaProposer: self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens self.hidden_size = self.spec_config.draft_model_config.get_hidden_size() self.dtype = vllm_config.model_config.dtype + self.num_speculative_tokens = self.spec_config.num_speculative_tokens def propose( self, + num_speculative_tokens: int, target_hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata, slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, # unused ) -> torch.Tensor: + assert num_speculative_tokens == self.num_speculative_tokens # Generate blocks and compute logits blocks = self.model(target_hidden_states) logits = self.model.compute_logits(blocks) diff --git a/vllm/v1/spec_decode/metrics.py b/vllm/v1/spec_decode/metrics.py index 5da41510b4..a3ccfb29e7 100644 --- a/vllm/v1/spec_decode/metrics.py +++ b/vllm/v1/spec_decode/metrics.py @@ -28,12 +28,14 @@ class SpecDecodingStats: num_draft_tokens: int = 0 num_accepted_tokens: int = 0 num_accepted_tokens_per_pos: list[int] = field(default_factory=list) + num_draft_tokens_per_pos: list[int] = field(default_factory=list) @classmethod def new(cls, num_spec_tokens: int) -> "SpecDecodingStats": return cls( num_spec_tokens=num_spec_tokens, num_accepted_tokens_per_pos=[0] * num_spec_tokens, + num_draft_tokens_per_pos=[0] * num_spec_tokens, ) def observe_draft(self, num_draft_tokens: int, num_accepted_tokens: int): @@ -43,6 +45,8 @@ class SpecDecodingStats: assert num_accepted_tokens <= self.num_spec_tokens for i in range(num_accepted_tokens): self.num_accepted_tokens_per_pos[i] += 1 + for i in range(num_draft_tokens): + self.num_draft_tokens_per_pos[i] += 1 class SpecDecodingLogging: diff --git a/vllm/v1/spec_decode/ngram_proposer.py b/vllm/v1/spec_decode/ngram_proposer.py index 53199d0ce2..e0240d0e66 100644 --- a/vllm/v1/spec_decode/ngram_proposer.py +++ b/vllm/v1/spec_decode/ngram_proposer.py @@ -55,6 +55,7 @@ class NgramProposer: # Trigger Numba JIT compilation for N-gram proposer. # This usually takes less than 1 second. self.propose( + self.k, [[]] * 1024, np.zeros(1024, dtype=np.int32), np.zeros((1024, self.max_model_len), dtype=np.int32), @@ -66,6 +67,7 @@ class NgramProposer: valid_ngram_requests: list, num_tokens_no_spec: np.ndarray, token_ids_cpu: np.ndarray, + k: int, ) -> list[list[int]]: """Batch version of ngram proposer using numba for acceleration. @@ -78,6 +80,8 @@ class NgramProposer: token_ids_cpu: Numpy array of shape (batch_size, max_model_len) representing the token IDs for each request. + k: + Number of speculative tokens to propose. Returns: list[list[int]]: @@ -110,7 +114,7 @@ class NgramProposer: self.min_n, self.max_n, self.max_model_len, - self.k, + k, self.valid_ngram_draft, self.valid_ngram_num_drafts, ) @@ -130,6 +134,7 @@ class NgramProposer: def propose( self, + num_speculative_tokens: int, sampled_token_ids: list[list[int]], num_tokens_no_spec: np.ndarray, token_ids_cpu: np.ndarray, @@ -137,6 +142,8 @@ class NgramProposer: | list[dict[str, torch.Tensor]] | None = None, # unused ) -> list[list[int]]: + assert num_speculative_tokens <= self.k + # find which requests need ngram proposals valid_ngram_requests = [] for i, sampled_ids in enumerate(sampled_token_ids): @@ -157,6 +164,7 @@ class NgramProposer: valid_ngram_requests, num_tokens_no_spec, token_ids_cpu, + num_speculative_tokens, ) return draft_token_ids diff --git a/vllm/v1/spec_decode/ngram_proposer_gpu.py b/vllm/v1/spec_decode/ngram_proposer_gpu.py index 7759d5c32f..b8a0116ede 100644 --- a/vllm/v1/spec_decode/ngram_proposer_gpu.py +++ b/vllm/v1/spec_decode/ngram_proposer_gpu.py @@ -314,6 +314,7 @@ class NgramProposerGPU: def propose( self, + num_speculative_tokens: int, num_tokens_no_spec: torch.Tensor, # [batch_size] token_ids_gpu: torch.Tensor, # [batch_size, max_len] valid_sampled_token_ids_gpu: torch.Tensor, # [batch_size, num_spec_tokens + 1] @@ -326,6 +327,7 @@ class NgramProposerGPU: updated lengths, then run the kernel. Args: + num_speculative_tokens: Number of speculative tokens to propose. num_tokens_no_spec: Number of tokens per sequence (read-only) token_ids_gpu: Token IDs tensor (modified in-place with new tokens) valid_sampled_token_ids_gpu: Newly sampled tokens to scatter @@ -336,6 +338,7 @@ class NgramProposerGPU: num_valid_draft_tokens: Count of leading valid draft tokens per request [batch_size] """ + assert num_speculative_tokens == self.k assert token_ids_gpu.device == self.device assert num_tokens_no_spec.device == self.device diff --git a/vllm/v1/spec_decode/step3p5.py b/vllm/v1/spec_decode/step3p5.py index ccca17a318..043f3f2be2 100644 --- a/vllm/v1/spec_decode/step3p5.py +++ b/vllm/v1/spec_decode/step3p5.py @@ -273,6 +273,7 @@ class Step3p5MTPProposer(EagleProposer): def propose( self, + num_speculative_tokens: int, target_token_ids: torch.Tensor, target_positions: torch.Tensor, target_hidden_states: torch.Tensor, @@ -286,6 +287,7 @@ class Step3p5MTPProposer(EagleProposer): | list[dict[str, torch.Tensor]] | None = None, ) -> torch.Tensor: + self.num_speculative_tokens = num_speculative_tokens self._last_draft_probs = None batch_size = common_attn_metadata.batch_size() diff --git a/vllm/v1/spec_decode/suffix_decoding.py b/vllm/v1/spec_decode/suffix_decoding.py index fee5d97468..66137a0063 100644 --- a/vllm/v1/spec_decode/suffix_decoding.py +++ b/vllm/v1/spec_decode/suffix_decoding.py @@ -34,12 +34,14 @@ class SuffixDecodingProposer: def propose( self, + num_speculative_tokens: int, input_batch: InputBatch, sampled_token_ids: list[list[int]], slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, # unused ) -> list[list[int]]: + assert num_speculative_tokens == self.num_speculative_tokens """ Propose speculative tokens for each request in the input batch. Suffix Decoding will speculate a dynamic number of tokens for each request every decoding step, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index afda4ec0bb..4e3842d108 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -620,9 +620,11 @@ class GPUModelRunner( ) self.num_spec_tokens = 0 + self.prev_num_spec_tokens = 0 self.valid_sampled_token_count_gpu: torch.Tensor | None = None if self.speculative_config: self.num_spec_tokens = self.speculative_config.num_speculative_tokens + self.prev_num_spec_tokens = self.num_spec_tokens draft_config = self.speculative_config.draft_model_config if draft_config is not None and draft_config.max_model_len is not None: self.effective_drafter_max_model_len = draft_config.max_model_len @@ -1764,7 +1766,7 @@ class GPUModelRunner( spec_flattened_indices.extend( range(flattened_index - draft_len + 1, flattened_index + 1) ) - start = prev_index * self.num_spec_tokens + start = prev_index * self.prev_num_spec_tokens # prev_draft_token_indices is used to find which draft_tokens_id # should be copied to input_ids # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]] @@ -4704,6 +4706,9 @@ class GPUModelRunner( def _copy_draft_token_ids_to_cpu( self, scheduler_output: "SchedulerOutput", zeros_only: bool = False ) -> None: + if torch.is_tensor(self._draft_token_ids): + assert isinstance(self._draft_token_ids, torch.Tensor) + self.prev_num_spec_tokens = self._draft_token_ids.shape[1] # Check if we need to copy draft tokens to CPU. In async scheduling, # we only copy when needed for structured output, penalties or bad_words. if self.use_async_scheduling and not ( @@ -4722,16 +4727,17 @@ class GPUModelRunner( assert self.draft_token_ids_cpu is not None default_stream = torch.cuda.current_stream() num_reqs = draft_token_ids.shape[0] + num_spec_tokens = draft_token_ids.shape[1] with torch.cuda.stream(self.draft_token_ids_copy_stream): if not zeros_only: # Trigger async copy of draft token ids to cpu. self.draft_token_ids_copy_stream.wait_stream(default_stream) - self.draft_token_ids_cpu[:num_reqs].copy_( + self.draft_token_ids_cpu[:num_reqs, :num_spec_tokens].copy_( draft_token_ids, non_blocking=True ) else: # No copy needed, just zero-out cpu tensor. - self.draft_token_ids_cpu[:num_reqs] = 0 + self.draft_token_ids_cpu[:num_reqs, :num_spec_tokens] = 0 self.draft_token_ids_event.record() def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]: @@ -4743,7 +4749,11 @@ class GPUModelRunner( assert self.draft_token_ids_event is not None assert self.draft_token_ids_cpu is not None self.draft_token_ids_event.synchronize() - return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids + assert isinstance(self._draft_token_ids, torch.Tensor) + num_spec_tokens = self._draft_token_ids.shape[1] + return self.draft_token_ids_cpu[ + : len(req_ids), :num_spec_tokens + ].tolist(), req_ids def _copy_valid_sampled_token_count( self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor @@ -4823,6 +4833,7 @@ class GPUModelRunner( num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens spec_config = self.speculative_config assert spec_config is not None + num_spec_tokens_to_schedule = scheduler_output.num_spec_tokens_to_schedule self._draft_probs = None self._draft_prob_req_ids = None if spec_config.method == "ngram": @@ -4831,6 +4842,7 @@ class GPUModelRunner( assert isinstance(sampled_token_ids, list) assert isinstance(self.drafter, NgramProposer) draft_token_ids = self.drafter.propose( + num_spec_tokens_to_schedule, sampled_token_ids, self.input_batch.num_tokens_no_spec, self.input_batch.token_ids_cpu, @@ -4864,6 +4876,7 @@ class GPUModelRunner( batch_size = next_token_ids.shape[0] draft_token_ids, num_valid_draft_tokens = self.drafter.propose( + num_spec_tokens_to_schedule, self.num_tokens_no_spec_gpu[:batch_size], self.token_ids_gpu_tensor[:batch_size], valid_sampled_token_ids_gpu, @@ -4885,7 +4898,10 @@ class GPUModelRunner( assert isinstance(sampled_token_ids, list) assert isinstance(self.drafter, SuffixDecodingProposer) draft_token_ids = self.drafter.propose( - self.input_batch, sampled_token_ids, slot_mappings=slot_mappings + num_spec_tokens_to_schedule, + self.input_batch, + sampled_token_ids, + slot_mappings=slot_mappings, ) elif spec_config.method == "medusa": assert isinstance(sampled_token_ids, list) @@ -4909,6 +4925,7 @@ class GPUModelRunner( hidden_states = sample_hidden_states[indices] draft_token_ids = self.drafter.propose( + num_speculative_tokens=num_spec_tokens_to_schedule, target_hidden_states=hidden_states, sampling_metadata=sampling_metadata, slot_mappings=slot_mappings, @@ -4926,6 +4943,7 @@ class GPUModelRunner( target_hidden_states = [h[:num_scheduled_tokens] for h in aux_hidden_states] draft_token_ids = self.drafter.propose( + num_speculative_tokens=num_spec_tokens_to_schedule, sampled_token_ids=sampled_token_ids, target_hidden_states=target_hidden_states, common_attn_metadata=common_attn_metadata, @@ -5059,6 +5077,7 @@ class GPUModelRunner( mm_embed_inputs = None draft_token_ids = self.drafter.propose( + num_speculative_tokens=num_spec_tokens_to_schedule, target_token_ids=target_token_ids, target_positions=target_positions, target_hidden_states=target_hidden_states,