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, 16:00 UTC
[Bugfix][TPU] Fix V1 TPU worker for sliding window (#16059)
Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
Apr 4, 2025 at 17:31 UTC
70ad3f9e98687776772c446530a091c4e2019e7b
1 file changed
+4
-3
vllm/v1/worker/tpu_worker.py
+4
-3
@@ -18,7 +18,7 @@ from vllm.logger import init_logger
18
from vllm.model_executor import set_random_seed
19
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
20
from vllm.v1.core.sched.output import SchedulerOutput
21
-from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
21
+from vllm.v1.kv_cache_interface import (AttentionSpec, KVCacheConfig,
22
KVCacheSpec)
23
from vllm.v1.outputs import ModelRunnerOutput
24
from vllm.v1.utils import bind_kv_cache
@@ -137,7 +137,7 @@ class TPUWorker:
137
kv_caches: dict[str, torch.Tensor] = {}
138
kv_cache_spec = self.model_runner.get_kv_cache_spec()
139
for layer_name, layer_spec in kv_cache_spec.items():
140
- if isinstance(layer_spec, FullAttentionSpec):
140
+ if isinstance(layer_spec, AttentionSpec):
141
dtype = layer_spec.dtype
142
143
# Use an empty tensor instead of `None`` to force Dynamo to pass
@@ -147,7 +147,8 @@ class TPUWorker:
147
device=self.device)
148
kv_caches[layer_name] = tpu_kv_cache
149
else:
150
- raise NotImplementedError
150
+ raise NotImplementedError(
151
+ f"Unsupported KV cache spec '{type(layer_spec)}'")
152
153
runner_kv_caches: list[torch.Tensor] = []
154
bind_kv_cache(