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 11, 10:00 UTC
[feature] extend DBO to XBO (#30120)
Signed-off-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: root <root@hk01dgx028.cm.cluster>
jiangkuaixue123 committed
Dec 16, 2025 at 13:04 UTC
b9ff4f2a8dffc84b2ce226e7e98c33756caf098f
10 files changed
+133
-73
tests/v1/attention/test_attention_splitting.py
+1
@@ -323,6 +323,7 @@ def test_prefill_split_across_ubatches(
323
num_tokens,
324
batch_spec.batch_size,
325
split_point=split_point,
326
+ num_ubatches=2,
327
)
328
assert ubatch_slices is not None and len(ubatch_slices) == 2
329
vllm/config/parallel.py
+10
@@ -156,6 +156,8 @@ class ParallelConfig:
156
157
enable_dbo: bool = False
158
"""Enable dual batch overlap for the model executor."""
159
+ ubatch_size: int = 0
160
+ """Number of ubatch size."""
161
162
dbo_decode_token_threshold: int = 32
163
"""The threshold for dual batch overlap for batches only containing decodes.
@@ -325,6 +327,14 @@ class ParallelConfig:
327
including data parallelism."""
328
return self.world_size * self.data_parallel_size
329
330
+ @property
331
+ def use_ubatching(self) -> bool:
332
+ return self.enable_dbo or self.ubatch_size > 1
333
+
334
+ @property
335
+ def num_ubatches(self) -> int:
336
+ return 2 if self.enable_dbo else self.ubatch_size
337
+
338
def get_next_dp_init_port(self) -> int:
339
"""
340
We might need to initialize process groups in multiple
vllm/config/vllm.py
+5
-2
@@ -870,9 +870,12 @@ class VllmConfig:
870
f"cudagraph_mode={self.compilation_config.cudagraph_mode}"
871
)
872
873
- if self.parallel_config.enable_dbo:
873
+ if self.parallel_config.use_ubatching:
874
a2a_backend = self.parallel_config.all2all_backend
875
- assert a2a_backend in ["deepep_low_latency", "deepep_high_throughput"], (
875
+ assert a2a_backend in [
876
+ "deepep_low_latency",
877
+ "deepep_high_throughput",
878
+ ], (
879
"Microbatching currently only supports the deepep_low_latency and "
880
f"deepep_high_throughput all2all backend. {a2a_backend} is not "
881
"supported. To fix use --all2all-backend=deepep_low_latency or "
vllm/engine/arg_utils.py
+6
@@ -408,6 +408,7 @@ class EngineArgs:
408
enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
409
all2all_backend: str | None = ParallelConfig.all2all_backend
410
enable_dbo: bool = ParallelConfig.enable_dbo
411
+ ubatch_size: int = ParallelConfig.ubatch_size
412
dbo_decode_token_threshold: int = ParallelConfig.dbo_decode_token_threshold
413
dbo_prefill_token_threshold: int = ParallelConfig.dbo_prefill_token_threshold
414
disable_nccl_for_dp_synchronization: bool = (
@@ -841,6 +842,10 @@ class EngineArgs:
842
"--all2all-backend", **parallel_kwargs["all2all_backend"]
843
)
844
parallel_group.add_argument("--enable-dbo", **parallel_kwargs["enable_dbo"])
845
+ parallel_group.add_argument(
846
+ "--ubatch-size",
847
+ **parallel_kwargs["ubatch_size"],
848
+ )
849
parallel_group.add_argument(
850
"--dbo-decode-token-threshold",
851
**parallel_kwargs["dbo_decode_token_threshold"],
@@ -1557,6 +1562,7 @@ class EngineArgs:
1562
enable_expert_parallel=self.enable_expert_parallel,
1563
all2all_backend=self.all2all_backend,
1564
enable_dbo=self.enable_dbo,
1565
+ ubatch_size=self.ubatch_size,
1566
dbo_decode_token_threshold=self.dbo_decode_token_threshold,
1567
dbo_prefill_token_threshold=self.dbo_prefill_token_threshold,
1568
disable_nccl_for_dp_synchronization=self.disable_nccl_for_dp_synchronization,
vllm/v1/attention/backends/utils.py
+9
-5
@@ -201,10 +201,11 @@ def _make_metadata_with_slice(
201
)
202
# NOTE: last token can be outside of the last request if we have CG padding.
203
204
- # If the "middle" request has tokens in both ubatches, we have to split it.
205
- # If ubatch_slice is the first ubatch then we will be splitting the last
206
- # request. If it's the second microbatch, then we will be splitting the
207
- # first request
204
+ # If the request is split across ubatches, we have to adjust the metadata.
205
+ # splits_first_request: The first request in this slice is the continuation of
206
+ # a request that started in a previous slice.
207
+ # splits_last_request: The last request in this slice continues into the
208
+ # next slice.
209
splits_first_request = first_tok > start_locs[first_req]
210
splits_last_request = last_tok < start_locs[last_req + 1] - 1
211
@@ -225,7 +226,10 @@ def _make_metadata_with_slice(
226
seq_lens_cpu = attn_metadata.seq_lens_cpu[request_slice]
227
228
if splits_last_request:
228
- tokens_skipped = query_start_loc_cpu[-1] - token_slice.stop
229
+ # NOTE: We use start_locs (the original query_start_loc_cpu) to calculate
230
+ # the tokens skipped because query_start_loc_cpu might have been modified
231
+ # if splits_first_request is True.
232
+ tokens_skipped = start_locs[last_req + 1] - token_slice.stop
233
query_start_loc[-1] -= tokens_skipped
234
query_start_loc_cpu[-1] -= tokens_skipped
235
vllm/v1/worker/dp_utils.py
+4
-4
@@ -11,7 +11,7 @@ from vllm.distributed.parallel_state import get_dp_group
11
from vllm.logger import init_logger
12
from vllm.v1.worker.ubatch_utils import (
13
check_ubatch_thresholds,
14
- is_second_ubatch_empty,
14
+ is_last_ubatch_empty,
15
)
16
17
logger = init_logger(__name__)
@@ -56,7 +56,7 @@ def _run_ar(
56
return tensor
57
58
59
-def _post_process_ubatch(tensor: torch.Tensor) -> bool:
59
+def _post_process_ubatch(tensor: torch.Tensor, num_ubatches: int) -> bool:
60
orig_num_tokens_tensor = tensor[0, :]
61
padded_num_tokens_tensor = tensor[1, :]
62
@@ -68,7 +68,7 @@ def _post_process_ubatch(tensor: torch.Tensor) -> bool:
68
# there are no "empty" second ubatches
69
orig_min_num_tokens = int(orig_num_tokens_tensor.min().item())
70
padded_max_num_tokens = int(padded_num_tokens_tensor.max().item())
71
- if is_second_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens):
71
+ if is_last_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens, num_ubatches):
72
logger.debug(
73
"Aborting ubatching %s %s", orig_min_num_tokens, padded_max_num_tokens
74
)
@@ -146,7 +146,7 @@ def _synchronize_dp_ranks(
146
assert should_attempt_dp_padding == should_dp_pad
147
148
# Check conditions for microbatching
149
- should_ubatch = _post_process_ubatch(tensor)
149
+ should_ubatch = _post_process_ubatch(tensor, parallel_config.num_ubatches)
150
151
if should_ubatch and not should_dp_pad:
152
logger.debug_once(
vllm/v1/worker/gpu_model_runner.py
+26
-7
@@ -2987,7 +2987,7 @@ class GPUModelRunner(
2987
2988
cascade_attn_prefix_lens = None
2989
# Disable cascade attention when using microbatching (DBO)
2990
- if self.cascade_attn_enabled and not self.parallel_config.enable_dbo:
2990
+ if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
2991
# Pre-compute cascade attention prefix lengths
2992
cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens(
2993
num_scheduled_tokens_np,
@@ -3028,6 +3028,13 @@ class GPUModelRunner(
3028
num_scheduled_tokens_np,
3029
num_tokens_padded,
3030
num_reqs_padded,
3031
+ self.parallel_config.num_ubatches,
3032
+ )
3033
+
3034
+ logger.debug(
3035
+ "ubatch_slices: %s, ubatch_slices_padded: %s",
3036
+ ubatch_slices,
3037
+ ubatch_slices_padded,
3038
)
3039
3040
pad_attn = cudagraph_mode == CUDAGraphMode.FULL
@@ -3710,11 +3717,14 @@ class GPUModelRunner(
3717
# wrap the model with full cudagraph wrapper if needed.
3718
cudagraph_mode = self.compilation_config.cudagraph_mode
3719
assert cudagraph_mode is not None
3713
- if cudagraph_mode.has_full_cudagraphs() and not self.parallel_config.enable_dbo:
3720
+ if (
3721
+ cudagraph_mode.has_full_cudagraphs()
3722
+ and not self.parallel_config.use_ubatching
3723
+ ):
3724
self.model = CUDAGraphWrapper(
3725
self.model, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
3726
)
3717
- elif self.parallel_config.enable_dbo:
3727
+ elif self.parallel_config.use_ubatching:
3728
if cudagraph_mode.has_full_cudagraphs():
3729
self.model = UBatchWrapper(
3730
self.model, self.vllm_config, CUDAGraphMode.FULL, self.device
@@ -4095,7 +4105,16 @@ class GPUModelRunner(
4105
batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
4106
)
4107
ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
4098
- should_ubatch, num_scheduled_tokens, num_tokens_padded, num_reqs_padded
4108
+ should_ubatch,
4109
+ num_scheduled_tokens,
4110
+ num_tokens_padded,
4111
+ num_reqs_padded,
4112
+ self.vllm_config.parallel_config.num_ubatches,
4113
+ )
4114
+ logger.debug(
4115
+ "ubatch_slices: %s, ubatch_slices_padded: %s",
4116
+ ubatch_slices,
4117
+ ubatch_slices_padded,
4118
)
4119
4120
attn_metadata: PerLayerAttnMetadata | None = None
@@ -4644,7 +4663,7 @@ class GPUModelRunner(
4663
# is above the threshold. Otherwise we just capture a non-ubatched
4664
# version of the graph
4665
allow_microbatching = (
4647
- self.parallel_config.enable_dbo
4666
+ self.parallel_config.use_ubatching
4667
and cudagraph_runtime_mode == CUDAGraphMode.FULL
4668
and uniform_decode
4669
and check_ubatch_thresholds(
@@ -4779,8 +4798,8 @@ class GPUModelRunner(
4798
if kv_cache_group_id < len(kernel_block_sizes)
4799
else None,
4800
num_metadata_builders=1
4782
- if not self.parallel_config.enable_dbo
4783
- else 2,
4801
+ if not self.parallel_config.use_ubatching
4802
+ else self.parallel_config.num_ubatches,
4803
)
4804
# Calculate reorder batch threshold (if needed)
4805
# Note (tdoublep): do this *after* constructing builders,
vllm/v1/worker/gpu_ubatch_wrapper.py
+19
-16
@@ -103,8 +103,10 @@ class UBatchWrapper:
103
self.vllm_config = vllm_config
104
self.compilation_config = vllm_config.compilation_config
105
self.comm_stream = torch.cuda.Stream(device=device)
106
- # Two ubatch threads plus the main thread
107
- self.ready_barrier = threading.Barrier(3)
106
+ # Ubatch threads plus the main thread
107
+ self.ready_barrier = threading.Barrier(
108
+ self.vllm_config.parallel_config.num_ubatches + 1
109
+ )
110
111
self.cudagraphs: dict[int, CUDAGraphMetaData] = {}
112
@@ -309,7 +311,7 @@ class UBatchWrapper:
311
create_forward_context(
312
attn_metadata[i] if attn_metadata is not None else None,
313
self.vllm_config,
312
- dp_metadata=dp_metadata,
314
+ dp_metadata=dp_metadata[i],
315
batch_descriptor=batch_descriptor,
316
cudagraph_runtime_mode=cudagraph_runtime_mode,
317
)
@@ -417,18 +419,19 @@ class UBatchWrapper:
419
420
# We shouldn't be here unless we are running with multiple DP ranks
421
assert dp_metadata is not None
420
- num_tokens_per_ubatch = (
421
- ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start
422
- )
423
- dp_size = self.vllm_config.parallel_config.data_parallel_size
424
- ubatch_num_tokens_across_dp = torch.tensor(
425
- [num_tokens_per_ubatch] * dp_size, device="cpu", dtype=torch.int32
426
- )
427
- ubatch_dp_metadata = DPMetadata.make(
428
- self.vllm_config.parallel_config,
429
- num_tokens_per_ubatch,
430
- ubatch_num_tokens_across_dp,
431
- )
422
+ ubatch_dp_metadata = []
423
+ for ubatch_slice in ubatch_slices:
424
+ dp_size = self.vllm_config.parallel_config.data_parallel_size
425
+ ubatch_num_tokens_across_dp = torch.tensor(
426
+ [ubatch_slice.num_tokens] * dp_size, device="cpu", dtype=torch.int32
427
+ )
428
+ ubatch_dp_metadata.append(
429
+ DPMetadata.make(
430
+ self.vllm_config.parallel_config,
431
+ ubatch_slice.num_tokens,
432
+ ubatch_num_tokens_across_dp,
433
+ )
434
+ )
435
436
if (
437
num_tokens not in self.cudagraphs
@@ -464,7 +467,7 @@ class UBatchWrapper:
467
intermediate_tensors=intermediate_tensors,
468
inputs_embeds=inputs_embeds,
469
compute_stream=compute_stream,
467
- dp_metadata=dp_metadata,
470
+ dp_metadata=ubatch_dp_metadata,
471
batch_descriptor=batch_descriptor,
472
cudagraph_runtime_mode=CUDAGraphMode.NONE,
473
)
vllm/v1/worker/ubatch_utils.py
+37
-34
@@ -27,14 +27,16 @@ class UBatchSlice:
27
UBatchSlices: TypeAlias = list[UBatchSlice]
28
29
30
-def is_second_ubatch_empty(orig_num_tokens: int, padded_num_tokens: int) -> bool:
31
- return (padded_num_tokens // 2) >= orig_num_tokens
30
+def is_last_ubatch_empty(
31
+ orig_num_tokens: int, padded_num_tokens: int, num_ubatches: int
32
+) -> bool:
33
+ return (padded_num_tokens // num_ubatches) * (num_ubatches - 1) >= orig_num_tokens
34
35
36
def check_ubatch_thresholds(
37
config: ParallelConfig, num_tokens: int, uniform_decode: bool
38
) -> bool:
37
- if not config.enable_dbo:
39
+ if not config.use_ubatching:
40
return False
41
if uniform_decode:
42
return num_tokens >= config.dbo_decode_token_threshold
@@ -42,21 +44,17 @@ def check_ubatch_thresholds(
44
return num_tokens >= config.dbo_prefill_token_threshold
45
46
45
-# This just pads the second ubatch slice out to the total number of tokens
47
+# This pads the last ubatch slice out to the total number of tokens
48
# (num_tokens + padding) since we do `create_ubatch_slices` before applying DP padding.
49
def _pad_out_ubatch_slices(
50
ubatch_slices: UBatchSlices, num_total_tokens: int, num_reqs_padded: int
51
) -> 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),
52
+ last_slice = ubatch_slices[-1]
53
+ padded_last_request_slice = slice(last_slice.request_slice.start, num_reqs_padded)
54
+ padded_last_token_slice = slice(last_slice.token_slice.start, num_total_tokens)
55
+
56
+ return ubatch_slices[:-1] + [
57
+ UBatchSlice(padded_last_request_slice, padded_last_token_slice)
58
]
59
60
@@ -65,40 +63,45 @@ def maybe_create_ubatch_slices(
63
num_scheduled_tokens: np.ndarray,
64
num_tokens_padded: int,
65
num_reqs_padded: int,
68
- split_point: int | None = None,
66
+ num_ubatches: int,
67
+ split_point: list[int] | int | None = None,
68
) -> tuple[UBatchSlices | None, UBatchSlices | None]:
69
if not should_ubatch:
70
return None, None
71
72
if split_point is None:
74
- split_point = int(num_tokens_padded) // 2
73
+ split_point = int(num_tokens_padded) // num_ubatches
74
+
75
+ token_split_points = [split_point * i for i in range(1, num_ubatches)]
76
77
# TODO(lucas): Refactor the gpu_model_runner.py so we can pass
78
# in cu_num_tokens directly (i.e. query_start_loc)
79
cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32)
80
np.cumsum(num_scheduled_tokens, dtype=np.int32, out=cu_num_tokens[1:])
81
81
- first_ubatch_token_slice = slice(0, split_point)
82
- second_ubatch_token_slice = slice(split_point, cu_num_tokens[-1])
82
+ ubatch_slices = []
83
+ start_token = 0
84
84
- # Determine request slices using exclusive stop semantics
85
- # First ubatch includes requests whose tokens overlap [0, split_point)
86
- first_ubatch_req_stop = int(
87
- np.searchsorted(cu_num_tokens, split_point, side="left")
88
- )
89
- first_ubatch_req_slice = slice(0, first_ubatch_req_stop)
85
+ # Add the end point to the split points to make iteration easier
86
+ all_points = token_split_points + [cu_num_tokens[-1]]
87
91
- # Second ubatch starts at the request that contains the split_point
92
- # or the request starting exactly at split_point (if on boundary)
93
- second_ubatch_req_start = int(
94
- np.searchsorted(cu_num_tokens, split_point, side="right") - 1
95
- )
96
- second_ubatch_req_slice = slice(second_ubatch_req_start, len(cu_num_tokens) - 1)
88
+ for end_token in all_points:
89
+ token_slice = slice(start_token, end_token)
90
98
- ubatch_slices = [
99
- UBatchSlice(first_ubatch_req_slice, first_ubatch_token_slice),
100
- UBatchSlice(second_ubatch_req_slice, second_ubatch_token_slice),
101
- ]
91
+ # Determine request slices using exclusive stop semantics
92
+ # Ubatch includes requests whose tokens overlap [start_token, end_token)
93
+
94
+ # Start at the request that contains the start_token
95
+ # or the request starting exactly at start_token (if on boundary)
96
+ req_start = int(np.searchsorted(cu_num_tokens, start_token, side="right") - 1)
97
+
98
+ # Stop at the request that starts at or after end_token
99
+ req_stop = int(np.searchsorted(cu_num_tokens, end_token, side="left"))
100
+
101
+ req_slice = slice(req_start, req_stop)
102
+ ubatch_slices.append(UBatchSlice(req_slice, token_slice))
103
+
104
+ start_token = end_token
105
106
ubatch_slices_padded = _pad_out_ubatch_slices(
107
ubatch_slices, num_tokens_padded, num_reqs_padded
vllm/v1/worker/ubatching.py
+16
-5
@@ -7,10 +7,15 @@ import torch
7
8
from vllm import forward_context
9
from vllm.forward_context import ForwardContext
10
+from vllm.logger import init_logger
11
from vllm.utils.torch_utils import current_stream
12
13
+logger = init_logger(__name__)
14
+
15
_THREAD_ID_TO_CONTEXT: dict = {}
13
-_CURRENT_CONTEXTS: list[Optional["UBatchContext"]] = [None, None]
16
+# Here we hardcode the number of microbatches to 2 for default.
17
+_NUM_UBATCHES: int = 2
18
+_CURRENT_CONTEXTS: list[Optional["UBatchContext"]] = []
19
20
21
class UBatchContext:
@@ -48,6 +53,7 @@ class UBatchContext:
53
global _CURRENT_CONTEXTS, _THREAD_ID_TO_CONTEXT
54
_THREAD_ID_TO_CONTEXT[threading.get_ident()] = self.id
55
_CURRENT_CONTEXTS[self.id] = self
56
+ # _NUM_UBATCHES is set in make_ubatch_contexts
57
self.ready_barrier.wait()
58
59
self.cpu_wait_event.wait()
@@ -181,7 +187,7 @@ dbo_switch_to_compute_sync = _register_ubatch_function(
187
def dbo_register_recv_hook(recv_hook):
188
if len(_THREAD_ID_TO_CONTEXT) > 0:
189
ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()]
184
- next_ctx = _CURRENT_CONTEXTS[(ctx_idx + 1) % 2]
190
+ next_ctx = _CURRENT_CONTEXTS[(ctx_idx + 1) % _NUM_UBATCHES]
191
next_ctx.recv_hook = recv_hook
192
193
@@ -202,7 +208,14 @@ def make_ubatch_contexts(
208
ready_barrier: threading.Barrier,
209
schedule: str = "default",
210
) -> list[UBatchContext]:
205
- assert num_micro_batches == 2, "only been tested with 2 micro-batches"
211
+ global _NUM_UBATCHES, _CURRENT_CONTEXTS
212
+ assert num_micro_batches > 1, "num_micro_batches must be greater than 1"
213
+
214
+ _NUM_UBATCHES = num_micro_batches
215
+ # Ensure the global context list is large enough
216
+ if len(_CURRENT_CONTEXTS) < num_micro_batches:
217
+ _CURRENT_CONTEXTS.extend([None] * (num_micro_batches - len(_CURRENT_CONTEXTS)))
218
+
219
"""
220
Create a context manager for micro-batching synchronization.
221
"""
@@ -210,8 +223,6 @@ def make_ubatch_contexts(
223
gpu_comm_done_events = [torch.Event() for _ in range(num_micro_batches)]
224
gpu_compute_done_events = [torch.Event() for _ in range(num_micro_batches)]
225
213
- assert len(forward_contexts) == 2
214
-
226
ctxs = []
227
for i in range(num_micro_batches):
228
ctx = UBatchContext(