[TPU] Skip creating empty tensor (#7630)

ce143353 · Woosuk Kwon · GitHub · bbf55c48 · ce143353
Unverified Commit ce143353 authored Aug 17, 2024 by Woosuk Kwon Committed by GitHub Aug 17, 2024
Show whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

vllm/worker/tpu_worker.py vllm/worker/tpu_worker.py +4 -1

No files found.
--- a/vllm/worker/tpu_worker.py
+++ b/vllm/worker/tpu_worker.py
@@ -271,7 +271,10 @@ def _make_src_to_dst(
    mapping: List[Tuple[int, int]],
    src_device: Union[torch.device, str],
    dst_device: Union[torch.device, str],
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> Optional[Tuple[torch.Tensor, torch.Tensor]]:
+    if not mapping:
+        return None
+
    src_indices = [i for i, _ in mapping]
    dst_indices = [i for _, i in mapping]
    src_indices = torch.tensor(src_indices,