[V1][Spec Decode] Add Dynamic SD (#32374)

Signed-off-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com>
Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

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)

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.

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

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(
             [

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,

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,

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,

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"
 

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 (

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:

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(

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:

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

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

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)

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(

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)

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:

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

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
 

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()
 

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,

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,