biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V1] Do not store None in self.generators (#11038)

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon
2024-12-09 15:08:19 -08:00
committed by GitHub
parent 5ed5d5f128
commit 6faec54505
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/v1/worker/gpu_input_batch.py
View File

@@ -102,6 +102,8 @@ class InputBatch:
self.top_k_reqs: Set[str] = set() self.top_k_reqs: Set[str] = set()
# req_index -> generator # 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.generators: Dict[int, torch.Generator] = {}
self.num_logprobs: Dict[str, int] = {} self.num_logprobs: Dict[str, int] = {}
@@ -147,7 +149,10 @@ class InputBatch:
if sampling_params.top_k > 0: if sampling_params.top_k > 0:
self.top_k_reqs.add(req_id) 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 num_logprobs = sampling_params.logprobs
if num_logprobs is not None and num_logprobs > 0: if num_logprobs is not None and num_logprobs > 0: