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 13, 22:00 UTC
[BugFix] Fix mixed penalties batch with async scheduling (#27910)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Nov 1, 2025 at 10:51 UTC
c2ed069b32e2805c05a858c6157f4c6393b145a8
1 file changed
+8
vllm/v1/sample/ops/penalties.py
+8
@@ -21,6 +21,14 @@ def apply_all_penalties(
21
"""
22
_, vocab_size = logits.shape
23
output_tokens_t = _convert_to_tensors(output_token_ids, vocab_size, logits.device)
24
+
25
+ # In the async scheduling case, rows that won't have penalties applied may contain
26
+ # -1 placeholder token ids. We must replace these with valid token ids so that the
27
+ # scatter done in apply_penalties is valid.
28
+ # NOTE(nick): The penalties implementation is currently quite inefficient and
29
+ # will be reworked anyhow.
30
+ output_tokens_t.masked_fill_(output_tokens_t == -1, vocab_size)
31
+
32
return apply_penalties(
33
logits,
34
prompt_token_ids,