[V1] Decouple GPU and TPU InputBatch (#19778)

Signed-off-by: Andrew Feldman <afeldman@redhat.com>

Author: afeldman-nm

vllm/v1/sample/tpu/metadata.py

@@ -5,7 +5,7 @@
 
 import torch
 
-from vllm.v1.worker.gpu_input_batch import InputBatch
+from vllm.v1.worker.tpu_input_batch import InputBatch
 
 DEFAULT_SAMPLING_PARAMS = dict(
     temperature=-1.0,

vllm/v1/worker/gpu_input_batch.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# Datastructures defining an input batch
+# Datastructures defining a GPU input batch
 
 from dataclasses import dataclass
 from typing import Optional, cast
@@ -453,6 +453,11 @@ def swap_states(self, i1: int, i2: int) -> None:
         self.block_table.swap_row(i1, i2)
 
     def condense(self, empty_req_indices: list[int]) -> None:
+        """Move non-empty requests down into lower, empty indices.
+        
+        Args:
+          empty_req_indices: empty batch indices, sorted descending.
+        """
         num_reqs = self.num_reqs
         if num_reqs == 0:
             # The batched states are empty.

vllm/v1/worker/lora_model_runner_mixin.py

@@ -5,6 +5,7 @@
 """
 
 from contextlib import contextmanager
+from typing import Union
 
 import numpy as np
 import torch.nn as nn
@@ -15,7 +16,10 @@
 from vllm.lora.request import LoRARequest
 from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
 from vllm.model_executor.models import supports_lora, supports_multimodal
-from vllm.v1.worker.gpu_input_batch import InputBatch
+from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch
+from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch
+
+InputBatch = Union[TPUInputBatch, GPUInputBatch]
 
 logger = init_logger(__name__)