[TPU] Async output processing for TPU (#8011)

2024-08-29 19:35:29 -07:00
parent 428dd1445e
commit 80c7b089b1
2 changed files with 10 additions and 4 deletions
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -347,10 +347,10 @@ class ModelConfig:
            self.use_async_output_proc = False
            return

-        if device_config.device_type != "cuda":
+        if device_config.device_type not in ("cuda", "tpu"):
            logger.warning(
-                "Async output processing is only supported for CUDA."
-                " Disabling it for other platforms.")
+                "Async output processing is only supported for CUDA or TPU. "
+                "Disabling it for other platforms.")
            self.use_async_output_proc = False
            return