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, 11:00 UTC
[Perf] Use small max_num_batched_tokens for A100 (#17885)
Signed-off-by: KuntaiDu <kuntai@uchicago.edu>
Kuntai Du committed
May 11, 2025 at 00:53 UTC
911215528324a52d74a729335506aa5ec0a7cc65
1 file changed
+5
-1
vllm/engine/arg_utils.py
+5
-1
@@ -1438,11 +1438,15 @@ class EngineArgs:
1438
from vllm.platforms import current_platform
1439
try:
1440
device_memory = current_platform.get_device_total_memory()
1441
+ device_name = current_platform.get_device_name().lower()
1442
except Exception:
1443
# This is only used to set default_max_num_batched_tokens
1444
device_memory = 0
1445
1445
- if device_memory >= 70 * GiB_bytes:
1446
+ # NOTE(Kuntai): Setting large `max_num_batched_tokens` for A100 reduces
1447
+ # throughput, see PR #17885 for more details.
1448
+ # So here we do an extra device name check to prevent such regression.
1449
+ if device_memory >= 70 * GiB_bytes and "a100" not in device_name:
1450
# For GPUs like H100 and MI300x, use larger default values.
1451
default_max_num_batched_tokens = {
1452
UsageContext.LLM_CLASS: 16384,