Skip to content

Commit

Permalink
[V1] Do not store None in self.generators (vllm-project#11038)
Browse files Browse the repository at this point in the history
Signed-off-by: Woosuk Kwon <[email protected]>
  • Loading branch information
WoosukKwon authored and weilong.yu committed Dec 13, 2024
1 parent 056b59a commit 1b7fbc5
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion vllm/v1/worker/gpu_input_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,8 @@ def __init__(
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] = {}
Expand Down Expand Up @@ -147,7 +149,10 @@ def add_request(
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:
Expand Down

0 comments on commit 1b7fbc5

Please sign in to comment.