@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 349 lines 12.3 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from typing import TYPE_CHECKING
4
5 import numpy as np
6 import torch
7
8 from vllm.sampling_params import SamplingParams
9 from vllm.triton_utils import tl, triton
10 from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
11 from vllm.v1.worker.gpu.sample.logits_processor.interface import (
12 LogitsContext,
13 LogitsProcessor,
14 LogitsProcRequestState,
15 )
16
17 if TYPE_CHECKING:
18 from vllm.config import VllmConfig
19
20 MAX_NUM_ALLOWED_TOKEN_IDS = 1024
21 MAX_NUM_LOGIT_BIAS_TOKENS = 1024
22 MAX_NUM_STOP_TOKEN_IDS = 128
23
24
25 class LogitBiasState(LogitsProcessor):
26 def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState):
27 self.prompt_len_np = req_states.prompt_len.np
28 max_num_reqs = req_states.max_num_reqs
29 device = req_states.device
30
31 # Allowed token IDs.
32 self.num_allowed_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
33 self.allowed_token_ids = StagedWriteTensor(
34 (max_num_reqs, MAX_NUM_ALLOWED_TOKEN_IDS), dtype=torch.int32, device=device
35 )
36 # Logit bias.
37 self.num_logit_bias = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
38 self.logit_bias_token_ids = StagedWriteTensor(
39 (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS), dtype=torch.int32, device=device
40 )
41 self.logit_bias = StagedWriteTensor(
42 (max_num_reqs, MAX_NUM_LOGIT_BIAS_TOKENS),
43 dtype=torch.float32,
44 device=device,
45 )
46 # Min tokens.
47 self.min_lens = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
48 self.num_stop_token_ids = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
49 self.restore_when_all_masked = UvaBackedTensor(max_num_reqs, dtype=torch.int32)
50 self.stop_token_ids = StagedWriteTensor(
51 (max_num_reqs, MAX_NUM_STOP_TOKEN_IDS), dtype=torch.int32, device=device
52 )
53
54 # Using any of the above.
55 self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool)
56
57 def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
58 prompt_len = int(self.prompt_len_np[req_idx])
59 # Using any logit bias.
60 use_logit_bias = False
61
62 # Allowed token IDs.
63 allowed_token_ids = sampling_params.allowed_token_ids
64 if allowed_token_ids:
65 num_allowed_token_ids = len(allowed_token_ids)
66 if num_allowed_token_ids > MAX_NUM_ALLOWED_TOKEN_IDS:
67 raise ValueError(
68 f"Too many allowed token IDs: {num_allowed_token_ids}. "
69 f"The max size is {MAX_NUM_ALLOWED_TOKEN_IDS}."
70 )
71 self.num_allowed_token_ids.np[req_idx] = num_allowed_token_ids
72 self.allowed_token_ids.stage_write(req_idx, 0, allowed_token_ids)
73 use_logit_bias = True
74 else:
75 self.num_allowed_token_ids.np[req_idx] = 0
76
77 # Logit bias.
78 logit_bias = sampling_params.logit_bias
79 if logit_bias:
80 num_logit_bias = len(logit_bias)
81 if num_logit_bias > MAX_NUM_LOGIT_BIAS_TOKENS:
82 raise ValueError(
83 f"Too many logit bias tokens: {num_logit_bias}. "
84 f"The max size is {MAX_NUM_LOGIT_BIAS_TOKENS}."
85 )
86 self.num_logit_bias.np[req_idx] = num_logit_bias
87 self.logit_bias_token_ids.stage_write(req_idx, 0, logit_bias.keys())
88 self.logit_bias.stage_write(req_idx, 0, logit_bias.values())
89 use_logit_bias = True
90 else:
91 self.num_logit_bias.np[req_idx] = 0
92
93 # Min tokens.
94 min_tokens = sampling_params.min_tokens
95 min_len = prompt_len + min_tokens
96 self.min_lens.np[req_idx] = min_len
97 stop_token_ids = sampling_params.all_stop_token_ids
98 if min_tokens > 0 and stop_token_ids:
99 num_stop_token_ids = len(stop_token_ids)
100 if num_stop_token_ids > MAX_NUM_STOP_TOKEN_IDS:
101 raise ValueError(
102 f"Too many stop tokens: {num_stop_token_ids}. "
103 f"The max size is {MAX_NUM_STOP_TOKEN_IDS}."
104 )
105 self.num_stop_token_ids.np[req_idx] = num_stop_token_ids
106 self.stop_token_ids.stage_write(req_idx, 0, stop_token_ids)
107 self.restore_when_all_masked.np[req_idx] = int(
108 sampling_params.structured_outputs is not None
109 )
110 use_logit_bias = True
111 else:
112 self.num_stop_token_ids.np[req_idx] = 0
113 self.restore_when_all_masked.np[req_idx] = 0
114
115 self.use_logit_bias[req_idx] = use_logit_bias
116 return use_logit_bias
117
118 def apply_staged_writes(self) -> None:
119 self.num_allowed_token_ids.copy_to_uva()
120 self.allowed_token_ids.apply_write()
121
122 self.num_logit_bias.copy_to_uva()
123 self.logit_bias_token_ids.apply_write()
124 self.logit_bias.apply_write()
125
126 self.min_lens.copy_to_uva()
127 self.num_stop_token_ids.copy_to_uva()
128 self.restore_when_all_masked.copy_to_uva()
129 self.stop_token_ids.apply_write()
130
131 def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
132 if not np.any(self.use_logit_bias[ctx.idx_mapping_np]):
133 # No request uses logit bias. Skip the kernel launch.
134 return logits
135
136 enable_stop_token_restore = bool(
137 np.any(self.restore_when_all_masked.np[ctx.idx_mapping_np])
138 )
139
140 apply_logit_bias(
141 logits,
142 ctx.expanded_idx_mapping,
143 ctx.pos,
144 self.num_allowed_token_ids.gpu,
145 self.allowed_token_ids.gpu,
146 self.num_logit_bias.gpu,
147 self.logit_bias_token_ids.gpu,
148 self.logit_bias.gpu,
149 self.min_lens.gpu,
150 self.num_stop_token_ids.gpu,
151 self.restore_when_all_masked.gpu,
152 self.stop_token_ids.gpu,
153 enable_stop_token_restore,
154 )
155 return logits
156
157
158 @triton.jit
159 def _bias_kernel(
160 logits_ptr,
161 logits_stride,
162 vocab_size,
163 expanded_idx_mapping_ptr,
164 # Allowed token IDs.
165 num_allowed_token_ids_ptr,
166 allowed_token_ids_ptr,
167 allowed_token_ids_stride,
168 # Logit bias.
169 num_logit_bias_ptr,
170 bias_token_ids_ptr,
171 bias_token_ids_stride,
172 bias_ptr,
173 bias_stride,
174 # Min tokens.
175 pos_ptr,
176 min_lens_ptr,
177 num_stop_token_ids_ptr,
178 restore_when_all_masked_ptr,
179 stop_token_ids_ptr,
180 stop_token_ids_stride,
181 BLOCK_SIZE: tl.constexpr,
182 LOGITS_BLOCK_SIZE: tl.constexpr,
183 CHECK_ALL_MASKED_ROWS: tl.constexpr,
184 ):
185 token_idx = tl.program_id(0).to(tl.int64)
186 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)
187
188 block = tl.arange(0, BLOCK_SIZE)
189
190 # Allowed token IDs.
191 num_allowed_token_ids = tl.load(num_allowed_token_ids_ptr + req_state_idx)
192 if num_allowed_token_ids > 0:
193 block = tl.arange(0, BLOCK_SIZE)
194 mask = block < num_allowed_token_ids
195
196 # Save logits for allowed token IDs.
197 allowed_token_ids = tl.load(
198 allowed_token_ids_ptr + req_state_idx * allowed_token_ids_stride + block,
199 mask=mask,
200 )
201 logits = tl.load(
202 logits_ptr + token_idx * logits_stride + allowed_token_ids, mask=mask
203 )
204
205 tl.debug_barrier() # save must read original logits before the -inf overwrite
206
207 # Set logits to -inf for all tokens.
208 for i in range(0, vocab_size, LOGITS_BLOCK_SIZE):
209 offset = i + tl.arange(0, LOGITS_BLOCK_SIZE)
210 tl.store(
211 logits_ptr + token_idx * logits_stride + offset,
212 -float("inf"),
213 mask=offset < vocab_size,
214 )
215
216 tl.debug_barrier() # -inf overwrite must finish before restoring saved logits
217
218 # Restore logits for allowed token IDs.
219 tl.store(
220 logits_ptr + token_idx * logits_stride + allowed_token_ids,
221 logits,
222 mask=mask,
223 )
224
225 # Logit bias.
226 num_logit_bias = tl.load(num_logit_bias_ptr + req_state_idx)
227 if num_logit_bias > 0:
228 mask = block < num_logit_bias
229 token_ids = tl.load(
230 bias_token_ids_ptr + req_state_idx * bias_token_ids_stride + block,
231 mask=mask,
232 )
233 bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)
234 logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask)
235 logits += bias
236 tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)
237
238 # Apply min tokens.
239 num_stop_token_ids = tl.load(num_stop_token_ids_ptr + req_state_idx)
240 pos = tl.load(pos_ptr + token_idx)
241 min_len = tl.load(min_lens_ptr + req_state_idx)
242 if num_stop_token_ids > 0 and pos + 1 < min_len:
243 mask = block < num_stop_token_ids
244 stop_token_ids = tl.load(
245 stop_token_ids_ptr + req_state_idx * stop_token_ids_stride + block,
246 mask=mask,
247 )
248 if CHECK_ALL_MASKED_ROWS:
249 should_restore_stop_logits = tl.load(
250 restore_when_all_masked_ptr + req_state_idx
251 )
252 if should_restore_stop_logits:
253 stop_logits = tl.load(
254 logits_ptr + token_idx * logits_stride + stop_token_ids,
255 mask=mask,
256 other=-float("inf"),
257 )
258
259 # Save must read original logits before the -inf overwrite.
260 tl.debug_barrier()
261
262 tl.store(
263 logits_ptr + token_idx * logits_stride + stop_token_ids,
264 -float("inf"),
265 mask=mask,
266 )
267
268 # The -inf overwrite must be visible to the row scan below.
269 tl.debug_barrier()
270
271 row_max = tl.full((), -float("inf"), tl.float32)
272 for i in range(0, vocab_size, LOGITS_BLOCK_SIZE):
273 offset = i + tl.arange(0, LOGITS_BLOCK_SIZE)
274 logits = tl.load(
275 logits_ptr + token_idx * logits_stride + offset,
276 mask=offset < vocab_size,
277 other=-float("inf"),
278 )
279 row_max = tl.maximum(row_max, tl.max(logits, axis=0))
280
281 if row_max == -float("inf"):
282 tl.store(
283 logits_ptr + token_idx * logits_stride + stop_token_ids,
284 stop_logits,
285 mask=mask
286 & (stop_logits > -float("inf"))
287 & (stop_logits < float("inf")),
288 )
289 else:
290 tl.store(
291 logits_ptr + token_idx * logits_stride + stop_token_ids,
292 -float("inf"),
293 mask=mask,
294 )
295 else:
296 tl.store(
297 logits_ptr + token_idx * logits_stride + stop_token_ids,
298 -float("inf"),
299 mask=mask,
300 )
301
302
303 def apply_logit_bias(
304 logits: torch.Tensor,
305 expanded_idx_mapping: torch.Tensor,
306 pos: torch.Tensor,
307 num_allowed_token_ids: torch.Tensor,
308 allowed_token_ids: torch.Tensor,
309 num_logit_bias: torch.Tensor,
310 logit_bias_token_ids: torch.Tensor,
311 logit_bias: torch.Tensor,
312 min_lens: torch.Tensor,
313 num_stop_token_ids: torch.Tensor,
314 restore_when_all_masked: torch.Tensor,
315 stop_token_ids: torch.Tensor,
316 check_all_masked_rows: bool = False,
317 ) -> None:
318 num_tokens, vocab_size = logits.shape
319 BLOCK_SIZE = triton.next_power_of_2(
320 max(
321 allowed_token_ids.shape[-1],
322 logit_bias_token_ids.shape[-1],
323 stop_token_ids.shape[-1],
324 )
325 )
326 LOGITS_BLOCK_SIZE = 8192
327 _bias_kernel[(num_tokens,)](
328 logits,
329 logits.stride(0),
330 vocab_size,
331 expanded_idx_mapping,
332 num_allowed_token_ids,
333 allowed_token_ids,
334 allowed_token_ids.stride(0),
335 num_logit_bias,
336 logit_bias_token_ids,
337 logit_bias_token_ids.stride(0),
338 logit_bias,
339 logit_bias.stride(0),
340 pos,
341 min_lens,
342 num_stop_token_ids,
343 restore_when_all_masked,
344 stop_token_ids,
345 stop_token_ids.stride(0),
346 BLOCK_SIZE=BLOCK_SIZE,
347 LOGITS_BLOCK_SIZE=LOGITS_BLOCK_SIZE,
348 CHECK_ALL_MASKED_ROWS=check_all_masked_rows,
349 )