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 9, 23:00 UTC
[Model Runner V2] Skip kernel launch for penalties & logit_bias (#32634)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 19, 2026 at 22:20 UTC
e9c83cdc51f7d9fc642599057ec490e291ea7be3
3 files changed
+35
-5
vllm/v1/worker/gpu/sample/logit_bias.py
+18
-1
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
import torch
5
6
from vllm.sampling_params import SamplingParams
@@ -49,12 +50,18 @@ class LogitBiasState:
50
device=device,
51
)
52
53
+ # Using any of the above.
54
+ self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool)
55
+
56
def add_request(
57
self,
58
req_idx: int,
59
prompt_len: int,
60
sampling_params: SamplingParams,
61
) -> None:
62
+ # Using any logit bias.
63
+ use_logit_bias = False
64
+
65
# Allowed token IDs.
66
allowed_token_ids = sampling_params.allowed_token_ids
67
if allowed_token_ids:
@@ -66,6 +73,7 @@ class LogitBiasState:
73
)
74
self.num_allowed_token_ids.np[req_idx] = num_allowed_token_ids
75
self.allowed_token_ids.stage_write(req_idx, 0, allowed_token_ids)
76
+ use_logit_bias = True
77
else:
78
self.num_allowed_token_ids.np[req_idx] = 0
79
@@ -81,6 +89,7 @@ class LogitBiasState:
89
self.num_logit_bias.np[req_idx] = num_logit_bias
90
self.logit_bias_token_ids.stage_write(req_idx, 0, logit_bias.keys())
91
self.logit_bias.stage_write(req_idx, 0, logit_bias.values())
92
+ use_logit_bias = True
93
else:
94
self.num_logit_bias.np[req_idx] = 0
95
@@ -89,7 +98,7 @@ class LogitBiasState:
98
min_len = prompt_len + min_tokens
99
self.min_lens.np[req_idx] = min_len
100
stop_token_ids = sampling_params.all_stop_token_ids
92
- if stop_token_ids:
101
+ if min_tokens > 0 and stop_token_ids:
102
num_stop_token_ids = len(stop_token_ids)
103
if num_stop_token_ids > MAX_NUM_STOP_TOKEN_IDS:
104
raise ValueError(
@@ -98,9 +107,12 @@ class LogitBiasState:
107
)
108
self.num_stop_token_ids.np[req_idx] = num_stop_token_ids
109
self.stop_token_ids.stage_write(req_idx, 0, stop_token_ids)
110
+ use_logit_bias = True
111
else:
112
self.num_stop_token_ids.np[req_idx] = 0
113
114
+ self.use_logit_bias[req_idx] = use_logit_bias
115
+
116
def apply_staged_writes(self) -> None:
117
self.num_allowed_token_ids.copy_to_uva()
118
self.allowed_token_ids.apply_write()
@@ -117,8 +129,13 @@ class LogitBiasState:
129
self,
130
logits: torch.Tensor,
131
idx_mapping: torch.Tensor,
132
+ idx_mapping_np: np.ndarray,
133
pos: torch.Tensor,
134
) -> None:
135
+ if not np.any(self.use_logit_bias[idx_mapping_np]):
136
+ # No request uses logit bias. Skip the kernel launch.
137
+ return
138
+
139
apply_logit_bias(
140
logits,
141
idx_mapping,
vllm/v1/worker/gpu/sample/penalties.py
+15
-2
@@ -18,6 +18,7 @@ class PenaltiesState:
18
self.repetition_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
19
self.frequency_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
20
self.presence_penalty = UvaBackedTensor(max_num_reqs, dtype=torch.float32)
21
+ self.use_penalty = np.zeros(max_num_reqs, dtype=bool)
22
23
# Initialize repetition penalty manually because 0 is an invalid value for it.
24
self.repetition_penalty.np.fill(1.0)
@@ -42,7 +43,10 @@ class PenaltiesState:
43
self.repetition_penalty.np[req_idx] = sampling_params.repetition_penalty
44
self.frequency_penalty.np[req_idx] = sampling_params.frequency_penalty
45
self.presence_penalty.np[req_idx] = sampling_params.presence_penalty
45
- if use_penalty(sampling_params):
46
+
47
+ do_penalty = use_penalty(sampling_params)
48
+ self.use_penalty[req_idx] = do_penalty
49
+ if do_penalty:
50
self._penalties_reqs.append(req_idx)
51
52
def apply_staged_writes(
@@ -66,7 +70,16 @@ class PenaltiesState:
70
self.frequency_penalty.copy_to_uva()
71
self.presence_penalty.copy_to_uva()
72
69
- def apply_penalties(self, logits: torch.Tensor, idx_mapping: torch.Tensor) -> None:
73
+ def apply_penalties(
74
+ self,
75
+ logits: torch.Tensor,
76
+ idx_mapping: torch.Tensor,
77
+ idx_mapping_np: np.ndarray,
78
+ ) -> None:
79
+ if not np.any(self.use_penalty[idx_mapping_np]):
80
+ # No request uses penalties. Skip the kernel launch.
81
+ return
82
+
83
apply_penalties(
84
logits,
85
idx_mapping,
vllm/v1/worker/gpu/sample/sampler.py
+2
-2
@@ -104,10 +104,10 @@ class Sampler:
104
logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)
105
106
# Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.
107
- self.logit_bias_state.apply_logit_bias(logits, idx_mapping, pos)
107
+ self.logit_bias_state.apply_logit_bias(logits, idx_mapping, idx_mapping_np, pos)
108
109
# Apply penalties in place.
110
- self.penalties_state.apply_penalties(logits, idx_mapping)
110
+ self.penalties_state.apply_penalties(logits, idx_mapping, idx_mapping_np)
111
112
# Apply temperature in place.
113
apply_temperature(logits, idx_mapping, self.sampling_states.temperature.gpu)