From 6faec545057e6152e92e8ab619fc018e20864943 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 9 Dec 2024 15:08:19 -0800 Subject: [PATCH] [V1] Do not store `None` in self.generators (#11038) Signed-off-by: Woosuk Kwon --- vllm/v1/worker/gpu_input_batch.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu_input_batch.py b/vllm/v1/worker/gpu_input_batch.py index 457784bb0287c..25d95ac6e26af 100644 --- a/vllm/v1/worker/gpu_input_batch.py +++ b/vllm/v1/worker/gpu_input_batch.py @@ -102,6 +102,8 @@ class InputBatch: self.top_k_reqs: Set[str] = set() # req_index -> generator + # NOTE(woosuk): The indices of the requests that do not have their own + # generator should not be included in the dictionary. self.generators: Dict[int, torch.Generator] = {} self.num_logprobs: Dict[str, int] = {} @@ -147,7 +149,10 @@ class InputBatch: if sampling_params.top_k > 0: self.top_k_reqs.add(req_id) - self.generators[req_index] = request.generator + # NOTE(woosuk): self.generators should not include the requests that + # do not have their own generator. + if request.generator is not None: + self.generators[req_index] = request.generator num_logprobs = sampling_params.logprobs if num_logprobs is not None and num_logprobs > 0: