fix: improve token_ids_cpu swap to copy only valid indices (#36255)
Signed-off-by: Jaebok Lee <jaebok9541@naver.com>
zzaebok committed
Aug 29, 2026 at 08:22 UTC
68c52b5a937fe63c7a315739ad468f9d30fd483b
1 file changed
+4
-4
vllm/v1/worker/tpu_input_batch.py
+4
-4
@@ -377,10 +377,10 @@ class InputBatch:
377
# self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
378
# self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
379
# instead, we need to temporarily copy the data for one of the indices
380
- # TODO(lucas): optimize this by only copying valid indices
381
- tmp = self.token_ids_cpu[i1, ...].copy()
382
- self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]
383
- self.token_ids_cpu[i2, ...] = tmp
380
+ n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2])
381
+ tmp = self.token_ids_cpu[i1, :n].copy()
382
+ self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n]
383
+ self.token_ids_cpu[i2, :n] = tmp
384
385
swap_dict_values(self.generators, i1, i2)
386
swap_dict_values(self.min_tokens, i1, i2)