A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 10:15 UTC
[Async][Feat] support apply penalty or bad_words for async + spec (#30495)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com> Signed-off-by: izhuhaoran <izhuhaoran@qq.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
zhrrr committed
Jan 9, 2026 at 10:31 UTC
8ff4a9956693450a37843bde88934ebcea440402
4 files changed
+70
-34
tests/v1/e2e/test_async_scheduling.py
+12
@@ -51,6 +51,14 @@ def test_without_spec_decoding(
51
dict(logprobs=2),
52
dict(logprobs=2, presence_penalty=-1.0),
53
dict(structured_outputs=struct_outputs),
54
+ dict(
55
+ structured_outputs=struct_outputs,
56
+ logprobs=2,
57
+ ),
58
+ dict(
59
+ structured_outputs=struct_outputs,
60
+ presence_penalty=-1.0,
61
+ ),
62
dict(
63
structured_outputs=struct_outputs,
64
logprobs=2,
@@ -105,11 +113,15 @@ def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch)
113
114
test_sampling_params = [
115
dict(),
116
+ dict(presence_penalty=-1.0),
117
+ dict(bad_words=["the", " the"]),
118
dict(logprobs=2),
119
+ dict(logprobs=2, presence_penalty=-1.0),
120
dict(structured_outputs=struct_outputs),
121
dict(
122
structured_outputs=struct_outputs,
123
logprobs=2,
124
+ presence_penalty=-1.0,
125
),
126
]
127
vllm/v1/engine/input_processor.py
-17
@@ -165,23 +165,6 @@ class InputProcessor:
165
"are not yet supported with speculative decoding."
166
)
167
168
- # Async scheduling + spec decode currently incompatible with some
169
- # sampling parameters.
170
- if (
171
- self.vllm_config.speculative_config is not None
172
- and self.vllm_config.scheduler_config.async_scheduling
173
- and (
174
- params.frequency_penalty != 0.0
175
- or params.presence_penalty != 0.0
176
- or params.repetition_penalty != 1.0
177
- or params.bad_words_token_ids
178
- )
179
- ):
180
- raise ValueError(
181
- "async scheduling with spec decoding doesn't yet support "
182
- "penalties or bad words in sampling parameters."
183
- )
184
-
168
def _validate_params(
169
self,
170
params: SamplingParams | PoolingParams,
vllm/v1/worker/gpu_input_batch.py
+34
-3
@@ -965,9 +965,40 @@ class InputBatch:
965
if sampled_token_ids is None:
966
assert self.async_copy_ready_event is not None
967
self.async_copy_ready_event.synchronize()
968
- sampled_token_ids = self.sampled_token_ids_cpu.squeeze(-1).tolist()
969
- # Replace placeholder token id with actual sampled id.
970
- req_output_token_ids[-1] = sampled_token_ids[prev_index]
968
+ sampled_token_ids = self.sampled_token_ids_cpu.tolist()
969
+ # Replace placeholder token id(s) with actual sampled id(s).
970
+ new_ids: list[int] = sampled_token_ids[prev_index]
971
+ if not new_ids:
972
+ continue
973
+ num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
974
+ # Also account for case where there may be a smaller number of
975
+ # output placeholders (tokens can be discarded after a kv-load failure).
976
+ first_placeholder = req_output_token_ids.index(-1)
977
+ num_placeholders = len(req_output_token_ids) - first_placeholder
978
+ num_to_replace = min(num_sampled_ids, num_placeholders)
979
+ del new_ids[num_to_replace:]
980
+ end_index = first_placeholder + num_to_replace
981
+ req_output_token_ids[first_placeholder:end_index] = new_ids
982
+
983
+ def update_async_spec_token_ids(self, draft_token_ids: list[list[int]]) -> None:
984
+ """
985
+ In async scheduling case, update spec_token_ids in sampling metadata with
986
+ real draft token ids from prior step. This is called right before they are
987
+ needed by the rejection sampler for penalty/bad_words computation.
988
+ """
989
+ if not draft_token_ids or not self.prev_req_id_to_index:
990
+ return
991
+
992
+ if (spec_token_ids := self.sampling_metadata.spec_token_ids) is not None:
993
+ for req_id, spec_ids in zip(self.req_ids, spec_token_ids):
994
+ if spec_ids:
995
+ prev_index = self.prev_req_id_to_index.get(req_id)
996
+ if prev_index is not None:
997
+ draft_ids = draft_token_ids[prev_index]
998
+ if draft_ids:
999
+ del draft_ids[len(spec_ids) :]
1000
+ spec_ids.clear()
1001
+ spec_ids.extend(draft_ids)
1002
1003
@property
1004
def num_reqs(self) -> int:
vllm/v1/worker/gpu_model_runner.py
+24
-14
@@ -2721,15 +2721,21 @@ class GPUModelRunner(
2721
) -> SamplerOutput:
2722
# Sample the next token and get logprobs if needed.
2723
sampling_metadata = self.input_batch.sampling_metadata
2724
+ # Update output token ids with tokens sampled in last step
2725
+ # if async scheduling and required by current sampling params.
2726
+ self.input_batch.update_async_output_token_ids()
2727
if spec_decode_metadata is None:
2725
- # Update output token ids with tokens sampled in last step
2726
- # if async scheduling and required by current sampling params.
2727
- self.input_batch.update_async_output_token_ids()
2728
return self.sampler(
2729
logits=logits,
2730
sampling_metadata=sampling_metadata,
2731
)
2732
2733
+ # Update spec_token_ids with real draft tokens from pre step only when
2734
+ # output_token_ids is needed (penalties or bad_words are in use).
2735
+ if self.use_async_scheduling and self._draft_token_req_ids is not None:
2736
+ draft_token_ids_cpu, _ = self._get_draft_token_ids_cpu()
2737
+ self.input_batch.update_async_spec_token_ids(draft_token_ids_cpu)
2738
+
2739
sampler_output = self.rejection_sampler(
2740
spec_decode_metadata,
2741
None, # draft_probs
@@ -3352,8 +3358,6 @@ class GPUModelRunner(
3358
) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
3359
kv_connector_output = self.kv_connector_output
3360
self.kv_connector_output = None
3355
- self._draft_token_ids = None
3356
- self._draft_token_req_ids = None
3361
3362
if self.execute_model_state is None:
3363
# Nothing to do (PP non-final rank case), output isn't used.
@@ -3393,6 +3397,8 @@ class GPUModelRunner(
3397
with record_function_or_nullcontext("gpu_model_runner: sample"):
3398
sampler_output = self._sample(logits, spec_decode_metadata)
3399
3400
+ self._draft_token_ids = None
3401
+ self._draft_token_req_ids = None
3402
self.input_batch.prev_sampled_token_ids = None
3403
3404
def propose_draft_token_ids(sampled_token_ids):
@@ -3517,17 +3523,18 @@ class GPUModelRunner(
3523
def take_draft_token_ids(self) -> DraftTokenIds | None:
3524
if not self.num_spec_tokens or not self._draft_token_req_ids:
3525
return None
3520
- req_ids = self._draft_token_req_ids
3521
- draft_token_ids = self._get_draft_token_ids_cpu(len(req_ids))
3526
+ draft_token_ids, req_ids = self._get_draft_token_ids_cpu()
3527
return DraftTokenIds(req_ids, draft_token_ids)
3528
3529
def _copy_draft_token_ids_to_cpu(
3530
self, scheduler_output: "SchedulerOutput", zeros_only: bool = False
3531
) -> None:
3527
- struct_output = scheduler_output.has_structured_output_requests
3528
- if self.use_async_scheduling and not struct_output:
3529
- # Draft tokens don't need to be copied to the CPU if async
3530
- # scheduling is in use and there are no structured output reqs.
3532
+ # Check if we need to copy draft tokens to CPU. In async scheduling,
3533
+ # we only copy when needed for structured output, penalties or bad_words.
3534
+ if self.use_async_scheduling and not (
3535
+ scheduler_output.has_structured_output_requests
3536
+ or self.input_batch.sampling_metadata.output_token_ids
3537
+ ):
3538
return
3539
# We must also set the corresponding request ids.
3540
self._draft_token_req_ids = self.input_batch.req_ids.copy()
@@ -3552,13 +3559,16 @@ class GPUModelRunner(
3559
self.draft_token_ids_cpu[:num_reqs] = 0
3560
self.draft_token_ids_event.record()
3561
3555
- def _get_draft_token_ids_cpu(self, num_reqs: int) -> list[list[int]]:
3562
+ def _get_draft_token_ids_cpu(self) -> tuple[list[list[int]], list[str]]:
3563
if isinstance(self._draft_token_ids, list):
3557
- return self._draft_token_ids
3564
+ return self._draft_token_ids, self.input_batch.req_ids
3565
+ req_ids = self._draft_token_req_ids
3566
+ if req_ids is None:
3567
+ return [], []
3568
assert self.draft_token_ids_event is not None
3569
assert self.draft_token_ids_cpu is not None
3570
self.draft_token_ids_event.synchronize()
3561
- return self.draft_token_ids_cpu[:num_reqs].tolist()
3571
+ return self.draft_token_ids_cpu[: len(req_ids)].tolist(), req_ids
3572
3573
def _copy_valid_sampled_token_count(
3574
self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor