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, 06:30 UTC
[Model Runner V2] Minor refactor for logit_bias (#32209)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 12, 2026 at 13:08 UTC
dec28688c5060f921594f9a7ce1647cf7a87a39e
1 file changed
+54
-26
vllm/v1/worker/gpu/sample/logit_bias.py
+54
-26
@@ -119,35 +119,18 @@ class LogitBiasState:
119
idx_mapping: torch.Tensor,
120
pos: torch.Tensor,
121
) -> None:
122
- num_reqs, vocab_size = logits.shape
123
- BLOCK_SIZE = triton.next_power_of_2(
124
- max(
125
- MAX_NUM_ALLOWED_TOKEN_IDS,
126
- MAX_NUM_LOGIT_BIAS_TOKENS,
127
- MAX_NUM_STOP_TOKEN_IDS,
128
- )
129
- )
130
- LOGITS_BLOCK_SIZE = 8192
131
- _bias_kernel[(num_reqs,)](
122
+ apply_logit_bias(
123
logits,
133
- logits.stride(0),
134
- vocab_size,
124
idx_mapping,
136
- self.num_allowed_token_ids,
137
- self.allowed_token_ids,
138
- self.allowed_token_ids.gpu.stride(0),
139
- self.num_logit_bias,
140
- self.logit_bias_token_ids,
141
- self.logit_bias_token_ids.gpu.stride(0),
142
- self.logit_bias,
143
- self.logit_bias.gpu.stride(0),
125
pos,
145
- self.min_lens,
146
- self.num_stop_token_ids,
147
- self.stop_token_ids,
148
- self.stop_token_ids.gpu.stride(0),
149
- BLOCK_SIZE=BLOCK_SIZE,
150
- LOGITS_BLOCK_SIZE=LOGITS_BLOCK_SIZE,
126
+ self.num_allowed_token_ids.gpu,
127
+ self.allowed_token_ids.gpu,
128
+ self.num_logit_bias.gpu,
129
+ self.logit_bias_token_ids.gpu,
130
+ self.logit_bias.gpu,
131
+ self.min_lens.gpu,
132
+ self.num_stop_token_ids.gpu,
133
+ self.stop_token_ids.gpu,
134
)
135
136
@@ -240,3 +223,48 @@ def _bias_kernel(
223
-float("inf"),
224
mask=mask,
225
)
226
+
227
+
228
+def apply_logit_bias(
229
+ logits: torch.Tensor,
230
+ idx_mapping: torch.Tensor,
231
+ pos: torch.Tensor,
232
+ num_allowed_token_ids: torch.Tensor,
233
+ allowed_token_ids: torch.Tensor,
234
+ num_logit_bias: torch.Tensor,
235
+ logit_bias_token_ids: torch.Tensor,
236
+ logit_bias: torch.Tensor,
237
+ min_lens: torch.Tensor,
238
+ num_stop_token_ids: torch.Tensor,
239
+ stop_token_ids: torch.Tensor,
240
+) -> None:
241
+ num_reqs, vocab_size = logits.shape
242
+ BLOCK_SIZE = triton.next_power_of_2(
243
+ max(
244
+ allowed_token_ids.shape[-1],
245
+ logit_bias_token_ids.shape[-1],
246
+ stop_token_ids.shape[-1],
247
+ )
248
+ )
249
+ LOGITS_BLOCK_SIZE = 8192
250
+ _bias_kernel[(num_reqs,)](
251
+ logits,
252
+ logits.stride(0),
253
+ vocab_size,
254
+ idx_mapping,
255
+ num_allowed_token_ids,
256
+ allowed_token_ids,
257
+ allowed_token_ids.stride(0),
258
+ num_logit_bias,
259
+ logit_bias_token_ids,
260
+ logit_bias_token_ids.stride(0),
261
+ logit_bias,
262
+ logit_bias.stride(0),
263
+ pos,
264
+ min_lens,
265
+ num_stop_token_ids,
266
+ stop_token_ids,
267
+ stop_token_ids.stride(0),
268
+ BLOCK_SIZE=BLOCK_SIZE,
269
+ LOGITS_BLOCK_SIZE=LOGITS_BLOCK_SIZE,
270
+ )