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
Fix handling of `max_num_batched_tokens` for pooling tasks (#23004)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com>
Maximilien de Bayser committed
Aug 16, 2025 at 14:36 UTC
52ce1420e9f6f52308f49a2898433a52674a4a8b
2 files changed
+5
-8
vllm/config/__init__.py
-3
@@ -3600,9 +3600,6 @@ class VllmConfig:
3600
logger.info(reason)
3601
self.scheduler_config.chunked_prefill_enabled = False
3602
self.scheduler_config.long_prefill_token_threshold = 0
3603
- self.scheduler_config.max_num_batched_tokens = max(
3604
- self.scheduler_config.max_model_len,
3605
- DEFAULT_MAX_NUM_BATCHED_TOKENS)
3603
3604
if self.cache_config is not None:
3605
self.cache_config.enable_prefix_caching = False
vllm/engine/arg_utils.py
+5
-5
@@ -1602,9 +1602,6 @@ class EngineArgs:
1602
self.enable_prefix_caching = incremental_prefill_supported
1603
logger.info("(%s) prefix caching by default", action)
1604
1605
- if not self.enable_chunked_prefill:
1606
- self.max_num_batched_tokens = model_config.max_model_len
1607
-
1605
# V1 should use the new scheduler by default.
1606
# Swap it only if this arg is set to the original V0 default
1607
if self.scheduler_cls == EngineArgs.scheduler_cls:
@@ -1692,8 +1689,11 @@ class EngineArgs:
1689
self.max_num_batched_tokens = \
1690
default_max_num_batched_tokens[usage_context]
1691
else:
1695
- self.max_num_batched_tokens = default_max_num_batched_tokens[
1696
- usage_context]
1692
+ if not self.enable_chunked_prefill:
1693
+ self.max_num_batched_tokens = model_config.max_model_len
1694
+ else:
1695
+ self.max_num_batched_tokens = \
1696
+ default_max_num_batched_tokens[usage_context]
1697
logger.debug(
1698
"Setting max_num_batched_tokens to %d for %s usage context.",
1699
self.max_num_batched_tokens, use_context_value)