A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
Remove extra tensor on CPU (#20693)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
Maximilien de Bayser committed
Jul 12, 2025 at 03:06 UTC
5de8d9f111242671beb6fdc28367df8796ce638f
1 file changed
+13
-5
vllm/v1/sample/logits_processor.py
+13
-5
@@ -234,10 +234,16 @@ class MinPLogitsProcessor(LogitsProcessor):
234
device="cpu",
235
pin_memory=pin_memory)
236
self.min_p_cpu = self.min_p_cpu_tensor.numpy()
237
- # Pre-allocated device tensor
238
- self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
239
- dtype=torch.float32,
240
- device=device)
237
+
238
+ self.use_double_tensor = torch.device("cpu") != torch.device(device)
239
+
240
+ if self.use_double_tensor:
241
+ # Pre-allocated device tensor
242
+ self.min_p_device: torch.Tensor = torch.empty((max_num_reqs, ),
243
+ dtype=torch.float32,
244
+ device=device)
245
+ else:
246
+ self.min_p_device = self.min_p_cpu_tensor
247
# Current slice of the device tensor
248
self.min_p: torch.Tensor = self.min_p_device[:0]
249
@@ -284,7 +290,9 @@ class MinPLogitsProcessor(LogitsProcessor):
290
size = batch_update.batch_size
291
if self.min_p_count and (needs_update or self.min_p.shape[0] != size):
292
self.min_p = self.min_p_device[:size]
287
- self.min_p.copy_(self.min_p_cpu_tensor[:size], non_blocking=True)
293
+ if self.use_double_tensor:
294
+ self.min_p.copy_(self.min_p_cpu_tensor[:size],
295
+ non_blocking=True)
296
self.min_p.unsqueeze_(1)
297
298
def apply(self, logits: torch.Tensor) -> torch.Tensor: