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 14, 05:30 UTC
[TPU][Bugfix] fix the missing apply_model in tpu worker (#25526)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Sep 23, 2025 at 22:18 UTC
190c45a6af666ce2be25a73aa68301e2b1de309e
2 files changed
+8
-6
tests/v1/tpu/test_tpu_int8.py
+1
-5
@@ -48,13 +48,9 @@ def test_model_tpu_int8(vllm_runner, model: str, dtype: str, max_tokens: int,
48
49
prompts = [
50
"A robot may not injure a human being",
51
- "It is only with the heart that one can see rightly;",
52
- "The greatest glory in living lies not in never falling,",
51
]
52
answers = [
55
- "or, being injured, not kill, except in",
56
- "without the heart, one can only see wrongly.",
57
- "but in rising every time we fall. - Nelson"
53
+ "or kill a human being",
54
]
55
56
with vllm_runner(model, dtype=dtype, hf_overrides=hf_overrides) as vllm:
vllm/v1/worker/tpu_worker.py
+7
-1
@@ -3,7 +3,7 @@
3
"""A TPU worker class."""
4
5
import os
6
-from typing import Any, Optional
6
+from typing import Any, Callable, Optional, TypeVar
7
8
import torch
9
import torch.distributed
@@ -31,6 +31,8 @@ from vllm.v1.worker.utils import bind_kv_cache
31
32
logger = init_logger(__name__)
33
34
+_R = TypeVar("_R")
35
+
36
if not USE_TPU_COMMONS:
37
logger.info("tpu_commons not found, using vLLM's TPUWorker.")
38
import torch_xla.core.xla_model as xm
@@ -333,6 +335,10 @@ class TPUWorker:
335
def shutdown(self) -> None:
336
self.model_runner.ensure_kv_transfer_shutdown()
337
338
+ def apply_model(self, fn: Callable[[nn.Module], _R]) -> _R:
339
+ """Apply a function on the model inside this worker."""
340
+ return fn(self.get_model())
341
+
342
343
if USE_TPU_COMMONS:
344
from tpu_commons.worker import TPUWorker as TPUCommonsWorker