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] Support logit_bias, allowed_token_ids, min_tokens (#32163)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Jan 12, 2026 at 11:31 UTC
ca81811bfeca05f3104f2de7c58dd6d57d54472d
2 files changed
+253
-5
vllm/v1/worker/gpu/buffer_utils.py
+11
-5
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Sequence
3
+from collections.abc import Iterable, Sequence
4
5
import numpy as np
6
import torch
@@ -92,9 +92,10 @@ class StagedWriteTensor:
92
max_concurrency: int = 2,
93
uva_instead_of_gpu: bool = False,
94
):
95
- if dtype not in [torch.int32, torch.int64]:
95
+ supported_dtypes = [torch.int32, torch.int64, torch.float32]
96
+ if dtype not in supported_dtypes:
97
raise ValueError(
97
- f"Unsupported dtype {dtype}: should be either int32 or int64"
98
+ f"Unsupported dtype {dtype}: should be one of {supported_dtypes}"
99
)
100
self.num_rows = size if isinstance(size, int) else size[0]
101
self.dtype = dtype
@@ -111,7 +112,7 @@ class StagedWriteTensor:
112
113
self._staged_write_indices: list[int] = []
114
self._staged_write_starts: list[int] = []
114
- self._staged_write_contents: list[int] = []
115
+ self._staged_write_contents: list[int | float] = []
116
self._staged_write_cu_lens: list[int] = []
117
118
self.write_indices = UvaBufferPool(
@@ -128,7 +129,12 @@ class StagedWriteTensor:
129
self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency
130
)
131
131
- def stage_write(self, index: int, start: int, x: list[int]) -> None:
132
+ def stage_write(
133
+ self,
134
+ index: int,
135
+ start: int,
136
+ x: Iterable[int] | Iterable[float],
137
+ ) -> None:
138
assert index >= 0
139
assert start >= 0
140
if not x:
vllm/v1/worker/gpu/sample/logit_bias.py
new
+242
@@ -0,0 +1,242 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+from vllm.sampling_params import SamplingParams
6
+from vllm.triton_utils import tl, triton
7
+from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
8
+
9
+MAX_NUM_ALLOWED_TOKEN_IDS = 1024
10
+MAX_NUM_LOGIT_BIAS_TOKENS = 1024
11
+MAX_NUM_STOP_TOKEN_IDS = 128
12
+
13
+
14
+class LogitBiasState:
15
+ def __init__(
16
+ self,
17
+ max_num_reqs: int,
18
+ device: torch.device,
19
+ ):
20
+ self.max_num_reqs = max_num_reqs
21
+
22
+ # Allowed token IDs.
23
+ self.num_allowed_token_ids = UvaBackedTensor(
24
+ self.max_num_reqs, dtype=torch.int32
25
+ )
26
+ self.allowed_token_ids = StagedWriteTensor(
27
+ (self.max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS),
28
+ dtype=torch.int32,
29
+ device=device,
30
+ )
31
+ # Logit bias.
32
+ self.num_logit_bias = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
33
+ self.logit_bias_token_ids = StagedWriteTensor(
34
+ (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
35
+ dtype=torch.int32,
36
+ device=device,
37
+ )
38
+ self.logit_bias = StagedWriteTensor(
39
+ (self.max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
40
+ dtype=torch.float32,
41
+ device=device,
42
+ )
43
+ # Min tokens.
44
+ self.min_lens = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
45
+ self.num_stop_token_ids = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
46
+ self.stop_token_ids = StagedWriteTensor(
47
+ (self.max_num_reqs, MAX_NUM_STOP_TOKEN_IDS),
48
+ dtype=torch.int32,
49
+ device=device,
50
+ )
51
+
52
+ def add_request(
53
+ self,
54
+ req_idx: int,
55
+ prompt_len: int,
56
+ sampling_params: SamplingParams,
57
+ ) -> None:
58
+ # Allowed token IDs.
59
+ allowed_token_ids = sampling_params.allowed_token_ids
60
+ if allowed_token_ids:
61
+ num_allowed_token_ids = len(allowed_token_ids)
62
+ if num_allowed_token_ids > MAX_NUM_ALLOWED_TOKEN_IDS:
63
+ raise ValueError(
64
+ f"Too many allowed token IDs: {num_allowed_token_ids}. "
65
+ f"The max size is {MAX_NUM_ALLOWED_TOKEN_IDS}."
66
+ )
67
+ self.num_allowed_token_ids.np[req_idx] = num_allowed_token_ids
68
+ self.allowed_token_ids.stage_write(req_idx, 0, allowed_token_ids)
69
+ else:
70
+ self.num_allowed_token_ids.np[req_idx] = 0
71
+
72
+ # Logit bias.
73
+ logit_bias = sampling_params.logit_bias
74
+ if logit_bias:
75
+ num_logit_bias = len(logit_bias)
76
+ if num_logit_bias > MAX_NUM_LOGIT_BIAS_TOKENS:
77
+ raise ValueError(
78
+ f"Too many logit bias tokens: {num_logit_bias}. "
79
+ f"The max size is {MAX_NUM_LOGIT_BIAS_TOKENS}."
80
+ )
81
+ self.num_logit_bias.np[req_idx] = num_logit_bias
82
+ self.logit_bias_token_ids.stage_write(req_idx, 0, logit_bias.keys())
83
+ self.logit_bias.stage_write(req_idx, 0, logit_bias.values())
84
+ else:
85
+ self.num_logit_bias.np[req_idx] = 0
86
+
87
+ # Min tokens.
88
+ min_tokens = sampling_params.min_tokens
89
+ min_len = prompt_len + min_tokens
90
+ self.min_lens.np[req_idx] = min_len
91
+ stop_token_ids = sampling_params.all_stop_token_ids
92
+ if stop_token_ids:
93
+ num_stop_token_ids = len(stop_token_ids)
94
+ if num_stop_token_ids > MAX_NUM_STOP_TOKEN_IDS:
95
+ raise ValueError(
96
+ f"Too many stop tokens: {num_stop_token_ids}. "
97
+ f"The max size is {MAX_NUM_STOP_TOKEN_IDS}."
98
+ )
99
+ self.num_stop_token_ids.np[req_idx] = num_stop_token_ids
100
+ self.stop_token_ids.stage_write(req_idx, 0, stop_token_ids)
101
+ else:
102
+ self.num_stop_token_ids.np[req_idx] = 0
103
+
104
+ def apply_staged_writes(self) -> None:
105
+ self.num_allowed_token_ids.copy_to_uva()
106
+ self.allowed_token_ids.apply_write()
107
+
108
+ self.num_logit_bias.copy_to_uva()
109
+ self.logit_bias_token_ids.apply_write()
110
+ self.logit_bias.apply_write()
111
+
112
+ self.min_lens.copy_to_uva()
113
+ self.num_stop_token_ids.copy_to_uva()
114
+ self.stop_token_ids.apply_write()
115
+
116
+ def apply_logit_bias(
117
+ self,
118
+ logits: torch.Tensor,
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,)](
132
+ logits,
133
+ logits.stride(0),
134
+ vocab_size,
135
+ 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),
144
+ 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,
151
+ )
152
+
153
+
154
+@triton.jit
155
+def _bias_kernel(
156
+ logits_ptr,
157
+ logits_stride,
158
+ vocab_size,
159
+ idx_mapping_ptr,
160
+ # Allowed token IDs.
161
+ num_allowed_token_ids_ptr,
162
+ allowed_token_ids_ptr,
163
+ allowed_token_ids_stride,
164
+ # Logit bias.
165
+ num_logit_bias_ptr,
166
+ bias_token_ids_ptr,
167
+ bias_token_ids_stride,
168
+ bias_ptr,
169
+ bias_stride,
170
+ # Min tokens.
171
+ pos_ptr,
172
+ min_lens_ptr,
173
+ num_stop_token_ids_ptr,
174
+ stop_token_ids_ptr,
175
+ stop_token_ids_stride,
176
+ BLOCK_SIZE: tl.constexpr,
177
+ LOGITS_BLOCK_SIZE: tl.constexpr,
178
+):
179
+ batch_idx = tl.program_id(0)
180
+ req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
181
+
182
+ block = tl.arange(0, BLOCK_SIZE)
183
+
184
+ # Allowed token IDs.
185
+ num_allowed_token_ids = tl.load(num_allowed_token_ids_ptr + req_state_idx)
186
+ if num_allowed_token_ids > 0:
187
+ block = tl.arange(0, BLOCK_SIZE)
188
+ mask = block < num_allowed_token_ids
189
+
190
+ # Save logits for allowed token IDs.
191
+ allowed_token_ids = tl.load(
192
+ allowed_token_ids_ptr + req_state_idx * allowed_token_ids_stride + block,
193
+ mask=mask,
194
+ )
195
+ logits = tl.load(
196
+ logits_ptr + batch_idx * logits_stride + allowed_token_ids, mask=mask
197
+ )
198
+
199
+ # Set logits to -inf for all tokens.
200
+ for i in range(0, vocab_size, LOGITS_BLOCK_SIZE):
201
+ offset = i + tl.arange(0, LOGITS_BLOCK_SIZE)
202
+ tl.store(
203
+ logits_ptr + batch_idx * logits_stride + offset,
204
+ -float("inf"),
205
+ mask=offset < vocab_size,
206
+ )
207
+
208
+ # Restore logits for allowed token IDs.
209
+ tl.store(
210
+ logits_ptr + batch_idx * logits_stride + allowed_token_ids,
211
+ logits,
212
+ mask=mask,
213
+ )
214
+
215
+ # Logit bias.
216
+ num_logit_bias = tl.load(num_logit_bias_ptr + req_state_idx)
217
+ if num_logit_bias > 0:
218
+ mask = block < num_logit_bias
219
+ token_ids = tl.load(
220
+ bias_token_ids_ptr + req_state_idx * bias_token_ids_stride + block,
221
+ mask=mask,
222
+ )
223
+ bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
224
+ logits = tl.load(logits_ptr + batch_idx * logits_stride + token_ids, mask=mask)
225
+ logits += bias
226
+ tl.store(logits_ptr + batch_idx * logits_stride + token_ids, logits, mask=mask)
227
+
228
+ # Apply min tokens.
229
+ num_stop_token_ids = tl.load(num_stop_token_ids_ptr + req_state_idx)
230
+ pos = tl.load(pos_ptr + batch_idx)
231
+ min_len = tl.load(min_lens_ptr + req_state_idx)
232
+ if num_stop_token_ids > 0 and pos < min_len:
233
+ mask = block < num_stop_token_ids
234
+ stop_token_ids = tl.load(
235
+ stop_token_ids_ptr + req_state_idx * stop_token_ids_stride + block,
236
+ mask=mask,
237
+ )
238
+ tl.store(
239
+ logits_ptr + batch_idx * logits_stride + stop_token_ids,
240
+ -float("inf"),
241
+ mask=mask,
242
+ )