diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 0ce0288f31..fe423d6e38 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1128,7 +1128,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.pooling_runner is not None: assert new_req_data.pooling_params is not None - assert new_req_data.prompt_token_ids is not None self.pooling_runner.add_request( req_id, req_index, diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index 839edc8b8c..ce94764851 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/vllm/v1/worker/gpu/pool/pooling_runner.py @@ -64,7 +64,7 @@ class PoolingRunner: req_id: str, req_index: int, pooling_params: PoolingParams, - prompt_token_ids: list[int], + prompt_token_ids: list[int] | None, ) -> None: task = pooling_params.task if task not in self.supported_tasks: @@ -77,6 +77,7 @@ class PoolingRunner: self.pooling_params[req_index] = pooling_params self.pooling_states[req_index] = PoolingStates() if pooling_params.requires_token_ids: + assert prompt_token_ids is not None self.prompt_token_ids[req_index] = torch.tensor( prompt_token_ids, dtype=torch.int64 )