Skip to content

GitLab

Unverified Commit 5a0a6dfd authored by Nick Hill's avatar Nick Hill Committed by GitHub
Browse files

[BugFix] Fix incorrect preallocated sampled_token_ids tensor size (#28025) 


Signed-off-by: default avatarNick Hill <nhill@redhat.com>
parent 938772af
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
vllm/v1/worker/gpu_model_runner.py
View file @ 5a0a6dfd
...@@ -524,7 +524,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin): ...@@ -524,7 +524,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
self._draft_token_ids: list[list[int]] | torch.Tensor | None = None self._draft_token_ids: list[list[int]] | torch.Tensor | None = None
self.transfer_event = torch.cuda.Event() self.transfer_event = torch.cuda.Event()
self.sampled_token_ids_pinned_cpu = torch.empty( self.sampled_token_ids_pinned_cpu = torch.empty(
(self.max_model_len, 1), (self.max_num_reqs, 1),
dtype=torch.int64, dtype=torch.int64,
device="cpu", device="cpu",
pin_memory=self.pin_memory, pin_memory=self.pin_memory,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment