xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-25 06:45:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Remove extra tensor on CPU (#20693)

Browse Source 
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
This commit is contained in:
Maximilien de Bayser 2025-07-12 03:06:34 -03:00 committed by GitHub
parent c1c8ca57ff
commit 5de8d9f111
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 13 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
vllm/v1/sample/logits_processor.py
View File

@ -234,10 +234,16 @@ class MinPLogitsProcessor(LogitsProcessor):
device="cpu",
pin_memory=pin_memory)
self.min_p_cpu = self.min_p_cpu_tensor.numpy()
# Pre-allocated device tensor
self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
dtype=torch.float32,
device=device)
self.use_double_tensor = torch.device("cpu") != torch.device(device)
if self.use_double_tensor:
# Pre-allocated device tensor
self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
dtype=torch.float32,
device=device)
else:
self.min_p_device = self.min_p_cpu_tensor
# Current slice of the device tensor
self.min_p: torch.Tensor = self.min_p_device[:0]
@ -284,7 +290,9 @@ class MinPLogitsProcessor(LogitsProcessor):
size = batch_update.batch_size
if self.min_p_count and (needs_update or self.min_p.shape[0] != size):
self.min_p = self.min_p_device[:size]
self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True)
if self.use_double_tensor:
self.min_p.copy_(self.min_p_cpu_tensor[:size],
non_blocking=True)
self.min_p.unsqueeze_(1)
def apply(self, logits: torch.Tensor) -> torch.Tensor: