iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
vLLM-4x-RTX-3090/patches/0021-spec-decode-structured-output-under-concurrency-on-a.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

1177 lines
47 KiB
Diff
Raw Permalink Blame History

From 53ba1efcf7012a9b847bfcc0de3cbdd17e3303ce Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Tue, 22 Sep 2026 04:17:29 -0400
Subject: [PATCH 021/103] spec-decode: structured output under concurrency on a
PP shape
Ports upstream vLLM PRs #54442 and #56802 (issue #54437): never sample from an
unmasked row, and key draft-token hand-off snapshots by scheduler step so an
alternating PP microbatch cannot overwrite the slot.
---
tests/v1/core/test_scheduler.py | 19 ++
tests/v1/worker/test_draft_tokens_handler.py | 210 ++++++++++++++++++
tests/v1/worker/test_grammar_fail_closed.py | 151 +++++++++++++
.../v1/worker/test_grammar_invalid_drafts.py | 76 +++++++
vllm/v1/core/sched/output.py | 9 +
vllm/v1/core/sched/scheduler.py | 13 +-
vllm/v1/engine/core.py | 4 +-
vllm/v1/executor/abstract.py | 9 +-
vllm/v1/executor/multiproc_executor.py | 7 +-
vllm/v1/executor/uniproc_executor.py | 9 +-
vllm/v1/structured_output/utils.py | 13 ++
vllm/v1/worker/gpu/model_runner.py | 59 ++++-
vllm/v1/worker/gpu/sample/batch_shard.py | 7 +
.../gpu/spec_decode/rejection_sampler.py | 18 +-
vllm/v1/worker/gpu/spec_decode/utils.py | 148 +++++++++---
vllm/v1/worker/gpu_model_runner.py | 3 +-
vllm/v1/worker/gpu_worker.py | 4 +-
vllm/v1/worker/worker_base.py | 33 ++-
18 files changed, 748 insertions(+), 44 deletions(-)
create mode 100644 tests/v1/worker/test_draft_tokens_handler.py
create mode 100644 tests/v1/worker/test_grammar_fail_closed.py
create mode 100644 tests/v1/worker/test_grammar_invalid_drafts.py
diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py
index 061c825..07ced9e 100644
--- a/tests/v1/core/test_scheduler.py
+++ b/tests/v1/core/test_scheduler.py
@@ -30,6 +30,7 @@ from vllm.multimodal.inputs import (
MultiModalKwargsItem,
PlaceholderRange,
)
+from vllm.platforms import current_platform
from vllm.sampling_params import SamplingParams, StructuredOutputsParams
from vllm.utils.hashing import sha256
from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
@@ -249,6 +250,24 @@ def test_schedule(enable_prefix_caching: bool, prompt_logprobs: int | None):
assert scheduler.running[i] == request
+def test_scheduler_output_step_identity_defaults_to_zero():
+ output = SchedulerOutput.make_empty()
+ assert output.scheduler_step == 0
+
+
+def test_scheduler_output_carries_step_identity():
+ if not current_platform.device_type:
+ pytest.skip("No inferred device type in this environment")
+
+ scheduler = create_scheduler()
+ request = create_requests(num_requests=1)[0]
+ scheduler.add_request(request)
+
+ output = scheduler.schedule()
+ assert output.scheduler_step == scheduler.current_step
+ assert output.scheduler_step == 1
+
+
def test_scheduler_stats_route_to_existing_output_client():
scheduler = create_scheduler()
request = create_requests(num_requests=1)[0]
diff --git a/tests/v1/worker/test_draft_tokens_handler.py b/tests/v1/worker/test_draft_tokens_handler.py
new file mode 100644
index 0000000..bf016eb
--- /dev/null
+++ b/tests/v1/worker/test_draft_tokens_handler.py
@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DraftTokensHandler step identity (#54437 piece 2 / Link A)."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
+from vllm.v1.worker.worker_base import WorkerWrapperBase
+
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
+
+K = 4
+
+
+def _device() -> torch.device:
+ if torch.cuda.is_available():
+ return torch.device("cuda:0")
+ return torch.device("cpu")
+
+
+def _batch(req_ids: list[str], structured: bool) -> SimpleNamespace:
+ return SimpleNamespace(req_ids=list(req_ids), has_structured_output_reqs=structured)
+
+
+def _drafts(start: int, num_reqs: int, device: torch.device) -> torch.Tensor:
+ return torch.arange(
+ start, start + num_reqs * K, dtype=torch.int32, device=device
+ ).reshape(num_reqs, K)
+
+
+def test_consumed_snapshot_survives_later_batch():
+ """Link A: a later execute_model must not erase an earlier step's drafts."""
+ device = _device()
+ handler = DraftTokensHandler(device, max_snapshots=4)
+
+ drafts_a = _drafts(1, 2, device)
+ drafts_b = _drafts(101, 2, device)
+
+ handler.snapshot_consumed_drafts(
+ step_id=1,
+ input_batch=_batch(["gA", "pA"], True),
+ draft_tokens=drafts_a,
+ )
+ got = handler.get_draft_tokens(step_id=1)
+ assert got is not None
+ assert got.req_ids == ["gA", "pA"]
+ assert got.draft_token_ids[0] == [1, 2, 3, 4]
+
+ handler.snapshot_consumed_drafts(
+ step_id=2,
+ input_batch=_batch(["gB", "pB"], True),
+ draft_tokens=drafts_b,
+ )
+ got_b = handler.get_draft_tokens(step_id=2)
+ assert got_b is not None
+ assert got_b.req_ids == ["gB", "pB"]
+
+ got_a = handler.get_draft_tokens(step_id=1)
+ assert got_a is not None
+ assert got_a.req_ids == ["gA", "pA"]
+ assert got_a.draft_token_ids[0] == [1, 2, 3, 4]
+
+
+def test_plain_batch_does_not_evict_structured_snapshot():
+ device = _device()
+ handler = DraftTokensHandler(device, max_snapshots=4)
+ drafts_a = _drafts(1, 2, device)
+ drafts_plain = _drafts(201, 2, device)
+
+ handler.snapshot_consumed_drafts(
+ step_id=1,
+ input_batch=_batch(["gA", "pA"], True),
+ draft_tokens=drafts_a,
+ )
+ handler.snapshot_consumed_drafts(
+ step_id=2,
+ input_batch=_batch(["p1", "p2"], False),
+ draft_tokens=drafts_plain,
+ )
+
+ got = handler.get_draft_tokens(step_id=1)
+ assert got is not None
+ assert got.req_ids == ["gA", "pA"]
+ assert handler.get_draft_tokens(step_id=2) is None
+
+
+def test_unknown_step_returns_none_for_fail_closed_fallback():
+ device = _device()
+ handler = DraftTokensHandler(device)
+ handler.snapshot_consumed_drafts(
+ step_id=3,
+ input_batch=_batch(["gA"], True),
+ draft_tokens=_drafts(1, 1, device),
+ )
+
+ assert handler.get_draft_tokens(step_id=99) is None
+
+
+def test_proposal_slot_without_step_id_stays_latest():
+ """Non-async post_step still reads the latest proposed drafts."""
+ device = _device()
+ handler = DraftTokensHandler(device)
+ handler.set_draft_tokens(_batch(["gA", "pA"], True), _drafts(1, 2, device))
+ handler.set_draft_tokens(_batch(["gB", "pB"], True), _drafts(101, 2, device))
+
+ got = handler.get_draft_tokens()
+ assert got is not None
+ assert got.req_ids == ["gB", "pB"]
+ assert got.draft_token_ids[0] == [101, 102, 103, 104]
+
+
+def test_proposal_plain_batch_returns_placeholders():
+ device = _device()
+ handler = DraftTokensHandler(device)
+ handler.set_draft_tokens(_batch(["p1", "p2"], False), _drafts(1, 2, device))
+
+ got = handler.get_draft_tokens()
+ assert got is not None
+ assert got.req_ids == ["p1", "p2"]
+ assert got.draft_token_ids == [[-1] * K, [-1] * K]
+
+
+def test_consumed_lookup_ignores_proposal_slot():
+ device = _device()
+ handler = DraftTokensHandler(device)
+ handler.snapshot_consumed_drafts(
+ step_id=5,
+ input_batch=_batch(["gA"], True),
+ draft_tokens=_drafts(1, 1, device),
+ )
+ handler.set_draft_tokens(_batch(["gB"], True), _drafts(101, 1, device))
+
+ got = handler.get_draft_tokens(step_id=5)
+ assert got is not None
+ assert got.req_ids == ["gA"]
+ assert got.draft_token_ids[0] == [1, 2, 3, 4]
+
+
+def test_ring_buffer_evicts_oldest_and_misses():
+ device = _device()
+ handler = DraftTokensHandler(device, max_snapshots=2)
+ handler.snapshot_consumed_drafts(
+ step_id=1,
+ input_batch=_batch(["a"], True),
+ draft_tokens=_drafts(1, 1, device),
+ )
+ handler.snapshot_consumed_drafts(
+ step_id=2,
+ input_batch=_batch(["b"], True),
+ draft_tokens=_drafts(11, 1, device),
+ )
+ handler.snapshot_consumed_drafts(
+ step_id=3,
+ input_batch=_batch(["c"], True),
+ draft_tokens=_drafts(21, 1, device),
+ )
+
+ assert handler.get_draft_tokens(step_id=1) is None
+ got = handler.get_draft_tokens(step_id=2)
+ assert got is not None
+ assert got.req_ids == ["b"]
+ got = handler.get_draft_tokens(step_id=3)
+ assert got is not None
+ assert got.req_ids == ["c"]
+
+
+def test_worker_wrapper_drops_unsupported_step_id():
+ wrapper = WorkerWrapperBase(rpc_rank=0)
+ seen: list[tuple] = []
+
+ class PluginWorker:
+ def take_draft_token_ids(self):
+ seen.append(())
+ return "legacy"
+
+ wrapper.worker = PluginWorker()
+ assert wrapper.take_draft_token_ids(step_id=7) == "legacy"
+ assert seen == [()]
+
+
+def test_worker_wrapper_forwards_step_id_when_supported():
+ wrapper = WorkerWrapperBase(rpc_rank=0)
+ seen: list[int | None] = []
+
+ class Worker:
+ def take_draft_token_ids(self, step_id: int | None = None):
+ seen.append(step_id)
+ return "ok"
+
+ wrapper.worker = Worker()
+ assert wrapper.take_draft_token_ids(step_id=7) == "ok"
+ assert wrapper.take_draft_token_ids() == "ok"
+ assert seen == [7, None]
+
+
+def test_worker_wrapper_drops_positional_only_step_id():
+ wrapper = WorkerWrapperBase(rpc_rank=0)
+ seen: list[tuple] = []
+
+ class PluginWorker:
+ def take_draft_token_ids(self, step_id: int | None = None, /):
+ seen.append((step_id,))
+ return "legacy"
+
+ wrapper.worker = PluginWorker()
+ assert wrapper.take_draft_token_ids(step_id=7) == "legacy"
+ assert seen == [(None,)]
diff --git a/tests/v1/worker/test_grammar_fail_closed.py b/tests/v1/worker/test_grammar_fail_closed.py
new file mode 100644
index 0000000..869cb98
--- /dev/null
+++ b/tests/v1/worker/test_grammar_fail_closed.py
@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Acceptance must stop at the last bitmask row that carries a real mask.
+
+`grammar_bitmask` stops constraining rows once it meets a -1 placeholder in the
+scheduled draft window, and the worker verifies the drafts it holds on the GPU
+rather than the scheduler's padded copy. When the two disagree -- the hand-off
+lost a step's drafts, see #54437 -- acceptance can walk into a row where every
+token is allowed.
+
+The grammar here admits only LEGAL. The scheduled window is back-filled one
+draft deep, so rows 0 and 1 carry a real mask and rows 2.. are `_full_mask`,
+while the drafts on the GPU would carry acceptance into row 2 and emit ILLEGAL.
+"""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.sched.output import GrammarOutput
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS
+from vllm.v1.worker.gpu.spec_decode.rejection_sampler import RejectionSampler
+from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
+
+VOCAB = 128
+K = 3 # drafts per step
+NUM_LOGITS = K + 1 # + the bonus row
+LEGAL = 7
+ILLEGAL = 99
+# Index of the first -1 in the scheduled window: rows 0..1 are masked.
+NUM_ACCEPTABLE = 1
+
+
+def _bitmask() -> np.ndarray:
+ """Rows 0..NUM_ACCEPTABLE constrained to LEGAL, the rest fully permissive."""
+ mask = np.zeros((NUM_LOGITS, cdiv(VOCAB, 32)), dtype=np.int32)
+ for row in range(NUM_LOGITS):
+ if row <= NUM_ACCEPTABLE:
+ mask[row, LEGAL // 32] = 1 << (LEGAL % 32)
+ else:
+ mask[row, :] = -1 # every token allowed
+ return mask
+
+
+def _input_batch(device: torch.device) -> SimpleNamespace:
+ cu_np = np.array([0, NUM_LOGITS], dtype=np.int32)
+ # draft_sampled = input_ids[logits_indices]; entry i + 1 is draft i, so the
+ # drafts the GPU holds are [LEGAL, LEGAL, ILLEGAL].
+ input_ids = torch.tensor(
+ [LEGAL, LEGAL, LEGAL, ILLEGAL], dtype=torch.int32, device=device
+ )
+ return SimpleNamespace(
+ num_reqs=1,
+ req_ids=["g"],
+ num_draft_tokens=K,
+ num_draft_tokens_per_req=np.array([K], dtype=np.int32),
+ input_ids=input_ids,
+ logits_indices=torch.arange(NUM_LOGITS, dtype=torch.int32, device=device),
+ positions=torch.arange(NUM_LOGITS, dtype=torch.int64, device=device),
+ cu_num_logits_np=cu_np,
+ cu_num_logits=torch.from_numpy(cu_np).to(device),
+ idx_mapping_np=np.array([0], dtype=np.int32),
+ idx_mapping=torch.zeros(1, dtype=torch.int32, device=device),
+ expanded_idx_mapping=torch.zeros(NUM_LOGITS, dtype=torch.int32, device=device),
+ expanded_local_pos=torch.arange(NUM_LOGITS, dtype=torch.int32, device=device),
+ seq_lens=torch.tensor([64], dtype=torch.int32, device=device),
+ )
+
+
+def _logits(device: torch.device) -> torch.Tensor:
+ """Every row prefers its own draft, so nothing but the mask can reject one."""
+ logits = torch.zeros(NUM_LOGITS, VOCAB, dtype=torch.float32, device=device)
+ logits[0, LEGAL] = 10.0
+ logits[1, LEGAL] = 10.0
+ logits[2, ILLEGAL] = 10.0 # row 2 is permissive, so this wins there
+ logits[3, LEGAL] = 10.0
+ return logits
+
+
+def _rejection_sampler(device: torch.device) -> RejectionSampler:
+ sampler = object.__new__(RejectionSampler)
+ sampler.num_speculative_steps = K
+ sampler.enable_adaptive_verification = False
+ sampler.use_block_verification = False
+ sampler.synthetic_conditional_rates = None
+ sampler.watermark_key = None
+ zeros_f = torch.zeros(1, dtype=torch.float32, device=device) # greedy
+ sampler.sampler = SimpleNamespace(
+ compute_nans=False,
+ logprobs_mode="raw_logprobs",
+ use_fp64_gumbel=False,
+ apply_sampling_params=lambda logits, *args, **kwargs: logits,
+ sampling_states=SimpleNamespace(
+ temperature=SimpleNamespace(gpu=zeros_f),
+ seeds=SimpleNamespace(gpu=torch.zeros(1, dtype=torch.int64, device=device)),
+ max_num_logprobs=lambda _: NO_LOGPROBS,
+ ),
+ req_states=SimpleNamespace(
+ prefill_len=SimpleNamespace(
+ gpu=torch.zeros(1, dtype=torch.int32, device=device)
+ )
+ ),
+ )
+ return sampler
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+def test_acceptance_stops_at_the_last_masked_row():
+ device = torch.device("cuda")
+ input_batch = _input_batch(device)
+ logits = _logits(device)
+
+ runner = object.__new__(GPUModelRunner)
+ runner.device = device
+ runner.vocab_size = VOCAB
+ runner.batch_sharder = None
+ runner.sampler = None
+ runner.model = SimpleNamespace(compute_logits=lambda hidden: logits)
+ runner.rejection_sampler = _rejection_sampler(device)
+ runner.speculator = SimpleNamespace(draft_logits=None)
+ runner.structured_outputs_worker = StructuredOutputsWorker(
+ max_num_logits=NUM_LOGITS,
+ vocab_size=VOCAB,
+ device=device,
+ mask_stride=NUM_LOGITS,
+ num_bonus_tokens=1,
+ )
+
+ grammar_output = GrammarOutput(["g"], _bitmask())
+ # The scheduler reports how many leading drafts the bitmask constrained.
+ # Absent on an older scheduler, where the whole window is invalidated.
+ grammar_output.num_acceptable_drafts = [NUM_ACCEPTABLE]
+
+ hidden = torch.zeros(NUM_LOGITS, 8, dtype=torch.float32, device=device)
+ sampler_output, num_sampled, _ = runner.sample(hidden, input_batch, grammar_output)
+
+ accepted = int(num_sampled[0].item())
+ emitted = sampler_output.sampled_token_ids[0, :accepted].tolist()
+ assert ILLEGAL not in emitted, (
+ f"emitted {emitted}: a draft verified against a permissive row was "
+ f"accepted, so the request sampled with no grammar constraint"
+ )
+ # Draft 0 may be accepted; draft 1 must not be, because accepting it moves
+ # sampling into row 2, which constrains nothing.
+ assert accepted == NUM_ACCEPTABLE + 1
+ assert emitted == [LEGAL] * accepted
diff --git a/tests/v1/worker/test_grammar_invalid_drafts.py b/tests/v1/worker/test_grammar_invalid_drafts.py
new file mode 100644
index 0000000..476fe76
--- /dev/null
+++ b/tests/v1/worker/test_grammar_invalid_drafts.py
@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Which speculative drafts a structured-output request may accept.
+
+`StructuredOutputManager.grammar_bitmask` fills row i before it inspects
+`req_tokens[i]`, so with the first -1 placeholder at index j rows 0..j carry a
+real mask and rows j+1.. carry the all-permissive `_full_mask`. Drafts j..K-1
+must therefore be rejected, or the request samples with no grammar constraint
+at a position the model really does sample.
+"""
+
+from unittest.mock import Mock
+
+import numpy as np
+import torch
+
+from vllm.v1.worker.gpu.model_runner import grammar_invalid_draft_positions
+
+
+def _input_batch(req_ids, num_logits_per_req, num_draft_tokens=6):
+ cu = np.concatenate([[0], np.cumsum(num_logits_per_req)]).astype(np.int32)
+ return Mock(
+ req_ids=list(req_ids),
+ cu_num_logits_np=cu,
+ num_draft_tokens=num_draft_tokens,
+ )
+
+
+def _positions(batch, grammar_req_ids, num_acceptable):
+ out = grammar_invalid_draft_positions(
+ batch, grammar_req_ids, num_acceptable, torch.device("cpu")
+ )
+ return None if out is None else out.tolist()
+
+
+def test_positions_reject_whole_window_when_nothing_was_backfilled():
+ # Two requests, 4 logits rows each (1 + 3 drafts). Request "g" is at rows
+ # 0..3, so its drafts sit at draft_sampled[1..3].
+ batch = _input_batch(["g", "p"], [4, 4])
+ assert _positions(batch, ["g"], [0]) == [1, 2, 3]
+
+
+def test_positions_keep_the_drafts_the_bitmask_could_see():
+ batch = _input_batch(["g", "p"], [4, 4])
+ # Two real drafts: they were masked correctly and may be accepted.
+ assert _positions(batch, ["g"], [2]) == [3]
+
+
+def test_positions_reject_nothing_when_fully_backfilled():
+ batch = _input_batch(["g", "p"], [4, 4])
+ assert _positions(batch, ["g"], [3]) is None
+
+
+def test_positions_are_offset_per_request():
+ batch = _input_batch(["p", "g"], [4, 4])
+ # "g" starts at logits row 4, so its drafts are draft_sampled[5..7].
+ assert _positions(batch, ["g"], [0]) == [5, 6, 7]
+
+
+def test_positions_fall_back_to_the_whole_window_without_the_field():
+ # An older scheduler, or warmup, supplies no num_acceptable_drafts; be
+ # conservative rather than accepting drafts whose mask is unknown.
+ batch = _input_batch(["g"], [4])
+ assert _positions(batch, ["g"], None) == [1, 2, 3]
+
+
+def test_no_positions_without_drafts_or_grammar_requests():
+ batch = _input_batch(["g"], [4], num_draft_tokens=0)
+ assert _positions(batch, ["g"], [0]) is None
+ batch = _input_batch(["g"], [4])
+ assert _positions(batch, [], []) is None
+
+
+def test_unknown_request_is_skipped():
+ batch = _input_batch(["p"], [4])
+ assert _positions(batch, ["g"], [0]) is None
diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py
index 3a46fa4..03bb2cd 100644
--- a/vllm/v1/core/sched/output.py
+++ b/vllm/v1/core/sched/output.py
@@ -304,6 +304,10 @@ class SchedulerOutput:
# Number of spec tokens to schedule for the next step.
num_spec_tokens_to_schedule: int = 0
+ # Identity of this schedule() call (`Scheduler.current_step`). Used to
+ # retrieve the matching consumption-time draft-token snapshot.
+ scheduler_step: int = 0
+
@classmethod
def make_empty(cls) -> "SchedulerOutput":
return cls(
@@ -325,3 +329,8 @@ class GrammarOutput:
structured_output_request_ids: list[str]
# Bitmask ordered as structured_output_request_ids.
grammar_bitmask: "npt.NDArray[np.int32]"
+ # Per request, ordered as structured_output_request_ids: how many leading
+ # drafts the bitmask constrained. `grammar_bitmask` fills every row after
+ # the first -1 placeholder with the all-permissive `_full_mask`, so drafts
+ # from this index on must not be accepted. None means nothing to invalidate.
+ num_acceptable_drafts: list[int] | None = None
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 25af1b6..b16ae9d 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -74,7 +74,8 @@ from vllm.v1.core.sched import e2_alloc_log as _e2_alloc_log
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.structured_output import StructuredOutputGrammar, StructuredOutputManager
+from vllm.v1.structured_output.utils import strip_speculative_padding
from vllm.v1.utils import record_function_or_nullcontext
logger = init_logger(__name__)
@@ -1465,6 +1466,7 @@ class Scheduler(SchedulerInterface):
kv_connector_block_state=kv_connector_block_state,
num_spec_tokens_to_schedule=num_spec_tokens_to_schedule,
ec_manager_metadata=self.encoder_cache_manager.get_manager_metadata(),
+ scheduler_step=self.current_step,
)
# NOTE(Kuntai): this function is designed for multiple purposes:
@@ -1941,7 +1943,14 @@ class Scheduler(SchedulerInterface):
structured_output_request_ids,
scheduler_output.scheduled_spec_decode_tokens,
)
- return GrammarOutput(structured_output_request_ids, bitmask)
+ spec_tokens = scheduler_output.scheduled_spec_decode_tokens
+ num_acceptable_drafts = [
+ len(strip_speculative_padding(spec_tokens.get(req_id, [])))
+ for req_id in structured_output_request_ids
+ ]
+ return GrammarOutput(
+ structured_output_request_ids, bitmask, num_acceptable_drafts
+ )
def update_from_output(
self,
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
index ffc9a5d..110f56d 100644
--- a/vllm/v1/engine/core.py
+++ b/vllm/v1/engine/core.py
@@ -725,7 +725,9 @@ class EngineCore:
# When draft tokens are used with structured output, validate them
# before computing the grammar bitmask for the deferred request.
if self.check_for_draft_tokens:
- draft_token_ids = self.model_executor.take_draft_token_ids()
+ draft_token_ids = self.model_executor.take_draft_token_ids(
+ step_id=deferred_scheduler_output.scheduler_step
+ )
if draft_token_ids is not None:
# Update the draft token ids in the scheduler output to
# filter out the invalid spec tokens, which will be padded
diff --git a/vllm/v1/executor/abstract.py b/vllm/v1/executor/abstract.py
index 964f16e..a24507b 100644
--- a/vllm/v1/executor/abstract.py
+++ b/vllm/v1/executor/abstract.py
@@ -259,8 +259,13 @@ class Executor(ABC):
def execute_dummy_batch(self) -> None:
self.collective_rpc("execute_dummy_batch")
- def take_draft_token_ids(self) -> DraftTokenIds | None:
- output: list[DraftTokenIds] = self.collective_rpc("take_draft_token_ids")
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
+ # Omit step_id when unset so plugin workers with the documented
+ # no-argument take_draft_token_ids() keep working.
+ kwargs = {} if step_id is None else {"step_id": step_id}
+ output: list[DraftTokenIds] = self.collective_rpc(
+ "take_draft_token_ids", kwargs=kwargs
+ )
return output[0]
def profile(self, is_start: bool = True, profile_prefix: str | None = None):
diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py
index aae270f..1ccd23a 100644
--- a/vllm/v1/executor/multiproc_executor.py
+++ b/vllm/v1/executor/multiproc_executor.py
@@ -405,10 +405,13 @@ class MultiprocExecutor(Executor):
def execute_dummy_batch(self) -> None:
self.collective_rpc("execute_dummy_batch", unique_reply_rank=self.output_rank)
- def take_draft_token_ids(self) -> DraftTokenIds | None:
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
# OPTIMIZATION: Get output only from a single worker (output_rank)
+ kwargs = {} if step_id is None else {"step_id": step_id}
return self.collective_rpc(
- "take_draft_token_ids", unique_reply_rank=self.output_rank
+ "take_draft_token_ids",
+ unique_reply_rank=self.output_rank,
+ kwargs=kwargs,
)
def collective_rpc( # type: ignore[override]
diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py
index da1f2c5..dc798cd 100644
--- a/vllm/v1/executor/uniproc_executor.py
+++ b/vllm/v1/executor/uniproc_executor.py
@@ -141,8 +141,13 @@ class UniProcExecutor(Executor):
single_value=True,
)
- def take_draft_token_ids(self) -> DraftTokenIds | None:
- return self.collective_rpc("take_draft_token_ids", single_value=True)
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
+ kwargs = {} if step_id is None else {"step_id": step_id}
+ return self.collective_rpc(
+ "take_draft_token_ids",
+ kwargs=kwargs,
+ single_value=True,
+ )
def check_health(self) -> None:
# UniProcExecutor will always be healthy as long as
diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py
index 92122f8..7c4af7a 100644
--- a/vllm/v1/structured_output/utils.py
+++ b/vllm/v1/structured_output/utils.py
@@ -573,3 +573,16 @@ def choice_as_grammar(choice: list[str]) -> str:
escaped_choices = (escape_ebnf_string(c) for c in choice)
grammar = "root ::= " + " | ".join(f'"{c}"' for c in escaped_choices)
return grammar
+
+def strip_speculative_padding(token_ids: list[int]) -> list[int]:
+ """Drop speculative-decoding padding from a token block.
+
+ ngram and other speculative backends pad rejected draft positions with a
+ -1 sentinel. Structured-output grammars treat every entry as a real token
+ id, so the sentinels (and everything after the first one) are removed here,
+ before tokens reach any backend, rather than inside a single backend.
+ """
+ for i, token_id in enumerate(token_ids):
+ if token_id < 0:
+ return token_ids[:i]
+ return token_ids
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index e3479c2..744fe5b 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -185,6 +185,39 @@ from vllm.v1.worker.workspace import lock_workspace, use_workspace_lane
logger = init_logger(__name__)
+def grammar_invalid_draft_positions(
+ input_batch: InputBatch,
+ grammar_req_ids: list[str],
+ num_acceptable_drafts: list[int] | None,
+ device: torch.device,
+) -> torch.Tensor | None:
+ """Indices of `draft_sampled` entries whose bitmask row is permissive.
+
+ Drafts from `num_acceptable_drafts` on were verified against rows carrying
+ `_full_mask`, so accepting one would sample with no grammar constraint.
+ Returns None when there is nothing to invalidate.
+ """
+ if not grammar_req_ids or input_batch.num_draft_tokens == 0:
+ return None
+ cu_num_logits = input_batch.cu_num_logits_np.tolist()
+ req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
+ positions: list[int] = []
+ for i, req_id in enumerate(grammar_req_ids):
+ req_idx = req_id_to_idx.get(req_id)
+ if req_idx is None:
+ continue
+ # Without the field (an older scheduler, or warmup) invalidate the whole
+ # window, which is the conservative choice.
+ num_acceptable = (
+ num_acceptable_drafts[i] if num_acceptable_drafts is not None else 0
+ )
+ start = cu_num_logits[req_idx] + 1 + num_acceptable
+ positions.extend(range(start, cu_num_logits[req_idx + 1]))
+ if not positions:
+ return None
+ return torch.tensor(positions, dtype=torch.int64, device=device)
+
+
class GPUModelRunner(LoRAModelRunnerMixin):
def __init__(self, vllm_config: VllmConfig, device: torch.device):
self.vllm_config = vllm_config
@@ -293,7 +326,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
self.use_aux_hidden_state_outputs = True
# Draft tokens propagation - for spec-dec + struct outputs.
- self.draft_tokens_handler = DraftTokensHandler(self.device)
+ self.draft_tokens_handler = DraftTokensHandler(
+ self.device,
+ max_snapshots=max(8, vllm_config.max_concurrent_batches + 1),
+ )
self.pcp_manager: pcp.PCPManager | None = None
@@ -1471,6 +1507,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
else None
),
)
+ if self.num_speculative_steps > 0:
+ # Snapshot the drafts that were just written into input_ids, tagged
+ # with this scheduler step. Proposal-time drafts (set later in
+ # sample_tokens) are for the next step and must not overwrite this.
+ self.draft_tokens_handler.snapshot_consumed_drafts(
+ scheduler_output.scheduler_step,
+ input_batch,
+ self.req_states.draft_tokens[input_batch.idx_mapping],
+ )
input_batch = pcp.maybe_partition_pcp_batch(
self.pcp_manager,
input_batch,
@@ -1543,6 +1588,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
sample_hidden_states = hidden_states[input_batch.logits_indices]
logits = self.model.compute_logits(sample_hidden_states)
+ invalid_draft_positions = None
if grammar_output is not None:
# Apply grammar bitmask to the logits in-place.
assert self.structured_outputs_worker is not None
@@ -1552,6 +1598,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
grammar_output.structured_output_request_ids,
grammar_output.grammar_bitmask,
)
+ invalid_draft_positions = grammar_invalid_draft_positions(
+ input_batch,
+ grammar_output.structured_output_request_ids,
+ grammar_output.num_acceptable_drafts,
+ self.device,
+ )
sampler_output: SamplerOutput | None
if input_batch.num_reqs == 0:
@@ -1570,6 +1622,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
input_batch,
# Draft logits are needed for probabilistic rejection sampling.
self.speculator.draft_logits,
+ invalid_draft_positions,
)
if shard_metadata is not None:
@@ -2193,8 +2246,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
return async_output
- def take_draft_token_ids(self) -> DraftTokenIds | None:
- return self.draft_tokens_handler.get_draft_tokens()
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
+ return self.draft_tokens_handler.get_draft_tokens(step_id=step_id)
@torch.inference_mode()
@step_eplb_after()
diff --git a/vllm/v1/worker/gpu/sample/batch_shard.py b/vllm/v1/worker/gpu/sample/batch_shard.py
index 0f484f1..07f7987 100644
--- a/vllm/v1/worker/gpu/sample/batch_shard.py
+++ b/vllm/v1/worker/gpu/sample/batch_shard.py
@@ -335,9 +335,16 @@ def _shard_grammar_output(
cursor += num_req_logits
if not local_ids:
return None
+ # num_acceptable_drafts is ordered with structured_output_request_ids.
+ num_acceptable = grammar_output.num_acceptable_drafts
+ if num_acceptable is not None:
+ owned = set(local_ids)
+ ids = grammar_output.structured_output_request_ids
+ num_acceptable = [n for i, n in zip(ids, num_acceptable) if i in owned]
return GrammarOutput(
structured_output_request_ids=local_ids,
grammar_bitmask=grammar_output.grammar_bitmask[keep_indices],
+ num_acceptable_drafts=num_acceptable,
)
diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
index 609924e..6c07e6a 100644
--- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
+++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py
@@ -176,6 +176,7 @@ class RejectionSampler:
idx_mapping_np: np.ndarray,
expanded_idx_mapping: torch.Tensor,
expanded_local_pos: torch.Tensor,
+ verify_draft_sampled: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
processed_logits = self.sampler.apply_sampling_params(
logits,
@@ -189,7 +190,7 @@ class RejectionSampler:
sampled, num_sampled = rejection_sample(
processed_logits,
draft_logits,
- draft_sampled,
+ verify_draft_sampled,
cu_num_logits,
pos,
idx_mapping,
@@ -216,7 +217,10 @@ class RejectionSampler:
pos: torch.Tensor,
max_chunk_logits: int,
max_num_logprobs: int,
+ verify_draft_sampled: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]:
+ if verify_draft_sampled is None:
+ verify_draft_sampled = draft_sampled
cu_num_logits_np = input_batch.cu_num_logits_np
use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES
num_reqs = input_batch.num_reqs
@@ -251,6 +255,7 @@ class RejectionSampler:
input_batch.idx_mapping_np[start:end],
input_batch.expanded_idx_mapping[lo:hi],
input_batch.expanded_local_pos[lo:hi],
+ verify_draft_sampled[lo:hi],
)
chunk_logprobs = self._get_logprobs_tensors(
sampled,
@@ -289,6 +294,7 @@ class RejectionSampler:
logits: torch.Tensor,
input_batch: InputBatch,
draft_logits: torch.Tensor | None = None,
+ invalid_draft_positions: torch.Tensor | None = None,
) -> SamplerOutput:
# NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
# that num_nans is computed before applying penalties and temperature.
@@ -297,6 +303,15 @@ class RejectionSampler:
draft_sampled = input_batch.input_ids[input_batch.logits_indices]
pos = input_batch.positions[input_batch.logits_indices]
+ # Marking a draft invalid (`is_valid_draft = draft_sampled >= 0`) pins the
+ # accepted length before it reaches a permissive bitmask row. Only
+ # verification sees the copy: `apply_sampling_params` and watermarking
+ # keep the real draft ids.
+ verify_draft_sampled = draft_sampled
+ if invalid_draft_positions is not None:
+ verify_draft_sampled = draft_sampled.clone()
+ verify_draft_sampled[invalid_draft_positions] = -1
+
max_num_logprobs = self.sampler.sampling_states.max_num_logprobs(
input_batch.idx_mapping_np
)
@@ -309,6 +324,7 @@ class RejectionSampler:
pos,
chunk_logit_limit,
max_num_logprobs,
+ verify_draft_sampled,
)
num_sampled, num_rejected = get_num_sampled_and_rejected(
diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py
index f73acb2..65ad160 100644
--- a/vllm/v1/worker/gpu/spec_decode/utils.py
+++ b/vllm/v1/worker/gpu/spec_decode/utils.py
@@ -1,5 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections import deque
+from dataclasses import dataclass
+
import numpy as np
import torch
@@ -10,6 +13,10 @@ from vllm.v1.outputs import DraftTokenIds
from vllm.v1.worker.gpu.async_utils import async_copy_to_np
from vllm.v1.worker.gpu.input_batch import InputBatch
+# Keep enough consumption-time snapshots to cover async PP in-flight
+# batches (`max_concurrent_batches` is `pp_size + 1`) with headroom.
+_DEFAULT_MAX_SNAPSHOTS = 8
+
logger = init_logger(__name__)
@@ -26,48 +33,135 @@ def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig:
return load_config
+@dataclass
+class _DraftSnapshot:
+ step_id: int | None
+ req_ids: list[str]
+ num_draft_tokens: int
+ draft_tokens_np: np.ndarray | None
+ copy_event: torch.cuda.Event | None
+
+ def to_draft_token_ids(self) -> DraftTokenIds:
+ if self.draft_tokens_np is not None:
+ if self.copy_event is not None:
+ self.copy_event.synchronize()
+ draft_token_ids = self.draft_tokens_np.tolist()
+ else:
+ draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids]
+ return DraftTokenIds(self.req_ids, draft_token_ids)
+
+
class DraftTokensHandler:
- def __init__(self, device: torch.device | None = None):
+ def __init__(
+ self,
+ device: torch.device | None = None,
+ max_snapshots: int = _DEFAULT_MAX_SNAPSHOTS,
+ ):
self.device = device
- 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.copy_stream: torch.cuda.Stream | None = None
+ if device is not None and device.type == "cuda":
+ self.copy_stream = torch.cuda.Stream(device)
+
+ self._max_snapshots = max(1, max_snapshots)
+ # Consumption-time snapshots tagged with scheduler step identity.
+ self._consumed: deque[_DraftSnapshot] = deque(maxlen=self._max_snapshots)
+ # Latest proposal-time drafts for the non-async post_step path.
+ self._proposed: _DraftSnapshot | None = None
+
+ def _copy_drafts(
+ self, draft_tokens: torch.Tensor
+ ) -> tuple[np.ndarray, torch.cuda.Event | None]:
+ if self.copy_stream is None:
+ return draft_tokens.detach().cpu().numpy(), None
- self.req_ids: list[str] = []
- self.draft_tokens_np: np.ndarray | None = None
- self.num_draft_tokens: int = 0
+ current_stream = torch.cuda.current_stream(self.device)
+ self.copy_stream.wait_stream(current_stream)
+ copy_event = torch.cuda.Event(blocking=True)
+ with torch.cuda.stream(self.copy_stream):
+ draft_tokens_np = async_copy_to_np(draft_tokens)
+ # draft_tokens is a temporary allocation on the main stream and
+ # read here on copy_stream; without record_stream, the caching
+ # allocator may reuse its memory before the async copy executes.
+ draft_tokens.record_stream(self.copy_stream)
+ copy_event.record()
+ return draft_tokens_np, copy_event
+
+ def snapshot_consumed_drafts(
+ self,
+ step_id: int,
+ input_batch: InputBatch,
+ draft_tokens: torch.Tensor,
+ ) -> None:
+ """Record drafts that prepare_inputs fed into input_ids for this step.
+
+ Tagged with the scheduler step so the engine can retrieve the matching
+ snapshot later by equality. Batches without structured-output requests
+ are not stored, so they cannot evict a useful snapshot.
+ """
+ if not input_batch.has_structured_output_reqs:
+ return
+
+ draft_tokens_np, copy_event = self._copy_drafts(draft_tokens)
+ self._consumed.append(
+ _DraftSnapshot(
+ step_id=step_id,
+ req_ids=list(input_batch.req_ids),
+ num_draft_tokens=draft_tokens.shape[1],
+ draft_tokens_np=draft_tokens_np,
+ copy_event=copy_event,
+ )
+ )
def set_draft_tokens(
self, input_batch: InputBatch, draft_tokens: torch.Tensor
) -> None:
- self.req_ids = input_batch.req_ids
- self.num_draft_tokens = draft_tokens.shape[1]
+ req_ids = list(input_batch.req_ids)
+ num_draft_tokens = draft_tokens.shape[1]
if not input_batch.has_structured_output_reqs:
# No draft token validation needs to be performed by
# the scheduler for this batch.
- self.draft_tokens_np = None
+ self._proposed = _DraftSnapshot(
+ step_id=None,
+ req_ids=req_ids,
+ num_draft_tokens=num_draft_tokens,
+ draft_tokens_np=None,
+ copy_event=None,
+ )
return
# For spec decoding + structured outputs, we must transfer the
# draft tokens back to the scheduler for grammar validation.
- current_stream = torch.cuda.current_stream(self.device)
- self.copy_stream.wait_stream(current_stream)
- with torch.cuda.stream(self.copy_stream):
- self.draft_tokens_np = async_copy_to_np(draft_tokens)
- # draft_tokens is a temporary allocation on the main stream and read here on
- # copy_stream; without record_stream, the caching allocator may reuse its
- # memory before the async copy executes.
- draft_tokens.record_stream(self.copy_stream)
- self.copy_event.record()
+ draft_tokens_np, copy_event = self._copy_drafts(draft_tokens)
+ self._proposed = _DraftSnapshot(
+ step_id=None,
+ req_ids=req_ids,
+ num_draft_tokens=num_draft_tokens,
+ draft_tokens_np=draft_tokens_np,
+ copy_event=copy_event,
+ )
- def get_draft_tokens(self) -> DraftTokenIds | None:
- if self.draft_tokens_np is not None:
- self.copy_event.synchronize()
- draft_token_ids = self.draft_tokens_np.tolist()
- else:
- # This case only happens when async scheduling is disabled.
- draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids]
- return DraftTokenIds(self.req_ids, draft_token_ids)
+ def get_draft_tokens(self, step_id: int | None = None) -> DraftTokenIds | None:
+ """Return draft tokens for grammar validation.
+
+ If ``step_id`` is set, look up the consumption-time snapshot for that
+ scheduler step by equality. A miss returns None so placeholders stay
+ -1 and fail-closed invalidation can pin acceptance via is_valid_draft.
+ If ``step_id`` is None, return the latest proposal-time drafts (the
+ non-async post_step path).
+ """
+ if step_id is not None:
+ for snapshot in reversed(self._consumed):
+ if snapshot.step_id == step_id:
+ return snapshot.to_draft_token_ids()
+ logger.debug(
+ "Draft-token snapshot miss for scheduler step %s",
+ step_id,
+ )
+ return None
+
+ if self._proposed is None:
+ return DraftTokenIds([], [])
+ return self._proposed.to_draft_token_ids()
def get_parallel_drafting_token_id(hf_config) -> int:
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index 97a59bd..743c593 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -4895,7 +4895,8 @@ class GPUModelRunner(
self.input_batch.num_tokens_no_spec[i] = pos + 1
self.input_batch.prev_req_id_to_index = prev_req_id_to_index
- def take_draft_token_ids(self) -> DraftTokenIds | None:
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
+ del step_id # V2-only identity; V1 keeps a single proposal slot.
if not self.num_spec_tokens or not self._draft_token_req_ids:
return None
draft_token_ids, req_ids = self._get_draft_token_ids_cpu()
diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py
index eead9c9..949f299 100644
--- a/vllm/v1/worker/gpu_worker.py
+++ b/vllm/v1/worker/gpu_worker.py
@@ -1362,8 +1362,8 @@ class Worker(WorkerBase):
return self.model_runner.pool() # type: ignore
return None
- def take_draft_token_ids(self) -> DraftTokenIds | None:
- return self.model_runner.take_draft_token_ids()
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
+ return self.model_runner.take_draft_token_ids(step_id=step_id)
def profile(self, is_start: bool = True, profile_prefix: str | None = None):
# Check if profiling is enabled
diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py
index cf5b0e7..3bf7ca6 100644
--- a/vllm/v1/worker/worker_base.py
+++ b/vllm/v1/worker/worker_base.py
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import inspect
from collections.abc import Callable
from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
@@ -24,18 +25,34 @@ from vllm.v1.kv_cache_interface import KVCacheSpec
if TYPE_CHECKING:
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
- from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput
+ from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
else:
SchedulerOutput = object
GrammarOutput = object
AsyncModelRunnerOutput = object
ModelRunnerOutput = object
+ DraftTokenIds = object
logger = init_logger(__name__)
_R = TypeVar("_R")
+def _callable_accepts_kwarg(func: Callable[..., Any], name: str) -> bool:
+ """Return True if ``func`` can be called with ``name`` as a keyword."""
+ try:
+ params = inspect.signature(func).parameters
+ except (TypeError, ValueError):
+ return False
+ param = params.get(name)
+ if param is not None:
+ return param.kind in (
+ inspect.Parameter.POSITIONAL_OR_KEYWORD,
+ inspect.Parameter.KEYWORD_ONLY,
+ )
+ return any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+
+
class CompilationTimes(NamedTuple):
language_model: float
encoder: float
@@ -357,6 +374,20 @@ class WorkerWrapperBase:
def __getattr__(self, attr: str):
return getattr(self.worker, attr)
+ def take_draft_token_ids(self, step_id: int | None = None) -> DraftTokenIds | None:
+ """Forward draft-token lookup, dropping ``step_id`` if unsupported.
+
+ Plugin workers are documented to implement no-argument
+ ``take_draft_token_ids()``. Passing ``step_id`` to those methods would
+ raise TypeError; fall back to the legacy call instead.
+ """
+ method = getattr(self.worker, "take_draft_token_ids", None)
+ if method is None:
+ return None
+ if step_id is not None and _callable_accepts_kwarg(method, "step_id"):
+ return method(step_id=step_id)
+ return method()
+
def _apply_mm_cache(self, scheduler_output: SchedulerOutput) -> None:
mm_cache = self.mm_receiver_cache
if mm_cache is None:
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink