[ModelRunner V2] Minor indexing optimizations (#51210)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Aug 6, 2026 at 14:44 UTC
c5d470ac4ccd17ac9663db0d1c0e2060e5ae15ad
3 files changed
+6
-6
vllm/v1/worker/gpu/input_batch.py
+2
-2
@@ -110,8 +110,8 @@ class InputBatch:
110
device = input_buffers.device
111
112
req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
113
- idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
114
- idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
113
+ idx_mapping_np = np.arange(num_reqs, dtype=np.intp)
114
+ idx_mapping = torch.arange(num_reqs, dtype=torch.int64, device=device)
115
expanded_idx_mapping = idx_mapping
116
expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
117
vllm/v1/worker/gpu/model_runner.py
+3
-3
@@ -978,11 +978,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
978
979
# batch_idx -> req_id
980
req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
981
- numtoks_iter = map(num_tokens_per_req.get, req_ids)
981
+ numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)
982
num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
983
984
- idx_mapping_iter = map(self.req_states.req_id_to_index.get, req_ids)
985
- idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.int32, count=num_reqs)
984
+ idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
985
+ idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
986
idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
987
988
# Get the number of draft tokens for each request.
vllm/v1/worker/gpu/pcp_manager.py
+1
-1
@@ -357,7 +357,7 @@ class PCPManager:
357
358
local_to_global_batch_req_idx_np = np.fromiter(
359
(segment.global_batch_req_idx for segment in local_segments),
360
- dtype=np.int32,
360
+ dtype=np.intp,
361
count=num_local_reqs,
362
)
363
local_start_pos_np = np.fromiter(