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, 22:15 UTC
[Bugfix] Fix max_num_batched_tokens for MLA (#13620)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Feb 20, 2025 at 20:45 UTC
71face8540042e309b01a18e5f184cc700a8bc6d
1 file changed
+14
-6
vllm/config.py
+14
-6
@@ -51,6 +51,9 @@ else:
51
52
logger = init_logger(__name__)
53
54
+# This value is chosen to have a balance between ITL and TTFT. Note it is
55
+# not optimized for throughput.
56
+_DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048
57
_POOLING_MODEL_MAX_NUM_BATCHED_TOKENS = 32768
58
_MULTIMODAL_MODEL_MAX_NUM_BATCHED_TOKENS = 5120
59
@@ -1526,15 +1529,17 @@ class SchedulerConfig:
1529
# for now. Have max_num_batched_tokens set to max_model_len
1530
# so we don't reject sequences on account of a short
1531
# max_num_batched_tokens.
1529
- self.max_num_batched_tokens = max(self.max_model_len, 2048)
1532
+ self.max_num_batched_tokens = max(
1533
+ self.max_model_len, _DEFAULT_MAX_NUM_BATCHED_TOKENS)
1534
else:
1531
- # This value is chosen to have a balance between ITL
1532
- # and TTFT. Note it is not optimized for throughput.
1533
- self.max_num_batched_tokens = 2048
1535
+ self.max_num_batched_tokens = (
1536
+ _DEFAULT_MAX_NUM_BATCHED_TOKENS)
1537
else:
1535
- # If max_model_len is too short, use 2048 as the default value
1538
+ # If max_model_len is too short, use
1539
+ # _DEFAULT_MAX_NUM_BATCHED_TOKENS as the default value
1540
# for higher throughput.
1537
- self.max_num_batched_tokens = max(self.max_model_len, 2048)
1541
+ self.max_num_batched_tokens = max(
1542
+ self.max_model_len, _DEFAULT_MAX_NUM_BATCHED_TOKENS)
1543
1544
if self.runner_type == "pooling":
1545
# Choose specific value for higher throughput
@@ -3333,6 +3338,9 @@ class VllmConfig:
3338
"caching to be disabled.")
3339
self.scheduler_config.enable_chunked_prefill = False
3340
self.scheduler_config.chunked_prefill_enabled = False
3341
+ self.scheduler_config.max_num_batched_tokens = max(
3342
+ self.scheduler_config.max_model_len,
3343
+ _DEFAULT_MAX_NUM_BATCHED_TOKENS)
3344
3345
if self.cache_config is not None:
3346
self.cache_config.enable_prefix_caching = False