diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index a9ee523b0a..8a952aa043 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/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) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 6be4ab146c..640f3ea202 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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. diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 62bacf5314..587293471c 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/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(