[ModelRunner V2] Minor indexing optimizations (#51210)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

vllm/v1/worker/gpu/input_batch.py

@@ -110,8 +110,8 @@ class InputBatch:
         device = input_buffers.device
 
         req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
-        idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
-        idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
+        idx_mapping_np = np.arange(num_reqs, dtype=np.intp)
+        idx_mapping = torch.arange(num_reqs, dtype=torch.int64, device=device)
         expanded_idx_mapping = idx_mapping
         expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
 

vllm/v1/worker/gpu/model_runner.py

@@ -978,11 +978,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # batch_idx -> req_id
         req_ids = sort_batch_req_ids(num_tokens_per_req, self.decode_query_len)
-        numtoks_iter = map(num_tokens_per_req.get, req_ids)
+        numtoks_iter = map(num_tokens_per_req.__getitem__, req_ids)
         num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)
 
-        idx_mapping_iter = map(self.req_states.req_id_to_index.get, req_ids)
-        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.int32, count=num_reqs)
+        idx_mapping_iter = map(self.req_states.req_id_to_index.__getitem__, req_ids)
+        idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.intp, count=num_reqs)
         idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)
 
         # Get the number of draft tokens for each request.

vllm/v1/worker/gpu/pcp_manager.py

@@ -357,7 +357,7 @@ class PCPManager:
 
         local_to_global_batch_req_idx_np = np.fromiter(
             (segment.global_batch_req_idx for segment in local_segments),
-            dtype=np.int32,
+            dtype=np.intp,
             count=num_local_reqs,
         )
         local_start_pos_np = np.fromiter(