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
[BugFix] Fix using `dbo_decode_token_threshold` always (and ignoring `dbo_prefill_token_threshold`) (#25622)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Sep 26, 2025 at 12:22 UTC
984d18498a420d771c4a38c6e6e3f3aacb090157
2 files changed
+9
-3
vllm/v1/worker/gpu_model_runner.py
+7
-2
@@ -1045,11 +1045,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1045
num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
1046
num_tokens_padded = num_tokens_unpadded + self.get_local_padding(
1047
num_tokens_unpadded)
1048
+ uniform_decode = \
1049
+ (max_num_scheduled_tokens == self.uniform_decode_query_len) and \
1050
+ (total_num_scheduled_tokens == num_reqs * max_num_scheduled_tokens)
1051
ubatch_slices, num_tokens_after_padding = \
1052
ubatch_split(num_scheduled_tokens,
1053
num_tokens_unpadded,
1054
num_tokens_padded,
1052
- self.vllm_config)
1055
+ uniform_decode=uniform_decode,
1056
+ vllm_config=self.vllm_config)
1057
1058
self.seq_lens.np[:num_reqs] = (
1059
self.input_batch.num_computed_tokens_cpu[:num_reqs] +
@@ -2989,7 +2993,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2993
num_scheduled_tokens,
2994
total_num_scheduled_tokens,
2995
total_num_scheduled_tokens,
2992
- self.vllm_config,
2996
+ uniform_decode=uniform_decode,
2997
+ vllm_config=self.vllm_config,
2998
)
2999
3000
# If we failed to microbatch, currently need to resynchronize
vllm/v1/worker/ubatch_splitting.py
+2
-1
@@ -139,6 +139,7 @@ def ubatch_split(
139
num_scheduled_tokens_per_request: np.ndarray,
140
num_tokens_unpadded: int,
141
num_tokens_padded: int,
142
+ uniform_decode: bool,
143
vllm_config: VllmConfig,
144
) -> tuple[Optional[UBatchSlices], Optional[torch.Tensor]]:
145
"""
@@ -164,7 +165,7 @@ def ubatch_split(
165
should_attempt_ubatching = check_ubatch_thresholds(
166
parallel_config,
167
num_tokens_unpadded,
167
- vllm_config,
168
+ uniform_decode=uniform_decode,
169
)
170
171
# Don't microbatch unless every other DP worker is also microbatching