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 7, 22:45 UTC
[Bug] Refactor max_num_batched_tokens to account for drafting (#34898)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Benjamin Chislett committed
Feb 22, 2026 at 11:18 UTC
682566b18e69d12a1ee603906417f508d61ac7ea
5 files changed
+63
-20
vllm/config/scheduler.py
+8
-1
@@ -46,12 +46,19 @@ class SchedulerConfig:
46
"""The runner type to launch for the model."""
47
48
max_num_batched_tokens: int = Field(default=DEFAULT_MAX_NUM_BATCHED_TOKENS, ge=1)
49
- """Maximum number of tokens to be processed in a single iteration.
49
+ """Maximum number of tokens that can be processed in a single iteration.
50
51
The default value here is mainly for convenience when testing.
52
In real usage, this should be set in `EngineArgs.create_engine_config`.
53
"""
54
55
+ max_num_scheduled_tokens: int | None = Field(default=None)
56
+ """Maximum number of tokens that the scheduler may issue in a single iteration.
57
+
58
+ This is usually equal to max_num_batched_tokens, but can be smaller in cases
59
+ when the model might append tokens into the batch (such as speculative decoding).
60
+ Defaults to max_num_batched_tokens."""
61
+
62
max_num_seqs: int = Field(default=DEFAULT_MAX_NUM_SEQS, ge=1)
63
"""Maximum number of sequences to be processed in a single iteration.
64
vllm/config/speculative.py
+16
@@ -750,6 +750,22 @@ class SpeculativeConfig:
750
f"errors during speculative decoding."
751
)
752
753
+ @property
754
+ def max_num_new_slots_for_drafting(self) -> int:
755
+ """
756
+ Calculate the maximum number of new slots that might be added to the batch
757
+ when drafting.
758
+ """
759
+ slots_per_req = 0 # for serial non-draft-model methods, no change needed
760
+ if self.parallel_drafting:
761
+ # For parallel drafting, we need one new slot per 'masked' token
762
+ slots_per_req = self.num_speculative_tokens - 1
763
+ if self.uses_draft_model():
764
+ # For draft model-based speculation, we need one new slot per request
765
+ # Since we do not slice the draft tokens
766
+ slots_per_req += 1
767
+ return slots_per_req
768
+
769
def use_eagle(self) -> bool:
770
return self.method in ("eagle", "eagle3", "mtp")
771
vllm/config/vllm.py
+33
-14
@@ -822,6 +822,8 @@ class VllmConfig:
822
self.speculative_config is None
823
)
824
825
+ self._set_max_num_scheduled_tokens()
826
+
827
if current_platform.support_static_graph_mode():
828
# if cudagraph_mode has full cudagraphs, we need to check support
829
if model_config := self.model_config:
@@ -1185,6 +1187,37 @@ class VllmConfig:
1187
if size % self.parallel_config.tensor_parallel_size == 0
1188
]
1189
1190
+ def _set_max_num_scheduled_tokens(self):
1191
+ """
1192
+ In most cases, the scheduler may schedule a batch with as many tokens as the
1193
+ worker is configured to handle. However for some speculative decoding methods,
1194
+ the drafter model may insert additional slots into the batch when drafting.
1195
+ To account for this, we need to decrease the max_num_scheduled_tokens by an
1196
+ upper bound on the number of slots that can be added.
1197
+ """
1198
+ if self.speculative_config is not None:
1199
+ scheduled_token_delta = (
1200
+ self.speculative_config.max_num_new_slots_for_drafting
1201
+ * self.scheduler_config.max_num_seqs
1202
+ )
1203
+ max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
1204
+ if self.scheduler_config.max_num_scheduled_tokens is None:
1205
+ self.scheduler_config.max_num_scheduled_tokens = (
1206
+ max_num_batched_tokens - scheduled_token_delta
1207
+ )
1208
+
1209
+ max_num_scheduled_tokens = self.scheduler_config.max_num_scheduled_tokens
1210
+ if max_num_batched_tokens < max_num_scheduled_tokens + (
1211
+ self.speculative_config.max_num_new_slots_for_drafting
1212
+ * self.scheduler_config.max_num_seqs
1213
+ ):
1214
+ raise ValueError(
1215
+ f"VllmConfig received max_num_scheduled_tokens but it does not have"
1216
+ " enough slots to support the speculative decoding settings."
1217
+ f" It should be greater by at least {scheduled_token_delta}, but"
1218
+ f" got {max_num_batched_tokens=} and {max_num_scheduled_tokens=}."
1219
+ )
1220
+
1221
def _set_cudagraph_sizes(self):
1222
"""
1223
vLLM defines the default candidate list of batch sizes for CUDA graph
@@ -1347,22 +1380,8 @@ class VllmConfig:
1380
computed_compile_ranges_split_points = []
1381
1382
# The upper bound of the compile ranges is the max_num_batched_tokens.
1350
- # For speculative decoding, the compile range must be extended
1351
- # - Sequential: + 1 * max_num_seqs (one draft token per iteration)
1352
- # - Parallel draft: + num_speculative_tokens * max_num_seqs
1383
compile_range_end = self.scheduler_config.max_num_batched_tokens
1384
if compile_range_end is not None:
1355
- if self.speculative_config is not None and (
1356
- self.speculative_config.uses_draft_model()
1357
- or self.speculative_config.use_eagle()
1358
- ):
1359
- multiplier = (
1360
- self.speculative_config.num_speculative_tokens
1361
- if self.speculative_config.parallel_drafting
1362
- else 1
1363
- )
1364
- compile_range_end += multiplier * self.scheduler_config.max_num_seqs
1365
-
1385
computed_compile_ranges_split_points.append(compile_range_end)
1386
1387
# Add the compile ranges for flashinfer
vllm/v1/core/sched/scheduler.py
+5
-1
@@ -99,7 +99,11 @@ class Scheduler(SchedulerInterface):
99
100
# Scheduling constraints.
101
self.max_num_running_reqs = self.scheduler_config.max_num_seqs
102
- self.max_num_scheduled_tokens = self.scheduler_config.max_num_batched_tokens
102
+ self.max_num_scheduled_tokens = (
103
+ self.scheduler_config.max_num_scheduled_tokens
104
+ if self.scheduler_config.max_num_scheduled_tokens
105
+ else self.scheduler_config.max_num_batched_tokens
106
+ )
107
self.max_model_len = vllm_config.model_config.max_model_len
108
self.enable_kv_cache_events = (
109
self.kv_events_config is not None
vllm/v1/spec_decode/eagle.py
+1
-4
@@ -100,11 +100,8 @@ class SpecDecodeBaseProposer:
100
if self.parallel_drafting:
101
self._init_parallel_drafting_params()
102
103
- # The drafter can get longer sequences than the target model.
103
max_batch_size = vllm_config.scheduler_config.max_num_seqs
105
- self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + (
106
- self.net_num_new_slots_per_request * max_batch_size
107
- )
104
+ self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
105
self.token_arange_np = np.arange(self.max_num_tokens)
106
107
# Multi-modal data support