@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3f41d102c5

[1/2][Model Runner V2] DBO support, eager mode (#50945)

Signed-off-by: specture724 <specture724@gmail.com> Signed-off-by: Anjie Hou <149605198+specture724@users.noreply.github.com> Co-authored-by: jiangkuaixue123 <jiangxiaozhou111@163.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>

Anjie Hou committed Sep 5, 2026 at 05:17 UTC 3f41d102c5b8636f7206a5fd5949dd8ae0b30a0d
17 files changed +1346 -111
tests/model_executor/test_routed_experts_capture.py
+1
@@ -453,6 +453,7 @@ def test_v2_model_runner_accepts_routed_experts(monkeypatch):
453 distributed_executor_backend=None,
454 pipeline_parallel_size=1,
455 enable_dbo=False,
456 + use_ubatching=False,
457 enable_elastic_ep=False,
458 ),
459 compilation_config=SimpleNamespace(
tests/v1/worker/test_gpu_ubatch_slicing.py new
+583
@@ -0,0 +1,583 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Dual batch overlap (DBO) in the V2 model runner.
4 +
5 +The V2 runner slices the batch *before* building attention metadata, where V1
6 +builds metadata for the whole batch and slices it afterwards. Both must describe
7 +the same microbatches, so the main test here pins V2's slicing against V1's
8 +`split_attn_metadata`. The rest covers the decision to microbatch and the
9 +threaded execution of the microbatches.
10 +"""
11 +
12 +import threading
13 +from dataclasses import replace
14 +from types import SimpleNamespace
15 +from typing import Any, cast
16 +from unittest.mock import patch
17 +
18 +import numpy as np
19 +import pytest
20 +import torch
21 +
22 +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
23 +from vllm.config import CUDAGraphMode, ModelConfig, ParallelConfig, VllmConfig
24 +from vllm.forward_context import create_forward_context
25 +from vllm.v1.kv_cache_interface import KVCacheConfig
26 +from vllm.v1.worker.gpu import dp_utils
27 +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
28 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
29 +from vllm.v1.worker.gpu.model_states.interface import ModelState
30 +from vllm.v1.worker.gpu.ubatch_utils import (
31 + UBatchRunner,
32 + UBatchState,
33 + _slice_input_batch,
34 + create_ubatch_slices,
35 + slice_model_inputs,
36 +)
37 +from vllm.v1.worker.ubatch_utils import (
38 + UBatchSlice,
39 + maybe_create_ubatch_slices,
40 + split_attn_metadata,
41 +)
42 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id, dbo_yield
43 +
44 +MAX_NUM_REQS = 32
45 +MAX_NUM_TOKENS = 128
46 +
47 +
48 +def _make_buffers() -> InputBuffers:
49 + return InputBuffers(
50 + max_num_reqs=MAX_NUM_REQS,
51 + max_num_tokens=MAX_NUM_TOKENS,
52 + device=torch.device("cpu"),
53 + )
54 +
55 +
56 +def _make_ubatch_buffers(num_ubatches: int = 2) -> list[tuple[torch.Tensor, ...]]:
57 + """The per-microbatch (query_start_loc, seq_lens) buffers UBatchRunner owns."""
58 + return [
59 + (
60 + torch.zeros(MAX_NUM_REQS + 1, dtype=torch.int32),
61 + torch.zeros(MAX_NUM_REQS, dtype=torch.int32),
62 + )
63 + for _ in range(num_ubatches)
64 + ]
65 +
66 +
67 +def _make_input_batch(
68 + query_lens: list[int],
69 + seq_lens: list[int],
70 + buffers: InputBuffers,
71 + num_reqs_padded: int | None = None,
72 + num_tokens_padded: int | None = None,
73 +) -> InputBatch:
74 + """Build an InputBatch shaped like the one `prepare_inputs` produces."""
75 + num_reqs = len(query_lens)
76 + num_tokens = int(sum(query_lens))
77 + num_reqs_padded = num_reqs_padded or num_reqs
78 + num_tokens_padded = num_tokens_padded or num_tokens
79 +
80 + # make_dummy writes the shared buffers, so fill them afterwards.
81 + base = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
82 +
83 + query_start_loc_np = np.zeros(num_reqs_padded + 1, dtype=np.int32)
84 + np.cumsum(query_lens, out=query_start_loc_np[1 : num_reqs + 1])
85 + # Padded entries repeat the token count, as `prepare_inputs` does.
86 + query_start_loc_np[num_reqs + 1 :] = num_tokens
87 + buffers.query_start_loc[: num_reqs_padded + 1] = torch.from_numpy(
88 + query_start_loc_np
89 + )
90 +
91 + buffers.seq_lens[:num_reqs] = torch.tensor(seq_lens, dtype=torch.int32)
92 + buffers.seq_lens[num_reqs:num_reqs_padded] = 0
93 +
94 + seq_lens_upper_bound = np.zeros(num_reqs_padded, dtype=np.int32)
95 + seq_lens_upper_bound[:num_reqs] = seq_lens
96 +
97 + return replace(
98 + base,
99 + req_ids=[f"req_{i}" for i in range(num_reqs)],
100 + num_reqs=num_reqs,
101 + num_reqs_after_padding=num_reqs_padded,
102 + idx_mapping=torch.arange(num_reqs, dtype=torch.int32),
103 + idx_mapping_np=np.arange(num_reqs, dtype=np.int32),
104 + num_scheduled_tokens=np.array(query_lens, dtype=np.int32),
105 + num_tokens=num_tokens,
106 + num_tokens_after_padding=num_tokens_padded,
107 + query_start_loc=buffers.query_start_loc[: num_reqs_padded + 1],
108 + query_start_loc_np=query_start_loc_np,
109 + seq_lens=buffers.seq_lens[:num_reqs_padded],
110 + seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_upper_bound),
111 + num_computed_tokens_np=np.array(seq_lens, dtype=np.int32)
112 + - np.array(query_lens, dtype=np.int32),
113 + prefill_len_np=np.zeros(num_reqs, dtype=np.int32),
114 + num_computed_prefill_tokens_np=np.zeros(num_reqs, dtype=np.int32),
115 + is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
116 + input_ids=buffers.input_ids[:num_tokens_padded],
117 + positions=buffers.positions[:num_tokens_padded],
118 + is_padding=buffers.is_padding[:num_tokens_padded],
119 + )
120 +
121 +
122 +# (query_lens, seq_lens): uniform decode, a boundary that lands inside a
123 +# request, and a boundary that lands exactly on a request start.
124 +BATCHES = {
125 + "uniform_decode": ([1] * 8, [128 + 4 * i for i in range(8)]),
126 + "split_request": ([1, 1, 10, 1, 1], [64, 96, 512, 32, 48]),
127 + "aligned_prefills": ([6, 2, 4, 4], [6, 130, 260, 4]),
128 + "split_first_and_last": ([3, 9, 3], [70, 300, 41]),
129 + # A long prefill trailing a decode: it straddles the boundary and is also
130 + # the last request, so the second microbatch holds nothing else.
131 + "trailing_prefill_spans_boundary": ([2, 20], [64, 300]),
132 +}
133 +
134 +
135 +@pytest.mark.parametrize("batch_name", list(BATCHES))
136 +def test_slicing_matches_v1_split_attn_metadata(batch_name: str):
137 + """V2's pre-sliced inputs describe the same microbatches as V1's."""
138 + query_lens, seq_lens = BATCHES[batch_name]
139 + num_scheduled_tokens = np.array(query_lens, dtype=np.int32)
140 + num_tokens = int(num_scheduled_tokens.sum())
141 +
142 + ubatch_slices, _ = maybe_create_ubatch_slices(
143 + True, num_scheduled_tokens, num_tokens, len(query_lens), num_ubatches=2
144 + )
145 + assert ubatch_slices is not None
146 +
147 + v1_metadata = create_common_attn_metadata(
148 + BatchSpec(seq_lens=seq_lens, query_lens=query_lens),
149 + block_size=16,
150 + device=torch.device("cpu"),
151 + )
152 + v1_ubatches = split_attn_metadata(ubatch_slices, v1_metadata)
153 +
154 + buffers = _make_buffers()
155 + ubatch_buffers = _make_ubatch_buffers()
156 + input_batch = _make_input_batch(query_lens, seq_lens, buffers)
157 +
158 + for i, (ubatch_slice, v1_ubatch) in enumerate(zip(ubatch_slices, v1_ubatches)):
159 + v2_ubatch = _slice_input_batch(input_batch, ubatch_slice, *ubatch_buffers[i])
160 +
161 + assert v2_ubatch.num_reqs == v1_ubatch.num_reqs
162 + assert v2_ubatch.num_tokens == v1_ubatch.num_actual_tokens
163 + assert int(v2_ubatch.num_scheduled_tokens.max()) == v1_ubatch.max_query_len
164 + torch.testing.assert_close(v2_ubatch.query_start_loc, v1_ubatch.query_start_loc)
165 + torch.testing.assert_close(v2_ubatch.seq_lens, v1_ubatch.seq_lens)
166 + torch.testing.assert_close(
167 + torch.from_numpy(v2_ubatch.query_start_loc_np),
168 + v1_ubatch.query_start_loc_cpu,
169 + )
170 + # `prepare_attn` derives max_seq_len from the upper bound. Compare the
171 + # bound itself rather than V1's max_seq_len: V1 floors it with the full
172 + # batch's max_seq_len to keep the value CUDA-graph capture installed,
173 + # which V2 gets from `prepare_attn(for_capture=True)` instead.
174 + torch.testing.assert_close(
175 + v2_ubatch.seq_lens_cpu_upper_bound, v1_ubatch.seq_lens_cpu_upper_bound
176 + )
177 +
178 +
179 +def test_microbatches_do_not_share_buffers():
180 + """Both microbatches are live at once, so their buffers must be distinct."""
181 + query_lens, seq_lens = BATCHES["split_request"]
182 + num_scheduled_tokens = np.array(query_lens, dtype=np.int32)
183 + num_tokens = int(num_scheduled_tokens.sum())
184 +
185 + ubatch_slices, _ = maybe_create_ubatch_slices(
186 + True, num_scheduled_tokens, num_tokens, len(query_lens), num_ubatches=2
187 + )
188 + assert ubatch_slices is not None
189 +
190 + buffers = _make_buffers()
191 + ubatch_buffers = _make_ubatch_buffers()
192 + input_batch = _make_input_batch(query_lens, seq_lens, buffers)
193 +
194 + first = _slice_input_batch(input_batch, ubatch_slices[0], *ubatch_buffers[0])
195 + first_query_start_loc = first.query_start_loc.clone()
196 + first_seq_lens = first.seq_lens.clone()
197 +
198 + # Slicing the second microbatch must not disturb the first.
199 + _slice_input_batch(input_batch, ubatch_slices[1], *ubatch_buffers[1])
200 +
201 + torch.testing.assert_close(first.query_start_loc, first_query_start_loc)
202 + torch.testing.assert_close(first.seq_lens, first_seq_lens)
203 +
204 +
205 +def test_trailing_microbatch_absorbs_cudagraph_padding():
206 + """The padded microbatch keeps padded rows empty and query_start_loc flat."""
207 + query_lens, seq_lens = [1] * 6, [64] * 6
208 + num_tokens_padded = 8
209 +
210 + buffers = _make_buffers()
211 + input_batch = _make_input_batch(
212 + query_lens,
213 + seq_lens,
214 + buffers,
215 + num_reqs_padded=num_tokens_padded,
216 + num_tokens_padded=num_tokens_padded,
217 + )
218 + ubatch_slices_padded = create_ubatch_slices(input_batch, num_ubatches=2)
219 +
220 + last = _slice_input_batch(
221 + input_batch, ubatch_slices_padded[-1], *_make_ubatch_buffers()[1]
222 + )
223 +
224 + # Two real decodes plus two padded rows.
225 + assert last.num_reqs == 2
226 + assert last.num_tokens == 2
227 + assert last.num_reqs_after_padding == 4
228 + assert last.num_tokens_after_padding == 4
229 + torch.testing.assert_close(
230 + last.query_start_loc, torch.tensor([0, 1, 2, 2, 2], dtype=torch.int32)
231 + )
232 + torch.testing.assert_close(
233 + last.seq_lens, torch.tensor([64, 64, 0, 0], dtype=torch.int32)
234 + )
235 +
236 +
237 +DECODE_THRESHOLD = 32
238 +PREFILL_THRESHOLD = 128
239 +DECODE_QUERY_LEN = 1
240 +
241 +
242 +def test_dummy_batch_can_be_microbatched():
243 + """A DP rank with no work still has to split, so its dummy batch must slice.
244 +
245 + Microbatching is all-or-nothing across the group, so a rank whose step is a
246 + dummy batch runs it microbatched like everyone else.
247 + """
248 + num_reqs, num_tokens = 1, 32
249 + buffers = _make_buffers()
250 + input_batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
251 +
252 + ubatch_slices_padded = create_ubatch_slices(input_batch, num_ubatches=2)
253 +
254 + ubatch_buffers = _make_ubatch_buffers()
255 + ubatches = [
256 + _slice_input_batch(input_batch, ubatch_slice, *ubatch_buffers[i])
257 + for i, ubatch_slice in enumerate(ubatch_slices_padded)
258 + ]
259 +
260 + assert [u.num_tokens for u in ubatches] == [16, 16]
261 + # The one dummy request spans both microbatches, 16 of its tokens in each.
262 + assert all(u.num_reqs == 1 for u in ubatches)
263 + for ubatch in ubatches:
264 + torch.testing.assert_close(
265 + ubatch.query_start_loc, torch.tensor([0, 16], dtype=torch.int32)
266 + )
267 +
268 +
269 +def _sync_dp(
270 + num_tokens_per_rank: list[int],
271 + uniform_token_count_per_rank: list[int] | None = None,
272 + allow_ubatching: bool = True,
273 +) -> tuple[BatchExecutionDescriptor, dp_utils.DPSyncState | None]:
274 + """Run the DP handshake with the all-reduce stubbed out.
275 +
276 + The microbatching decision lives inside `sync_cudagraph_and_dp_padding`, so
277 + the cases below stand in for the collective rather than for the decision,
278 + and need no process group. Every rank asks for eager, which keeps the
279 + non-microbatched path off the cudagraph manager. This rank is rank 0.
280 + """
281 + dp_size = len(num_tokens_per_rank)
282 + uniform_token_counts = uniform_token_count_per_rank or [0] * dp_size
283 + reduced = torch.zeros(6, dp_size, dtype=torch.int32)
284 + reduced[0] = torch.tensor(num_tokens_per_rank, dtype=torch.int32)
285 + reduced[1] = CUDAGraphMode.NONE.value
286 + reduced[2] = torch.tensor(uniform_token_counts, dtype=torch.int32)
287 + reduced[3] = -1 # max_query_len, -1 means None
288 + reduced[4] = int(allow_ubatching)
289 + reduced[5] = 8 # num_reqs
290 +
291 + with (
292 + patch.object(dp_utils.dist, "all_reduce", lambda t, group: t.copy_(reduced)),
293 + patch.object(dp_utils, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)),
294 + ):
295 + return dp_utils.sync_cudagraph_and_dp_padding(
296 + cudagraph_manager=None,
297 + desired_batch_desc=BatchExecutionDescriptor(
298 + cg_mode=CUDAGraphMode.NONE,
299 + num_tokens=num_tokens_per_rank[0],
300 + num_reqs=8,
301 + ),
302 + num_tokens=num_tokens_per_rank[0],
303 + num_reqs=8,
304 + uniform_token_count=uniform_token_counts[0] or None,
305 + dp_size=dp_size,
306 + dp_rank=0,
307 + parallel_config=ParallelConfig(
308 + enable_dbo=True,
309 + dbo_decode_token_threshold=DECODE_THRESHOLD,
310 + dbo_prefill_token_threshold=PREFILL_THRESHOLD,
311 + ),
312 + allow_ubatching=allow_ubatching,
313 + uniform_decode=uniform_token_counts[0] == DECODE_QUERY_LEN,
314 + )
315 +
316 +
317 +def test_every_dp_rank_must_agree_to_microbatch():
318 + """One rank below the threshold is enough to keep the group on one batch.
319 +
320 + The thresholds are not communicated -- they are checked against the token
321 + counts the all-reduce already carries -- so this also pins that the whole
322 + vector is consulted, not just this rank's own entry.
323 + """
324 + assert _sync_dp([256, 256])[0].num_ubatches == 2
325 + assert _sync_dp([256, 100])[0].num_ubatches == 1
326 + assert _sync_dp([100, 256])[0].num_ubatches == 1
327 +
328 +
329 +def test_decode_threshold_only_applies_when_every_rank_is_a_uniform_decode():
330 + """Otherwise the group is held to the prefill threshold."""
331 + # 64 tokens clears the decode threshold but not the prefill one, so it only
332 + # microbatches when both ranks are uniform decodes.
333 + assert _sync_dp([64, 64], uniform_token_count_per_rank=[1, 1])[0].num_ubatches == 2
334 + assert _sync_dp([64, 64], uniform_token_count_per_rank=[1, 0])[0].num_ubatches == 1
335 +
336 +
337 +def test_microbatching_pads_all_ranks_to_the_largest():
338 + """Ranks must run the same token count so microbatch sizes line up."""
339 + desc, dp_sync = _sync_dp([200, 256])
340 + assert desc.num_tokens == 256
341 + assert dp_sync is not None
342 + assert dp_sync.num_tokens_across_dp.tolist() == [256, 256]
343 +
344 +
345 +def test_microbatching_survives_a_rank_that_cannot_fill_it():
346 + """A rank whose real tokens all land in the first microbatch still splits.
347 +
348 + Rank 0 has 128 real tokens but the split point is at 512 // 2 = 256, so its
349 + second microbatch is pure padding. That microbatch does no work, which is
350 + fine -- it still has to run so the expert all-to-all stays collective.
351 + """
352 + desc, _ = _sync_dp([128, 512])
353 + assert desc.num_tokens == 512
354 + assert desc.num_ubatches == 2
355 +
356 +
357 +def test_all_padding_microbatch_has_no_work_to_do():
358 + """The microbatch past the last real token is well-formed and empty.
359 +
360 + It carries the last request with none of its query tokens, so attention and
361 + sampling see nothing to do while the microbatch still runs.
362 + """
363 + query_lens, seq_lens = [1] * 6, [64] * 6
364 + num_tokens_padded = 16 # split at 8, past all 6 real tokens
365 +
366 + buffers = _make_buffers()
367 + input_batch = _make_input_batch(
368 + query_lens, seq_lens, buffers, num_tokens_padded=num_tokens_padded
369 + )
370 + ubatch_slices_padded = create_ubatch_slices(input_batch, num_ubatches=2)
371 + ubatch_buffers = _make_ubatch_buffers()
372 +
373 + first = _slice_input_batch(input_batch, ubatch_slices_padded[0], *ubatch_buffers[0])
374 + last = _slice_input_batch(input_batch, ubatch_slices_padded[1], *ubatch_buffers[1])
375 +
376 + # The first microbatch keeps every real request; the second gets none.
377 + assert first.num_reqs == 6
378 + assert first.num_tokens == 6
379 + assert last.num_tokens == 0
380 + assert last.num_reqs == 1
381 + assert last.num_reqs_after_padding == 1
382 + assert last.num_tokens_after_padding == 8
383 +
384 + # Zero query tokens for the one request it holds, so no attention work.
385 + torch.testing.assert_close(
386 + last.query_start_loc, torch.tensor([0, 0], dtype=torch.int32)
387 + )
388 + assert last.num_scheduled_tokens.tolist() == [0]
389 + # ...and its sequence still ends where its tokens do: they all live in the
390 + # first microbatch, which computes them before this one runs.
391 + torch.testing.assert_close(last.seq_lens, torch.tensor([64], dtype=torch.int32))
392 + torch.testing.assert_close(
393 + last.seq_lens_cpu_upper_bound, torch.tensor([64], dtype=torch.int32)
394 + )
395 +
396 +
397 +def test_microbatching_off_when_a_rank_does_not_allow_it():
398 + assert _sync_dp([256, 256], allow_ubatching=False)[0].num_ubatches == 1
399 +
400 +
401 +def test_slice_model_inputs_handles_mrope_positions():
402 + model_inputs = {
403 + "input_ids": torch.arange(8),
404 + # M-RoPE positions carry a leading section dim.
405 + "positions": torch.arange(24).view(3, 8),
406 + "inputs_embeds": None,
407 + "intermediate_tensors": None,
408 + }
409 + sliced = slice_model_inputs(model_inputs, slice(4, 8))
410 +
411 + torch.testing.assert_close(sliced["input_ids"], torch.arange(4, 8))
412 + torch.testing.assert_close(sliced["positions"], torch.arange(24).view(3, 8)[:, 4:8])
413 + assert sliced["inputs_embeds"] is None
414 + assert sliced["intermediate_tensors"] is None
415 +
416 +
417 +def _make_dbo_config() -> VllmConfig:
418 + return VllmConfig(
419 + model_config=ModelConfig(model="facebook/opt-125m", dtype="float16", seed=0),
420 + parallel_config=ParallelConfig(
421 + enable_dbo=True, all2all_backend="deepep_low_latency"
422 + ),
423 + )
424 +
425 +
426 +def _make_model_inputs(num_tokens: int, device: torch.device) -> dict[str, Any]:
427 + return {
428 + "input_ids": torch.arange(num_tokens, device=device),
429 + "positions": torch.arange(num_tokens, device=device),
430 + "inputs_embeds": None,
431 + "intermediate_tensors": None,
432 + }
433 +
434 +
435 +@pytest.fixture(autouse=True)
436 +def _no_comm_sm_reservation(monkeypatch):
437 + """CI runs on MIG slices with fewer SMs than the default comm reservation."""
438 + monkeypatch.setenv("VLLM_DBO_COMM_SMS", "0")
439 +
440 +
441 +def _make_execution_runner(vllm_config: VllmConfig) -> UBatchRunner:
442 + """A runner for the execution tests below, which never call `prepare()`.
443 +
444 + The attention-side dependencies are only read while building per-microbatch
445 + metadata, so they are left out here.
446 + """
447 + return UBatchRunner(
448 + vllm_config,
449 + torch.device("cuda:0"),
450 + model_state=cast(ModelState, None),
451 + attn_groups=[],
452 + kv_cache_config=cast(KVCacheConfig, None),
453 + max_num_reqs=MAX_NUM_REQS,
454 + )
455 +
456 +
457 +def _make_ubatch_state(
458 + vllm_config: VllmConfig, ubatch_slices: list[UBatchSlice]
459 +) -> UBatchState:
460 + return UBatchState(
461 + slices=ubatch_slices,
462 + forward_contexts=[
463 + create_forward_context(None, vllm_config) for _ in ubatch_slices
464 + ],
465 + )
466 +
467 +
468 +def test_thresholds_below_the_microbatch_count_are_rejected():
469 + """A batch with fewer tokens than microbatches cannot be split at all.
470 +
471 + Nothing downstream copes with an unsplittable batch, so the thresholds are
472 + what keep one out: a one-token decode split two ways leaves an empty
473 + microbatch.
474 + """
475 + with pytest.raises(ValueError, match="number of microbatches"):
476 + ParallelConfig(enable_dbo=True, dbo_decode_token_threshold=1)
477 +
478 +
479 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU")
480 +def test_runner_allocates_one_buffer_pair_per_microbatch():
481 + """All microbatches are live at once, so none may share rebase buffers."""
482 + vllm_config = _make_dbo_config()
483 + runner = _make_execution_runner(vllm_config)
484 +
485 + for buffers in (runner.ubatch_query_start_loc, runner.ubatch_seq_lens):
486 + assert len(buffers) == runner.num_ubatches
487 + storages = {b.untyped_storage().data_ptr() for b in buffers}
488 + assert len(storages) == runner.num_ubatches
489 +
490 +
491 +class _YieldingModel(torch.nn.Module):
492 + """Toy model that hands off to the other microbatch mid-forward."""
493 +
494 + def __init__(self, trace: list[tuple[str, int]]):
495 + super().__init__()
496 + self.trace = trace
497 +
498 + def forward(self, input_ids, positions, intermediate_tensors=None, **kwargs):
499 + self.trace.append(("enter", dbo_current_ubatch_id()))
500 + out = input_ids.float().unsqueeze(-1) * 2 + positions.float().unsqueeze(-1)
501 + # Stands in for the expert all-to-all handoff point.
502 + dbo_yield()
503 + self.trace.append(("exit", dbo_current_ubatch_id()))
504 + return out
505 +
506 +
507 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU")
508 +def test_ubatch_runner_overlaps_and_matches_single_batch():
509 + """Microbatches interleave at the yield point and produce the same output."""
510 + vllm_config = _make_dbo_config()
511 + device = torch.device("cuda:0")
512 + runner = _make_execution_runner(vllm_config)
513 +
514 + model_inputs = _make_model_inputs(16, device)
515 + ubatch_state = _make_ubatch_state(
516 + vllm_config,
517 + [
518 + UBatchSlice(slice(0, 4), slice(0, 8)),
519 + UBatchSlice(slice(4, 8), slice(8, 16)),
520 + ],
521 + )
522 +
523 + trace: list[tuple[str, int]] = []
524 + model = _YieldingModel(trace)
525 + output = runner.run(model, model_inputs, ubatch_state)
526 +
527 + expected = model_inputs["input_ids"].float().unsqueeze(-1) * 2 + model_inputs[
528 + "positions"
529 + ].float().unsqueeze(-1)
530 + torch.testing.assert_close(output, expected)
531 +
532 + # Both microbatches reach the handoff before either finishes.
533 + assert trace == [("enter", 0), ("enter", 1), ("exit", 0), ("exit", 1)]
534 +
535 +
536 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="DBO needs a GPU")
537 +def test_ubatch_runner_names_the_microbatch_that_failed():
538 + """A failing microbatch surfaces as a named error, not a downstream KeyError.
539 +
540 + The model here never yields, so the microbatches run back to back and no
541 + handoff is left outstanding when one of them dies. A microbatch that fails
542 + *while its sibling is parked at a yield* hangs the step instead -- the
543 + shared handoff protocol in `ubatching.py` has no way to unwind a parked
544 + microbatch, and the V1 runner has the same gap. Fixing that means changing
545 + `ubatching.py`, which is out of scope for the V2 runner.
546 + """
547 + vllm_config = _make_dbo_config()
548 + device = torch.device("cuda:0")
549 + runner = _make_execution_runner(vllm_config)
550 +
551 + class _FailingModel(torch.nn.Module):
552 + def forward(self, input_ids, positions, **kwargs):
553 + if dbo_current_ubatch_id() == 0:
554 + raise ValueError("boom")
555 + return input_ids.float().unsqueeze(-1)
556 +
557 + model_inputs = _make_model_inputs(8, device)
558 + ubatch_state = _make_ubatch_state(
559 + vllm_config,
560 + [
561 + UBatchSlice(slice(0, 2), slice(0, 4)),
562 + UBatchSlice(slice(2, 4), slice(4, 8)),
563 + ],
564 + )
565 +
566 + result: dict[str, BaseException] = {}
567 +
568 + def _call():
569 + try:
570 + runner.run(_FailingModel(), model_inputs, ubatch_state)
571 + except BaseException as e: # noqa: BLE001
572 + result["error"] = e
573 +
574 + # Run behind a watchdog: `UBatchRunner.run` joins without a timeout, so a
575 + # regression here would wedge the suite rather than fail it.
576 + caller = threading.Thread(target=_call, daemon=True)
577 + caller.start()
578 + caller.join(timeout=60.0)
579 + assert not caller.is_alive(), "UBatchRunner.run hung on a failing microbatch"
580 +
581 + assert isinstance(result["error"], RuntimeError)
582 + assert "Microbatch 0" in str(result["error"])
583 + assert isinstance(result["error"].__cause__, ValueError)
vllm/config/parallel.py
+11
@@ -1065,6 +1065,17 @@ class ParallelConfig:
1065 "Unable to use nsight profiling unless workers run with Ray."
1066 )
1067
1068 + # A batch below one token per microbatch cannot be split, so the
1069 + # thresholds have to keep it out rather than the split having to cope.
1070 + if self.use_ubatching and (
1071 + min(self.dbo_decode_token_threshold, self.dbo_prefill_token_threshold)
1072 + < self.num_ubatches
1073 + ):
1074 + raise ValueError(
1075 + "dbo_decode_token_threshold and dbo_prefill_token_threshold must "
1076 + f"be at least the number of microbatches ({self.num_ubatches})."
1077 + )
1078 +
1079 return self
1080
1081 def reconfigure_for_independent_dp_rank(self) -> None:
vllm/config/vllm.py
+43 -2
@@ -2596,8 +2596,8 @@ class VllmConfig:
2596 ):
2597 unsupported.append("EAGLE3 with pipeline parallelism")
2598
2599 - if self.parallel_config.enable_dbo:
2600 - unsupported.append("dual batch overlap")
2599 + if self.parallel_config.use_ubatching:
2600 + unsupported.extend(self._get_dbo_unsupported_features())
2601
2602 if self.parallel_config.enable_elastic_ep:
2603 unsupported.append("elastic expert parallelism")
@@ -2739,6 +2739,47 @@ class VllmConfig:
2739 f"{'; '.join(blockers)}."
2740 )
2741
2742 + def _get_dbo_unsupported_features(self) -> list[str]:
2743 + """Collect what the V2 model runner cannot combine with DBO.
2744 +
2745 + The V2 runner microbatches a plain decoder forward pass. Anything that
2746 + slices or replays the batch differently (drafting, adapters, pipeline
2747 + stages, context parallelism, encoders) is not handled yet.
2748 + """
2749 + # TODO: DBO with model runner V2 is under development.
2750 + # It should be enabled with explicit VLLM_USE_V2_MODEL_RUNNER environ.
2751 + # Remove it when stable.
2752 + if envs.VLLM_USE_V2_MODEL_RUNNER is None:
2753 + return ["dual batch overlap"]
2754 +
2755 + unsupported: list[str] = []
2756 + model_config = self.model_config
2757 + parallel_config = self.parallel_config
2758 +
2759 + if self.lora_config is not None:
2760 + unsupported.append("dual batch overlap with LoRA")
2761 + if self.speculative_config is not None:
2762 + unsupported.append("dual batch overlap with speculative decoding")
2763 + if parallel_config.pipeline_parallel_size > 1:
2764 + unsupported.append("dual batch overlap with pipeline parallelism")
2765 + if (
2766 + parallel_config.decode_context_parallel_size > 1
2767 + or parallel_config.prefill_context_parallel_size > 1
2768 + ):
2769 + unsupported.append("dual batch overlap with context parallelism")
2770 + if model_config is not None and (
2771 + model_config.is_multimodal_model or model_config.is_encoder_decoder
2772 + ):
2773 + unsupported.append("dual batch overlap with multimodal models")
2774 + if model_config is not None and model_config.is_hybrid:
2775 + unsupported.append("dual batch overlap with hybrid models")
2776 + if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
2777 + unsupported.append("dual batch overlap with CUDA graphs")
2778 + if self.is_mm_encoder_only:
2779 + unsupported.append("dual batch overlap with encoder only models")
2780 +
2781 + return unsupported
2782 +
2783 def _validate_v2_model_runner(self) -> None:
2784 """Check for features not yet supported by the V2 model runner."""
2785 if not HAS_TRITON:
vllm/model_executor/models/diffusion_gemma.py
+2
@@ -988,7 +988,9 @@ class DiffusionGemmaModelState(ModelState):
988 attn_groups,
989 kv_cache_config,
990 for_capture=False,
991 + ubatch_idx: int = 0,
992 ) -> dict[str, Any]:
993 + assert ubatch_idx == 0, "DBO is not supported"
994 if cudagraph_mode == CUDAGraphMode.FULL:
995 num_reqs = input_batch.num_reqs_after_padding
996 num_tokens = input_batch.num_tokens_after_padding
vllm/v1/worker/gpu/attn_utils.py
+15 -8
@@ -22,6 +22,7 @@ from vllm.v1.kv_cache_interface import (
22 UniformTypeKVCacheSpecs,
23 )
24 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
25 +from vllm.v1.worker.ubatch_utils import get_num_ubatches
26 from vllm.v1.worker.utils import (
27 AttentionGroup,
28 add_kv_sharing_layers_to_kv_cache_groups,
@@ -140,14 +141,19 @@ def init_attn_backend(
141 vllm_config=vllm_config,
142 device=device,
143 kernel_block_size=kernel_block_size,
143 - num_metadata_builders=1,
144 + # Microbatches build attention metadata concurrently, and some
145 + # builders keep the prepared metadata on themselves (MLA stores
146 + # it on the prefill backend), so each ubatch needs its own.
147 + num_metadata_builders=get_num_ubatches(vllm_config.parallel_config),
148 )
145 - builder = group.get_metadata_builder(0)
146 - if attn_backend_workspace is None:
147 - if hasattr(builder, "_get_workspace_buffer"):
148 - attn_backend_workspace = builder._get_workspace_buffer()
149 - else:
150 - if hasattr(builder, "set_workspace_buffer"):
149 + # The microbatches' builders share the workspace: they all issue
150 + # attention on the one compute stream the threads hand off, so the
151 + # buffer is written serially, as it already is across steps.
152 + for builder in group.metadata_builders:
153 + if attn_backend_workspace is None:
154 + if hasattr(builder, "_get_workspace_buffer"):
155 + attn_backend_workspace = builder._get_workspace_buffer()
156 + elif hasattr(builder, "set_workspace_buffer"):
157 builder.set_workspace_buffer(attn_backend_workspace)
158 attn_cg_support_info = get_attn_cg_support(attn_groups, vllm_config)
159 return attn_groups, attn_cg_support_info, kernel_block_sizes
@@ -271,6 +277,7 @@ def build_attn_metadata(
277 for_cudagraph_capture: bool = False,
278 causal: bool | torch.Tensor | Mapping[int, bool] = True,
279 rswa_prefix_lens: torch.Tensor | None = None,
280 + ubatch_idx: int = 0,
281 ) -> dict[str, Any]:
282 seq_lens = seq_lens[:num_reqs]
283 if dcp_local_seq_lens is not None:
@@ -319,7 +326,7 @@ def build_attn_metadata(
326 )
327
328 for attn_group in attn_groups[i]:
322 - attn_metadata_builder = attn_group.get_metadata_builder(0)
329 + attn_metadata_builder = attn_group.get_metadata_builder(ubatch_idx)
330 if for_cudagraph_capture:
331 metadata = attn_metadata_builder.build_for_cudagraph_capture(
332 common_attn_metadata
vllm/v1/worker/gpu/cudagraph_utils.py
+8
@@ -68,6 +68,8 @@ class BatchExecutionDescriptor:
68 # uniform_token_count unset, so this is what keeps a prefill batch out of one.
69 max_query_len: int | None = None
70 num_active_loras: int = 0
71 + # Number of microbatches the batch is split into (DBO). 1 means no splitting.
72 + num_ubatches: int = 1
73
74
75 def make_cudagraph_stats(
@@ -100,11 +102,13 @@ def _is_compatible(
102 uniform_token_count: int | None,
103 num_active_loras: int,
104 max_query_len: int | None,
105 + num_ubatches: int,
106 ) -> bool:
107 # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count
108 # desc.num_reqs=None means no request padding needed (PIECEWISE)
109 # desc.max_query_len=None means the graph does not constrain query length; a
110 # caller that does not track max_query_len must not match one that does
111 + # A graph captured for N microbatches can only serve a batch split N ways.
112 return (
113 (
114 desc.uniform_token_count is None
@@ -117,6 +121,7 @@ def _is_compatible(
121 and (desc.num_reqs is None or desc.num_reqs >= num_reqs)
122 and desc.num_tokens >= num_tokens
123 and desc.num_active_loras == num_active_loras
124 + and desc.num_ubatches == num_ubatches
125 )
126
127
@@ -434,6 +439,7 @@ class CudaGraphManager:
439 uniform_token_count: int | None,
440 num_active_loras: int,
441 max_query_len: int | None = None,
442 + num_ubatches: int = 1,
443 ) -> BatchExecutionDescriptor:
444 """Find matching cudagraph descriptor from priority-ordered candidates."""
445
@@ -448,6 +454,7 @@ class CudaGraphManager:
454 uniform_token_count,
455 effective_loras,
456 max_query_len,
457 + num_ubatches,
458 ):
459 return desc
460 return BatchExecutionDescriptor(
@@ -455,6 +462,7 @@ class CudaGraphManager:
462 num_tokens=num_tokens,
463 num_reqs=num_reqs,
464 num_active_loras=effective_loras,
465 + num_ubatches=num_ubatches,
466 )
467
468 def run_fullgraph(self, desc: BatchExecutionDescriptor):
vllm/v1/worker/gpu/dp_utils.py
+55 -5
@@ -7,12 +7,14 @@ from dataclasses import dataclass, replace
7 import torch
8 import torch.distributed as dist
9
10 +from vllm.config import ParallelConfig
11 from vllm.config.compilation import CUDAGraphMode
12 from vllm.distributed.parallel_state import get_dp_group
13 from vllm.v1.worker.gpu.cudagraph_utils import (
14 BatchExecutionDescriptor,
15 CudaGraphManager,
16 )
17 +from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds, get_num_ubatches
18
19
20 @dataclass(frozen=True)
@@ -48,27 +50,35 @@ def sync_cudagraph_and_dp_padding(
50 dp_rank: int,
51 max_query_len: int | None = None,
52 num_active_loras: int = 0,
53 + parallel_config: ParallelConfig | None = None,
54 + allow_ubatching: bool = False,
55 + uniform_decode: bool = False,
56 ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:
57 """
58 Coordinates the batch descriptor and DP padding across all ranks.
59
60 + `parallel_config` is only needed to decide whether to microbatch, so callers
61 + that never do (`allow_ubatching=False`) can leave it out.
62 +
63 Returns (synced_batch_desc, sync). `sync` is None when no rank has work.
64 """
65 assert dp_size > 1, "DP size must be greater than 1"
66 group = get_dp_group().cpu_group
59 - tensor = torch.zeros(5, dp_size, dtype=torch.int32, device="cpu")
67 + tensor = torch.zeros(6, dp_size, dtype=torch.int32, device="cpu")
68 tensor[0][dp_rank] = num_tokens
69 tensor[1][dp_rank] = desired_batch_desc.cg_mode.value
70 tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None)
71 tensor[3][dp_rank] = max_query_len or -1 # (-1 means None)
64 - tensor[4][dp_rank] = num_reqs
72 + tensor[4][dp_rank] = int(allow_ubatching)
73 + tensor[5][dp_rank] = num_reqs
74 dist.all_reduce(tensor, group=group)
75
76 num_tokens_across_dp = tensor[0]
77 cg_mode_across_dp = tensor[1]
78 uniform_token_counts_across_dp = tensor[2]
79 max_query_lens_across_dp = tensor[3]
71 - num_reqs_across_dp = tensor[4]
80 + allow_ubatching_across_dp = tensor[4]
81 + num_reqs_across_dp = tensor[5]
82
83 # If ranks disagree on the uniform token count, or its 0 (means None) set to None
84 synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0])
@@ -83,6 +93,40 @@ def sync_cudagraph_and_dp_padding(
93 )
94 return synced_desc, None
95
96 + if torch.all(allow_ubatching_across_dp == 1):
97 + # This rank voted too, so its caller passed a config.
98 + assert parallel_config is not None
99 + # A uniform decode only if every rank runs the same uniform query
100 + # length, so every rank reaches the same answer here.
101 + uniform_decode_across_dp = uniform_decode and torch.all(
102 + uniform_token_counts_across_dp == (uniform_token_count or 0)
103 + )
104 + if check_ubatch_thresholds(
105 + parallel_config,
106 + # Thresholds only grow with the token count, so holding the smallest
107 + # rank to them holds every rank to them.
108 + int(num_tokens_across_dp.min()),
109 + uniform_decode=uniform_decode_across_dp,
110 + ):
111 + # Microbatching is all-or-nothing: the expert all-to-all is
112 + # collective, so every rank splits and pads to the same token count.
113 + # A rank too small to fill a microbatch does no work in it, like a
114 + # dummy run. Microbatched steps run eager; nothing is captured yet.
115 + ubatch_num_tokens = int(num_tokens_across_dp.max())
116 + return BatchExecutionDescriptor(
117 + cg_mode=CUDAGraphMode.NONE,
118 + num_tokens=ubatch_num_tokens,
119 + num_reqs=num_reqs,
120 + num_ubatches=get_num_ubatches(parallel_config),
121 + ), DPSyncState(
122 + num_tokens_across_dp=torch.full_like(
123 + num_tokens_across_dp, ubatch_num_tokens
124 + ),
125 + uniform_token_count=synced_uniform_token_count,
126 + eager=True,
127 + num_reqs=int(num_reqs_across_dp.max()),
128 + )
129 +
130 synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min().item()))
131
132 # If any rank wants to run eager, all ranks run eager
@@ -98,7 +142,7 @@ def sync_cudagraph_and_dp_padding(
142 num_tokens_across_dp=num_tokens_across_dp,
143 uniform_token_count=synced_uniform_token_count,
144 eager=True,
101 - num_reqs=int(num_reqs_across_dp.max().item()),
145 + num_reqs=int(num_reqs_across_dp.max()),
146 ),
147 )
148
@@ -136,7 +180,7 @@ def sync_cudagraph_and_dp_padding(
180 synced_desc.num_reqs
181 if synced_desc.cg_mode == CUDAGraphMode.FULL
182 and synced_desc.num_reqs is not None
139 - else int(num_reqs_across_dp.max().item())
183 + else int(num_reqs_across_dp.max())
184 ),
185 )
186
@@ -151,6 +195,9 @@ def dispatch_cg_and_sync_dp(
195 max_query_len: int | None = None,
196 need_eager: bool = False,
197 num_active_loras: int = 0,
198 + parallel_config: ParallelConfig | None = None,
199 + allow_ubatching: bool = False,
200 + uniform_decode: bool = False,
201 dp_sync: DPSyncState | None = None,
202 ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]:
203 """Pick a cudagraph descriptor for this batch, agreeing it across DP ranks.
@@ -239,4 +286,7 @@ def dispatch_cg_and_sync_dp(
286 dp_rank,
287 max_query_len=max_query_len,
288 num_active_loras=num_active_loras,
289 + parallel_config=parallel_config,
290 + allow_ubatching=allow_ubatching,
291 + uniform_decode=uniform_decode,
292 )
vllm/v1/worker/gpu/model_runner.py
+41 -3
@@ -161,6 +161,11 @@ from vllm.v1.worker.gpu.spec_decode.speculator import DraftModelSpeculator
161 from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
162 from vllm.v1.worker.gpu.states import RequestState
163 from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
164 +from vllm.v1.worker.gpu.ubatch_utils import (
165 + UBatchRunner,
166 + UBatchState,
167 + maybe_build_ubatch_runner,
168 +)
169 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
170 from vllm.v1.worker.utils import (
171 KVBlockZeroer,
@@ -230,6 +235,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
235 self.dp_size = self.parallel_config.data_parallel_size
236 self.dp_rank = self.parallel_config.data_parallel_rank
237
238 + # Dual batch overlap. Created in initialize_kv_cache(), once everything
239 + # it runs the microbatched forward with exists.
240 + self.ubatch_runner: UBatchRunner | None = None
241 +
242 # Detect EP all2all peer faults to prevent emitting corrupted output.
243 # Only meaningful for MoE + DP with an FT-capable all2all backend.
244 self.check_ep_fault = False
@@ -634,6 +643,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
643 self.block_tables,
644 cls=self.pcp_manager_cls,
645 )
646 + self.ubatch_runner = maybe_build_ubatch_runner(
647 + self.vllm_config,
648 + self.device,
649 + self.model_state,
650 + self.attn_groups,
651 + self.kv_cache_config,
652 + self.max_num_reqs,
653 + )
654 initialize_mamba_ssu_backend(
655 self.vllm_config.mamba_config,
656 self.kv_cache_config,
@@ -1597,6 +1614,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1614 max_query_len=max_query_len,
1615 need_eager=is_profile or skip_compiled,
1616 num_active_loras=num_active_loras,
1617 + parallel_config=self.parallel_config,
1618 + allow_ubatching=(
1619 + self.ubatch_runner is not None and not skip_attn_for_dummy_run
1620 + ),
1621 + uniform_decode=uniform_tok_count == self.decode_query_len,
1622 )
1623
1624 if batch_desc.num_tokens == 0:
@@ -1665,7 +1687,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1687
1688 attn_metadata = None
1689 slot_mappings_by_layer = None
1668 - if not (dummy_run and skip_attn_for_dummy_run):
1690 + ubatch_state: UBatchState | None = None
1691 + if batch_desc.num_ubatches > 1:
1692 + assert self.ubatch_runner is not None
1693 + assert block_tables is not None and slot_mappings is not None
1694 + ubatch_state = self.ubatch_runner.prepare(
1695 + input_batch, block_tables, slot_mappings
1696 + )
1697 + elif not (dummy_run and skip_attn_for_dummy_run):
1698 assert slot_mappings is not None
1699 slot_mappings_by_layer = build_slot_mappings_by_layer(
1700 slot_mappings, self.kv_cache_config
@@ -1751,7 +1780,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1780 del intermediate_tensors
1781
1782 # Update the EPLB meta.
1754 - self.eplb.prepare_forward(self.model_config, input_batch.num_tokens)
1783 + ubatch_slices = ubatch_state.slices if ubatch_state is not None else None
1784 + self.eplb.prepare_forward(
1785 + self.model_config, input_batch.num_tokens, ubatch_slices
1786 + )
1787
1788 self.step_timing.record_batch(
1789 input_batch, batch_desc.cg_mode == CUDAGraphMode.FULL
@@ -1783,12 +1815,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1815 dp_sync.num_tokens_across_dp if dp_sync is not None else None
1816 ),
1817 batch_descriptor=batch_descriptor,
1818 + ubatch_slices=ubatch_slices,
1819 slot_mapping=slot_mappings_by_layer,
1820 skip_compiled=skip_compiled,
1821 is_padding=input_batch.is_padding,
1822 ):
1823 self.kv_connector.pre_forward(scheduler_output)
1791 - if batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
1824 + if ubatch_state is not None:
1825 + assert self.ubatch_runner is not None
1826 + model_output = self.ubatch_runner.run(
1827 + self.model, model_inputs, ubatch_state
1828 + )
1829 + elif batch_desc.cg_mode == CUDAGraphMode.PIECEWISE:
1830 # Run the PIECEWISE graph (compiled PW cudagraph or breakable
1831 # cudagraph, chosen inside run_pw_graph). cg_mode is only
1832 # PIECEWISE after the cudagraph manager exists.
vllm/v1/worker/gpu/model_states/default.py
+2
@@ -172,6 +172,7 @@ class DefaultModelState(ModelState):
172 attn_groups: list[list[AttentionGroup]],
173 kv_cache_config: KVCacheConfig,
174 for_capture: bool = False,
175 + ubatch_idx: int = 0,
176 ) -> dict[str, Any]:
177 if cudagraph_mode == CUDAGraphMode.FULL:
178 # Use padded sizes - padding is handled by model_runner.prepare_attn.
@@ -224,5 +225,6 @@ class DefaultModelState(ModelState):
225 mm_req_doc_ranges=req_doc_ranges,
226 for_cudagraph_capture=for_capture,
227 rswa_prefix_lens=input_batch.prompt_lens,
228 + ubatch_idx=ubatch_idx,
229 )
230 return attn_metadata
vllm/v1/worker/gpu/model_states/encoder_decoder.py
+2
@@ -116,7 +116,9 @@ class EncoderDecoderModelState(ModelState):
116 attn_groups: list[list[AttentionGroup]],
117 kv_cache_config: KVCacheConfig,
118 for_capture: bool = False,
119 + ubatch_idx: int = 0,
120 ) -> dict[str, Any]:
121 + assert ubatch_idx == 0, "DBO is not supported"
122 if cudagraph_mode == CUDAGraphMode.FULL:
123 num_reqs = input_batch.num_reqs_after_padding
124 num_tokens = input_batch.num_tokens_after_padding
vllm/v1/worker/gpu/model_states/encoder_only.py
+2
@@ -167,7 +167,9 @@ class EncoderOnlyModelState(DefaultModelState):
167 attn_groups: list[list[AttentionGroup]],
168 kv_cache_config: KVCacheConfig,
169 for_capture: bool = False,
170 + ubatch_idx: int = 0,
171 ) -> dict[str, Any]:
172 + assert ubatch_idx == 0, "DBO is not supported"
173 attn_metadata = super().prepare_attn(
174 input_batch,
175 cudagraph_mode,
vllm/v1/worker/gpu/model_states/interface.py
+1
@@ -221,6 +221,7 @@ class ModelState(ABC):
221 attn_groups: list[list[AttentionGroup]],
222 kv_cache_config: KVCacheConfig,
223 for_capture: bool = False,
224 + ubatch_idx: int = 0,
225 ) -> dict[str, Any]:
226 raise NotImplementedError
227
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+2
@@ -238,7 +238,9 @@ class MambaHybridModelState(DefaultModelState):
238 attn_groups: list[list[AttentionGroup]],
239 kv_cache_config: KVCacheConfig,
240 for_capture: bool = False,
241 + ubatch_idx: int = 0,
242 ) -> dict[str, Any]:
243 + assert ubatch_idx == 0, "DBO is not supported"
244 if cudagraph_mode == CUDAGraphMode.FULL:
245 num_reqs = input_batch.num_reqs_after_padding
246 num_tokens = input_batch.num_tokens_after_padding
vllm/v1/worker/gpu/ubatch_utils.py new
+473
@@ -0,0 +1,473 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Microbatching (DBO) helpers for the V2 GPU model runner."""
4 +
5 +import threading
6 +from dataclasses import replace
7 +from typing import Any, NamedTuple
8 +
9 +import numpy as np
10 +import torch
11 +
12 +from vllm.config import CUDAGraphMode, VllmConfig
13 +from vllm.forward_context import (
14 + DPMetadata,
15 + ForwardContext,
16 + create_forward_context,
17 + override_forward_context,
18 +)
19 +from vllm.logger import init_logger
20 +from vllm.sequence import IntermediateTensors
21 +from vllm.utils.torch_utils import current_stream
22 +from vllm.v1.kv_cache_interface import KVCacheConfig
23 +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
24 +from vllm.v1.worker.gpu.input_batch import InputBatch
25 +from vllm.v1.worker.gpu.model_states.interface import ModelState
26 +from vllm.v1.worker.ubatch_utils import (
27 + UBatchSlice,
28 + UBatchSlices,
29 + create_sm_control_context,
30 + maybe_create_ubatch_slices,
31 +)
32 +from vllm.v1.worker.ubatching import make_ubatch_contexts
33 +from vllm.v1.worker.utils import AttentionGroup
34 +
35 +logger = init_logger(__name__)
36 +
37 +
38 +class UBatchState(NamedTuple):
39 + """Per-microbatch inputs for one dual-batch-overlap step."""
40 +
41 + slices: UBatchSlices
42 + forward_contexts: list[ForwardContext]
43 +
44 +
45 +def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices:
46 + """Split a DP-padded batch into the slices its microbatches run on.
47 +
48 + Splitting at the midpoint of the DP-padded token count leaves the trailing
49 + microbatch anywhere from full to entirely inside the padding. Clamping its
50 + request slice onto the last request leaves it holding that request with
51 + zero query tokens -- the shape a straddling request already takes -- so it
52 + stays well-formed with no work to do, like a dummy run.
53 + """
54 + _, ubatch_slices = maybe_create_ubatch_slices(
55 + True,
56 + input_batch.num_scheduled_tokens,
57 + input_batch.num_tokens_after_padding,
58 + input_batch.num_reqs_after_padding,
59 + num_ubatches,
60 + )
61 + assert ubatch_slices is not None
62 + return [
63 + UBatchSlice(
64 + slice(
65 + min(s.request_slice.start, input_batch.num_reqs - 1),
66 + min(s.request_slice.stop, input_batch.num_reqs_after_padding),
67 + ),
68 + s.token_slice,
69 + )
70 + for s in ubatch_slices
71 + ]
72 +
73 +
74 +def _slice_input_batch(
75 + input_batch: InputBatch,
76 + ubatch_slice: UBatchSlice,
77 + query_start_loc_buf: torch.Tensor,
78 + seq_lens_buf: torch.Tensor,
79 +) -> InputBatch:
80 + """Build the sub-`InputBatch` a single microbatch runs on.
81 +
82 + A request straddling the boundary appears in both microbatches, its query
83 + truncated to the tokens each one owns.
84 +
85 + Everything except `query_start_loc` and `seq_lens` is a view of the full
86 + batch's buffers, so slicing costs no allocation. Those two are rebased onto
87 + the microbatch's token range and written to the caller's buffers, which
88 + keeps them valid under CUDA graph capture and replay.
89 +
90 + The sub-batch describes the forward pass only. Sampling runs once over the
91 + merged batch, so `logits_indices`, `cu_num_logits` and the draft-token
92 + fields are carried over unchanged and must not be read off a sub-batch.
93 + """
94 + assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty"
95 +
96 + req_start = ubatch_slice.request_slice.start
97 + req_stop = ubatch_slice.request_slice.stop
98 + tok_start = ubatch_slice.token_slice.start
99 + tok_stop = ubatch_slice.token_slice.stop
100 +
101 + num_reqs_padded = req_stop - req_start
102 + num_tokens_padded = tok_stop - tok_start
103 + # The trailing microbatch may extend past the real batch into CUDA graph
104 + # padding; the unpadded counts stop at the last real request/token.
105 + num_reqs = max(0, min(req_stop, input_batch.num_reqs) - req_start)
106 + num_tokens = max(0, min(tok_stop, input_batch.num_tokens) - tok_start)
107 +
108 + query_start_loc = _slice_query_start_loc(
109 + input_batch.query_start_loc,
110 + query_start_loc_buf[: num_reqs_padded + 1],
111 + req_start,
112 + req_stop,
113 + tok_start,
114 + num_tokens_padded,
115 + )
116 + query_start_loc_np = np.clip(
117 + input_batch.query_start_loc_np[req_start : req_stop + 1] - tok_start,
118 + 0,
119 + num_tokens_padded,
120 + ).astype(np.int32)
121 +
122 + seq_lens = _slice_seq_lens(
123 + input_batch.seq_lens,
124 + seq_lens_buf[:num_reqs_padded],
125 + input_batch.query_start_loc,
126 + req_start,
127 + req_stop,
128 + tok_stop,
129 + )
130 + # Same truncation as `_slice_seq_lens`, on the host-side upper bound.
131 + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[
132 + req_start:req_stop
133 + ].clone()
134 + tokens_truncated = max(0, int(input_batch.query_start_loc_np[req_stop]) - tok_stop)
135 + if tokens_truncated:
136 + seq_lens_cpu_upper_bound[-1] -= tokens_truncated
137 +
138 + # Query lengths of the truncated requests, so consumers that derive
139 + # max_query_len from this array see the microbatch's own lengths.
140 + num_scheduled_tokens = np.diff(query_start_loc_np)[:num_reqs]
141 +
142 + dcp_local_seq_lens = input_batch.dcp_local_seq_lens
143 + if dcp_local_seq_lens is not None:
144 + # NOTE: a request split across microbatches keeps its full local
145 + # seq_len here. DCP is not yet validated with DBO.
146 + dcp_local_seq_lens = dcp_local_seq_lens[req_start:req_stop]
147 +
148 + return replace(
149 + input_batch,
150 + req_ids=input_batch.req_ids[req_start : min(req_stop, input_batch.num_reqs)],
151 + num_reqs=num_reqs,
152 + num_reqs_after_padding=num_reqs_padded,
153 + idx_mapping=input_batch.idx_mapping[req_start:req_stop],
154 + idx_mapping_np=input_batch.idx_mapping_np[req_start:req_stop],
155 + num_scheduled_tokens=num_scheduled_tokens,
156 + num_tokens=num_tokens,
157 + num_tokens_after_padding=num_tokens_padded,
158 + query_start_loc=query_start_loc,
159 + query_start_loc_np=query_start_loc_np,
160 + seq_lens=seq_lens,
161 + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
162 + dcp_local_seq_lens=dcp_local_seq_lens,
163 + num_computed_tokens_np=input_batch.num_computed_tokens_np[req_start:req_stop],
164 + prefill_len_np=input_batch.prefill_len_np[req_start:req_stop],
165 + num_computed_prefill_tokens_np=input_batch.num_computed_prefill_tokens_np[
166 + req_start:req_stop
167 + ],
168 + is_prefilling_np=input_batch.is_prefilling_np[req_start:req_stop],
169 + input_ids=input_batch.input_ids[tok_start:tok_stop],
170 + positions=input_batch.positions[tok_start:tok_stop],
171 + is_padding=input_batch.is_padding[tok_start:tok_stop],
172 + prompt_lens=(
173 + None
174 + if input_batch.prompt_lens is None
175 + else input_batch.prompt_lens[req_start:req_stop]
176 + ),
177 + )
178 +
179 +
180 +def _slice_query_start_loc(
181 + query_start_loc: torch.Tensor,
182 + out: torch.Tensor,
183 + req_start: int,
184 + req_stop: int,
185 + tok_start: int,
186 + num_tokens_padded: int,
187 +) -> torch.Tensor:
188 + """Rebase query_start_loc onto the microbatch's token range.
189 +
190 + The clamp truncates the requests straddling either boundary to the tokens
191 + this microbatch owns.
192 + """
193 + torch.sub(query_start_loc[req_start : req_stop + 1], tok_start, out=out)
194 + return out.clamp_(0, num_tokens_padded)
195 +
196 +
197 +def _slice_seq_lens(
198 + seq_lens: torch.Tensor,
199 + out: torch.Tensor,
200 + query_start_loc: torch.Tensor,
201 + req_start: int,
202 + req_stop: int,
203 + tok_stop: int,
204 +) -> torch.Tensor:
205 + """Slice seq_lens, shortening the request truncated at the boundary.
206 +
207 + Only tokens past the end of this microbatch are missing; tokens it owns in
208 + an *earlier* microbatch are computed before this one reads them, so they
209 + still count. Deriving the truncation from this microbatch's own query
210 + lengths would subtract them twice. It is computed from tensors, not Python
211 + ints, so this stays valid inside a CUDA graph.
212 + """
213 + out.copy_(seq_lens[req_start:req_stop])
214 + last = req_stop - req_start - 1
215 + out[last] -= (query_start_loc[req_stop] - tok_stop).clamp_(min=0)
216 + return out
217 +
218 +
219 +def slice_model_inputs(
220 + model_inputs: dict[str, Any], token_slice: slice
221 +) -> dict[str, Any]:
222 + """Narrow the model's per-token inputs to one microbatch."""
223 + sliced = dict(model_inputs)
224 + for key in ("input_ids", "inputs_embeds"):
225 + value = model_inputs.get(key)
226 + if value is not None:
227 + sliced[key] = value[token_slice]
228 +
229 + positions = model_inputs["positions"]
230 + # M-RoPE carries a leading section dim.
231 + sliced["positions"] = (
232 + positions[:, token_slice] if positions.ndim == 2 else positions[token_slice]
233 + )
234 +
235 + intermediate_tensors = model_inputs.get("intermediate_tensors")
236 + if intermediate_tensors is not None:
237 + sliced["intermediate_tensors"] = intermediate_tensors[token_slice]
238 + return sliced
239 +
240 +
241 +def merge_ubatch_outputs(outputs: list[Any]) -> Any:
242 + """Reassemble the full-batch model output from the per-microbatch ones.
243 +
244 + Models return hidden states, a tuple of those plus auxiliary hidden states
245 + (EAGLE3), or IntermediateTensors on non-final pipeline ranks.
246 + """
247 + first = outputs[0]
248 + if isinstance(first, IntermediateTensors):
249 + return IntermediateTensors(
250 + {
251 + key: torch.cat([out.tensors[key] for out in outputs], dim=0)
252 + for key in first.tensors
253 + }
254 + )
255 + if isinstance(first, tuple):
256 + return tuple(torch.cat(parts, dim=0) for parts in zip(*outputs))
257 + return torch.cat(outputs, dim=0)
258 +
259 +
260 +class UBatchRunner:
261 + """Runs the model on two or more microbatches that overlap each other.
262 +
263 + Each microbatch gets a thread and its own forward context. Only one thread
264 + holds the GPU at a time; they hand off to each other at the communication
265 + points inside the model (see `vllm.v1.worker.ubatching`), so one
266 + microbatch's expert all-to-all overlaps the other's compute.
267 + """
268 +
269 + def __init__(
270 + self,
271 + vllm_config: VllmConfig,
272 + device: torch.device,
273 + model_state: ModelState,
274 + attn_groups: list[list[AttentionGroup]],
275 + kv_cache_config: KVCacheConfig,
276 + max_num_reqs: int,
277 + ):
278 + self.vllm_config = vllm_config
279 + self.parallel_config = vllm_config.parallel_config
280 + self.num_ubatches = self.parallel_config.num_ubatches
281 + self.device = device
282 + self.model_state = model_state
283 + self.attn_groups = attn_groups
284 + self.kv_cache_config = kv_cache_config
285 + # `query_start_loc` and `seq_lens` are rebased onto each microbatch's
286 + # own token range, so they cannot be views of the full batch's buffers.
287 + # Allocating up front keeps their addresses stable across replays.
288 + self.ubatch_query_start_loc = [
289 + torch.zeros(max_num_reqs + 1, dtype=torch.int32, device=device)
290 + for _ in range(self.num_ubatches)
291 + ]
292 + self.ubatch_seq_lens = [
293 + torch.zeros(max_num_reqs, dtype=torch.int32, device=device)
294 + for _ in range(self.num_ubatches)
295 + ]
296 + self.comm_stream = torch.cuda.Stream(device=device)
297 + # The microbatch threads plus the thread that starts them.
298 + self.ready_barrier = threading.Barrier(self.num_ubatches + 1)
299 + self.sm_control = create_sm_control_context(self.parallel_config)
300 +
301 + def prepare(
302 + self,
303 + input_batch: InputBatch,
304 + block_tables: tuple[torch.Tensor, ...],
305 + slot_mappings: torch.Tensor,
306 + ) -> UBatchState:
307 + """Split the batch into the microbatches the step will run on.
308 +
309 + Attention metadata is built per microbatch and carried in the forward
310 + contexts the threads install, not in the caller's.
311 + """
312 + ubatch_slices = create_ubatch_slices(input_batch, self.num_ubatches)
313 +
314 + attn_metadata = []
315 + slot_mappings_by_layer = []
316 + is_padding = []
317 + for i, ubatch_slice in enumerate(ubatch_slices):
318 + ubatch = _slice_input_batch(
319 + input_batch,
320 + ubatch_slice,
321 + self.ubatch_query_start_loc[i],
322 + self.ubatch_seq_lens[i],
323 + )
324 + ubatch_slot_mappings = slot_mappings[:, ubatch_slice.token_slice]
325 + ubatch_block_tables = tuple(
326 + block_table[ubatch_slice.request_slice] for block_table in block_tables
327 + )
328 + attn_metadata.append(
329 + self.model_state.prepare_attn(
330 + ubatch,
331 + CUDAGraphMode.NONE,
332 + ubatch_block_tables,
333 + ubatch_slot_mappings,
334 + self.attn_groups,
335 + self.kv_cache_config,
336 + ubatch_idx=i,
337 + )
338 + )
339 + slot_mappings_by_layer.append(
340 + build_slot_mappings_by_layer(ubatch_slot_mappings, self.kv_cache_config)
341 + )
342 + is_padding.append(ubatch.is_padding)
343 +
344 + return UBatchState(
345 + slices=ubatch_slices,
346 + forward_contexts=self._make_forward_contexts(
347 + ubatch_slices, attn_metadata, slot_mappings_by_layer, is_padding
348 + ),
349 + )
350 +
351 + def _make_forward_contexts(
352 + self,
353 + ubatch_slices: UBatchSlices,
354 + attn_metadata: list[dict[str, Any]],
355 + slot_mappings_by_layer: list[dict[str, torch.Tensor]],
356 + is_padding: list[torch.Tensor | None],
357 + ) -> list[ForwardContext]:
358 + """Build one forward context per microbatch.
359 +
360 + The expert all-to-all runs once per microbatch, so each context's DP
361 + metadata carries that microbatch's token count. Microbatching pads every
362 + DP rank alike, so each rank can assume that count for the others too.
363 + """
364 + dp_size = self.parallel_config.data_parallel_size
365 + forward_contexts = []
366 + for i, ubatch_slice in enumerate(ubatch_slices):
367 + num_tokens = ubatch_slice.num_tokens
368 + num_tokens_across_dp = torch.full(
369 + (dp_size,), num_tokens, dtype=torch.int32, device="cpu"
370 + )
371 + forward_contexts.append(
372 + create_forward_context(
373 + attn_metadata[i],
374 + self.vllm_config,
375 + dp_metadata=DPMetadata.make(
376 + self.parallel_config, num_tokens, num_tokens_across_dp
377 + ),
378 + cudagraph_runtime_mode=CUDAGraphMode.NONE,
379 + slot_mapping=slot_mappings_by_layer[i],
380 + is_padding=is_padding[i],
381 + )
382 + )
383 + return forward_contexts
384 +
385 + def run(
386 + self,
387 + model: Any,
388 + model_inputs: dict[str, Any],
389 + ubatch_state: UBatchState,
390 + ) -> Any:
391 + ubatch_slices = ubatch_state.slices
392 + assert len(ubatch_slices) == len(ubatch_state.forward_contexts)
393 + assert len(ubatch_slices) == self.num_ubatches
394 +
395 + ubatch_contexts = make_ubatch_contexts(
396 + num_micro_batches=self.num_ubatches,
397 + comm_stream=self.comm_stream,
398 + compute_stream=current_stream(),
399 + forward_contexts=ubatch_state.forward_contexts,
400 + ready_barrier=self.ready_barrier,
401 + )
402 +
403 + outputs: dict[int, Any] = {}
404 + errors: dict[int, BaseException] = {}
405 +
406 + @torch.inference_mode()
407 + def run_ubatch(ubatch_context, inputs: dict[str, Any]) -> None:
408 + try:
409 + with ubatch_context:
410 + outputs[ubatch_context.id] = model(**inputs)
411 + except BaseException as e: # noqa: BLE001
412 + # Recorded so the join below can report which microbatch broke.
413 + # Note the siblings are not told to unwind: microbatches hand
414 + # control to each other, so one that dies mid-forward leaves
415 + # the others parked on a handoff that never comes and the step
416 + # hangs. The V1 runner behaves the same way; fixing it means
417 + # changing the shared handoff protocol in ubatching.py, which
418 + # is deliberately out of scope here.
419 + errors[ubatch_context.id] = e
420 +
421 + # The threads manage the forward context themselves; clear it here so
422 + # it is restored correctly once they are done.
423 + with override_forward_context(None), self.sm_control:
424 + threads = []
425 + for ubatch_context, ubatch_slice in zip(ubatch_contexts, ubatch_slices):
426 + thread = threading.Thread(
427 + target=run_ubatch,
428 + args=(
429 + ubatch_context,
430 + slice_model_inputs(model_inputs, ubatch_slice.token_slice),
431 + ),
432 + )
433 + threads.append(thread)
434 + thread.start()
435 +
436 + # Wait for every thread to reach its context, then start the first.
437 + self.ready_barrier.wait()
438 + ubatch_contexts[0].cpu_wait_event.set()
439 + for thread in threads:
440 + thread.join()
441 +
442 + if errors:
443 + failed = min(errors)
444 + raise RuntimeError(
445 + f"Microbatch {failed} of {self.num_ubatches} failed"
446 + ) from errors[failed]
447 + return merge_ubatch_outputs([outputs[i] for i in range(self.num_ubatches)])
448 +
449 +
450 +def maybe_build_ubatch_runner(
451 + vllm_config: VllmConfig,
452 + device: torch.device,
453 + model_state: ModelState,
454 + attn_groups: list[list[AttentionGroup]],
455 + kv_cache_config: KVCacheConfig,
456 + max_num_reqs: int,
457 +) -> UBatchRunner | None:
458 + """Build the microbatch runner, or None when DBO is not in use.
459 +
460 + Microbatching needs the DP handshake to agree on it, so it is only
461 + available with more than one DP rank (as in the V1 runner).
462 + """
463 + parallel_config = vllm_config.parallel_config
464 + if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1:
465 + return None
466 +
467 + logger.info_once(
468 + "Dual batch overlap is enabled. Microbatched steps run without "
469 + "CUDA graphs on the V2 model runner."
470 + )
471 + return UBatchRunner(
472 + vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs
473 + )
vllm/v1/worker/gpu_ubatch_wrapper.py
+2 -93
@@ -11,7 +11,6 @@ import torch
11 import vllm.envs as envs
12 from vllm.compilation.cuda_graph import CUDAGraphWrapper
13 from vllm.config import CUDAGraphMode, VllmConfig
14 -from vllm.distributed import get_ep_group
14 from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
15 from vllm.forward_context import (
16 DPMetadata,
@@ -23,9 +22,7 @@ from vllm.logger import init_logger
22 from vllm.model_executor.offloader.base import get_offloader
23 from vllm.platforms import current_platform
24 from vllm.sequence import IntermediateTensors
26 -from vllm.utils.deep_gemm import set_num_sms as deep_gemm_set_num_sms
27 -from vllm.utils.import_utils import has_deep_gemm
28 -from vllm.utils.platform_utils import num_compute_units
25 +from vllm.v1.worker.ubatch_utils import create_sm_control_context
26 from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts
27
28 logger = init_logger(__name__)
@@ -65,51 +62,6 @@ class CUDAGraphMetaData:
62 outputs: Any | None = None
63
64
68 -class SMControlContextManager:
69 - def __init__(
70 - self,
71 - comm_sms: int,
72 - set_comm_sms: Callable[[int], None],
73 - set_compute_sms: Callable[[int], None],
74 - ):
75 - """
76 - Context manager for controlling SM (Streaming Multiprocessor)
77 - allocation. Upon entering the context, it sets the number of SMs
78 - allocated for communication and computation to comm_sms and
79 - total_sms - comm_sms respectively. Upon exiting, it restores the
80 - allocation to use all available SMs (i.e. total_sms).
81 -
82 - Args:
83 - comm_sms (int): The number of SMs to allocate for communication.
84 - (The remainder will be used for computation.)
85 - set_comm_sms (Callable[[int], None]):
86 - A function that sets the number of SMs for communication.
87 - set_compute_sms (Callable[[int], None]):
88 - A function that sets the number of SMs for computation.
89 - """
90 -
91 - assert current_platform.is_cuda() or current_platform.is_rocm(), (
92 - "SM/CU control is supported on CUDA and ROCm platforms"
93 - )
94 - device = torch.accelerator.current_device_index()
95 - total_sms = num_compute_units(device)
96 -
97 - assert comm_sms < total_sms
98 - self.total_sms = total_sms
99 - self.compute_sms = total_sms - comm_sms
100 - self.comm_sms = comm_sms
101 - self.set_comm_sms = set_comm_sms
102 - self.set_compute_sms = set_compute_sms
103 -
104 - def __enter__(self):
105 - self.set_comm_sms(self.comm_sms)
106 - self.set_compute_sms(self.compute_sms)
107 -
108 - def __exit__(self, exc_type, exc_value, traceback):
109 - self.set_comm_sms(self.total_sms)
110 - self.set_compute_sms(self.total_sms)
111 -
112 -
65 class UBatchWrapper:
66 def __init__(
67 self,
@@ -135,7 +87,7 @@ class UBatchWrapper:
87 runnable, vllm_config, runtime_mode=runtime_mode
88 )
89
138 - self.sm_control = self._create_sm_control_context(vllm_config)
90 + self.sm_control = create_sm_control_context(vllm_config.parallel_config)
91 self.device = device
92 self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
93 self._runnable_str = str(runnable) if self.is_debugging_mode else None
@@ -151,49 +103,6 @@ class UBatchWrapper:
103 if self.cudagraph_wrapper is not None:
104 self.cudagraph_wrapper.clear_graphs()
105
154 - @staticmethod
155 - def _create_sm_control_context(vllm_config: VllmConfig):
156 - comm_sms: int = envs.VLLM_DBO_COMM_SMS
157 - rocm_deepep_ht_dbo = (
158 - current_platform.is_rocm()
159 - and vllm_config.parallel_config.enable_dbo
160 - and vllm_config.parallel_config.all2all_backend == "deepep_high_throughput"
161 - )
162 - if rocm_deepep_ht_dbo:
163 - # On ROCm, reserving CUs for DeepEP HT communication under DBO
164 - # corrupts DP+EP generation accuracy. Keep the backend active, but
165 - # leave all CUs visible to the compute and communication kernels.
166 - comm_sms = 0
167 -
168 - set_comm_sms = lambda sms: None
169 - if vllm_config.parallel_config.enable_expert_parallel:
170 - # Currently only DeepEP highthroughput supports SM control so this
171 - # only affects that case.
172 - ep_group = get_ep_group()
173 - device_communicator = ep_group.device_communicator
174 - all2all_manager = None
175 - if device_communicator is not None:
176 - all2all_manager = device_communicator.all2all_manager
177 -
178 - if all2all_manager is not None:
179 - max_sms_used = all2all_manager.max_sms_used()
180 - if max_sms_used is not None:
181 - comm_sms = min(comm_sms, max_sms_used)
182 -
183 - if comm_sms > 0 and all2all_manager is not None:
184 - set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms)
185 -
186 - # TODO(lucas): support other kernels besides DeepGEMM
187 - set_compute_sms = lambda sms: None
188 - if has_deep_gemm() and comm_sms > 0:
189 - set_compute_sms = lambda sms: deep_gemm_set_num_sms(sms)
190 -
191 - return SMControlContextManager(
192 - comm_sms=comm_sms,
193 - set_comm_sms=set_comm_sms,
194 - set_compute_sms=set_compute_sms,
195 - )
196 -
106 def __getattr__(self, key: str):
107 # allow accessing the attributes of the runnable.
108 if hasattr(self.runnable, key):
vllm/v1/worker/ubatch_utils.py
+103
@@ -1,12 +1,19 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from collections.abc import Callable
4 from dataclasses import dataclass
5 from typing import TypeAlias
6
7 import numpy as np
8 import torch
9
10 +import vllm.envs as envs
11 from vllm.config import ParallelConfig
12 +from vllm.distributed import get_ep_group
13 +from vllm.platforms import current_platform
14 +from vllm.utils.deep_gemm import set_num_sms as deep_gemm_set_num_sms
15 +from vllm.utils.import_utils import has_deep_gemm
16 +from vllm.utils.platform_utils import num_compute_units
17 from vllm.v1.attention.backend import CommonAttentionMetadata
18
19
@@ -29,12 +36,108 @@ class UBatchSlice:
36 UBatchSlices: TypeAlias = list[UBatchSlice]
37
38
39 +class SMControlContextManager:
40 + def __init__(
41 + self,
42 + comm_sms: int,
43 + set_comm_sms: Callable[[int], None],
44 + set_compute_sms: Callable[[int], None],
45 + ):
46 + """
47 + Context manager for controlling SM (Streaming Multiprocessor)
48 + allocation. Upon entering the context, it sets the number of SMs
49 + allocated for communication and computation to comm_sms and
50 + total_sms - comm_sms respectively. Upon exiting, it restores the
51 + allocation to use all available SMs (i.e. total_sms).
52 +
53 + Args:
54 + comm_sms (int): The number of SMs to allocate for communication.
55 + (The remainder will be used for computation.)
56 + set_comm_sms (Callable[[int], None]):
57 + A function that sets the number of SMs for communication.
58 + set_compute_sms (Callable[[int], None]):
59 + A function that sets the number of SMs for computation.
60 + """
61 +
62 + assert current_platform.is_cuda() or current_platform.is_rocm(), (
63 + "SM/CU control is supported on CUDA and ROCm platforms"
64 + )
65 + device = torch.accelerator.current_device_index()
66 + total_sms = num_compute_units(device)
67 +
68 + assert comm_sms < total_sms
69 + self.total_sms = total_sms
70 + self.compute_sms = total_sms - comm_sms
71 + self.comm_sms = comm_sms
72 + self.set_comm_sms = set_comm_sms
73 + self.set_compute_sms = set_compute_sms
74 +
75 + def __enter__(self):
76 + self.set_comm_sms(self.comm_sms)
77 + self.set_compute_sms(self.compute_sms)
78 +
79 + def __exit__(self, exc_type, exc_value, traceback):
80 + self.set_comm_sms(self.total_sms)
81 + self.set_compute_sms(self.total_sms)
82 +
83 +
84 +def create_sm_control_context(
85 + parallel_config: ParallelConfig,
86 +) -> SMControlContextManager:
87 + """Reserve SMs for communication kernels while microbatches overlap."""
88 + comm_sms: int = envs.VLLM_DBO_COMM_SMS
89 + rocm_deepep_ht_dbo = (
90 + current_platform.is_rocm()
91 + and parallel_config.enable_dbo
92 + and parallel_config.all2all_backend == "deepep_high_throughput"
93 + )
94 + if rocm_deepep_ht_dbo:
95 + # On ROCm, reserving CUs for DeepEP HT communication under DBO
96 + # corrupts DP+EP generation accuracy. Keep the backend active, but
97 + # leave all CUs visible to the compute and communication kernels.
98 + comm_sms = 0
99 +
100 + set_comm_sms = lambda sms: None
101 + if parallel_config.enable_expert_parallel:
102 + # Currently only DeepEP highthroughput supports SM control so this
103 + # only affects that case.
104 + ep_group = get_ep_group()
105 + device_communicator = ep_group.device_communicator
106 + all2all_manager = None
107 + if device_communicator is not None:
108 + all2all_manager = device_communicator.all2all_manager
109 +
110 + if all2all_manager is not None:
111 + max_sms_used = all2all_manager.max_sms_used()
112 + if max_sms_used is not None:
113 + comm_sms = min(comm_sms, max_sms_used)
114 +
115 + if comm_sms > 0 and all2all_manager is not None:
116 + set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms)
117 +
118 + # TODO(lucas): support other kernels besides DeepGEMM
119 + set_compute_sms = lambda sms: None
120 + if has_deep_gemm() and comm_sms > 0:
121 + set_compute_sms = lambda sms: deep_gemm_set_num_sms(sms)
122 +
123 + return SMControlContextManager(
124 + comm_sms=comm_sms,
125 + set_comm_sms=set_comm_sms,
126 + set_compute_sms=set_compute_sms,
127 + )
128 +
129 +
130 def is_last_ubatch_empty(
131 orig_num_tokens: int, padded_num_tokens: int, num_ubatches: int
132 ) -> bool:
133 return (padded_num_tokens // num_ubatches) * (num_ubatches - 1) >= orig_num_tokens
134
135
136 +def get_num_ubatches(parallel_config: ParallelConfig) -> int:
137 + """How many microbatches a step is split into; 1 when microbatching is off."""
138 + return parallel_config.num_ubatches if parallel_config.use_ubatching else 1
139 +
140 +
141 def check_ubatch_thresholds(
142 config: ParallelConfig, num_tokens: int, uniform_decode: bool
143 ) -> bool: