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, 12:15 UTC
[Frontend][TPU] Add TPU default max-num-batched-tokens based on device name (#17508)
Signed-off-by: Chenyaaang <chenyangli@google.com>
Chenyaaang committed
May 2, 2025 at 21:42 UTC
87baebebd8c960a36e2e778d4fd97ff84843d5e3
3 files changed
+40
-4
vllm/config.py
+8
@@ -3140,6 +3140,14 @@ def _get_and_verify_max_len(
3140
# derived length from the HF model config.
3141
if max_model_len is None:
3142
max_model_len = int(derived_max_model_len)
3143
+ if current_platform.is_tpu():
3144
+ logger.warning(
3145
+ "--max-model-len is not specified, "
3146
+ "it's currently using model's default length %s, "
3147
+ "which might be too large."
3148
+ "Please input with --max-model-len based on your "
3149
+ "request input length and output length, to avoid "
3150
+ "unnecessary degradation.", max_model_len)
3151
elif max_model_len > derived_max_model_len:
3152
# Some models might have a separate key for specifying model_max_length
3153
# that will be bigger than derived_max_model_len. We compare user input
vllm/engine/arg_utils.py
+29
-3
@@ -1441,8 +1441,8 @@ class EngineArgs:
1441
# as the platform that vLLM is running on (e.g. the case of scaling
1442
# vLLM with Ray) and has no GPUs. In this case we use the default
1443
# values for non-H100/H200 GPUs.
1444
+ from vllm.platforms import current_platform
1445
try:
1445
- from vllm.platforms import current_platform
1446
device_memory = current_platform.get_device_total_memory()
1447
except Exception:
1448
# This is only used to set default_max_num_batched_tokens
@@ -1463,11 +1463,37 @@ class EngineArgs:
1463
}
1464
default_max_num_seqs = 256
1465
1466
+ # tpu specific default values.
1467
+ if current_platform.is_tpu():
1468
+ default_max_num_batched_tokens_tpu = {
1469
+ UsageContext.LLM_CLASS: {
1470
+ 'V6E': 2048,
1471
+ 'V5E': 1024,
1472
+ 'V5P': 512,
1473
+ },
1474
+ UsageContext.OPENAI_API_SERVER: {
1475
+ 'V6E': 1024,
1476
+ 'V5E': 512,
1477
+ 'V5P': 256,
1478
+ }
1479
+ }
1480
+
1481
use_context_value = usage_context.value if usage_context else None
1482
if (self.max_num_batched_tokens is None
1483
and usage_context in default_max_num_batched_tokens):
1469
- self.max_num_batched_tokens = default_max_num_batched_tokens[
1470
- usage_context]
1484
+ if current_platform.is_tpu():
1485
+ chip_name = current_platform.get_device_name()
1486
+ if chip_name in default_max_num_batched_tokens_tpu[
1487
+ usage_context]:
1488
+ self.max_num_batched_tokens = \
1489
+ default_max_num_batched_tokens_tpu[
1490
+ usage_context][chip_name]
1491
+ else:
1492
+ self.max_num_batched_tokens = \
1493
+ default_max_num_batched_tokens[usage_context]
1494
+ else:
1495
+ self.max_num_batched_tokens = default_max_num_batched_tokens[
1496
+ usage_context]
1497
logger.debug(
1498
"Setting max_num_batched_tokens to %d for %s usage context.",
1499
self.max_num_batched_tokens, use_context_value)
vllm/platforms/tpu.py
+3
-1
@@ -3,6 +3,7 @@
3
from typing import TYPE_CHECKING, Optional, Union
4
5
import torch
6
+from tpu_info import device
7
8
import vllm.envs as envs
9
from vllm.inputs import ProcessorInputs, PromptType
@@ -54,7 +55,8 @@ class TpuPlatform(Platform):
55
56
@classmethod
57
def get_device_name(cls, device_id: int = 0) -> str:
57
- return "tpu"
58
+ chip_type, _ = device.get_local_chips()
59
+ return f"TPU {chip_type.name}"
60
61
@classmethod
62
def get_device_total_memory(cls, device_id: int = 0) -> int: