@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 48fb076cbc
Fork 0

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 15, 03:30 UTC

[V1] LogitsProcessor programming model (#16728)

Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com> Signed-off-by: Andrew Feldman <afeldman@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>

afeldman-nm committed Jul 2, 2025 at 12:10 UTC 48fb076cbc651f655aae8ffdff6930907309b2f7
13 files changed +1403 -395
tests/v1/sample/test_logits_processors.py new
+626
@@ -0,0 +1,626 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import random
4 +from collections.abc import Callable
5 +from typing import NamedTuple, Optional, Union
6 +
7 +import numpy as np
8 +import pytest
9 +import torch
10 +
11 +from tests.v1.sample.utils import (LogitsprocsTestFakes, create_fake_logits,
12 + create_penalty_tensor,
13 + create_prompt_tokens_tensor,
14 + fake_apply_logitsprocs,
15 + fake_update_logitsprocs_state)
16 +from vllm.platforms import current_platform
17 +from vllm.sampling_params import SamplingParams
18 +from vllm.utils import is_pin_memory_available
19 +# yapf: disable
20 +from vllm.v1.sample.logits_processor import (BatchUpdate, BatchUpdateBuilder,
21 + LogitBiasLogitsProcessor,
22 + LogitsProcessor,
23 + MinPLogitsProcessor,
24 + MinTokensLogitsProcessor,
25 + MoveDirectionality,
26 + init_builtin_logitsprocs)
27 +# yapf: enable
28 +from vllm.v1.sample.metadata import SamplingMetadata
29 +
30 +PIN_MEMORY_AVAILABLE = is_pin_memory_available()
31 +MAX_NUM_REQS = 256
32 +VOCAB_SIZE = 1024
33 +NUM_OUTPUT_TOKENS = 20
34 +CUDA_DEVICES = [
35 + f"{current_platform.device_type}:{i}"
36 + for i in range(1 if current_platform.device_count() == 1 else 2)
37 +]
38 +MAX_NUM_PROMPT_TOKENS = 64
39 +MIN_TOKENS_LEN_THRESHOLD = 5
40 +REQS_PER_LOGITPROC = 50
41 +STR_NO_LOGITPROC = "none"
42 +
43 +# LogitsProcessor subclass or "none"
44 +LogitprocType = Union[type[LogitsProcessor], str]
45 +
46 +
47 +class LogitsProcsRequestParams:
48 + """Encapsulates key params for a single request in a batch.
49 +
50 + Params can be customized based on the enabled logitproc
51 + """
52 + workload_index: int
53 + logitproc_type: LogitprocType # Logitproc enabled, specified by str id
54 + out_tokens: list[int] # Output tokens required for min tokens test
55 + params: SamplingParams # Settings customized for logitproc
56 +
57 + def __init__(self, workload_index: int, logitproc_type: LogitprocType):
58 + self.workload_index = workload_index
59 + self.logitproc_type = logitproc_type
60 + # Number of output tokens is randomly 0 or twice the min-tokens
61 + # threshold which will be used in testing. Output token values
62 + # don't matter *for these tests* so use 0 as a dummy value
63 + self.out_tokens = ([0] *
64 + (MIN_TOKENS_LEN_THRESHOLD * random.randint(0, 2)))
65 + self.params = _sampling_params_from_logitproc(logitproc_type)
66 +
67 + def __str__(self):
68 + """For debugging"""
69 + summ = ', '.join(f'{k}={v}' for k, v in vars(self).items())
70 + return f"MyClass({summ})"
71 +
72 +
73 +def _generate_fake_sampling_metadata(
74 + num_output_tokens: int,
75 + batch_size: int,
76 + vocab_size: int,
77 + device: torch.device,
78 +) -> SamplingMetadata:
79 + """Generate fake sampling metadata with fake logitsprocs"""
80 + output_token_ids: list[list[int]] = []
81 + prompt_token_ids: list[list[int]] = []
82 + for _ in range(batch_size):
83 + output_token_ids.append(
84 + np.random.randint(0, vocab_size, size=num_output_tokens).tolist())
85 + prompt_token_ids.append(
86 + np.random.randint(0,
87 + vocab_size,
88 + size=np.random.randint(
89 + 1, MAX_NUM_PROMPT_TOKENS)).tolist())
90 + logitsprocs = init_builtin_logitsprocs(
91 + pin_memory_available=PIN_MEMORY_AVAILABLE,
92 + max_num_reqs=MAX_NUM_REQS + 1,
93 + device=device)
94 +
95 + fake_sampling_metadata = SamplingMetadata(
96 + temperature=torch.full((batch_size, ), 0.0),
97 + all_greedy=True,
98 + all_random=False,
99 + top_p=None,
100 + top_k=None,
101 + generators={},
102 + max_num_logprobs=0,
103 + prompt_token_ids=create_prompt_tokens_tensor(prompt_token_ids,
104 + vocab_size, device),
105 + output_token_ids=output_token_ids,
106 + frequency_penalties=create_penalty_tensor(batch_size, 0.0, device),
107 + presence_penalties=create_penalty_tensor(batch_size, 0.0, device),
108 + repetition_penalties=create_penalty_tensor(batch_size, 1.0, device),
109 + no_penalties=True,
110 + allowed_token_ids_mask=None,
111 + bad_words_token_ids={},
112 + logitsprocs=logitsprocs)
113 + return fake_sampling_metadata
114 +
115 +
116 +def _generate_test_fakes(batch_size: int, device: str) -> LogitsprocsTestFakes:
117 + """Generate fake logits and sampling metadata"""
118 + fake_logits = create_fake_logits(batch_size, VOCAB_SIZE)
119 + # Create one dominant token per batch, to support min-p test
120 + for i in range(batch_size):
121 + fake_logits[i, 0] = 10.0 # High logit for first token
122 + fake_logits[i, 1:] = 1e-2 # Others remain low
123 + sampling_metadata = _generate_fake_sampling_metadata(
124 + NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))
125 + return LogitsprocsTestFakes(
126 + logits=fake_logits,
127 + sampling_metadata=sampling_metadata,
128 + )
129 +
130 +
131 +def _sampling_params_from_logitproc(
132 + logitproc_type: LogitprocType) -> SamplingParams:
133 + """Customize request SamplingParams for a specified logitproc"""
134 + # SamplingParams for req with no logitproc
135 + kwargs = {"min_p": 0.0, "logit_bias": None, "min_tokens": 0}
136 + if fxn := logitsprocs_test_mapping[logitproc_type].gen_request_fxn:
137 + fxn(kwargs)
138 + return SamplingParams(**kwargs)
139 +
140 +
141 +def _generate_mixed_logitsprocs_batch_params(
142 + reqs_per_logitproc: int,
143 + logitsprocs_types: list[str],
144 +) -> list[LogitsProcsRequestParams]:
145 + """Define key params for a batch of requests with a different
146 + logitproc enabled per request.
147 +
148 + The batch will have `reqs_per_logitproc` repeats for all
149 + `logitsprocs_types` under test, including the case where
150 + no logitsproc is enabled. The batch is randomly shuffled. The
151 + size of the batch is `reqs_per_logitproc` times
152 + `n = len(logitsprocs_types)`
153 +
154 + Args:
155 + reqs_per_logitproc: number of requests using each logitproc
156 + logitsprocs_types: logitsprocs under test
157 +
158 + Returns:
159 + List of per-request params which configure the engine for that request's
160 + enabled logitproc
161 + """
162 + batch_size = len(logitsprocs_types) * reqs_per_logitproc
163 + # Generate multiple repeats of key params for each logitproc;
164 + # apply random inverse permutation to the iteration
165 + # over logitsprocs, such that logitsprocs are shuffled.
166 + batch_perm = random.sample(range(batch_size), k=batch_size)
167 + return [
168 + LogitsProcsRequestParams(
169 + workload_index=idx,
170 + logitproc_type=logitsprocs_types[pdx // reqs_per_logitproc])
171 + for idx, pdx in enumerate(batch_perm)
172 + ]
173 +
174 +
175 +def _raise_error_invalid(
176 + msg_suffix: str,
177 + batch_index: int,
178 + request_params: LogitsProcsRequestParams,
179 + step_idx: int,
180 + err_cls: type[Exception] = ValueError,
181 +) -> None:
182 + raise err_cls(f"Validation failed for step={step_idx}, "
183 + f"batch_index={batch_index}, "
184 + f"workload_index={request_params.workload_index}, "
185 + f"req_params={request_params}. Reason: {msg_suffix}")
186 +
187 +
188 +def _logit_bias_params(kwargs: dict) -> None:
189 + """Logit bias config"""
190 + kwargs["logit_bias"] = {
191 + random.randint(0, VOCAB_SIZE - 1): random.choice([-0.1, 0.2])
192 + }
193 +
194 +
195 +def _logit_bias_validate(
196 + test_fakes: LogitsprocsTestFakes,
197 + persistent_batch: list[LogitsProcsRequestParams],
198 + logits_new: torch.Tensor,
199 + batch_index: int,
200 + request_params: LogitsProcsRequestParams,
201 + step_idx: int,
202 +) -> None:
203 + """Validate logit bias logitproc applied correctly"""
204 + logit_bias = request_params.params.logit_bias
205 + logits_old = (
206 + test_fakes.logits[persistent_batch[batch_index].workload_index].cpu())
207 + logits_new = logits_new[batch_index].cpu()
208 + for token_id in range(VOCAB_SIZE):
209 + logit_old_value = logits_old[token_id]
210 + logit_new_value = logits_new[token_id]
211 + if token_id in logit_bias:
212 + bias_value = logit_bias[token_id]
213 + exp_value = bias_value + logit_old_value
214 + if logit_new_value != pytest.approx(exp_value):
215 + _raise_error_invalid(msg_suffix=(
216 + f"Biased token {token_id} logit value {logit_new_value} "
217 + f"does not match expected value {exp_value} "
218 + f"given bias {bias_value}"),
219 + batch_index=batch_index,
220 + request_params=request_params,
221 + step_idx=step_idx)
222 +
223 + else:
224 + if logit_new_value != pytest.approx(logit_old_value):
225 + _raise_error_invalid(msg_suffix=(
226 + f"Unbiased token {token_id} logit value {logit_new_value} "
227 + f"does not match expected value {logit_old_value}"),
228 + batch_index=batch_index,
229 + request_params=request_params,
230 + step_idx=step_idx)
231 +
232 +
233 +def _min_p_params(kwargs: dict) -> None:
234 + """Min-p logitproc config"""
235 + kwargs["min_p"] = 0.1
236 +
237 +
238 +def _min_p_validate(
239 + test_fakes: LogitsprocsTestFakes,
240 + persistent_batch: list[LogitsProcsRequestParams],
241 + logits_new: torch.Tensor,
242 + batch_index: int,
243 + request_params: LogitsProcsRequestParams,
244 + step_idx: int,
245 +) -> None:
246 + """Validate min-p logitproc applied correctly"""
247 + for token_id in range(VOCAB_SIZE):
248 + logits_for_token = logits_new[batch_index][token_id]
249 + if token_id == 0:
250 + # Dominant token should always be unmasked
251 + if logits_for_token == -float("inf"):
252 + _raise_error_invalid(
253 + msg_suffix="Invalid: dominant token 0 masked (-inf)",
254 + batch_index=batch_index,
255 + request_params=request_params,
256 + step_idx=step_idx)
257 + else:
258 + if request_params.params.min_p > 0.0:
259 + # Non-dominant tokens should be masked when min_p > 0
260 + if logits_for_token != -float("inf"):
261 + _raise_error_invalid(
262 + msg_suffix=
263 + f"Invalid: non-dominant token {token_id} not masked",
264 + batch_index=batch_index,
265 + request_params=request_params,
266 + step_idx=step_idx)
267 + else:
268 + # No masking when min_p is 0
269 + if logits_for_token == -float("inf"):
270 + _raise_error_invalid(
271 + msg_suffix=
272 + f"Invalid: token {token_id} masked when min_p=0.0",
273 + batch_index=batch_index,
274 + request_params=request_params,
275 + step_idx=step_idx)
276 +
277 +
278 +def _min_tokens_params(kwargs: dict) -> None:
279 + """Min-tokens logitproc config"""
280 + kwargs["min_tokens"] = MIN_TOKENS_LEN_THRESHOLD
281 + kwargs["stop_token_ids"] = [
282 + np.random.randint(0, VOCAB_SIZE - 1)
283 + for _ in range(np.random.randint(0, VOCAB_SIZE))
284 + ]
285 +
286 +
287 +def _min_tokens_validate(
288 + test_fakes: LogitsprocsTestFakes,
289 + persistent_batch: list[LogitsProcsRequestParams],
290 + logits_new: torch.Tensor,
291 + batch_index: int,
292 + request_params: LogitsProcsRequestParams,
293 + step_idx: int,
294 +) -> None:
295 + """Validate min-tokens logitsproc applied correctly"""
296 + ref_num_out_tokens = len(request_params.out_tokens)
297 + min_reached = ref_num_out_tokens >= MIN_TOKENS_LEN_THRESHOLD
298 + ref_all_stop_token_ids = request_params.params.all_stop_token_ids
299 + mt_lp: MinTokensLogitsProcessor = next(
300 + test_fakes.get_logitsprocs_by_cls(MinTokensLogitsProcessor))
301 + assert isinstance(mt_lp, MinTokensLogitsProcessor)
302 + min_tok = mt_lp.min_toks.get(batch_index, None)
303 +
304 + # Validate min-token logits processor state
305 + if min_tok:
306 + (_, out_tok, all_stop_token_ids) = min_tok
307 + num_out_tokens = len(out_tok)
308 + if num_out_tokens != ref_num_out_tokens:
309 + _raise_error_invalid(msg_suffix=(
310 + "Number of output tokens in min-token logit processor "
311 + f"request metadata ({num_out_tokens}) does not match "
312 + f"reference ({ref_num_out_tokens})."),
313 + batch_index=batch_index,
314 + request_params=request_params,
315 + step_idx=step_idx)
316 + if ref_all_stop_token_ids != all_stop_token_ids:
317 + _raise_error_invalid(msg_suffix=(
318 + "Stop token ids do not match reference; all_stop_token_ids: "
319 + f"{sorted(all_stop_token_ids)}, ref_all_stop_token_ids: "
320 + f"{sorted(ref_all_stop_token_ids)}"),
321 + batch_index=batch_index,
322 + request_params=request_params,
323 + step_idx=step_idx)
324 + if min_reached:
325 + _raise_error_invalid(msg_suffix=(
326 + "Expected min-tokens request with min reached, but batch "
327 + "index is recognized by min-tokens logits processor."),
328 + batch_index=batch_index,
329 + request_params=request_params,
330 + step_idx=step_idx,
331 + err_cls=RuntimeError)
332 +
333 + elif not min_reached:
334 + _raise_error_invalid(msg_suffix=(
335 + "Expected min-tokens request with min not reached, but batch "
336 + "index is not recognized by min-tokens logits processor."),
337 + batch_index=batch_index,
338 + request_params=request_params,
339 + step_idx=step_idx,
340 + err_cls=RuntimeError)
341 +
342 + # Validate min-token logits
343 + for token_id in range(VOCAB_SIZE):
344 + logits_for_token = logits_new[batch_index][token_id]
345 + if token_id in ref_all_stop_token_ids and not min_reached:
346 + if logits_for_token != -float("inf"):
347 + _raise_error_invalid(
348 + msg_suffix=(f"Token {token_id} is a stop token and "
349 + "the sequence has not reached min length, "
350 + "but the token is not masked "
351 + f"(logit={logits_for_token})"),
352 + batch_index=batch_index,
353 + request_params=request_params,
354 + step_idx=step_idx)
355 + else:
356 + if logits_for_token == -float("inf"):
357 + _raise_error_invalid(
358 + msg_suffix=(f"Token {token_id} should not be masked but "
359 + f"is (output len={ref_num_out_tokens})"),
360 + batch_index=batch_index,
361 + request_params=request_params,
362 + step_idx=step_idx)
363 +
364 +
365 +def _none_validate(
366 + test_fakes: LogitsprocsTestFakes,
367 + persistent_batch: list[LogitsProcsRequestParams],
368 + logits_new: torch.Tensor,
369 + batch_index: int,
370 + request_params: LogitsProcsRequestParams,
371 + step_idx: int,
372 +) -> None:
373 + """Validate that no logits processors are applied"""
374 + logits = (
375 + test_fakes.logits[persistent_batch[batch_index].workload_index].cpu())
376 + ref_logits = logits_new[batch_index]
377 + if not torch.all(ref_logits == logits):
378 + mismatch_toks = (ref_logits
379 + != logits).nonzero(as_tuple=True)[0].tolist()
380 + mismatch_strs = []
381 + for token in mismatch_toks:
382 + val = float(logits[token])
383 + ref_val = float(ref_logits[token])
384 + mismatch_strs.append(f"({token=},{val=},{ref_val=})")
385 + _raise_error_invalid(msg_suffix=(
386 + f"Unexpected modification of logits: {','.join(mismatch_strs)}"),
387 + batch_index=batch_index,
388 + request_params=request_params,
389 + step_idx=step_idx)
390 +
391 +
392 +class LogitsprocTestHelpers(NamedTuple):
393 + """Supports setting up and validating logitsprocs unit tests."""
394 + eval_fxn: Callable
395 + gen_request_fxn: Optional[Callable] = None
396 +
397 +
398 +logitsprocs_test_mapping = {
399 + STR_NO_LOGITPROC:
400 + LogitsprocTestHelpers(eval_fxn=_none_validate),
401 + LogitBiasLogitsProcessor:
402 + LogitsprocTestHelpers(gen_request_fxn=_logit_bias_params,
403 + eval_fxn=_logit_bias_validate),
404 + MinPLogitsProcessor:
405 + LogitsprocTestHelpers(gen_request_fxn=_min_p_params,
406 + eval_fxn=_min_p_validate),
407 + MinTokensLogitsProcessor:
408 + LogitsprocTestHelpers(gen_request_fxn=_min_tokens_params,
409 + eval_fxn=_min_tokens_validate),
410 +}
411 +
412 +
413 +def _get_test_cases() -> list[list[str]]:
414 + """Each test case is a set of logitsprocs"""
415 + logitsprocs_types = list(logitsprocs_test_mapping.keys())
416 + return [[STR_NO_LOGITPROC]] + [[logitproc_type, STR_NO_LOGITPROC]
417 + for logitproc_type in logitsprocs_types
418 + if logitproc_type != STR_NO_LOGITPROC
419 + ] + [logitsprocs_types]
420 +
421 +
422 +def _generate_fake_step_update(
423 + persistent_batch: list[LogitsProcsRequestParams],
424 + workload_params: list[LogitsProcsRequestParams],
425 + wdx: int,
426 + batch_update_builder: BatchUpdateBuilder,
427 +) -> tuple[Optional[BatchUpdate], int, int]:
428 + batch_size = len(persistent_batch)
429 + workload_size = len(workload_params)
430 + workload_reqs_remaining = workload_size - wdx
431 + max_add_remove_per_step = max(1, int(0.2 * workload_size))
432 +
433 + # 50% of steps: add no reqs
434 + # Other 50%: add a limited number of reqs (less than the number
435 + # of workload reqs remaining, less than an arbitrary max)
436 + # If no workload reqs remain: 100% of steps have 0 adds
437 + num_step_add = random.choice([
438 + 0,
439 + random.randint(1, min(max_add_remove_per_step,
440 + workload_reqs_remaining))
441 + ]) if workload_reqs_remaining else 0
442 +
443 + # 50% of steps: remove no requests
444 + # Other 50%: remove a limited number of reqs (less than the number
445 + # persistent batch reqs remaining, less than an arbitrary max)
446 + # If persistent batch is empty: 100% of steps have 0 removals until
447 + # more requests are added. Assume that removed requests are always
448 + # drawn from the current batch, before new adds
449 + num_step_remove = random.choice([
450 + 0, random.randint(1, min(max_add_remove_per_step, batch_size))
451 + ]) if batch_size else 0
452 +
453 + num_step_add_replace = min(num_step_add, num_step_remove)
454 +
455 + # Generate fake removed request indices drawn from persistent batch indices
456 + for removal in random.sample(range(batch_size), num_step_remove):
457 + batch_update_builder.removed_append(removal)
458 +
459 + # Get added requests from workload
460 + for add_req_params in workload_params[wdx:(wdx + num_step_add_replace)]:
461 + # Replace as many removed requests as possible with added requests
462 + add_remove_idx = batch_update_builder.pop_removed()
463 + batch_update_builder.added.append(
464 + (add_remove_idx, add_req_params.params, add_req_params.out_tokens))
465 + persistent_batch[add_remove_idx] = add_req_params
466 +
467 + # Append remaining added requests to end of batch
468 + add_reqs_append = workload_params[(wdx +
469 + num_step_add_replace):(wdx +
470 + num_step_add)]
471 + batch_update_builder.added.extend([
472 + (adx + batch_size, add_req_params.params, add_req_params.out_tokens)
473 + for adx, add_req_params in enumerate(add_reqs_append)
474 + ])
475 + persistent_batch.extend(add_reqs_append)
476 + pre_condense_batch_size = len(persistent_batch)
477 + wdx += num_step_add # Update workload offset
478 +
479 + # Simulate condensing persistent batch
480 + last_nonempty_index = pre_condense_batch_size - 1
481 + condensed_to_idxs = set()
482 + while batch_update_builder.removed:
483 + if (last_nonempty_index in batch_update_builder.removed
484 + or last_nonempty_index in condensed_to_idxs):
485 + last_nonempty_index -= 1
486 + continue
487 + # last_nonempty_index is the highest persistent batch index that was
488 + # not removed
489 + first_empty_index = batch_update_builder.peek_removed()
490 + assert first_empty_index is not None
491 + if first_empty_index > last_nonempty_index:
492 + break
493 + # first_empty_index is the lowest removed persistent batch index
494 + # that is less than last_nonempty_index
495 + #
496 + # move last_nonempty_index -> first_empty_index
497 + batch_update_builder.pop_removed()
498 + condensed_to_idxs.add(first_empty_index)
499 + persistent_batch[first_empty_index] = persistent_batch[
500 + last_nonempty_index]
501 + batch_update_builder.moved.append(
502 + (last_nonempty_index, first_empty_index,
503 + MoveDirectionality.UNIDIRECTIONAL))
504 +
505 + last_nonempty_index -= 1
506 +
507 + # Now removed requests & gaps left by non-removed requests that got
508 + # moved downward are grouped consecutively in the upper indices of
509 + # the persistent batch. Truncate them to get condensed persistent batch
510 + condensed_batch_size = batch_size + num_step_add - num_step_remove
511 + persistent_batch[:] = persistent_batch[0:condensed_batch_size]
512 +
513 + if condensed_batch_size > 1:
514 + # Simulate arbitrary reorder_batch() in the kernel backend
515 + # Generate a random number k of non-overlapping swap tuples
516 + k = random.randint(0, condensed_batch_size // 2)
517 + idxs = list(range(condensed_batch_size))
518 + random.shuffle(idxs)
519 + swaps = [
520 + tuple(sorted([idxs[2 * i], idxs[2 * i + 1]])) for i in range(k)
521 + ]
522 + batch_update_builder.moved.extend([
523 + (sw[0], sw[1], MoveDirectionality.SWAP) for sw in swaps
524 + ])
525 + for adx, bdx in swaps:
526 + persistent_batch[adx], persistent_batch[bdx] = persistent_batch[
527 + bdx], persistent_batch[adx]
528 +
529 + return (batch_update_builder.get_and_reset(condensed_batch_size), wdx,
530 + workload_size - wdx)
531 +
532 +
533 +def _assert_valid(
534 + batch_size: int,
535 + persistent_batch: list[LogitsProcsRequestParams],
536 + test_fakes: LogitsprocsTestFakes,
537 + slice_idxs: list[int],
538 + logits_w_lp: torch.Tensor,
539 + step_idx: int,
540 +) -> None:
541 + if not slice_idxs:
542 + # Trivial case of empty persistent batch
543 + assert len(persistent_batch) == 0
544 + if logits_w_lp.shape[0] != 0:
545 + raise ValueError("Fake persistent batch is empty but logitsprocs "
546 + f"output batch has shape {logits_w_lp.shape}")
547 + return
548 +
549 + # Validate logits for each fake request
550 + for batch_index in range(batch_size):
551 + request_params = persistent_batch[batch_index]
552 + # Invoke the appropriate validation function for
553 + # the logitproc employed by this request
554 + fxn = logitsprocs_test_mapping[request_params.logitproc_type].eval_fxn
555 + fxn(test_fakes=test_fakes,
556 + persistent_batch=persistent_batch,
557 + logits_new=logits_w_lp,
558 + batch_index=batch_index,
559 + request_params=request_params,
560 + step_idx=step_idx)
561 +
562 +
563 +@pytest.mark.parametrize("device", CUDA_DEVICES)
564 +@pytest.mark.parametrize("reqs_per_logitproc", [REQS_PER_LOGITPROC])
565 +@pytest.mark.parametrize("logitsprocs_under_test", _get_test_cases())
566 +def test_logitsprocs(device: str, reqs_per_logitproc: int,
567 + logitsprocs_under_test: list[str]):
568 + random.seed(40)
569 + torch.set_default_device(device)
570 +
571 + # Define a shuffled batch of requests which individually use a different
572 + # logitproc, or no logitproc at all
573 + workload_params = _generate_mixed_logitsprocs_batch_params(
574 + reqs_per_logitproc=reqs_per_logitproc,
575 + logitsprocs_types=logitsprocs_under_test)
576 + workload_size = len(workload_params)
577 +
578 + # Create fake test data structures for testing.
579 + test_fakes = _generate_test_fakes(workload_size, device)
580 +
581 + wdx = 0 # Next request index in workload to add
582 + persistent_batch: list[LogitsProcsRequestParams] = [
583 + ] # Persistent batch state, as list of workload indices
584 +
585 + # Generate fake removed request indices from current persistent
586 + # batch before adds
587 + batch_update_builder = BatchUpdateBuilder()
588 +
589 + # Break when entire workload has been added previously and persistent
590 + # batch is empty
591 + workload_reqs_remaining = workload_size
592 + batch_size = 0
593 + step_idx = 0
594 + while True:
595 + if not (workload_reqs_remaining or batch_size):
596 + break
597 +
598 + (
599 + batch_update,
600 + wdx,
601 + workload_reqs_remaining,
602 + ) = _generate_fake_step_update(
603 + persistent_batch=persistent_batch,
604 + workload_params=workload_params,
605 + wdx=wdx,
606 + batch_update_builder=batch_update_builder,
607 + )
608 + batch_size = len(persistent_batch)
609 +
610 + # Apply fake batch update to logitsprocs
611 + fake_update_logitsprocs_state(test_fakes, batch_update)
612 +
613 + # Emulate application of logits processors in engine
614 + slice_idxs = [req.workload_index for req in persistent_batch]
615 + logits_w_lp = fake_apply_logitsprocs(test_fakes, slice_idxs).cpu()
616 +
617 + _assert_valid(
618 + batch_size=batch_size,
619 + persistent_batch=persistent_batch,
620 + test_fakes=test_fakes,
621 + slice_idxs=slice_idxs,
622 + logits_w_lp=logits_w_lp,
623 + step_idx=step_idx,
624 + )
625 +
626 + step_idx += 1
tests/v1/sample/test_logprobs_e2e.py
+3 -2
@@ -13,9 +13,10 @@ EXPECTED_VALUE = 0.62
13
14 # FIXME(rob): enable prefix caching once supported.
15 MODEL = "meta-llama/Llama-3.2-1B-Instruct"
16 -MODEL_ARGS = f"pretrained={MODEL},enforce_eager=True,enable_prefix_caching=False" # noqa: E501
16 +MODEL_ARGS = f"pretrained={MODEL},enforce_eager=True,enable_prefix_caching=False,gpu_memory_utilization=0.8" # noqa: E501
17 SERVER_ARGS = [
18 - "--enforce_eager", "--no_enable_prefix_caching", "--disable-log-requests"
18 + "--enforce_eager", "--no_enable_prefix_caching", "--disable-log-requests",
19 + "--gpu-memory-utilization=0.8"
20 ]
21 NUM_CONCURRENT = 100
22
tests/v1/sample/test_rejection_sampler.py
+2 -3
@@ -7,6 +7,7 @@ import torch
7 import torch.nn.functional as F
8
9 from vllm.platforms import current_platform
10 +from vllm.v1.sample.logits_processor import LogitsProcessorManager
11 from vllm.v1.sample.metadata import SamplingMetadata
12 from vllm.v1.sample.rejection_sampler import (PLACEHOLDER_TOKEN_ID,
13 RejectionSampler)
@@ -58,7 +59,6 @@ def create_sampling_metadata(
59 all_random=not all_greedy,
60 top_p=top_p,
61 top_k=top_k,
61 - min_p=torch.empty(1, ),
62 generators=generators,
63 max_num_logprobs=0,
64 no_penalties=False,
@@ -67,10 +67,9 @@ def create_sampling_metadata(
67 presence_penalties=torch.tensor([]),
68 repetition_penalties=torch.tensor([]),
69 output_token_ids=[],
70 - min_tokens={},
71 - logit_bias=[None],
70 allowed_token_ids_mask=None,
71 bad_words_token_ids={},
72 + logitsprocs=LogitsProcessorManager(),
73 )
74
75
tests/v1/sample/test_sampler.py
+5 -149
@@ -8,10 +8,13 @@ import pytest
8 import torch
9
10 from vllm.platforms import current_platform
11 -from vllm.utils import make_tensor_with_pad
11 +from vllm.utils import is_pin_memory_available, make_tensor_with_pad
12 +from vllm.v1.sample.logits_processor import LogitsProcessorManager
13 from vllm.v1.sample.metadata import SamplingMetadata
14 from vllm.v1.sample.sampler import Sampler
15
16 +PIN_MEMORY_AVAILABLE = is_pin_memory_available()
17 +MAX_NUM_REQS = 256
18 VOCAB_SIZE = 1024
19 NUM_OUTPUT_TOKENS = 20
20 CUDA_DEVICES = [
@@ -48,18 +51,6 @@ def _create_prompt_tokens_tensor(
51 )
52
53
51 -def _create_logit_bias(
52 - batch_size: int,
53 - vocab_size: int,
54 - bias_value: float,
55 -) -> list[Optional[dict[int, float]]]:
56 - res: list[Optional[dict[int, float]]] = []
57 - for i in range(batch_size):
58 - logit_bias = {min(i, vocab_size - 1): bias_value}
59 - res.append(logit_bias)
60 - return res
61 -
62 -
54 def _create_allowed_token_ids(
55 batch_size: int,
56 vocab_size: int,
@@ -145,7 +136,6 @@ def _create_default_sampling_metadata(
136 all_random=False,
137 top_p=None,
138 top_k=None,
148 - min_p=None,
139 generators={},
140 max_num_logprobs=0,
141 prompt_token_ids=_create_prompt_tokens_tensor(prompt_token_ids,
@@ -155,43 +145,13 @@ def _create_default_sampling_metadata(
145 presence_penalties=_create_penalty_tensor(batch_size, 0.0, device),
146 repetition_penalties=_create_penalty_tensor(batch_size, 1.0, device),
147 no_penalties=True,
158 - min_tokens={},
159 - logit_bias=[None] * batch_size,
148 allowed_token_ids_mask=None,
149 bad_words_token_ids={},
150 + logitsprocs=LogitsProcessorManager(),
151 )
152 return fake_sampling_metadata
153
154
166 -def _generate_min_token_penalties_and_stop_tokens(
167 - num_output_tokens: int, batch_size: int, vocab_size: int,
168 - batch_indices_for_min_token_penalty: list[int]
169 -) -> dict[int, tuple[int, set[int]]]:
170 - """
171 - Generates and returns a dict of minimum token penalties and
172 - corresponding stop token IDs (`min_tokens`, `stop_token_ids`) for each
173 - batch.
174 -
175 - If a batch index is included in `batch_indices_for_min_token_penalty`,
176 - a higher `min_tokens` value is assigned (within a randomized range),
177 - and a random set of stop token IDs is created. Otherwise, a lower
178 - `min_tokens` value is assigned, and the stop token IDs set is empty.
179 - """
180 - min_tokens: dict[int, tuple[int, set[int]]] = {}
181 - for index in range(batch_size):
182 - if index in batch_indices_for_min_token_penalty:
183 - min_tokens[index] = (
184 - np.random.randint(num_output_tokens + 1,
185 - 2 * num_output_tokens),
186 - set(
187 - np.random.randint(0, vocab_size - 1)
188 - for _ in range(np.random.randint(0, vocab_size))))
189 - else:
190 - min_tokens[index] = (np.random.randint(0,
191 - num_output_tokens), set())
192 - return min_tokens
193 -
194 -
155 def _create_weighted_output_token_list(
156 batch_size: int,
157 vocab_size: int) -> tuple[list[list[int]], list[list[int]]]:
@@ -227,36 +187,6 @@ def _create_weighted_output_token_list(
187 return output_token_ids, sorted_token_ids_in_output
188
189
230 -@pytest.mark.parametrize("device", CUDA_DEVICES)
231 -@pytest.mark.parametrize("batch_size", [1, 2, 32])
232 -def test_sampler_min_tokens_penalty(device: str, batch_size: int):
233 - """
234 - Tests that if the number of output tokens is less than
235 - SamplingParams.min_tokens then we will set the logits for
236 - the stop token ids to -inf.
237 - """
238 - torch.set_default_device(device)
239 - fake_logits = _create_fake_logits(batch_size, VOCAB_SIZE)
240 - sampling_metadata = _create_default_sampling_metadata(
241 - NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))
242 - batch_indices_for_min_token_penalty = np.random.randint(
243 - 0, batch_size - 1, size=np.random.randint(0, batch_size)).tolist()
244 - min_tokens = _generate_min_token_penalties_and_stop_tokens(
245 - NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE,
246 - batch_indices_for_min_token_penalty)
247 - sampling_metadata.min_tokens = min_tokens
248 - sampler = Sampler()
249 - logits = sampler.apply_penalties(fake_logits, sampling_metadata)
250 - logits = logits.cpu()
251 - for batch_idx in range(batch_size):
252 - for token_id in range(VOCAB_SIZE):
253 - _, stop_token_ids = min_tokens.get(batch_idx, (0, set()))
254 - if token_id in stop_token_ids:
255 - assert logits[batch_idx][token_id] == -float("inf")
256 - else:
257 - assert logits[batch_idx][token_id] != -float("inf")
258 -
259 -
190 @pytest.mark.parametrize("device", CUDA_DEVICES)
191 @pytest.mark.parametrize("batch_size", [1, 2, 32])
192 @pytest.mark.parametrize("presence_penalty", [-2.0, 2.0])
@@ -401,80 +331,6 @@ def test_sampler_repetition_penalty(device: str, batch_size: int,
331 or non_penalized_token_id in output_tokens)
332
333
404 -@pytest.mark.parametrize("device", CUDA_DEVICES)
405 -@pytest.mark.parametrize("batch_size", [1, 2, 32])
406 -@pytest.mark.parametrize("min_p", [0.0, 0.1])
407 -def test_sampler_min_p(device: str, batch_size: int, min_p: float):
408 - """
409 - Tests that when min_p is applied, tokens with probability below
410 - min_p * max_prob are masked with -inf.
411 - """
412 - torch.set_default_device(device)
413 - fake_logits = _create_fake_logits(batch_size, VOCAB_SIZE)
414 -
415 - # Create one dominant token per batch
416 - for i in range(batch_size):
417 - fake_logits[i, 0] = 10.0 # High logit for first token
418 - fake_logits[i, 1:] = 1e-2 # Others remain low
419 -
420 - sampling_metadata = _create_default_sampling_metadata(
421 - NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))
422 -
423 - # Configure min_p parameters
424 - sampling_metadata.min_p = torch.full((batch_size, ), min_p, device=device)
425 -
426 - sampler = Sampler()
427 - logits = sampler.apply_min_p(fake_logits, sampling_metadata.min_p)
428 - logits = logits.cpu()
429 -
430 - for batch_idx in range(batch_size):
431 - for token_id in range(VOCAB_SIZE):
432 - if token_id == 0:
433 - # Dominant token should always be unmasked
434 - assert logits[batch_idx][token_id] != -float("inf")
435 - else:
436 - if min_p > 0.0:
437 - # Non-dominant tokens should be masked when min_p > 0
438 - assert logits[batch_idx][token_id] == -float("inf")
439 - else:
440 - # No masking when min_p is 0
441 - assert logits[batch_idx][token_id] != -float("inf")
442 -
443 -
444 -@pytest.mark.parametrize("device", CUDA_DEVICES)
445 -@pytest.mark.parametrize("batch_size", [1, 2, 32])
446 -@pytest.mark.parametrize("bias_value", [-0.1, 1.2])
447 -def test_sampler_logit_bias(device: str, batch_size: int, bias_value: float):
448 - """
449 - Test to verify that when the repetition penalty is enabled, tokens
450 - are penalized based on their presence in the prompt or the existing
451 - output.
452 - """
453 - torch.set_default_device(device)
454 - # Create fake logits where each token is assigned the same
455 - # logit value.
456 - fake_logits = _create_fake_logits(batch_size, VOCAB_SIZE)
457 - sampling_metadata = _create_default_sampling_metadata(
458 - NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))
459 - sampling_metadata.logit_bias = _create_logit_bias(
460 - batch_size=batch_size,
461 - vocab_size=VOCAB_SIZE,
462 - bias_value=bias_value,
463 - )
464 - sampler = Sampler()
465 - logits = sampler.apply_logits_bias(fake_logits, sampling_metadata)
466 - logits = logits.cpu()
467 - for batch_idx in range(batch_size):
468 - logits_for_req = logits[batch_idx]
469 - biased_index = min(batch_idx, VOCAB_SIZE - 1)
470 - for token_id in range(VOCAB_SIZE):
471 - if biased_index == token_id:
472 - assert logits_for_req[token_id] == pytest.approx(bias_value +
473 - 1e-2)
474 - else:
475 - assert logits_for_req[token_id] == pytest.approx(1e-2)
476 -
477 -
334 @pytest.mark.parametrize("device", CUDA_DEVICES)
335 @pytest.mark.parametrize("batch_size", [1, 2, 32])
336 @pytest.mark.parametrize("num_allowed_token_ids", [0, 1, 2])
tests/v1/sample/utils.py
+80 -1
@@ -1,12 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Iterator
5 from enum import Enum
5 -from typing import Optional
6 +from typing import NamedTuple, Optional
7
8 import regex as re
9 +import torch
10
11 from vllm import CompletionOutput
12 +from vllm.utils import make_tensor_with_pad
13 +from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor
14 +from vllm.v1.sample.metadata import SamplingMetadata
15
16
17 class BatchLogprobsComposition(Enum):
@@ -134,3 +139,77 @@ def compute_correct_cumulative_logprob(
139 logprobs = completion_output.logprobs
140 assert logprobs is not None
141 return sum([lp[tok_id].logprob for tok_id, lp in zip(token_ids, logprobs)])
142 +
143 +
144 +def create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor:
145 + fake_logits = torch.full((batch_size, vocab_size), 1e-2, dtype=torch.float)
146 + return fake_logits
147 +
148 +
149 +def create_penalty_tensor(batch_size: int, penalty_value: float,
150 + device: torch.device) -> torch.Tensor:
151 + return torch.full((batch_size, ),
152 + fill_value=penalty_value,
153 + dtype=torch.float,
154 + device=device)
155 +
156 +
157 +def create_prompt_tokens_tensor(
158 + prompt_token_ids: list[list[int]],
159 + vocab_size: int,
160 + device: torch.device,
161 +) -> torch.Tensor:
162 + return make_tensor_with_pad(
163 + prompt_token_ids,
164 + pad=vocab_size,
165 + device=device,
166 + dtype=torch.int64,
167 + pin_memory=False,
168 + )
169 +
170 +
171 +class LogitsprocsTestFakes(NamedTuple):
172 + """Wraps fake data structures to support testing"""
173 + logits: torch.Tensor
174 + sampling_metadata: SamplingMetadata
175 +
176 + def get_logitsprocs_by_cls(
177 + self,
178 + cls: type[LogitsProcessor],
179 + ) -> Iterator[LogitsProcessor]:
180 + """Yield logits processors of a specific class.
181 +
182 + Args:
183 + cls: :class:`LogitsProcessor` subclass
184 +
185 + Returns:
186 + Iterator over logits processors
187 + """
188 + return (lp for lp in self.sampling_metadata.logitsprocs.all
189 + if isinstance(lp, cls))
190 +
191 + def get_logitsprocs(self) -> Iterator[LogitsProcessor]:
192 + """Iterator over all logits processors."""
193 + return self.sampling_metadata.logitsprocs.all
194 +
195 +
196 +def fake_update_logitsprocs_state(
197 + test_fakes: LogitsprocsTestFakes,
198 + batch_update: BatchUpdate,
199 +) -> None:
200 + """Imitate logits processors persistent batch state update
201 + in engine core"""
202 + for logitproc in test_fakes.get_logitsprocs():
203 + logitproc.update_state(batch_update)
204 +
205 +
206 +def fake_apply_logitsprocs(
207 + test_fakes: LogitsprocsTestFakes,
208 + slice_indices: list[int],
209 +) -> torch.Tensor:
210 + """Imitate application of logits processors in engine core"""
211 + logits = test_fakes.logits[torch.tensor(slice_indices,
212 + dtype=torch.long)].clone()
213 + for processor in test_fakes.get_logitsprocs():
214 + logits = processor.apply(logits)
215 + return logits
tests/v1/worker/test_gpu_input_batch.py
+29 -26
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import inspect
5 +from collections.abc import Sequence
6 from typing import Optional
7
8 import numpy as np
@@ -12,6 +13,7 @@ from vllm.platforms import current_platform
13 from vllm.sampling_params import SamplingParams
14 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
15 from vllm.v1.pool.metadata import PoolingMetadata
16 +from vllm.v1.sample.logits_processor import LogitsProcessorManager
17 from vllm.v1.sample.metadata import SamplingMetadata
18 from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable
19 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -26,13 +28,18 @@ CUDA_DEVICES = [
28 MAX_NUM_PROMPT_TOKENS = 64
29
30
29 -def _compare_objs(obj1, obj2):
31 +def _compare_objs(obj1,
32 + obj2,
33 + skip: Sequence = ("logitsprocs", "batch_update_builder")):
34 attrs = inspect.getmembers(obj1, lambda a: not (inspect.isroutine(a)))
35 attr_names = set([
36 a[0] for a in attrs
37 if not (a[0].startswith('__') and a[0].endswith('__'))
38 ])
39 for attr_name in attr_names:
40 + if attr_name in skip:
41 + continue
42 +
43 a = getattr(obj1, attr_name)
44 b = getattr(obj2, attr_name)
45
@@ -58,13 +65,11 @@ def _compare_objs(obj1, obj2):
65 f" in {obj1} and {obj2}: {a} != {b}"
66
67
61 -def _remove_requests(
62 - input_batch: InputBatch, batch_size: int,
63 - reqs: list[CachedRequestState]) -> tuple[set[str], list[int]]:
68 +def _remove_requests(input_batch: InputBatch, batch_size: int,
69 + reqs: list[CachedRequestState]) -> set[str]:
70 """
65 - Remove some requests randomly from the batch and returns a tuple
66 - of 1) set of request removed 2) indices of the requests removed
67 - ordered in descending order
71 + Remove some requests randomly from the batch and returns
72 + set of request removed
73 """
74
75 num_reqs_to_remove = np.random.randint(0, batch_size)
@@ -73,13 +78,11 @@ def _remove_requests(
78 req_index_to_remove = np.random.randint(0, batch_size)
79 req_indices_to_remove.add(req_index_to_remove)
80
76 - req_indices_to_remove_list = list(req_indices_to_remove)
77 - req_indices_to_remove_list.sort(reverse=True)
81 req_ids_to_remove: set[str] = set()
82 for index in req_indices_to_remove:
83 input_batch.remove_request(reqs[index].req_id)
84 req_ids_to_remove.add(reqs[index].req_id)
82 - return req_ids_to_remove, req_indices_to_remove_list
85 + return req_ids_to_remove
86
87
88 def _construct_expected_sampling_metadata(
@@ -100,7 +103,6 @@ def _construct_expected_sampling_metadata(
103 repetition_penalties = [1.0 for _ in range(num_reqs)]
104 top_k = [0 for _ in range(num_reqs)]
105 top_p = [0.0 for _ in range(num_reqs)]
103 - min_p = [0.0 for _ in range(num_reqs)]
106 temperature = [0.0 for _ in range(num_reqs)]
107 min_tokens = {}
108 logit_bias = [None] * num_reqs
@@ -123,7 +125,6 @@ def _construct_expected_sampling_metadata(
125 req.sampling_params.repetition_penalty)
126 top_k[index_in_input_batch] = req.sampling_params.top_k
127 top_p[index_in_input_batch] = req.sampling_params.top_p
126 - min_p[index_in_input_batch] = req.sampling_params.min_p
128 temperature[index_in_input_batch] = req.sampling_params.temperature
129 min_tokens[index_in_input_batch] = (
130 req.sampling_params.min_tokens,
@@ -145,8 +146,6 @@ def _construct_expected_sampling_metadata(
146 top_p, dtype=torch.float, device=device),
147 top_k=None if all(x == 0 for x in top_k) else torch.tensor(
148 top_k, dtype=torch.int, device=device),
148 - min_p=None if all(x == 0.0 for x in min_p) else torch.tensor(
149 - min_p, dtype=torch.float, device=device),
149 generators={},
150 max_num_logprobs=0,
151 prompt_token_ids=make_tensor_with_pad(
@@ -165,13 +164,12 @@ def _construct_expected_sampling_metadata(
164 dtype=torch.float,
165 device=device),
166 output_token_ids=output_token_ids,
168 - min_tokens=min_tokens,
167 no_penalties=(all(x == 0 for x in presence_penalties)
168 and all(x == 0 for x in frequency_penalties)
169 and all(x == 1 for x in repetition_penalties)),
172 - logit_bias=logit_bias,
170 allowed_token_ids_mask=allowed_token_ids_mask,
171 bad_words_token_ids=bad_words_token_ids,
172 + logitsprocs=LogitsProcessorManager(),
173 )
174
175
@@ -225,6 +223,8 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
223 and the `make_sampling_metadata` method is invoked on the batch. The
224 output of `make_sampling_metadata` is then compared against the expected
225 results to ensure correctness.
226 +
227 + Note: Ignore logits processor logic, which is tested separately
228 """
229 input_batch: InputBatch = InputBatch(
230 max_num_reqs=batch_size,
@@ -238,21 +238,22 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
238 reqs: list[CachedRequestState] = []
239 req_id_reqs = {}
240 req_id_output_token_ids = {}
241 +
242 # Add requests
243 for req_index in range(batch_size):
244 req: CachedRequestState = _construct_cached_request_state(req_index)
244 - input_batch.add_request(req, req_index)
245 + assigned_req_index = input_batch.add_request(req)
246 + assert req_index == assigned_req_index
247 reqs.append(req)
248 req_id_reqs[req.req_id] = req
249 req_id_output_token_ids[req.req_id] = req.output_token_ids
250
251 # Remove some requests
250 - req_ids_to_remove, req_indices_to_remove = _remove_requests(
251 - input_batch, batch_size, reqs)
252 + req_ids_to_remove = _remove_requests(input_batch, batch_size, reqs)
253 req_ids_retained = set(req_id_reqs.keys()) - req_ids_to_remove
254
255 # Compact the input batch
255 - input_batch.condense(req_indices_to_remove)
256 + input_batch.condense()
257
258 # Generate the sampling metadata
259 sampling_metadata = input_batch._make_sampling_metadata()
@@ -290,10 +291,8 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):
291 sampling_metadata.prompt_token_ids)
292 assert (expected_sampling_metadata.output_token_ids ==
293 sampling_metadata.output_token_ids)
293 - assert expected_sampling_metadata.min_tokens == sampling_metadata.min_tokens
294 assert expected_sampling_metadata.no_penalties == \
295 sampling_metadata.no_penalties
296 - assert expected_sampling_metadata.logit_bias == sampling_metadata.logit_bias
296 if sampling_metadata.allowed_token_ids_mask:
297 assert torch.allclose(
298 expected_sampling_metadata.allowed_token_ids_mask,
@@ -315,6 +314,8 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,
314 and the `make_sampling_metadata` method is invoked on the batch. The
315 output of `make_sampling_metadata` is then compared against the expected
316 results to ensure correctness.
317 +
318 + Note: Ignore logits processor logic, which is tested separately
319 """
320 input_batch: InputBatch = InputBatch(
321 max_num_reqs=batch_size,
@@ -341,7 +342,8 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,
342 # Add requests
343 for req_index in range(batch_size):
344 req: CachedRequestState = _construct_cached_request_state(req_index)
344 - input_batch.add_request(req, req_index)
345 + assigned_req_index = input_batch.add_request(req)
346 + assert assigned_req_index == req_index
347 reqs.append(req)
348 req_id_reqs[req.req_id] = req
349 req_id_output_token_ids[req.req_id] = req.output_token_ids
@@ -354,9 +356,10 @@ def test_swap_states_in_input_batch(device: str, batch_size: int,
356
357 for req_index in range(batch_size):
358 req = reordered_reqs[req_index]
357 - ref_input_batch.add_request(req, req_index)
359 + assigned_req_index = ref_input_batch.add_request(req)
360 + assert assigned_req_index == req_index
361
359 - input_batch.refresh_sampling_metadata()
360 - ref_input_batch.refresh_sampling_metadata()
362 + input_batch.refresh_metadata()
363 + ref_input_batch.refresh_metadata()
364
365 _compare_objs(input_batch, ref_input_batch)
vllm/v1/sample/logits_processor.py new
+516
@@ -0,0 +1,516 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import dataclasses
3 +from abc import ABC, abstractmethod
4 +from collections.abc import Iterator, Sequence
5 +from dataclasses import dataclass, field
6 +from enum import Enum
7 +from itertools import chain
8 +from typing import Optional, Union
9 +
10 +import torch
11 +from torch._prims_common import DeviceLikeType
12 +
13 +from vllm import PoolingParams, SamplingParams
14 +from vllm.logger import init_logger
15 +
16 +logger = init_logger(__name__)
17 +
18 +
19 +class MoveDirectionality(Enum):
20 + # One-way i1->i2 req move within batch
21 + UNIDIRECTIONAL = 0
22 + # Two-way i1<->i2 req swap within batch
23 + SWAP = 1
24 +
25 +
26 +# (index, params, output_tok_ids) tuples for new
27 +# requests added to the batch.
28 +AddedRequest = tuple[int, Union[SamplingParams, PoolingParams], list[int]]
29 +# (index 1, index 2, directionality) tuples representing
30 +# one-way moves or two-way swaps of requests in batch
31 +MovedRequest = tuple[int, int, MoveDirectionality]
32 +# Batch indices of any removed requests.
33 +RemovedRequest = int
34 +
35 +
36 +@dataclasses.dataclass(frozen=True)
37 +class BatchUpdate:
38 + """Persistent batch state change info for logitsprocs"""
39 + batch_size: int # Current num reqs in batch
40 +
41 + # Metadata for requests added to, removed from, and moved
42 + # within the persistent batch.
43 + #
44 + # Note: each added request is represented as
45 + # (index, params, output_tok_ids)
46 + # Key assumption: output_tok_ids is a reference to the
47 + # request's running output tokens list; in this way
48 + # the logits processors always see the latest list of
49 + # generated tokens
50 + removed: Sequence[RemovedRequest]
51 + moved: Sequence[MovedRequest]
52 + added: Sequence[AddedRequest]
53 +
54 +
55 +class BatchUpdateBuilder:
56 + """Helps track persistent batch state changes and build
57 + a batch update data structure for logitsprocs
58 +
59 + Assumptions:
60 + * All information about requests removed from persistent batch
61 + during a step is aggregated in self._removed through calls to
62 + self.removed_append() at the beginning of a step. This must happen
63 + before the first time that self.removed, self.pop_removed()
64 + or self.peek_removed() are invoked in a given step
65 + * After the first time that self.removed, self.pop_removed()
66 + or self.peek_removed() are read in a step, no new removals
67 + are registered using self.removed_append()
68 + * Elements of self._removed are never directly modified, added or
69 + removed (i.e. modification is only via self.removed_append() and
70 + self.pop_removed())
71 +
72 + Guarantees under above assumptions:
73 + * self.removed is always sorted in descending order
74 + * self.pop_removed() and self.peek_removed() both return
75 + the lowest removed request index in the current step
76 + """
77 +
78 + _removed: list[RemovedRequest]
79 + _is_removed_sorted: bool
80 + moved: list[MovedRequest]
81 + added: list[AddedRequest]
82 +
83 + def __init__(
84 + self,
85 + removed: Optional[list[RemovedRequest]] = None,
86 + moved: Optional[list[MovedRequest]] = None,
87 + added: Optional[list[AddedRequest]] = None,
88 + ) -> None:
89 + self._removed = removed or []
90 + self.moved = moved or []
91 + self.added = added or []
92 + self._is_removed_sorted = False
93 +
94 + def _ensure_removed_sorted(self) -> None:
95 + """Sort removed request indices in
96 + descending order.
97 +
98 + Idempotent after first call in a
99 + given step, until reset.
100 + """
101 + if not self._is_removed_sorted:
102 + self._removed.sort(reverse=True)
103 + self._is_removed_sorted = True
104 +
105 + @property
106 + def removed(self) -> list[RemovedRequest]:
107 + """Removed request indices sorted in
108 + descending order"""
109 + self._ensure_removed_sorted()
110 + return self._removed
111 +
112 + def removed_append(self, index: int) -> None:
113 + """Register the removal of a request from
114 + the persistent batch.
115 +
116 + Must not be called after the first time
117 + self.removed, self.pop_removed() or
118 + self.peek_removed() are invoked.
119 +
120 + Args:
121 + index: request index
122 + """
123 + if self._is_removed_sorted:
124 + raise RuntimeError("Cannot register new removed request after"
125 + " self.removed has been read.")
126 + self._removed.append(index)
127 +
128 + def has_removed(self) -> bool:
129 + return bool(self._removed)
130 +
131 + def peek_removed(self) -> Optional[int]:
132 + """Return lowest removed request index"""
133 + if self.has_removed():
134 + self._ensure_removed_sorted()
135 + return self._removed[-1]
136 + return None
137 +
138 + def pop_removed(self) -> Optional[int]:
139 + """Pop lowest removed request index"""
140 + if self.has_removed():
141 + self._ensure_removed_sorted()
142 + return self._removed.pop()
143 + return None
144 +
145 + def get_and_reset(self, batch_size: int) -> Optional[BatchUpdate]:
146 + """Generate a logitsprocs batch update data structure
147 + and reset internal batch update builder state.
148 +
149 + Args:
150 + batch_size: current persistent batch size
151 +
152 + Returns:
153 + Frozen logitsprocs batch update instance; `None` if no updates
154 + """
155 + # Reset removal-sorting logic
156 + self._is_removed_sorted = False
157 + if not any((self._removed, self.moved, self.added)):
158 + # No update; short-circuit
159 + return None
160 + # Build batch state update
161 + batch_update = BatchUpdate(
162 + batch_size=batch_size,
163 + removed=self._removed,
164 + moved=self.moved,
165 + added=self.added,
166 + )
167 + # Reset removed/moved/added update lists
168 + self._removed = []
169 + self.moved = []
170 + self.added = []
171 + return batch_update
172 +
173 +
174 +class LogitsProcessor(ABC):
175 +
176 + @abstractmethod
177 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
178 + raise NotImplementedError
179 +
180 + @abstractmethod
181 + def is_argmax_invariant(self) -> bool:
182 + """True if logits processor has no impact on the
183 + argmax computation in greedy sampling.
184 + NOTE: may or may not have the same value for all
185 + instances of a given LogitsProcessor subclass,
186 + depending on subclass implementation.
187 + TODO(andy): won't be utilized until logits
188 + processors are user-extensible
189 + """
190 + raise NotImplementedError
191 +
192 + @abstractmethod
193 + def update_state(
194 + self,
195 + batch_update: Optional[BatchUpdate],
196 + ) -> None:
197 + """Called when there are new output tokens, prior
198 + to each forward pass.
199 +
200 + Args:
201 + batch_update is non-None iff there have been
202 + changes to the batch makeup.
203 + """
204 + raise NotImplementedError
205 +
206 +
207 +@dataclass
208 +class LogitsProcessorManager:
209 + """Encapsulates initialized logitsproc objects."""
210 + argmax_invariant: list[LogitsProcessor] = field(
211 + default_factory=list) # argmax-invariant logitsprocs
212 + non_argmax_invariant: list[LogitsProcessor] = field(
213 + default_factory=list) # non-argmax-invariant logitsprocs
214 +
215 + @property
216 + def all(self) -> Iterator[LogitsProcessor]:
217 + """Iterator over all logits processors."""
218 + return chain(self.argmax_invariant, self.non_argmax_invariant)
219 +
220 +
221 +###### ----- Built-in LogitsProcessor impls below here
222 +
223 +
224 +class MinPLogitsProcessor(LogitsProcessor):
225 +
226 + def __init__(self, max_num_reqs: int, pin_memory: bool,
227 + device: DeviceLikeType):
228 + super().__init__()
229 + self.min_p_count: int = 0
230 +
231 + self.min_p_cpu_tensor = torch.zeros((max_num_reqs, ),
232 + dtype=torch.float32,
233 + device="cpu",
234 + pin_memory=pin_memory)
235 + self.min_p_cpu = self.min_p_cpu_tensor.numpy()
236 + # Pre-allocated device tensor
237 + self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
238 + dtype=torch.float32,
239 + device=device)
240 + # Current slice of the device tensor
241 + self.min_p: torch.Tensor = self.min_p_device[:0]
242 +
243 + def is_argmax_invariant(self) -> bool:
244 + """Min-p never impacts greedy sampling"""
245 + return True
246 +
247 + def get_min_p_by_index(self, index: int) -> float:
248 + return float(self.min_p_cpu[index])
249 +
250 + def update_state(self, batch_update: Optional[BatchUpdate]):
251 + if not batch_update:
252 + return
253 +
254 + needs_update = False
255 + # Process added requests.
256 + for index, params, _ in batch_update.added:
257 + min_p = params.min_p if isinstance(params, SamplingParams) else 0.0
258 + if self.min_p_cpu[index] != min_p:
259 + needs_update = True
260 + self.min_p_cpu[index] = min_p
261 + if min_p:
262 + self.min_p_count += 1
263 +
264 + if self.min_p_count:
265 + # Process removed requests.
266 + needs_update |= bool(batch_update.removed)
267 + for index in batch_update.removed:
268 + if self.min_p_cpu[index]:
269 + self.min_p_count -= 1
270 +
271 + # Process moved requests, unidirectional (a->b) and swap (a<->b)
272 + for adx, bdx, direct in batch_update.moved:
273 + change = (min_p_a :=
274 + self.min_p_cpu[adx]) != (min_p_b :=
275 + self.min_p_cpu[bdx])
276 + needs_update |= change
277 + if change:
278 + self.min_p_cpu[bdx] = min_p_a
279 + if direct == MoveDirectionality.SWAP:
280 + self.min_p_cpu[adx] = min_p_b
281 +
282 + # Update tensors if needed.
283 + size = batch_update.batch_size
284 + if self.min_p_count and (needs_update or self.min_p.shape[0] != size):
285 + self.min_p = self.min_p_device[:size]
286 + self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True)
287 + self.min_p.unsqueeze_(1)
288 +
289 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
290 + if not self.min_p_count:
291 + return logits
292 +
293 + # Convert logits to probability distribution
294 + probability_values = torch.nn.functional.softmax(logits, dim=-1)
295 + # Calculate maximum probabilities per sequence
296 + max_probabilities = torch.amax(probability_values,
297 + dim=-1,
298 + keepdim=True)
299 + # Adjust min_p
300 + adjusted_min_p = max_probabilities.mul_(self.min_p)
301 + # Identify valid tokens using threshold comparison
302 + invalid_token_mask = probability_values < adjusted_min_p
303 + # Apply mask using boolean indexing
304 + logits[invalid_token_mask] = -float('inf')
305 + return logits
306 +
307 +
308 +class LogitBiasLogitsProcessor(LogitsProcessor):
309 +
310 + def __init__(self, pin_memory: bool, device: torch.device):
311 + super().__init__()
312 + self.biases: dict[int, dict[int, float]] = {}
313 + self.device = device
314 + self.pin_memory = pin_memory
315 +
316 + self.bias_tensor: torch.Tensor = torch.tensor(())
317 + self.logits_slice = (self._device_tensor([], torch.int32),
318 + self._device_tensor([], torch.int32))
319 +
320 + def is_argmax_invariant(self) -> bool:
321 + """Logit bias can rebalance token probabilities and change the
322 + outcome of argmax in greedy sampling."""
323 + return False
324 +
325 + def update_state(self, batch_update: Optional[BatchUpdate]):
326 + if not batch_update:
327 + return
328 +
329 + # Process added requests.
330 + needs_update = bool(batch_update.added)
331 + for index, params, _ in batch_update.added:
332 + if isinstance(params, SamplingParams) and (lb :=
333 + params.logit_bias):
334 + self.biases[index] = lb
335 + else:
336 + self.biases.pop(index, None)
337 +
338 + if self.biases:
339 + # Process removed requests.
340 + for index in batch_update.removed:
341 + if self.biases.pop(index, None):
342 + needs_update = True
343 +
344 + # Process moved requests, unidirectional (a->b) and swap (a<->b)
345 + for a_index, b_index, direct in batch_update.moved:
346 + if direct == MoveDirectionality.UNIDIRECTIONAL:
347 + if (a_entry := self.biases.pop(a_index, None)) is None:
348 + if self.biases.pop(b_index, None) is not None:
349 + needs_update = True
350 + else:
351 + self.biases[b_index] = a_entry
352 + needs_update = True
353 + else:
354 + a_entry = self.biases.pop(a_index, None)
355 + if (b_entry := self.biases.pop(b_index, None)) is not None:
356 + self.biases[a_index] = b_entry
357 + needs_update = True
358 + if a_entry is not None:
359 + self.biases[b_index] = a_entry
360 + needs_update = True
361 +
362 + # Update tensors if needed.
363 + if needs_update:
364 + reqs, tok_ids, biases = [], [], []
365 + for req, lb in self.biases.items():
366 + reqs.extend([req] * len(lb))
367 + tok_ids.extend(lb.keys())
368 + biases.extend(lb.values())
369 +
370 + self.bias_tensor = self._device_tensor(biases, torch.float32)
371 + self.logits_slice = (self._device_tensor(reqs, torch.int32),
372 + self._device_tensor(tok_ids, torch.int32))
373 +
374 + def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:
375 + return (torch.tensor(data,
376 + device="cpu",
377 + dtype=dtype,
378 + pin_memory=self.pin_memory).to(device=self.device,
379 + non_blocking=True))
380 +
381 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
382 + if self.biases:
383 + logits[self.logits_slice] += self.bias_tensor
384 + return logits
385 +
386 +
387 +class MinTokensLogitsProcessor(LogitsProcessor):
388 +
389 + def __init__(self, pin_memory: bool, device: torch.device):
390 + # index -> (min_toks, output_token_ids, stop_token_ids)
391 + super().__init__()
392 + self.min_toks: dict[int, tuple[int, Sequence[int], set[int]]] = {}
393 + self.device = device
394 + self.pin_memory = pin_memory
395 +
396 + # (req_idx_tensor,eos_tok_id_tensor)
397 + self.logits_slice: tuple[torch.Tensor,
398 + torch.Tensor] = (self._device_tensor(
399 + [], torch.int32),
400 + self._device_tensor(
401 + [], torch.int32))
402 +
403 + def is_argmax_invariant(self) -> bool:
404 + """By censoring stop tokens, min-tokens can change the outcome
405 + of the argmax operation in greedy sampling."""
406 + return False
407 +
408 + def update_state(self, batch_update: Optional[BatchUpdate]):
409 + needs_update = False
410 +
411 + if batch_update:
412 + # Process added requests.
413 + needs_update |= bool(batch_update.added)
414 + for index, params, output_tok_ids in batch_update.added:
415 + if (isinstance(params, SamplingParams)
416 + and (min_tokens := params.min_tokens)
417 + and len(output_tok_ids) < min_tokens):
418 + # Replace request metadata at batch index
419 + self.min_toks[index] = (min_tokens, output_tok_ids,
420 + params.all_stop_token_ids)
421 + else:
422 + # Drop request metadata at batch index
423 + self.min_toks.pop(index, None)
424 +
425 + if self.min_toks:
426 + # Process removed requests.
427 + for index in batch_update.removed:
428 + if self.min_toks.pop(index, None):
429 + needs_update = True
430 +
431 + # Process moved requests, unidirectional (a->b) and
432 + # swapped (a<->b)
433 + for a_index, b_index, direct in batch_update.moved:
434 + if direct == MoveDirectionality.UNIDIRECTIONAL:
435 + if (a_entry := self.min_toks.pop(a_index,
436 + None)) is None:
437 + if self.min_toks.pop(b_index, None) is not None:
438 + needs_update = True
439 + else:
440 + self.min_toks[b_index] = a_entry
441 + needs_update = True
442 + else:
443 + a_entry = self.min_toks.pop(a_index, None)
444 + if (b_entry := self.min_toks.pop(b_index,
445 + None)) is not None:
446 + self.min_toks[a_index] = b_entry
447 + needs_update = True
448 + if a_entry is not None:
449 + self.min_toks[b_index] = a_entry
450 + needs_update = True
451 +
452 + if self.min_toks:
453 + # Check for any requests that have attained their min tokens.
454 + to_remove = tuple(index for index, (min_toks, out_tok_ids,
455 + _) in self.min_toks.items()
456 + if len(out_tok_ids) >= min_toks)
457 + if to_remove:
458 + needs_update = True
459 + for index in to_remove:
460 + del self.min_toks[index]
461 +
462 + # Update tensors if needed.
463 + if needs_update:
464 + reqs: list[int] = []
465 + tok_ids: list[int] = []
466 + for req, (_, _, stop_tok_ids) in self.min_toks.items():
467 + reqs.extend([req] * len(stop_tok_ids))
468 + tok_ids.extend(stop_tok_ids)
469 +
470 + self.logits_slice = (self._device_tensor(reqs, torch.int32),
471 + self._device_tensor(tok_ids, torch.int32))
472 +
473 + def _device_tensor(self, data: list, dtype: torch.dtype) -> torch.Tensor:
474 + return (torch.tensor(data,
475 + device="cpu",
476 + dtype=dtype,
477 + pin_memory=self.pin_memory).to(device=self.device,
478 + non_blocking=True))
479 +
480 + def apply(self, logits: torch.Tensor) -> torch.Tensor:
481 + if self.min_toks:
482 + # Inhibit EOS token for requests which have not reached min length
483 + logits[self.logits_slice] = -float("inf")
484 + return logits
485 +
486 +
487 +def init_builtin_logitsprocs(pin_memory_available: bool, max_num_reqs: int,
488 + device: torch.device) -> LogitsProcessorManager:
489 + """Construct 'builtin' vLLM logitsprocs which the engine
490 + loads by default.
491 +
492 + Args:
493 + pin_memory_available: pinned memory is available for use
494 + for use by logitsproc
495 + max_num_reqs: ceiling on request count in persistent batch
496 + device: inference device
497 +
498 + Returns:
499 + Data structure encapsulating loaded logitsprocs
500 + """
501 + min_tokens_logitproc = MinTokensLogitsProcessor(
502 + pin_memory=pin_memory_available, device=device)
503 + logit_bias_logitproc = LogitBiasLogitsProcessor(
504 + pin_memory=pin_memory_available, device=device)
505 + min_p_logitproc = MinPLogitsProcessor(
506 + pin_memory=pin_memory_available,
507 + device=device,
508 + # +1 for temporary swap space
509 + max_num_reqs=max_num_reqs + 1)
510 + return LogitsProcessorManager(
511 + non_argmax_invariant=[
512 + min_tokens_logitproc,
513 + logit_bias_logitproc,
514 + ],
515 + argmax_invariant=[min_p_logitproc],
516 + )
vllm/v1/sample/metadata.py
+5 -6
@@ -6,6 +6,8 @@ from typing import Optional
6
7 import torch
8
9 +from vllm.v1.sample.logits_processor import LogitsProcessorManager
10 +
11
12 @dataclass
13 class SamplingMetadata:
@@ -16,7 +18,6 @@ class SamplingMetadata:
18
19 top_p: Optional[torch.Tensor]
20 top_k: Optional[torch.Tensor]
19 - min_p: Optional[torch.Tensor]
21
22 generators: dict[int, torch.Generator]
23
@@ -31,14 +32,12 @@ class SamplingMetadata:
32
33 output_token_ids: list[list[int]]
34
34 - # req_index -> (min_tokens, stop_token_ids)
35 - min_tokens: dict[int, tuple[int, set[int]]]
36 -
37 - logit_bias: list[Optional[dict[int, float]]]
38 -
35 # `allowed_token_ids_mask` is a 2D bool tensor of shape (max batch size,
36 # vocab size).
37 allowed_token_ids_mask: Optional[torch.Tensor]
38
39 # req_index -> bad_words_token_ids
40 bad_words_token_ids: dict[int, list[list[int]]]
41 +
42 + # Loaded logits processors
43 + logitsprocs: LogitsProcessorManager
vllm/v1/sample/ops/penalties.py
-16
@@ -7,22 +7,6 @@ from vllm.model_executor.layers.utils import apply_penalties
7 from vllm.utils import is_pin_memory_available, make_tensor_with_pad
8
9
10 -def apply_min_token_penalties(
11 - logits: torch.Tensor, output_token_ids: list[list[int]],
12 - min_tokens: dict[int, tuple[int, set[int]]]) -> None:
13 - """
14 - Applies minimum token penalty by setting the logits of the stop tokens
15 - to -inf.
16 - """
17 - min_tokens_logits_to_penalize: list[tuple[int, int]] = []
18 - for index, (min_token, stop_token_ids) in min_tokens.items():
19 - if len(output_token_ids[index]) < min_token:
20 - for stop_token_id in stop_token_ids:
21 - min_tokens_logits_to_penalize.append((index, stop_token_id))
22 - if min_tokens_logits_to_penalize:
23 - logits[tuple(zip(*min_tokens_logits_to_penalize))] = -float("inf")
24 -
25 -
10 def apply_all_penalties(
11 logits: torch.Tensor,
12 prompt_token_ids: torch.Tensor,
vllm/v1/sample/sampler.py
+11 -71
@@ -5,12 +5,11 @@
5 import torch
6 import torch.nn as nn
7
8 -from vllm.utils import async_tensor_h2d, is_pin_memory_available
8 +from vllm.utils import is_pin_memory_available
9 from vllm.v1.outputs import LogprobsTensors, SamplerOutput
10 from vllm.v1.sample.metadata import SamplingMetadata
11 from vllm.v1.sample.ops.bad_words import apply_bad_words
12 -from vllm.v1.sample.ops.penalties import (apply_all_penalties,
13 - apply_min_token_penalties)
12 +from vllm.v1.sample.ops.penalties import apply_all_penalties
13 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
14
15 _SAMPLING_EPS = 1e-5
@@ -44,8 +43,11 @@ class Sampler(nn.Module):
43 logits = self.apply_allowed_token_ids(logits, sampling_metadata)
44 # Apply bad words exclusion.
45 logits = self.apply_bad_words(logits, sampling_metadata)
47 - # Apply logits bias.
48 - logits = self.apply_logits_bias(logits, sampling_metadata)
46 +
47 + # Apply logits processors which can impact greedy sampling
48 + for processor in (sampling_metadata.logitsprocs.non_argmax_invariant):
49 + logits = processor.apply(logits)
50 +
51 # Apply penalties (e.g., min_tokens, freq_penalties).
52 logits = self.apply_penalties(logits, sampling_metadata)
53 # Sample the next token.
@@ -110,9 +112,10 @@ class Sampler(nn.Module):
112 # Apply temperature.
113 logits = self.apply_temperature(logits, sampling_metadata.temperature)
114
113 - # Apply min_p.
114 - if sampling_metadata.min_p is not None:
115 - logits = self.apply_min_p(logits, sampling_metadata.min_p)
115 + # Apply logits processors that only apply to random sampling
116 + # (argmax invariant)
117 + for processor in sampling_metadata.logitsprocs.argmax_invariant:
118 + logits = processor.apply(logits)
119
120 # Apply top_k and/or top_p.
121 random_sampled = self.topk_topp_sampler(
@@ -187,10 +190,6 @@ class Sampler(nn.Module):
190 logits: torch.Tensor,
191 sampling_metadata: SamplingMetadata,
192 ) -> torch.Tensor:
190 - if sampling_metadata.min_tokens:
191 - apply_min_token_penalties(logits,
192 - sampling_metadata.output_token_ids,
193 - sampling_metadata.min_tokens)
193 if not sampling_metadata.no_penalties:
194 assert sampling_metadata.prompt_token_ids is not None
195 logits = apply_all_penalties(
@@ -203,65 +202,6 @@ class Sampler(nn.Module):
202 )
203 return logits
204
206 - def apply_min_p(
207 - self,
208 - logits: torch.Tensor,
209 - min_p: torch.Tensor,
210 - ) -> torch.Tensor:
211 - """
212 - Filters logits using adaptive probability thresholding.
213 - """
214 - # Convert logits to probability distribution
215 - probability_values = torch.nn.functional.softmax(logits, dim=-1)
216 - # Calculate maximum probabilities per sequence
217 - max_probabilities = torch.amax(probability_values,
218 - dim=-1,
219 - keepdim=True)
220 - # Reshape min_p for broadcasting
221 - adjusted_min_p = min_p.unsqueeze(1) * max_probabilities
222 - # Identify valid tokens using threshold comparison
223 - valid_token_mask = probability_values >= adjusted_min_p
224 - # Apply mask using boolean indexing
225 - logits[~valid_token_mask] = -float('inf')
226 - return logits
227 -
228 - def apply_logits_bias(
229 - self,
230 - logits: torch.Tensor,
231 - sampling_metadata: SamplingMetadata,
232 - ) -> torch.Tensor:
233 - # TODO(houseroad): this implementation is extremely inefficient.
234 - # One idea is implement this as a PyTorch C++ op, and we may
235 - # even optimize the logit_bias layout.
236 -
237 - rows: list[int] = []
238 - cols: list[int] = []
239 - vals: list[float] = []
240 -
241 - # Get vocabulary size from logits
242 - vocab_size = logits.shape[-1]
243 -
244 - for i, logit_bias in enumerate(sampling_metadata.logit_bias):
245 - if logit_bias:
246 - for token_id, bias in logit_bias.items():
247 - # Check token_id bounds to ensure within vocabulary
248 - if token_id < 0 or token_id >= vocab_size:
249 - raise ValueError(
250 - f"token_id {token_id} in logit_bias contains "
251 - f"out-of-vocab token id. Vocabulary size: "
252 - f"{vocab_size}")
253 - rows.append(i)
254 - cols.append(token_id)
255 - vals.append(bias)
256 -
257 - if rows:
258 - indices = async_tensor_h2d([rows, cols], torch.int64,
259 - logits.device, self.pin_memory)
260 - values = async_tensor_h2d(vals, torch.float, logits.device,
261 - self.pin_memory)
262 - logits.index_put_(tuple(indices), values=values, accumulate=True)
263 - return logits
264 -
205 def apply_allowed_token_ids(
206 self,
207 logits: torch.Tensor,
vllm/v1/spec_decode/utils.py
+12 -17
@@ -1,23 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.sampling_params import SamplingParams
4 from vllm.triton_utils import tl, triton
4 -from vllm.v1.worker.gpu_input_batch import InputBatch
5 -
6 -
7 -def is_spec_decode_supported(req_id: str, input_batch: InputBatch) -> bool:
8 - if req_id in input_batch.min_p_reqs:
9 - # Spec decode doesn't support min_p sampling.
10 - return False
11 - elif (req_id in input_batch.frequency_penalties_reqs
12 - or req_id in input_batch.presence_penalties_reqs
13 - or req_id in input_batch.repetition_penalties_reqs):
14 - # Spec decode doesn't support penalties.
15 - return False
16 - elif req_id in input_batch.num_logprobs:
17 - # Spec decode doesn't support logprobs.
18 - return False
19 -
20 - return True
5 +
6 +_SAMPLING_EPS = 1e-5
7 +
8 +
9 +def is_spec_decode_unsupported(sampling_params: SamplingParams) -> bool:
10 + """True if request is incompatible with speculative decoding"""
11 + return (sampling_params.frequency_penalty != 0.0
12 + or sampling_params.presence_penalty != 0.0
13 + or sampling_params.repetition_penalty != 1.0
14 + or sampling_params.min_p > _SAMPLING_EPS
15 + or sampling_params.logprobs is not None)
16
17
18 @triton.jit
vllm/v1/worker/gpu_input_batch.py
+95 -63
@@ -15,12 +15,14 @@ from vllm.sampling_params import SamplingParams, SamplingType
15 from vllm.utils import swap_dict_values
16 from vllm.v1.outputs import LogprobsTensors
17 from vllm.v1.pool.metadata import PoolingMetadata
18 +from vllm.v1.sample.logits_processor import (BatchUpdateBuilder,
19 + MoveDirectionality,
20 + init_builtin_logitsprocs)
21 from vllm.v1.sample.metadata import SamplingMetadata
22 +from vllm.v1.spec_decode.utils import is_spec_decode_unsupported
23 from vllm.v1.utils import copy_slice
24 from vllm.v1.worker.block_table import MultiGroupBlockTable
25
22 -_SAMPLING_EPS = 1e-5
23 -
26
27 @dataclass
28 class CachedRequestState:
@@ -67,8 +69,10 @@ class InputBatch:
69 pin_memory: bool,
70 vocab_size: int,
71 block_sizes: list[int], # The block_size of each kv cache group
72 + is_spec_decode: bool = False,
73 logits_processing_needs_token_ids: bool = False,
74 ):
75 + self.is_spec_decode = is_spec_decode
76 self.max_num_reqs = max_num_reqs
77 self.max_model_len = max_model_len
78 self.max_num_batched_tokens = max_num_batched_tokens
@@ -146,15 +150,8 @@ class InputBatch:
150 self.top_k_cpu = self.top_k_cpu_tensor.numpy()
151 self.top_k_reqs: set[str] = set()
152
149 - self.min_p = torch.empty((max_num_reqs, ),
150 - dtype=torch.float32,
151 - device=device)
152 - self.min_p_cpu_tensor = torch.empty((max_num_reqs, ),
153 - dtype=torch.float32,
154 - device="cpu",
155 - pin_memory=pin_memory)
156 - self.min_p_cpu = self.min_p_cpu_tensor.numpy()
157 - self.min_p_reqs: set[str] = set()
153 + # IDs of requests which do not support spec decoding
154 + self.spec_decode_unsupported_reqs: set[str] = set()
155
156 # Frequency penalty related data structures
157 self.frequency_penalties = torch.empty((max_num_reqs, ),
@@ -194,9 +191,6 @@ class InputBatch:
191 self.repetition_penalties_cpu_tensor.numpy()
192 self.repetition_penalties_reqs: set[str] = set()
193
197 - # req_index -> (min_tokens, stop_token_ids)
198 - self.min_tokens: dict[int, tuple[int, set[int]]] = {}
199 -
194 # lora related
195 self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
196 dtype=np.int32)
@@ -216,8 +210,20 @@ class InputBatch:
210 # To accumulate prompt logprobs tensor chunks across prefill steps.
211 self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
212
219 - self.logit_bias: list[Optional[dict[int,
220 - float]]] = [None] * max_num_reqs
213 + # Internal representation of per-step batch state changes, used for
214 + # reordering persistent batch and generating logitsprocs batch state
215 + # updates. Should reset each step.
216 + self.batch_update_builder = BatchUpdateBuilder()
217 +
218 + # Define logits processors.
219 + # TODO(andy): logits processor list should be extensible via engine
220 + # constructor argument; for now the list is fixed.
221 + self.logitsprocs = init_builtin_logitsprocs(
222 + pin_memory_available=pin_memory,
223 + max_num_reqs=max_num_reqs + 1,
224 + device=device)
225 +
226 + # TODO convert this to LogitsProcessor
227 self.has_allowed_token_ids: set[str] = set()
228 # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
229 # the value is False. Since we use masked_fill_ to set -inf.
@@ -240,14 +246,28 @@ class InputBatch:
246 # while performing state updates to the batch.
247 return cast(list[str], self._req_ids)
248
249 + def _get_next_add_index(self) -> int:
250 + if (req_index := self.batch_update_builder.pop_removed()) is not None:
251 + # Fill the empty index.
252 + return req_index
253 + # Append to end
254 + return self.num_reqs
255 +
256 + def _register_add_request(self, request: "CachedRequestState") -> int:
257 + """Track add-request operations"""
258 + req_index = self._get_next_add_index()
259 + assert req_index < self.max_num_reqs
260 + params = (request.sampling_params
261 + if request.sampling_params else request.pooling_params)
262 + self.batch_update_builder.added.append(
263 + (req_index, params, request.output_token_ids))
264 + return req_index
265 +
266 def add_request(
267 self,
268 request: "CachedRequestState",
246 - req_index: Optional[int] = None,
247 - ) -> None:
248 - if req_index is None:
249 - req_index = self.num_reqs
250 - assert req_index < self.max_num_reqs
269 + ) -> int:
270 + req_index = self._register_add_request(request)
271
272 req_id = request.req_id
273 if req_index == len(self._req_ids):
@@ -278,6 +298,9 @@ class InputBatch:
298 self.block_table.add_row(request.block_ids, req_index)
299
300 if sampling_params := request.sampling_params:
301 + if (self.is_spec_decode
302 + and is_spec_decode_unsupported(sampling_params)):
303 + self.spec_decode_unsupported_reqs.add(req_id)
304 if sampling_params.sampling_type == SamplingType.GREEDY:
305 # Avoid later division by zero.
306 self.temperature_cpu[req_index] = -1.0
@@ -295,11 +318,8 @@ class InputBatch:
318 else:
319 top_k = self.vocab_size
320 self.top_k_cpu[req_index] = top_k
298 - self.min_p_cpu[req_index] = sampling_params.min_p
321 self.frequency_penalties_cpu[
322 req_index] = sampling_params.frequency_penalty
301 - if sampling_params.min_p > _SAMPLING_EPS:
302 - self.min_p_reqs.add(req_id)
323 if sampling_params.frequency_penalty != 0.0:
324 self.frequency_penalties_reqs.add(req_id)
325 self.presence_penalties_cpu[
@@ -310,10 +330,6 @@ class InputBatch:
330 req_index] = sampling_params.repetition_penalty
331 if sampling_params.repetition_penalty != 1.0:
332 self.repetition_penalties_reqs.add(req_id)
313 - if sampling_params.min_tokens:
314 - self.min_tokens[req_index] = (
315 - sampling_params.min_tokens,
316 - sampling_params.all_stop_token_ids)
333
334 # NOTE(woosuk): self.generators should not include the requests that
335 # do not have their own generator.
@@ -325,8 +341,6 @@ class InputBatch:
341 if sampling_params.prompt_logprobs is not None:
342 self.num_prompt_logprobs[
343 req_id] = sampling_params.prompt_logprobs
328 - if sampling_params.logit_bias is not None:
329 - self.logit_bias[req_index] = sampling_params.logit_bias
344
345 if sampling_params.allowed_token_ids:
346 self.has_allowed_token_ids.add(req_id)
@@ -368,12 +382,22 @@ class InputBatch:
382 # No LoRA
383 self.request_lora_mapping[req_index] = 0
384
385 + return req_index
386 +
387 def remove_request(self, req_id: str) -> Optional[int]:
372 - """This method must always be followed by a call to condense()."""
388 + """This method must always be followed by a call to condense().
389 +
390 + Args:
391 + req_id: request to remove
392 +
393 + Returns:
394 + Removed request index, or `None` if `req_id` not recognized
395 + """
396
397 req_index = self.req_id_to_index.pop(req_id, None)
398 if req_index is None:
399 return None
400 + self.batch_update_builder.removed_append(req_index)
401 self._req_ids[req_index] = None
402 self.req_output_token_ids[req_index] = None
403
@@ -381,8 +405,7 @@ class InputBatch:
405 self.random_reqs.discard(req_id)
406 self.top_p_reqs.discard(req_id)
407 self.top_k_reqs.discard(req_id)
384 - self.min_p_reqs.discard(req_id)
385 - self.min_tokens.pop(req_index, None)
408 + self.spec_decode_unsupported_reqs.discard(req_id)
409 self.frequency_penalties_reqs.discard(req_id)
410 self.presence_penalties_reqs.discard(req_id)
411 self.repetition_penalties_reqs.discard(req_id)
@@ -400,7 +423,6 @@ class InputBatch:
423 self.lora_id_to_lora_request.pop(lora_id)
424 self.request_lora_mapping[req_index] = 0
425
403 - self.logit_bias[req_index] = None
426 self.has_allowed_token_ids.discard(req_id)
427 if self.allowed_token_ids_mask_cpu_tensor is not None:
428 # False means we don't fill with -inf.
@@ -410,6 +432,8 @@ class InputBatch:
432 return req_index
433
434 def swap_states(self, i1: int, i2: int) -> None:
435 + self.batch_update_builder.moved.append(
436 + (i1, i2, MoveDirectionality.SWAP))
437 old_id_i1 = self._req_ids[i1]
438 old_id_i2 = self._req_ids[i2]
439 self._req_ids[i1], self._req_ids[i2] =\
@@ -439,8 +463,6 @@ class InputBatch:
463 self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
464 self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] =\
465 self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]
442 - self.min_p_cpu[i1], self.min_p_cpu[i2] =\
443 - self.min_p_cpu[i2], self.min_p_cpu[i1]
466
467 # NOTE: the following is unsafe
468 # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
@@ -452,13 +474,10 @@ class InputBatch:
474 self.token_ids_cpu[i2, ...] = tmp
475
476 swap_dict_values(self.generators, i1, i2)
455 - swap_dict_values(self.min_tokens, i1, i2)
477 swap_dict_values(self.bad_words_token_ids, i1, i2)
478
479 self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
480 self.request_lora_mapping[i2], self.request_lora_mapping[i1]
460 - self.logit_bias[i1], self.logit_bias[i2] =\
461 - self.logit_bias[i2], self.logit_bias[i1]
481
482 if self.allowed_token_ids_mask_cpu_tensor is not None:
483 self.allowed_token_ids_mask_cpu_tensor[i1], \
@@ -467,12 +486,23 @@ class InputBatch:
486 self.allowed_token_ids_mask_cpu_tensor[i1]
487 self.block_table.swap_row(i1, i2)
488
470 - def condense(self, empty_req_indices: list[int]) -> None:
471 - """Move non-empty requests down into lower, empty indices.
472 -
489 + def condense(self) -> None:
490 + """Slide non-empty requests down into lower, empty indices.
491 +
492 + Any consecutive empty indices at the very end of the list are not
493 + filled.
494 +
495 Args:
474 - empty_req_indices: empty batch indices, sorted descending.
496 + empty_req_indices: empty indices which may be filled.
497 +
498 + Returns:
499 + swaps: list of (from,to) swap tuples for moved requests
500 + empty_req_indices: indices not filled by condensation
501 """
502 + if not (empty_req_indices := self.batch_update_builder.removed):
503 + # All removed requests were replaced by added requests, or else no
504 + # requests were removed at all. No condense() needed
505 + return
506 num_reqs = self.num_reqs
507 if num_reqs == 0:
508 # The batched states are empty.
@@ -489,11 +519,17 @@ class InputBatch:
519 last_req_index -= 1
520
521 # Find the smallest empty index.
492 - empty_index = empty_req_indices.pop()
522 + empty_index = self.batch_update_builder.peek_removed()
523 + assert empty_index is not None
524 if empty_index >= last_req_index:
525 break
526
496 - # Swap the states.
527 + # Move active request down into empty request
528 + # index.
529 + self.batch_update_builder.pop_removed()
530 + self.batch_update_builder.moved.append(
531 + (last_req_index, empty_index,
532 + MoveDirectionality.UNIDIRECTIONAL))
533 req_id = self._req_ids[last_req_index]
534 output_token_ids = self.req_output_token_ids[last_req_index]
535 assert req_id is not None
@@ -524,20 +560,14 @@ class InputBatch:
560 empty_index] = self.presence_penalties_cpu[last_req_index]
561 self.repetition_penalties_cpu[
562 empty_index] = self.repetition_penalties_cpu[last_req_index]
527 - self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]
563 generator = self.generators.pop(last_req_index, None)
564 if generator is not None:
565 self.generators[empty_index] = generator
566
532 - min_token = self.min_tokens.pop(last_req_index, None)
533 - if min_token is not None:
534 - self.min_tokens[empty_index] = min_token
535 -
567 self.request_lora_mapping[empty_index] = self.request_lora_mapping[
568 last_req_index]
569
539 - self.logit_bias[empty_index] = self.logit_bias[last_req_index]
540 -
570 + # TODO convert these to LogitsProcessors
571 if self.allowed_token_ids_mask_cpu_tensor is not None:
572 self.allowed_token_ids_mask_cpu_tensor[
573 empty_index] = self.allowed_token_ids_mask_cpu_tensor[
@@ -547,6 +577,7 @@ class InputBatch:
577 last_req_index, None)
578 if bad_words_token_ids is not None:
579 self.bad_words_token_ids[empty_index] = bad_words_token_ids
580 +
581 # Decrement last_req_index since it is now empty.
582 last_req_index -= 1
583
@@ -554,8 +585,17 @@ class InputBatch:
585 del self._req_ids[self.num_reqs:]
586 del self.req_output_token_ids[self.num_reqs:]
587
557 - def refresh_sampling_metadata(self):
558 - self.sampling_metadata = self._make_sampling_metadata()
588 + def refresh_metadata(self):
589 + """Apply batch updates, reset input batch at end of step
590 +
591 + * Apply batch add/remove/permute to logits procs' states
592 + * If batch state is modified, update sampling metadata
593 + """
594 + batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)
595 + for logit_proc in self.logitsprocs.all:
596 + logit_proc.update_state(batch_update)
597 + if batch_update:
598 + self.sampling_metadata = self._make_sampling_metadata()
599
600 def _make_sampling_metadata(self) -> SamplingMetadata:
601 num_reqs = self.num_reqs
@@ -568,8 +608,6 @@ class InputBatch:
608 copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs)
609 if not self.no_top_k:
610 copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs)
571 - if not self.no_min_p:
572 - copy_slice(self.min_p_cpu_tensor, self.min_p, num_reqs)
611
612 if not self.no_penalties:
613 # Since syncing these tensors is expensive only copy them
@@ -607,7 +645,6 @@ class InputBatch:
645 all_random=self.all_random,
646 top_p=None if self.no_top_p else self.top_p[:num_reqs],
647 top_k=None if self.no_top_k else self.top_k[:num_reqs],
610 - min_p=None if self.no_min_p else self.min_p[:num_reqs],
648 generators=self.generators,
649 max_num_logprobs=self.max_num_logprobs,
650 prompt_token_ids=prompt_token_ids,
@@ -615,11 +652,10 @@ class InputBatch:
652 presence_penalties=self.presence_penalties[:num_reqs],
653 repetition_penalties=self.repetition_penalties[:num_reqs],
654 output_token_ids=cast(list[list[int]], self.req_output_token_ids),
618 - min_tokens=self.min_tokens,
655 no_penalties=self.no_penalties,
620 - logit_bias=self.logit_bias[:num_reqs],
656 allowed_token_ids_mask=allowed_token_ids_mask,
657 bad_words_token_ids=self.bad_words_token_ids,
658 + logitsprocs=self.logitsprocs,
659 )
660
661 @property
@@ -702,10 +738,6 @@ class InputBatch:
738 def no_top_k(self) -> bool:
739 return len(self.top_k_reqs) == 0
740
705 - @property
706 - def no_min_p(self) -> bool:
707 - return len(self.min_p_reqs) == 0
708 -
741 @property
742 def no_penalties(self) -> bool:
743 return (len(self.presence_penalties_reqs) == 0
vllm/v1/worker/gpu_model_runner.py
+19 -41
@@ -63,12 +63,12 @@ from vllm.v1.spec_decode.eagle import EagleProposer
63 from vllm.v1.spec_decode.medusa import MedusaProposer
64 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
65 from vllm.v1.spec_decode.ngram_proposer import NgramProposer
66 -from vllm.v1.spec_decode.utils import is_spec_decode_supported
66 from vllm.v1.utils import bind_kv_cache
67 from vllm.v1.worker.block_table import BlockTable
68 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
69 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
70
71 +from ..sample.logits_processor import LogitsProcessorManager
72 from .utils import (gather_mm_placeholders, initialize_kv_cache_for_kv_sharing,
73 sanity_check_mm_encoder_outputs, scatter_mm_placeholders)
74
@@ -212,6 +212,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
212 pin_memory=self.pin_memory,
213 vocab_size=self.model_config.get_vocab_size(),
214 block_sizes=[self.cache_config.block_size],
215 + is_spec_decode=bool(self.vllm_config.speculative_config),
216 )
217
218 self.use_cuda_graph = (
@@ -316,7 +317,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
317 # from the KV cache of `shared_kv_cache_layers[layer_name]`.
318 self.shared_kv_cache_layers: dict[str, str] = {}
319
319 - def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> bool:
320 + def _may_reorder_batch(self, scheduler_output: "SchedulerOutput") -> None:
321 """
322 Update the order of requests in the batch based on the attention
323 backend's needs. For example, some attention backends (namely MLA) may
@@ -325,21 +326,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
326
327 Args:
328 scheduler_output: The scheduler output.
328 -
329 - Returns:
330 - True if the batch was reordered, False otherwise.
329 """
332 - batch_reordered = self.attn_metadata_builders[0].reorder_batch(
333 - self.input_batch, scheduler_output)
330 + self.attn_metadata_builders[0].reorder_batch(self.input_batch,
331 + scheduler_output)
332
333 # For models with multiple KV cache groups, the groups should agree on
334 # the same order of requests. We ensure this by only allowing the first
335 # group to reorder the batch and asserting that all other groups do not
336 # reorder the batch.
337 for i in range(1, len(self.kv_cache_config.kv_cache_groups)):
340 - assert not self.attn_metadata_builders[i].reorder_batch(
338 + batch_reordered = self.attn_metadata_builders[i].reorder_batch(
339 self.input_batch, scheduler_output)
342 - return batch_reordered
340 + assert not batch_reordered
341
342 # Note: used for model runner override.
343 def _init_device_properties(self) -> None:
@@ -372,11 +370,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
370 # then resubmitted with the same ID. In this case, we treat them as two
371 # distinct requests - clearing the cached states for the first request
372 # and handling the second as a new request.
375 - removed_req_indices: list[int] = []
373 for req_id in scheduler_output.finished_req_ids:
377 - req_index = self.input_batch.remove_request(req_id)
378 - if req_index is not None:
379 - removed_req_indices.append(req_index)
374 + self.input_batch.remove_request(req_id)
375
376 # Free the cached encoder outputs.
377 for req_id, input_id in scheduler_output.free_encoder_input_ids:
@@ -399,9 +394,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
394 # have low request overlap (e.g., alternating between two distinct
395 # sets of requests), this optimization becomes very inefficient.
396 for req_id in unscheduled_req_ids:
402 - req_index = self.input_batch.remove_request(req_id)
403 - assert req_index is not None
404 - removed_req_indices.append(req_index)
397 + self.input_batch.remove_request(req_id)
398
399 req_ids_to_add: list[str] = []
400 # Add new requests to the cached states.
@@ -545,31 +538,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
538 # NOTE(woosuk): `num_tokens` here may include spec tokens.
539 self.input_batch.num_tokens[req_index] = end_token_index
540
548 - # Check if the batch has changed. If not, we can skip copying the
549 - # sampling metadata from CPU to GPU.
550 - batch_changed = len(removed_req_indices) > 0 or len(req_ids_to_add) > 0
551 -
541 # Add the new or resumed requests to the persistent batch.
542 # The smaller empty indices are filled first.
554 - removed_req_indices.sort(reverse=True)
543 for req_id in req_ids_to_add:
544 req_state = self.requests[req_id]
557 - if removed_req_indices:
558 - # Fill the empty index.
559 - req_index = removed_req_indices.pop()
560 - else:
561 - # Append to the end.
562 - req_index = None
563 - self.input_batch.add_request(req_state, req_index)
545 + self.input_batch.add_request(req_state)
546
565 - # Condense the batched states if there are empty indices.
566 - if removed_req_indices:
567 - self.input_batch.condense(removed_req_indices)
568 -
569 - batch_reordered = self._may_reorder_batch(scheduler_output)
570 -
571 - if batch_changed or batch_reordered:
572 - self.input_batch.refresh_sampling_metadata()
547 + # Condense the batched states if there are gaps left by removed requests
548 + self.input_batch.condense()
549 + # Allow attention backend to reorder the batch, potentially
550 + self._may_reorder_batch(scheduler_output)
551 + # Refresh batch metadata with any pending updates.
552 + self.input_batch.refresh_metadata()
553
554 def _get_cumsum_and_arange(
555 self,
@@ -1296,7 +1276,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1276 scheduler_output: "SchedulerOutput",
1277 intermediate_tensors: Optional[IntermediateTensors] = None,
1278 ) -> Union[ModelRunnerOutput, IntermediateTensors]:
1299 -
1279 self._update_states(scheduler_output)
1280 if not scheduler_output.total_num_scheduled_tokens:
1281 if not has_kv_transfer_group():
@@ -1765,7 +1744,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1744 # Skip requests that require sampling parameters that are not
1745 # supported with speculative decoding.
1746 req_id = self.input_batch.req_ids[i]
1768 - if not is_spec_decode_supported(req_id, self.input_batch):
1747 + if req_id in self.input_batch.spec_decode_unsupported_reqs:
1748 draft_token_ids.append([])
1749 continue
1750
@@ -2121,7 +2100,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2100 all_random=False,
2101 top_p=dummy_tensors(0.9),
2102 top_k=dummy_tensors(logits.size(1) - 1),
2124 - min_p=None,
2103 generators={},
2104 max_num_logprobs=None,
2105 no_penalties=True,
@@ -2130,10 +2108,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2108 presence_penalties=dummy_tensors(0.1),
2109 repetition_penalties=dummy_tensors(0.1),
2110 output_token_ids=[[] for _ in range(num_reqs)],
2133 - min_tokens={},
2134 - logit_bias=[None for _ in range(num_reqs)],
2111 allowed_token_ids_mask=None,
2112 bad_words_token_ids={},
2113 + logitsprocs=LogitsProcessorManager(),
2114 )
2115 try:
2116 sampler_output = self.sampler(logits=logits,
@@ -2425,6 +2402,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2402 pin_memory=self.pin_memory,
2403 vocab_size=self.model_config.get_vocab_size(),
2404 block_sizes=block_sizes,
2405 + is_spec_decode=bool(self.vllm_config.speculative_config),
2406 )
2407
2408 def _allocate_kv_cache_tensors(