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 14, 05:30 UTC
[BugFix] Fix batch updates for pooling models (#23398)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Aug 22, 2025 at 17:20 UTC
c80c53a30ff7a9c074ec6a7d88021ebe8c19e2e9
3 files changed
+95
-79
vllm/v1/sample/logits_processor/state.py
+16
-4
@@ -50,6 +50,10 @@ class BatchUpdateBuilder:
50
self.added = added or []
51
self._is_removed_sorted = False
52
53
+ # Used to track changes in the pooling case
54
+ # where we don't populate the added list.
55
+ self.batch_changed = False
56
+
57
def _ensure_removed_sorted(self) -> None:
58
"""Sort removed request indices in
59
descending order.
@@ -80,6 +84,7 @@ class BatchUpdateBuilder:
84
raise RuntimeError("Cannot register new removed request after"
85
" self.removed has been read.")
86
self._removed.append(index)
87
+ self.batch_changed = True
88
89
def has_removed(self) -> bool:
90
return bool(self._removed)
@@ -98,9 +103,15 @@ class BatchUpdateBuilder:
103
return self._removed.pop()
104
return None
105
101
- def _is_update(self) -> bool:
102
- """True if there is a batch state change"""
103
- return any((self._removed, self.moved, self.added))
106
+ def reset(self) -> bool:
107
+ """Returns True if there were any changes to the batch."""
108
+ self._is_removed_sorted = False
109
+ self._removed.clear()
110
+ self.moved.clear()
111
+ self.added.clear()
112
+ batch_changed = self.batch_changed
113
+ self.batch_changed = False
114
+ return batch_changed
115
116
def get_and_reset(self, batch_size: int) -> Optional[BatchUpdate]:
117
"""Generate a logitsprocs batch update data structure and reset
@@ -114,7 +125,8 @@ class BatchUpdateBuilder:
125
"""
126
# Reset removal-sorting logic
127
self._is_removed_sorted = False
117
- if not self._is_update():
128
+ self.batch_changed = False
129
+ if not any((self._removed, self.moved, self.added)):
130
# No update; short-circuit
131
return None
132
# Build batch state update
vllm/v1/worker/gpu_input_batch.py
+76
-70
@@ -65,8 +65,7 @@ class CachedRequestState:
65
def get_token_id(self, idx: int) -> int:
66
if idx < self.num_prompt_tokens:
67
return self.prompt_token_ids[idx]
68
- else:
69
- return self.output_token_ids[idx - self.num_prompt_tokens]
68
+ return self.output_token_ids[idx - self.num_prompt_tokens]
69
70
71
class InputBatch:
@@ -261,30 +260,27 @@ class InputBatch:
260
Not applicable to pooling models.
261
"""
262
264
- # Detailed added request metadata is only required for non-pooling
265
- # models, to support logitsprocs
266
- assert request.sampling_params
267
-
263
# Fill the next empty index if there is one.
264
if (new_req_index := self.batch_update_builder.pop_removed()) is None:
265
# Append to end otherwise.
266
new_req_index = self.num_reqs
267
268
assert new_req_index < self.max_num_reqs
274
- self.batch_update_builder.added.append(
275
- (new_req_index, request.sampling_params, request.prompt_token_ids,
276
- request.output_token_ids))
269
+ self.batch_update_builder.batch_changed = True
270
+ if request.sampling_params:
271
+ # Detailed added request metadata is only required for non-pooling
272
+ # models, to support logitsprocs.
273
+ self.batch_update_builder.added.append(
274
+ (new_req_index, request.sampling_params,
275
+ request.prompt_token_ids, request.output_token_ids))
276
+
277
return new_req_index
278
279
def add_request(
280
self,
281
request: "CachedRequestState",
282
) -> int:
283
- if not self.is_pooling_model:
284
- # New request index bookkeeping for autoregressive models.
285
- req_index = self._register_add_request(request)
286
- else:
287
- req_index = self.num_reqs
283
+ req_index = self._register_add_request(request)
284
285
req_id = request.req_id
286
if req_index == len(self._req_ids):
@@ -389,7 +385,7 @@ class InputBatch:
385
self.logits_processing_needs_token_ids[req_index] = (
386
pooling_params.requires_token_ids)
387
else:
392
- raise NotImplementedError(request)
388
+ raise NotImplementedError("Unrecognized request type")
389
390
# Add request lora ID
391
if request.lora_request:
@@ -419,13 +415,25 @@ class InputBatch:
415
req_index = self.req_id_to_index.pop(req_id, None)
416
if req_index is None:
417
return None
422
- if not self.is_pooling_model:
423
- # Autoregressive models require bookkeeping of removed requests to
424
- # support logitsprocs.
425
- self.batch_update_builder.removed_append(req_index)
418
+
419
+ self.batch_update_builder.removed_append(req_index)
420
self._req_ids[req_index] = None
421
self.req_output_token_ids[req_index] = None
422
423
+ # LoRA
424
+ lora_id = self.request_lora_mapping[req_index]
425
+ if lora_id != 0:
426
+ lora_req_ids = self.lora_id_to_request_ids[lora_id]
427
+ lora_req_ids.discard(req_id)
428
+ if not lora_req_ids:
429
+ del self.lora_id_to_request_ids[lora_id]
430
+ del self.lora_id_to_lora_request[lora_id]
431
+ self.request_lora_mapping[req_index] = 0
432
+
433
+ if self.is_pooling_model:
434
+ self.pooling_params.pop(req_id, None)
435
+ return req_index
436
+
437
self.greedy_reqs.discard(req_id)
438
self.random_reqs.discard(req_id)
439
self.top_p_reqs.discard(req_id)
@@ -439,29 +447,14 @@ class InputBatch:
447
self.num_prompt_logprobs.pop(req_id, None)
448
self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
449
442
- # LoRA
443
- lora_id = self.request_lora_mapping[req_index]
444
- if lora_id != 0:
445
- lora_req_ids = self.lora_id_to_request_ids[lora_id]
446
- lora_req_ids.discard(req_id)
447
- if not lora_req_ids:
448
- del self.lora_id_to_request_ids[lora_id]
449
- del self.lora_id_to_lora_request[lora_id]
450
- self.request_lora_mapping[req_index] = 0
451
-
450
self.has_allowed_token_ids.discard(req_id)
451
if self.allowed_token_ids_mask_cpu_tensor is not None:
452
# False means we don't fill with -inf.
453
self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
454
self.bad_words_token_ids.pop(req_index, None)
457
- self.pooling_params.pop(req_id, None)
455
return req_index
456
457
def swap_states(self, i1: int, i2: int) -> None:
461
- # For autoregressive models, track detailed request reordering info
462
- # to support logitsprocs
463
- self.batch_update_builder.moved.append(
464
- (i1, i2, MoveDirectionality.SWAP))
458
old_id_i1 = self._req_ids[i1]
459
old_id_i2 = self._req_ids[i2]
460
self._req_ids[i1], self._req_ids[i2] =\
@@ -479,18 +472,6 @@ class InputBatch:
472
self.num_prompt_tokens[i2], self.num_prompt_tokens[i1]
473
self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] =\
474
self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1]
482
- self.temperature_cpu[i1], self.temperature_cpu[i2] =\
483
- self.temperature_cpu[i2], self.temperature_cpu[i1]
484
- self.top_p_cpu[i1], self.top_p_cpu[i2] =\
485
- self.top_p_cpu[i2], self.top_p_cpu[i1]
486
- self.top_k_cpu[i1], self.top_k_cpu[i2] =\
487
- self.top_k_cpu[i2], self.top_k_cpu[i1]
488
- self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] =\
489
- self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1]
490
- self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] =\
491
- self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
492
- self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] =\
493
- self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]
475
476
# NOTE: the following is unsafe
477
# self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
@@ -501,18 +482,41 @@ class InputBatch:
482
self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]
483
self.token_ids_cpu[i2, ...] = tmp
484
504
- swap_dict_values(self.generators, i1, i2)
505
- swap_dict_values(self.bad_words_token_ids, i1, i2)
485
+ self.block_table.swap_row(i1, i2)
486
507
- self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
487
+ self.request_lora_mapping[i1], self.request_lora_mapping[i2] = \
488
self.request_lora_mapping[i2], self.request_lora_mapping[i1]
489
490
+ if self.is_pooling_model:
491
+ # Sampling and logits parameters don't apply to pooling models.
492
+ return
493
+
494
+ # For autoregressive models, track detailed request reordering info
495
+ # to support logitsprocs.
496
+ self.batch_update_builder.moved.append(
497
+ (i1, i2, MoveDirectionality.SWAP))
498
+
499
+ self.temperature_cpu[i1], self.temperature_cpu[i2] = \
500
+ self.temperature_cpu[i2], self.temperature_cpu[i1]
501
+ self.top_p_cpu[i1], self.top_p_cpu[i2] = \
502
+ self.top_p_cpu[i2], self.top_p_cpu[i1]
503
+ self.top_k_cpu[i1], self.top_k_cpu[i2] = \
504
+ self.top_k_cpu[i2], self.top_k_cpu[i1]
505
+ self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = \
506
+ self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1]
507
+ self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = \
508
+ self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
509
+ self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = \
510
+ self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]
511
+
512
+ swap_dict_values(self.generators, i1, i2)
513
+ swap_dict_values(self.bad_words_token_ids, i1, i2)
514
+
515
if self.allowed_token_ids_mask_cpu_tensor is not None:
516
self.allowed_token_ids_mask_cpu_tensor[i1], \
517
self.allowed_token_ids_mask_cpu_tensor[i2] =\
518
self.allowed_token_ids_mask_cpu_tensor[i2], \
519
self.allowed_token_ids_mask_cpu_tensor[i1]
515
- self.block_table.swap_row(i1, i2)
520
521
def condense(self) -> None:
522
"""Slide non-empty requests down into lower, empty indices.
@@ -529,12 +533,6 @@ class InputBatch:
533
"""
534
num_reqs = self.num_reqs
535
532
- if self.is_pooling_model:
533
- # Will be contiguous in pooling case, just trim the lists.
534
- del self._req_ids[num_reqs:]
535
- del self.req_output_token_ids[num_reqs:]
536
- return
537
-
536
if not (empty_req_indices := self.batch_update_builder.removed):
537
# All removed requests were replaced by added requests, or else no
538
# requests were removed at all. No condense() needed
@@ -562,11 +560,6 @@ class InputBatch:
560
# Move active request down into empty request
561
# index.
562
self.batch_update_builder.pop_removed()
565
- # Autoregressive models require detailed tracking of condense
566
- # operations to support logitsprocs
567
- self.batch_update_builder.moved.append(
568
- (last_req_index, empty_index,
569
- MoveDirectionality.UNIDIRECTIONAL))
563
req_id = self._req_ids[last_req_index]
564
output_token_ids = self.req_output_token_ids[last_req_index]
565
assert req_id is not None
@@ -587,6 +580,21 @@ class InputBatch:
580
self.num_computed_tokens_cpu[
581
empty_index] = self.num_computed_tokens_cpu[last_req_index]
582
self.block_table.move_row(last_req_index, empty_index)
583
+
584
+ self.request_lora_mapping[empty_index] = self.request_lora_mapping[
585
+ last_req_index]
586
+
587
+ if self.is_pooling_model:
588
+ last_req_index -= 1
589
+ # Samping state not used by pooling models.
590
+ continue
591
+
592
+ # Autoregressive models require detailed tracking of condense
593
+ # operations to support logitsprocs
594
+ self.batch_update_builder.moved.append(
595
+ (last_req_index, empty_index,
596
+ MoveDirectionality.UNIDIRECTIONAL))
597
+
598
self.temperature_cpu[empty_index] = self.temperature_cpu[
599
last_req_index]
600
self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index]
@@ -601,9 +609,6 @@ class InputBatch:
609
if generator is not None:
610
self.generators[empty_index] = generator
611
604
- self.request_lora_mapping[empty_index] = self.request_lora_mapping[
605
- last_req_index]
606
-
612
# TODO convert these to LogitsProcessors
613
if self.allowed_token_ids_mask_cpu_tensor is not None:
614
self.allowed_token_ids_mask_cpu_tensor[
@@ -626,8 +631,9 @@ class InputBatch:
631
"""Apply any batch updates to sampling metadata."""
632
633
if self.is_pooling_model:
629
- # Batch changes every step for pooling models.
630
- self.sampling_metadata = self._make_sampling_metadata()
634
+ batch_changed = self.batch_update_builder.reset()
635
+ if batch_changed:
636
+ self.sampling_metadata = self._make_sampling_metadata()
637
return
638
639
# For non-pooling models - generate and apply logitsprocs update;
@@ -720,7 +726,8 @@ class InputBatch:
726
)
727
728
def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
723
- max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()
729
+ num_reqs = self.num_reqs
730
+ max_prompt_len = self.num_prompt_tokens[:num_reqs].max()
731
prompt_token_ids_cpu_tensor = torch.empty(
732
(self.num_reqs, max_prompt_len),
733
device="cpu",
@@ -728,11 +735,10 @@ class InputBatch:
735
pin_memory=self.pin_memory,
736
)
737
prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
731
- prompt_token_ids[:] = self.token_ids_cpu[:self.
732
- num_reqs, :max_prompt_len]
738
+ prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len]
739
# Use the value of vocab_size as a pad since we don't have a
740
# token_id of this value.
735
- for i in range(self.num_reqs):
741
+ for i in range(num_reqs):
742
prompt_token_ids[i, self.num_prompt_tokens[i]:] = self.vocab_size
743
return prompt_token_ids_cpu_tensor.to(device=self.device,
744
non_blocking=True)
vllm/v1/worker/gpu_model_runner.py
+3
-5
@@ -1489,10 +1489,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1489
for raw_output, seq_len, prompt_len in zip(
1490
raw_pooler_output, seq_lens_cpu, pooling_metadata.prompt_lens):
1491
1492
- if seq_len == prompt_len:
1493
- pooler_output.append(raw_output.data)
1494
- else:
1495
- pooler_output.append(None)
1492
+ output = raw_output.data if seq_len == prompt_len else None
1493
+ pooler_output.append(output)
1494
1495
return ModelRunnerOutput(
1496
req_ids=self.input_batch.req_ids,
@@ -1522,7 +1520,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1520
# Prepare the decoder inputs.
1521
(attn_metadata, logits_indices, spec_decode_metadata,
1522
num_scheduled_tokens_np, spec_decode_common_attn_metadata,
1525
- max_query_len) = (self._prepare_inputs(scheduler_output))
1523
+ max_query_len) = self._prepare_inputs(scheduler_output)
1524
1525
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1526
if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE