fixes

vllm-project · Feb 11, 2025 · 6f5b773 · 6f5b773
1 parent 9a0da81
commit 6f5b773
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/vllm/v1/worker/tpu_model_runner.py b/vllm/v1/worker/tpu_model_runner.py
@@ -76,7 +76,7 @@ class DecodeData:
     attn_metadata: Optional[PallasMetadata] = None
 
 
-class TPUModelRunner(ModelRunnerBase):
+class TPUModelRunner:
 
     def __init__(
         self,
@@ -692,7 +692,7 @@ def execute_model(
         scheduler_output: "SchedulerOutput",
     ) -> ModelRunnerOutput:
         # Update cached state
-        self.update_states(scheduler_output)
+        self._update_states(scheduler_output)
 
         # If necessary, swap decodes/prompts to have all decodes on the start
         ensure_decodes_first(self.input_batch)

diff --git a/vllm/v1/worker/tpu_worker.py b/vllm/v1/worker/tpu_worker.py
@@ -4,6 +4,7 @@
 
 import torch
 import torch.distributed
+import torch.nn as nn
 import torch_xla.core.xla_model as xm
 import torch_xla.runtime as xr
 
@@ -19,14 +20,13 @@
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, LayerBlockType, get_dtype_size
 from vllm.v1.worker.tpu_model_runner import ExecutionMode, TPUModelRunner
-from vllm.v1.worker.worker_base import WorkerBase
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 from vllm.v1.utils import bind_kv_cache
 
 logger = init_logger(__name__)
 
 
-class TPUWorker(WorkerBase):
+class TPUWorker:
 
     def __init__(
         self,
@@ -212,6 +212,7 @@ def check_health(self) -> None:
         # worker will always be healthy as long as it's running.
         return
 
+
 def init_tpu_worker_distributed_environment(
     parallel_config: ParallelConfig,
     rank: int,