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 6, 12:00 UTC
fix(worker): optimize swap_states to copy only active token prefixes (#34733)
Signed-off-by: Philip Ottesen <phiott256@gmail.com>
Philip Ottesen committed
Mar 18, 2026 at 22:59 UTC
0091017188ab26d4e4d146e0ec748d6ee34968d8
1 file changed
+20
-8
vllm/v1/worker/gpu_input_batch.py
+20
-8
@@ -529,6 +529,12 @@ class InputBatch:
529
def swap_states(self, i1: int, i2: int) -> None:
530
old_id_i1 = self._req_ids[i1]
531
old_id_i2 = self._req_ids[i2]
532
+ # Only swap the active token prefix for each request. Copying full
533
+ # max_model_len rows is expensive and unnecessary during reordering.
534
+ i1_active_token_count = self._get_active_token_count(i1)
535
+ i2_active_token_count = self._get_active_token_count(i2)
536
+ max_active_token_count = max(i1_active_token_count, i2_active_token_count)
537
+
538
self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa
539
self.req_output_token_ids[i1], self.req_output_token_ids[i2] = (
540
self.req_output_token_ids[i2],
@@ -560,12 +566,15 @@ class InputBatch:
566
# self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
567
# self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
568
# instead, we need to temporarily copy the data for one of the indices
563
- # TODO(lucas): optimize this by only copying valid indices
564
- tmp = self.token_ids_cpu[i1, ...].copy()
565
- self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]
566
- self.token_ids_cpu[i2, ...] = tmp
569
+ tmp_token_ids = self.token_ids_cpu[i1, :max_active_token_count].copy()
570
+ self.token_ids_cpu[i1, :max_active_token_count] = self.token_ids_cpu[
571
+ i2, :max_active_token_count
572
+ ]
573
+ self.token_ids_cpu[i2, :max_active_token_count] = tmp_token_ids
574
568
- self.is_token_ids[[i1, i2], ...] = self.is_token_ids[[i2, i1], ...]
575
+ self.is_token_ids[[i1, i2], :max_active_token_count] = self.is_token_ids[
576
+ [i2, i1], :max_active_token_count
577
+ ]
578
579
# Swap prompt embeddings if they exist
580
embeds_i1 = self.req_prompt_embeds.get(i1)
@@ -629,6 +638,11 @@ class InputBatch:
638
self.allowed_token_ids_mask_cpu_tensor[i1],
639
)
640
641
+ def _get_active_token_count(self, req_index: int) -> int:
642
+ return int(self.num_tokens_no_spec[req_index]) + len(
643
+ self.spec_token_ids[req_index]
644
+ )
645
+
646
def condense(self) -> None:
647
"""Slide non-empty requests down into lower, empty indices.
648
@@ -678,9 +692,7 @@ class InputBatch:
692
self.req_output_token_ids[last_req_index] = None
693
self.req_id_to_index[req_id] = empty_index
694
681
- num_tokens = self.num_tokens_no_spec[last_req_index] + len(
682
- self.spec_token_ids[last_req_index]
683
- )
695
+ num_tokens = self._get_active_token_count(last_req_index)
696
697
(self.spec_token_ids[last_req_index], self.spec_token_ids[empty_index]) = (
698
self.spec_token_ids[empty_index],