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
[Core][Bookkeeping Optimization] Update against numpy view of is_token_ids tensor (#27618)
Signed-off-by: Jialin Ouyang <Jialin.Ouyang@gmail.com>
Jialin Ouyang committed
Oct 28, 2025 at 01:13 UTC
b46e4a06f1aa8f243d3ed8d3b30d14ae56863d3d
2 files changed
+3
-2
vllm/v1/worker/gpu_input_batch.py
+2
-1
@@ -108,9 +108,10 @@ class InputBatch:
108
pin_memory=False,
109
)
110
self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
111
- self.is_token_ids = torch.zeros(
111
+ self.is_token_ids_tensor = torch.zeros(
112
(max_num_reqs, max_model_len), device="cpu", dtype=bool, pin_memory=False
113
)
114
+ self.is_token_ids = self.is_token_ids_tensor.numpy()
115
# Store prompt embeddings per request to avoid OOM from large upfront
116
# allocation if max_model_len is big.
117
# Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size)
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -1103,7 +1103,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1103
out=self.input_ids.cpu[:total_num_scheduled_tokens],
1104
)
1105
if self.enable_prompt_embeds:
1106
- is_token_ids = self.input_batch.is_token_ids.flatten()
1106
+ is_token_ids = self.input_batch.is_token_ids_tensor.flatten()
1107
torch.index_select(
1108
is_token_ids,
1109
0,