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 12, 01:00 UTC
[BugFix] Fix DBO assert `assert B_block_table == B_q` (#29933)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Dec 4, 2025 at 14:48 UTC
c8ab988b15af5e30e87c6eb27a0ededf0377ac9e
5 files changed
+83
-63
tests/v1/attention/test_attention_splitting.py
+9
-3
@@ -13,7 +13,7 @@ from vllm.v1.attention.backends.utils import (
13
split_attn_metadata,
14
split_decodes_and_prefills,
15
)
16
-from vllm.v1.worker.ubatch_utils import create_ubatch_slices
16
+from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices
17
18
19
@pytest.fixture
@@ -294,8 +294,14 @@ def test_prefill_split_across_ubatches(
294
qsl_np = common.query_start_loc_cpu.numpy()
295
num_tokens = common.num_actual_tokens
296
297
- ubatch_slices = create_ubatch_slices(num_scheduled_tokens, split_point)
298
- assert len(ubatch_slices) == 2
297
+ ubatch_slices, _ = maybe_create_ubatch_slices(
298
+ True,
299
+ num_scheduled_tokens,
300
+ num_tokens,
301
+ batch_spec.batch_size,
302
+ split_point=split_point,
303
+ )
304
+ assert ubatch_slices is not None and len(ubatch_slices) == 2
305
306
first_meta = _make_metadata_with_slice(ubatch_slices[0], common)
307
second_meta = _make_metadata_with_slice(ubatch_slices[1], common)
vllm/v1/spec_decode/eagle.py
+2
-2
@@ -1258,7 +1258,7 @@ class EagleProposer:
1258
num_tokens_padded: int,
1259
) -> tuple[int, torch.Tensor]:
1260
# TODO(Flechman): support DBO ubatching
1261
- ubatch_slices, num_toks_across_dp = coordinate_batch_across_dp(
1261
+ should_ubatch, num_toks_across_dp = coordinate_batch_across_dp(
1262
num_tokens_unpadded=num_tokens_unpadded,
1263
parallel_config=self.vllm_config.parallel_config,
1264
allow_microbatching=False,
@@ -1267,7 +1267,7 @@ class EagleProposer:
1267
uniform_decode=None,
1268
num_scheduled_tokens_per_request=None,
1269
)
1270
- assert ubatch_slices is None, "DBO ubatching not implemented for EAGLE"
1270
+ assert not should_ubatch, "DBO ubatching not implemented for EAGLE"
1271
1272
num_tokens_dp_padded = num_tokens_padded
1273
if num_toks_across_dp is not None:
vllm/v1/worker/dp_utils.py
+4
-39
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+
5
import numpy as np
6
import torch
7
import torch.distributed as dist
@@ -9,10 +10,7 @@ from vllm.config import ParallelConfig
10
from vllm.distributed.parallel_state import get_dp_group
11
from vllm.logger import init_logger
12
from vllm.v1.worker.ubatch_utils import (
12
- UBatchSlice,
13
- UBatchSlices,
13
check_ubatch_thresholds,
15
- create_ubatch_slices,
14
is_second_ubatch_empty,
15
)
16
@@ -91,20 +89,6 @@ def _post_process_dp_padding(tensor: torch.Tensor, should_dp_pad: bool) -> torch
89
return num_tokens_across_dp.cpu()
90
91
94
-# This just pads the second ubatch slice out to the total number of tokens
95
-# (num_tokens + padding) since we do `create_ubatch_slices` before applying DP padding.
96
-def _pad_out_ubatch_slice(
97
- ubatch_slices: UBatchSlices, num_total_tokens: int
98
-) -> UBatchSlices:
99
- padded_second_token_slice = slice(
100
- ubatch_slices[1].token_slice.start, num_total_tokens
101
- )
102
- ubatch_slices[1] = UBatchSlice(
103
- ubatch_slices[1].request_slice, padded_second_token_slice
104
- )
105
- return ubatch_slices
106
-
107
-
92
def _synchronize_dp_ranks(
93
num_tokens_unpadded: int,
94
num_tokens_padded: int,
@@ -175,7 +159,7 @@ def coordinate_batch_across_dp(
159
num_tokens_padded: int | None = None,
160
uniform_decode: bool | None = None,
161
num_scheduled_tokens_per_request: np.ndarray | None = None,
178
-) -> tuple[UBatchSlices | None, torch.Tensor | None]:
162
+) -> tuple[bool, torch.Tensor | None]:
163
"""
164
Coordinates amongst all DP ranks to determine if and how the full batch
165
should be split into microbatches.
@@ -204,7 +188,7 @@ def coordinate_batch_across_dp(
188
"""
189
if parallel_config.data_parallel_size == 1:
190
# Early exit.
207
- return None, None
191
+ return False, None
192
193
# If the caller has explicitly enabled microbatching.
194
should_attempt_ubatching = False
@@ -228,23 +212,4 @@ def coordinate_batch_across_dp(
212
parallel_config,
213
)
214
231
- # Don't microbatch unless every other DP worker is also microbatching
232
- if not should_ubatch:
233
- return (None, num_tokens_after_padding)
234
-
235
- # This doesn't actually pad the ubatch slices. It just initializes the
236
- # split point to the padded value so that padding can be applied
237
- # to the second ubatch in pad_out_ubatch_slice after attention
238
- # metadata creation
239
- assert num_tokens_after_padding is not None
240
- num_tokens_padded = int(num_tokens_after_padding[0].item())
241
- token_split_point = int(num_tokens_padded) // 2
242
-
243
- assert num_scheduled_tokens_per_request is not None
244
- ubatch_slices = create_ubatch_slices(
245
- num_scheduled_tokens_per_request, token_split_point
246
- )
247
- ubatch_slices = _pad_out_ubatch_slice(ubatch_slices, num_tokens_padded)
248
- assert sum(s.num_tokens for s in ubatch_slices) == num_tokens_padded
249
-
250
- return (ubatch_slices, num_tokens_after_padding)
215
+ return (should_ubatch, num_tokens_after_padding)
vllm/v1/worker/gpu_model_runner.py
+29
-16
@@ -153,6 +153,7 @@ from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
153
from vllm.v1.worker.ubatch_utils import (
154
UBatchSlices,
155
check_ubatch_thresholds,
156
+ maybe_create_ubatch_slices,
157
)
158
from vllm.v1.worker.utils import is_residual_scattered_for_sp
159
@@ -2743,7 +2744,7 @@ class GPUModelRunner(
2744
) -> tuple[
2745
CUDAGraphMode,
2746
BatchDescriptor,
2746
- UBatchSlices | None,
2747
+ bool,
2748
torch.Tensor | None,
2749
CUDAGraphStat | None,
2750
]:
@@ -2779,7 +2780,7 @@ class GPUModelRunner(
2780
2781
# Extra coordination when running data-parallel since we need to coordinate
2782
# across ranks
2782
- ubatch_slices, num_tokens_across_dp = None, None
2783
+ should_ubatch, num_tokens_across_dp = False, None
2784
if self.vllm_config.parallel_config.data_parallel_size > 1:
2785
# Disable DP padding when running eager to avoid excessive padding when
2786
# running prefills. This lets us set cudagraph_mode="NONE" on the prefiller
@@ -2789,8 +2790,8 @@ class GPUModelRunner(
2790
self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
2791
)
2792
2792
- ubatch_slices, num_tokens_across_dp = coordinate_batch_across_dp(
2793
- num_tokens_unpadded=num_tokens_padded,
2793
+ should_ubatch, num_tokens_across_dp = coordinate_batch_across_dp(
2794
+ num_tokens_unpadded=num_tokens,
2795
parallel_config=self.parallel_config,
2796
allow_microbatching=allow_microbatching,
2797
allow_dp_padding=allow_dp_padding,
@@ -2822,7 +2823,7 @@ class GPUModelRunner(
2823
return (
2824
cudagraph_mode,
2825
batch_descriptor,
2825
- ubatch_slices,
2826
+ should_ubatch,
2827
num_tokens_across_dp,
2828
cudagraph_stats,
2829
)
@@ -2921,7 +2922,7 @@ class GPUModelRunner(
2922
(
2923
cudagraph_mode,
2924
batch_desc,
2924
- ubatch_slices,
2925
+ should_ubatch,
2926
num_tokens_across_dp,
2927
cudagraph_stats,
2928
) = self._determine_batch_execution_and_padding(
@@ -2934,10 +2935,10 @@ class GPUModelRunner(
2935
2936
logger.debug(
2937
"Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
2937
- "ubatch_slices: %s, num_tokens_across_dp: %s",
2938
+ "should_ubatch: %s, num_tokens_across_dp: %s",
2939
cudagraph_mode,
2940
batch_desc,
2940
- ubatch_slices,
2941
+ should_ubatch,
2942
num_tokens_across_dp,
2943
)
2944
@@ -2945,10 +2946,18 @@ class GPUModelRunner(
2946
num_reqs_padded = (
2947
batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
2948
)
2949
+ ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
2950
+ should_ubatch,
2951
+ num_scheduled_tokens_np,
2952
+ num_tokens_padded,
2953
+ num_reqs_padded,
2954
+ )
2955
2949
- use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
2956
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
2957
2958
+ use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
2959
+ ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
2960
+
2961
(attn_metadata, spec_decode_common_attn_metadata) = (
2962
self._build_attention_metadata(
2963
num_tokens=num_tokens_unpadded,
@@ -2956,7 +2965,7 @@ class GPUModelRunner(
2965
num_reqs=num_reqs,
2966
num_reqs_padded=num_reqs_padded if pad_attn else None,
2967
max_query_len=max_num_scheduled_tokens,
2959
- ubatch_slices=ubatch_slices,
2968
+ ubatch_slices=ubatch_slices_attn,
2969
logits_indices=logits_indices,
2970
use_spec_decode=use_spec_decode,
2971
num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
@@ -2993,7 +3002,7 @@ class GPUModelRunner(
3002
num_tokens_across_dp=num_tokens_across_dp,
3003
cudagraph_runtime_mode=cudagraph_mode,
3004
batch_descriptor=batch_desc,
2996
- ubatch_slices=ubatch_slices,
3005
+ ubatch_slices=ubatch_slices_padded,
3006
),
3007
record_function_or_nullcontext("gpu_model_runner: forward"),
3008
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
@@ -3945,7 +3954,7 @@ class GPUModelRunner(
3954
3955
num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
3956
3948
- _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp, _ = (
3957
+ _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = (
3958
self._determine_batch_execution_and_padding(
3959
num_tokens=num_tokens_unpadded,
3960
num_reqs=num_reqs,
@@ -3979,6 +3988,9 @@ class GPUModelRunner(
3988
num_reqs_padded = (
3989
batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
3990
)
3991
+ ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
3992
+ should_ubatch, num_scheduled_tokens, num_tokens_padded, num_reqs_padded
3993
+ )
3994
3995
attn_metadata: PerLayerAttnMetadata | None = None
3996
@@ -4000,11 +4012,12 @@ class GPUModelRunner(
4012
self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
4013
self.query_start_loc.copy_to_gpu()
4014
4015
+ pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL
4016
attn_metadata, _ = self._build_attention_metadata(
4017
num_tokens=num_tokens_unpadded,
4018
num_reqs=num_reqs_padded,
4019
max_query_len=max_query_len,
4007
- ubatch_slices=ubatch_slices,
4020
+ ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
4021
for_cudagraph_capture=is_graph_capturing,
4022
)
4023
@@ -4056,11 +4069,11 @@ class GPUModelRunner(
4069
num_tokens_padded, None, False
4070
)
4071
4059
- if ubatch_slices is not None:
4072
+ if ubatch_slices_padded is not None:
4073
# Adjust values to reflect a single ubatch.
4074
# TODO(sage,lucas): this is cruft that should be addressed in
4075
# the padding refactor.
4063
- num_tokens_padded = ubatch_slices[0].num_tokens
4076
+ num_tokens_padded = ubatch_slices_padded[0].num_tokens
4077
if num_tokens_across_dp is not None:
4078
num_tokens_across_dp[:] = num_tokens_padded
4079
@@ -4073,7 +4086,7 @@ class GPUModelRunner(
4086
num_tokens_across_dp=num_tokens_across_dp,
4087
cudagraph_runtime_mode=cudagraph_runtime_mode,
4088
batch_descriptor=batch_desc,
4076
- ubatch_slices=ubatch_slices,
4089
+ ubatch_slices=ubatch_slices_padded,
4090
),
4091
):
4092
outputs = self.model(
vllm/v1/worker/ubatch_utils.py
+39
-3
@@ -42,9 +42,37 @@ def check_ubatch_thresholds(
42
return num_tokens >= config.dbo_prefill_token_threshold
43
44
45
-def create_ubatch_slices(
46
- num_scheduled_tokens: np.ndarray, split_point: int
45
+# This just pads the second ubatch slice out to the total number of tokens
46
+# (num_tokens + padding) since we do `create_ubatch_slices` before applying DP padding.
47
+def _pad_out_ubatch_slices(
48
+ ubatch_slices: UBatchSlices, num_total_tokens: int, num_reqs_padded: int
49
) -> UBatchSlices:
50
+ # TODO(lucas): handle empty second ubatch
51
+ padded_second_request_slice = slice(
52
+ ubatch_slices[1].request_slice.start, num_reqs_padded
53
+ )
54
+ padded_second_token_slice = slice(
55
+ ubatch_slices[1].token_slice.start, num_total_tokens
56
+ )
57
+ return [
58
+ ubatch_slices[0],
59
+ UBatchSlice(padded_second_request_slice, padded_second_token_slice),
60
+ ]
61
+
62
+
63
+def maybe_create_ubatch_slices(
64
+ should_ubatch: bool,
65
+ num_scheduled_tokens: np.ndarray,
66
+ num_tokens_padded: int,
67
+ num_reqs_padded: int,
68
+ split_point: int | None = None,
69
+) -> tuple[UBatchSlices | None, UBatchSlices | None]:
70
+ if not should_ubatch:
71
+ return None, None
72
+
73
+ if split_point is None:
74
+ split_point = int(num_tokens_padded) // 2
75
+
76
# TODO(lucas): Refactor the gpu_model_runner.py so we can pass
77
# in cu_num_tokens directly (i.e. query_start_loc)
78
cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32)
@@ -67,7 +95,15 @@ def create_ubatch_slices(
95
)
96
second_ubatch_req_slice = slice(second_ubatch_req_start, len(cu_num_tokens) - 1)
97
70
- return [
98
+ ubatch_slices = [
99
UBatchSlice(first_ubatch_req_slice, first_ubatch_token_slice),
100
UBatchSlice(second_ubatch_req_slice, second_ubatch_token_slice),
101
]
102
+
103
+ ubatch_slices_padded = _pad_out_ubatch_slices(
104
+ ubatch_slices, num_tokens_padded, num_reqs_padded
105
+ )
106
+
107
+ assert sum(s.num_tokens for s in ubatch_slices_padded) == num_tokens_padded
108
+
109
+ return ubatch_slices, ubatch_slices_padded