vLLM 源码完全拆解（二十七）：Dual Batch Overlap——两个线程、两条流与一次 yield，如何把专家通信藏进计算里

从「专家并行 + 数据并行」说起

DeepSeek 风格的稀疏 MoE 模型有两本书要算：一本是稀疏专家激活，另一本是 KV 缓存。以 DeepSeek-R1 为例，671B 总参数里每次前向只激活约 37B；同时 MLA 的低秩投影在张量并行下会被每个 rank 各自复制一份，TP 越大，有效 KV 缓存越小。

专家并行（Expert Parallel，EP）把专家摊到多个 rank 上，每个 rank 只持有自己那份专家。前向时，每个 token 按路由结果被发往持有对应专家的 rank，算完再收回来。这样每个 rank 都能把显存几乎全部用于 KV 缓存和专家权重。再叠加数据并行（DP），注意力层在 DP 组内复制、专家层在 EP 组内切分，就是 vLLM 文档里说的 Wide-EP：注意力计算按 DP rank 各自独立，MoE 的 token 在 EP 组里全局交换。

Wide-EP 的收益是显存与批次规模，代价是 MoE 层前后多出了通信：每一层都要做一次 dispatch（token 按专家发出去）和一次 combine（结果按 token 收回来）。EP 规模越大，参与 all-to-all 的 rank 越多，拓扑越复杂，这段通信的绝对时长就越可观。

通信为什么很难被「异步」掉

按理论，all-to-all 应该可以和计算重叠：一边发、一边算。但落到单卡上的执行序列里，这条路并不平坦：

1. dispatch 的结果是专家 GEMM 的输入。没有 dispatch 完成，本地专家无活可干。
2. combine 的输入是专家 GEMM 的输出。GEMM 没做完，combine 无从发起。
3. 同一时刻，一个微批的通信与另一个微批的计算可以并行，但单个微批内部这个链条无法拆开。

所以，真正能重叠的不是「一个微批的通信与它自己的计算」，而是「两个微批」——A 批在通信时，B 批在计算；B 批在通信时，A 批在计算。这就是双批重叠的全部动机。vLLM 的设计文档 docs/design/dbo.md 开篇也是这样写的：DBO 的核心动机是把 MoE 层稀疏的 all-to-all 与周围的计算重叠起来，目前只面向 DP+EP 部署。

上图是 vLLM 官方博客里一份 DeepSeek 解码负载的 profiler trace。横轴是时间，每个格子是一个流上的算子。可以看到 QKV、Attn、OProj、Select 这些计算算子都很短，但 MoE Dispatch/Combine 那一段拉得很长——即使这一批的计算量很小，通信仍按自己的节奏走完。在 EP 规模较大的部署里，这段通信往往就是把 GPU 利用率压在低位的主要因子。

关键数字：阈值与开关

DBO 的配置集中在 ParallelConfig 里，逻辑非常简单：

1
# vllm/config/parallel.py（节选）
2
enable_dbo: bool = False
3
ubatch_size: int = 0                       # 手动指定微批数
4
dbo_decode_token_threshold: int = 32       # 纯解码批的启用门槛
5
dbo_prefill_token_threshold: int = 512     # 含 prefill 批的启用门槛
6

7
@property
8
def use_ubatching(self):
9
    return self.enable_dbo or self.ubatch_size > 1
10

11
@property
12
def num_ubatches(self):
13
    return 2 if self.enable_dbo else self.ubatch_size

--enable-dbo 等价于「启用双批」，num_ubatches 固定为 2；--ubatch-size N 则允许把批次切成 N 份（ubatching.py 里有 _NUM_UBATCHES: int = 2 的默认值，运行时会被实际微批数覆盖，说明主线仍把双批作为默认形态）。两个阈值决定了「多小的批不值得切」：纯 decode 批至少 32 个 token、含 prefill 的批至少 512 个 token 才会被切开。配置校验里还有一条硬约束：两个阈值都不得小于微批数，否则切完会出现空批。

这组数字背后是一笔很朴素的账：切批本身有开销——多一次 DP 协调的 all-reduce、多一次元数据切分、多两个线程的调度、多一份流与事件的同步。批太小，重叠省下的时间还不够付这些账。

为什么必须全组一致

当 EP 组横跨多个 DP rank，每个 rank 的路由结果不同，dispatch 的目标 rank 也不同。一次 all-to-all 是一个集体操作，所有 rank 必须同时发起、同时完成。如果 rank 0 把批次切成两半、rank 1 不切，那么 rank 0 发出的两段 dispatch 会与 rank 1 的单次 dispatch 对不上号。

所以 DBO 的第一条铁律是：所有 DP rank 对「是否切批」必须给出一致答案。实现方式是一次 CPU 上的 all-reduce。

1
# vllm/v1/worker/gpu/dp_utils.py（节选）
2
tensor = torch.zeros(7, dp_size, dtype=torch.int32, device="cpu")
3
tensor[0][dp_rank] = num_tokens
4
tensor[1][dp_rank] = desired_batch_desc.cg_mode.value
5
tensor[2][dp_rank] = uniform_token_count or 0
6
tensor[3][dp_rank] = max_query_len or -1
7
tensor[4][dp_rank] = int(allow_ubatching)      # 本 rank 的投票
8
tensor[5][dp_rank] = num_reqs
9
tensor[6][dp_rank] = int(allow_skip_drafts)
10
get_dp_group().cpu_all_reduce_(tensor)

7 行 int32 的小张量在 CPU 上做一次求和型 all-reduce，把每个 rank 的意图直接抹成全局可见。求和之后，torch.all(allow_ubatching_across_dp == 1) 就是「所有 rank 都同意切批」；任何一个 rank 投反对票，结果就不为 1，全组都不切。这里选 CPU 而不是 NCCL，是因为这个决策直接决定后续要发多少个 kernel、走哪条执行路径，CPU 必须马上拿到结果；7 个数的一趟 CPU all-reduce 代价在微秒级，相对一步前向可以忽略。关于这类并行协调的细节，可参考 AllReduce 完全拆解：Ring、树算法与 NCCL——张量并行推理的通信账单从哪来。

谁在投反对票

allow_ubatching 来自 GPUModelRunner._allow_microbatching()。它做一件很具体的事：拒绝切开「读者正在等写者」的批次。

1
def _allow_microbatching(self, num_reqs, num_scheduled_tokens_np) -> bool:
2
    if not self.parallel_config.use_ubatching or num_reqs < 2:
3
        return True
4
    computed = self.input_batch.num_computed_tokens_cpu[:num_reqs]
5
    query_lens = num_scheduled_tokens_np[:num_reqs]
6
    readers = np.flatnonzero(
7
        (query_lens > (self.reorder_batch_threshold or 1)) & (computed > 0)
8
    )
9
    if readers.size == 0:
10
        return True
11

12
    for block_table in self.input_batch.block_table.block_tables:
13
        block_size = block_table.block_size
14
        table = block_table.get_numpy_array()
15
        start = computed // block_size
16
        stop = (computed + query_lens + block_size - 1) // block_size
17
        ...
18
        for reader in readers:
19
            whole = computed[reader] // block_size
20
            if np.isin(table[reader, :whole], written).any():
21
                return False
22
    return True

这段代码处理的是前缀缓存与写者的时间差。设想同一批里有两个请求：请求 A 是长 prefill，正在往一组 KV 块里写；请求 B 命中了前缀缓存，它的历史 KV 有一部分正好落在 A 这一拍才写的那组块里。整批一起跑时，这一步先写 KV、再做注意力，B 读到的是 A 刚刚写好的块，正确。但切成两半后，如果 B 在前半段、A 在后半段，B 的注意力就会去读一组尚未被填写的块。

代码的做法是：找出「读历史」请求（有 query 且已经算过 token，说明它不是首段 prefill），对每个读请求，把它已经完成的整块（whole）拿出来，和本批所有请求将要写的块集合 written 求交。交集非空就返回 False，全组在这一步不切批。注意只检查「整块」：读请求的末尾可能落在半满的块里，那是它自己的私有块，它会先写后读，不构成竞态。

这是一个典型的「正确性优先于性能」的 veto：切批只是机会主义优化，遇到写读耦合就退回整批执行，代价是一步的通信没有被藏住，收益是结果永远与不开 DBO 时一致。

投票通过之后的第二道闸：空微批不能进图

即便所有 rank 都同意切，还要处理一种边界：某个 rank 的第二个微批可能是空的。

ubatch_utils.py 里保留了一个 is_last_ubatch_empty() 辅助函数：

1
def is_last_ubatch_empty(orig_num_tokens, padded_num_tokens, num_ubatches):
2
    return (padded_num_tokens // num_ubatches) * (num_ubatches - 1) >= orig_num_tokens

按这个函数的定义，如果把全局 pad 后的 token 数均分，第一份已经装下了本 rank 的全部真实 token，第二个微批就是空的。设计文档的表述是「若某个 rank 在 pad 后第二个微批为空，则放弃切批，所有 rank 都不切」，当前主线的执行路径则不同：注释里写明「Empty microbatches run as dummy batches.」——空微批仍要参与 all-to-all，只是走 dummy 路径，不产生真实的专家负载；切分本身是全局约定，不能因为一个 rank 没货就临阵变卦。is_last_ubatch_empty() 这个函数还留在工具模块里，当前主线的 gpu_model_runner.py 已不再调用它，读源码时不要把它当作执行路径上的真正闸门。

真正被拦住的是 CUDA Graph：

1
if 2 * int(num_tokens_across_dp.min()) < ubatch_desc.num_tokens:
2
    # 如果有 rank 的第二个微批为空，就不用 CUDA Graph 跑
3
    ubatch_desc = None

最小值乘 2 若仍小于图桶的 token 数，说明 pad 之后必有一边是空的。图桶是按「每个微批都有真实 token」的形状捕获与回放的，空微批与图桶的假设不匹配，于是这一步退回 eager 的线程路径。切批继续、图回放放弃——这是当前代码的选择，而不是整体放弃切批。

空微批的代价确实存在：本 rank 的第二个微批没有计算，DP 组内的负载不均会被进一步放大，慢微批拖住快微批。所以它更像一个「允许但尽量避免」的兜底，而不是理想路径。

Pad 到最大，而非各扫门前雪

所有 rank 决定切批后，token 数不是各自为政，而是统一 pad 到全局最大：

1
# Expert all-to-all requires every rank to split and pad equally.
2
# Empty microbatches run as dummy batches.
3
ubatch_num_tokens = int(num_tokens_across_dp.max())

原因和投票的动机一样：all-to-all 是集体的，各 rank 的微批必须能对位。rank 0 有 900 个 token 切成 450+450，rank 1 有 600 个 token——它也要把 token 补到 450+450 对齐，第二段的 150 个 pad 行走 dummy 路径，dispatch 里不产生真实的专家负载。pad 本身是一种浪费，但比切不对位要多花一次集合通信便宜得多。

对 pad 的专门处理

最后一个微批要覆盖到 num_tokens_padded，因为切片是在 DP pad 之前计算的：

1
def _pad_out_ubatch_slices(ubatch_slices, num_total_tokens, num_reqs_padded):
2
    last_slice = ubatch_slices[-1]
3
    return ubatch_slices[:-1] + [UBatchSlice(
4
        slice(last_slice.request_slice.start, num_reqs_padded),
5
        slice(last_slice.token_slice.start, num_total_tokens),
6
    )]

于是存在两份切片：ubatch_slices（真实 token）和 ubatch_slices_padded（含 pad）。模型输入、slot mapping、DP 元数据都用 padded 版以保证各微批形状一致；注意力元数据则按图模式选择——FULL CUDA Graph 回放时用 padded 版（图里的形状必须与捕获时完全一致），否则用真实版（避免注意力去读不存在的 pad token）。

它持有什么

1
class UBatchContext:
2
    def __init__(self, id, comm_stream, compute_stream, forward_context,
3
                 ready_barrier, cpu_wait_event, cpu_signal_event,
4
                 gpu_comm_done_event, gpu_compute_done_event,
5
                 schedule="default", token_offset=0):
6
        ...
7
        self.current_stream = compute_stream
8
        self.recv_hook = None

四类同步原语，正好对应 CPU 与 GPU 两个世界：

注意 comm_stream 由 wrapper 在构造时创建一次、两个微批共用；compute_stream 是主线程当前的计算流，两个微批也共用。两个世界交错的那条线，就在于「哪个微批当前持有 current_stream」。

进入与退出

1
def __enter__(self):
2
    _THREAD_ID_TO_CONTEXT[threading.get_ident()] = self.id
3
    _CURRENT_CONTEXTS[self.id] = self
4
    self.ready_barrier.wait()          # 所有线程就位
5
    self.cpu_wait_event.wait()         # 等我被唤醒
6
    self.cpu_wait_event.clear()
7
    self._restore_context()            # 装上我自己的 ForwardContext
8
    self.update_stream(self.compute_stream)
9
    return self

ForwardContext 是 vLLM 的全局上下文（注意力元数据、slot mapping、DP 元数据都在里面）。两个线程共享一个 Python 进程，这个全局变量只能有一个值，所以谁在跑，谁就把自己的上下文装回去，跑完再交棒。

clear() 的时机也很讲究：等待用的 Event 在进入时清掉，所以在 yield 里「先 signal 对方、再 wait 自己」的顺序下，不会出现「对方还没睡就又被我唤醒」的丢失唤醒（lost wakeup）。

交接棒：_cpu_yield

1
def _cpu_yield(self):
2
    assert forward_context._forward_context == self.forward_context
3
    assert current_stream() == self.current_stream
4
    assert not self.cpu_wait_event.is_set()
5

6
    self.cpu_signal_event.set()        # 叫醒对方
7
    self.cpu_wait_event.wait()         # 自己睡下
8
    self.cpu_wait_event.clear()
9
    self._restore_context()

代码注释写得很直白：「正确性关键：同一时刻只能有一个线程在跑。」三句断言就是这条规则的守卫。

这里有一个值得停下来想的点：为什么使用 threading.Event 而不是把模型拆成独立进程？因为 ForwardContext、torch.cuda 的当前流、环境变量、CUDA Graph 的捕获都在同一个进程里；线程是最廉价的方式，可以共享这些状态，代价就是 GIL。

乒乓的完整形态

yield_() 保存当前流、叫醒对方、睡下、醒来后恢复流：

1
def yield_(self):
2
    self.current_stream = current_stream()
3
    self._cpu_yield()
4
    self.update_stream(self.current_stream)

而带流切换的版本才是 DBO 的关键——它让「计算 → 通信」的交接带着两件事：本微批的计算流记录一个完成事件，唤醒对方后切到通信流并等计算完成。

1
def yield_and_switch_from_compute_to_comm(self):
2
    assert current_stream() == self.compute_stream
3
    self._signal_compute_done()        # 在 compute_stream 上 record
4
    self._cpu_yield()                  # 叫醒另一个微批去算
5
    assert self.current_stream == self.compute_stream
6
    self.update_stream(self.comm_stream)
7
    self._wait_compute_done()          # comm_stream 等 compute_stream
8

9
def yield_and_switch_from_comm_to_compute(self):
10
    assert current_stream() == self.comm_stream
11
    self._signal_comm_done()
12
    self._cpu_yield()
13
    assert self.current_stream == self.comm_stream
14
    self.update_stream(self.compute_stream)
15
    self._wait_comm_done()

为什么 CPU 这边先交棒、再在 GPU 流上等待？因为 GPU 的命令是异步排进流的，CPU 不需要等 GPU。真正需要串行的是 CPU 上的发射顺序——哪些 kernel 先被排进哪条流。先叫醒对方，让对方在它自己的 compute_stream 上把计算 kernel 排进去，再回来在自己的流上等事件，就形成了 GPU 上「我的通信在跑、对方的计算在跑」的并发画面。

全局访问器与 recv hook

模型层的代码不持有 UBatchContext 引用，它通过一组全局函数访问「当前线程的上下文」：

1
_THREAD_ID_TO_CONTEXT: dict = {}
2

3
def _register_ubatch_function(func):
4
    def wrapper(*args, **kwargs):
5
        if len(_THREAD_ID_TO_CONTEXT) > 0:
6
            ctx = _CURRENT_CONTEXTS[_THREAD_ID_TO_CONTEXT[threading.get_ident()]]
7
            func(ctx, *args, **kwargs)
8
    return wrapper
9

10
dbo_yield = _register_ubatch_function(UBatchContext.yield_)
11
dbo_yield_and_switch_from_compute_to_comm = _register_ubatch_function(
12
    UBatchContext.yield_and_switch_from_compute_to_comm)

如果当前线程没有注册上下文（DBO 没开，或不在微批线程里），dbo_yield 就是一个空操作。这让模型层代码可以无条件地在关键位置写 dbo_yield()，而不需要到处加 if。

recv hook 是另一个关键的全局机制：

1
def dbo_register_recv_hook(recv_hook):
2
    ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()]
3
    next_ctx = _CURRENT_CONTEXTS[(ctx_idx + 1) % _NUM_UBATCHES]
4
    next_ctx.recv_hook = recv_hook
5

6
def maybe_run_recv_hook(self):
7
    if self.recv_hook is not None:
8
        self.recv_hook()
9
        self.recv_hook = None

当前线程把「接收我的 dispatch 结果」这个动作挂到对方的上下文上。对方在某处调用 dbo_maybe_run_recv_hook() 时，顺手替我执行这个等待。换句话说，hook 让「别人做别的事的时候，帮我把我自己的通信收掉」。

第 1 段：dispatch 前的交接

1
previous_event = dbo_get_previous_event(self.buffer.capture)
2

3
# 我们在这之前 yield：dispatch kernel 会阻塞 CPU，先把另一个微批的计算排进队列
4
dbo_yield_and_switch_from_compute_to_comm()
5
_dbo_record_stream(tokens, token_scales, rank_topk_ids, rank_topk_weights)
6

7
(..., handle, event) = self.buffer.dispatch(
8
    x=token_data, ...,
9
    async_finish=self.async_prepare and not dbo_enabled(),
10
    allocate_on_comm_stream=False,
11
)
12

13
a2a_idx = dbo_current_ubatch_id()
14
self.handles[a2a_idx] = handle
15

16
dbo_switch_to_compute_sync()
17
return lambda: self._receiver(event, ...)

三件事需要讲清：

previous_event = dbo_get_previous_event(self.buffer.capture) 在 yield 之前取。注释解释了为什么：DeepEP 需要一个「开始之前已经完成」的事件来做 stream 内部的依赖排序；如果放在 yield 之后取，CPU 上已经被对方插入了计算命令，这个事件会把对方的计算尾巴也算进来，导致重叠失效。

async_finish=self.async_prepare and not dbo_enabled()——DBO 开启时，dispatch 不走 DeepEP 的异步完成模式，而是走 DBO 自己的 hook 机制。两套异步不能叠加，否则「谁等谁、在哪条流上等」会乱。

self.handles[a2a_idx] = handle——DeepEP 的 dispatch 返回一个 handle，combine 时按它匹配。两个微批并发执行意味着会有两个 handle 同时在飞，所以要按微批 id 分开存。这是典型「单线程假设被打破后必须显式记账」的改动。

dbo_switch_to_compute_sync() 在 dispatch kernel 入队之后立刻切回计算流并等通信完成，这样后续在计算流上的量化、专家 GEMM 都正确地排在 dispatch 之后。

_dbo_record_stream：缓存分配器的跨流生命周期

这是最容易出错、也最容易在生产中被忽略的类型：PyTorch 的缓存分配器按流回收内存块。一个张量在通信流上分配（DeepEP 的输出），却在计算流上被读；分配器只记录「这个块属于通信流」，一旦通信流上的引用释放，就可能被分给下一个需求。下一个需求恰好是另一个微批在计算流上的一次分配——于是它在用一块被别的流读着的内存。

1
def _dbo_record_stream(*tensors):
2
    """Mark tensors allocated on the other DBO stream as used by this one."""
3
    if not dbo_enabled():
4
        return
5
    stream = current_stream()
6
    for t in tensors:
7
        if t is not None:
8
            t.record_stream(stream)

record_stream 告诉分配器「这块内存在这条流上也用过了，等这条流到了某个同步点再回收」。它的开销很小，但省掉它就会出现随机的、难以复现的结果错乱。

第 2 段：接收结果与挂 recv hook

dispatch 返回的 lambda 就是 receiver：它等 DeepEP 的事件、处理 topk ids 的 -1、按需在 dispatch 后量化。这个 receiver 在多线程路径下被挂到对面：

1
hook, receiver = (prepare_ret if isinstance(prepare_ret, tuple) else (None, prepare_ret))
2
if hook is not None:
3
    if dbo_enabled():
4
        dbo_register_recv_hook(hook)   # 挂到对方上下文
5
        dbo_yield()                    # 让对方先跑
6
    else:
7
        hook()

注意这里同时挂 hook 又 yield——挂完就睡，让对方在它的路径上 dbo_maybe_run_recv_hook() 时把我的等待执行掉。

DeepEP 低延迟后端用的是另一套：low_latency_combine 返回 recv_hook，由 dbo_maybe_run_recv_hook() 在同一微批路径上执行。也就是说，超低延迟模式下 vLLM 利用的是 DeepEP 内核自己已经分好的接收阶段，DBO 只负责把它们排进合适的时间点。

第 3 段：combine 前再交一次棒

1
previous_event = dbo_get_previous_event(self.buffer.capture)
2
dbo_yield_and_switch_from_compute_to_comm()
3
_dbo_record_stream(fused_expert_output)
4

5
combined_x, _, event = self.buffer.combine(
6
    x=fused_expert_output, handle=handle, topk_weights=None,
7
    previous_event=previous_event,
8
    async_finish=do_async and not dbo_enabled(),
9
)
10
dbo_switch_to_compute()
11

12
if do_async:
13
    def _receiver():
14
        if event.event is not None:
15
            event.current_stream_wait()
16
        dbo_switch_to_comm()
17
        output.copy_(combined_x, non_blocking=True)
18
        dbo_yield_and_switch_from_comm_to_compute()
19
    return _receiver

结构与第 1 段对称：交棒 → 发 combine → 切回计算流。最后那一段 receiver 里，拷贝回 output 发生在通信流上，拷贝完再交棒回计算流——这样「写回模型的输出张量」这件事留在通信流上，不占计算流的关键路径。

共享专家：在被让出的空档里干活

DeepSeek 的 MoE 层有共享专家，它的输入不经过 all-to-all，是纯本地计算。DBO 把它的执行安排在「通信正在飞」的那段时间里：

1
def _maybe_apply_shared_experts(self, shared_experts, shared_experts_input):
2
    if shared_experts is not None:
3
        assert self.prepare_finalize.supports_async()
4
        shared_experts(shared_experts_input, SharedExpertsOrder.MK_INTERNAL_OVERLAPPED)

在 _finalize 中，finalize_async() 返回后先立即调用共享专家，再处理 combine 的 hook/receiver——共享专家的 MLP 计算和 combine 通信在时间上重叠了。

__call__ 的分支

1
def __call__(self, *args, **kwargs):
2
    forward_context = get_forward_context()
3
    ubatch_slices = forward_context.ubatch_slices
4
    cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode
5

6
    if ubatch_slices is None:
7
        if cudagraph_runtime_mode is CUDAGraphMode.FULL:
8
            if forward_context.batch_descriptor.num_tokens in self.cudagraphs:
9
                cudagraph_runtime_mode = CUDAGraphMode.NONE
10
        if cudagraph_runtime_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE):
11
            return self.runnable(*args, **kwargs)
12
        return self.cudagraph_wrapper(*args, **kwargs)
13
    ...

这个 ubatch_slices is None 分支里的第一段逻辑非常细小但很重要：如果某个形状已经捕获过 ubatched 的图，那么同一形状 非 ubatched 的运行不能再走图捕获——因为 FULL 图按形状只保存一份，ubatched 版占了那个位置。此时强行 cudagraph_wrapper() 会尝试在正常运行时捕获一张新图，这是不能接受的（运行时捕获意味着停下来等编译）。处理方式是退回 eager。

每微批的输入切片与 DP 元数据

1
compute_stream = torch.cuda.current_stream()
2

3
ubatch_dp_metadata = []
4
for ubatch_slice in ubatch_slices:
5
    ubatch_num_tokens_across_dp = torch.tensor(
6
        [ubatch_slice.num_tokens] * dp_size, device="cpu", dtype=torch.int32)
7
    ubatch_dp_metadata.append(DPMetadata.make(
8
        self.vllm_config.parallel_config, ubatch_slice.num_tokens,
9
        ubatch_num_tokens_across_dp))

DPMetadata 记录的是「每个 DP rank 在这一步有多少 token」，TP 层的 all-reduce 融合与 DP 的负载均衡都读它。切批之后，每个微批要有自己的一份——因为在微批的视角里，DP 组内的 token 数就是「我这个微批的 token 数」，假设所有 rank 同样切分（前面已经 pad 到一致了）。

_slice_model_inputs() 再按 token_slice 切 input_ids、positions、inputs_embeds、intermediate_tensors。positions 有一个细节：MRoPE（多模态旋转位置编码）的 positions 是二维 [3, num_tokens]，第二种情况按列切，否则按第一维切。

跑微批

1
@torch.inference_mode()
2
def _ubatch_thread(results, model, ubatch_metadata):
3
    torch.accelerator.set_device_index(self.device)
4
    with ubatch_metadata.context:
5
        model_output = model(input_ids=..., positions=..., ...)
6
    results.append((ubatch_metadata.context.id, model_output))
7

8
with override_forward_context(None):
9
    threads = [threading.Thread(target=_ubatch_thread, args=(results, model, m))
10
               for m in ubatch_metadata]
11
    for t in threads: t.start()
12
    self.ready_barrier.wait()               # 等所有线程进入上下文
13
    ubatch_metadata[0].context.cpu_wait_event.set()   # 唤醒第一个
14
    for t in threads: t.join()

override_forward_context(None) 把主线程的全局上下文清掉——因为马上要由各微批线程各自装载。ready_barrier 在 __enter__ 里被所有线程等一次，确保两个线程都完成了上下文装载与流设置，主线程才开始「发第一棒」。主线程只会发第一棒；之后的每一棒都靠微批线程在 yield 里互相传递。

所有线程结束后，torch.cat(sorted_results, dim=0) 按微批 id 排序拼回完整输出；如果模型返回的是元组（例如 EAGLE-3 草稿需要收集 aux hidden states），会逐分量拼接，保持结构一致。

线程安全与顺序：为什么主线程要 join

看上去主线程只发一棒、然后 join，什么都没「管」。但它必须等所有线程结束才能返回输出，因为模型输出是按微批 id 拼起来的，而微批的执行顺序是动态的——谁先结束不一定。sorted(results) 按 id 排序正好解决了这一点。

捕获流程

_capture_ubatches() 的注释已经把四步说得很清楚，这里对着代码：

1
@torch.inference_mode()
2
def _capture_ubatch_thread(results, ubatch_metadata):
3
    torch.accelerator.set_device_index(self.device)
4
    with torch.cuda.stream(ctx.compute_stream):
5
        _ = torch.cuda.current_blas_handle()
6
    with torch.cuda.stream(ctx.comm_stream):
7
        _ = torch.cuda.current_blas_handle()
8
    with ctx:                                   # 进上下文、等障碍、被唤醒
9
        model_output = model(...)
10
    results.append((ctx.id, model_output))

第一步：每个线程先在自己两条流上都初始化 CUDA 上下文。current_blas_handle() 会触发 cuBLAS 的惰性初始化。这必须在捕获之前、在正确的设备上下文下完成——否则捕获图时第一个线程触发初始化会在图中插入初始化调用，或者更糟，在错误的流上初始化。

第二步与第三步在 torch.cuda.graph(...) 里完成：

1
with torch.cuda.graph(cudagraph.cudagraph, stream=compute_stream, pool=self.graph_pool):
2
    ubatch_metadata[0].context.cpu_wait_event.set()   # 发第一棒
3
    for thread in ubatch_threads:
4
        thread.join()
5
    sorted_results = [value for position, value in sorted(results)]
6
    result = _cat_ubatch_outputs(sorted_results)
7
    cudagraph_metadata.outputs = result
8
    get_offloader().join_after_forward()

第四步：把图与它的元数据存进 self.cudagraphs[num_tokens]。

捕获期间，两个线程照常执行 yield 乒乓——只是这时 GPU 上不实际执行 kernel，而是在记录。关键前提是乒乓的 CPU 交接不依赖 GPU 的返回值：yield 只操作 threading.Event 和 torch.cuda.Event 的 record/wait，不读回 GPU 数据，不做 item() 或 synchronize()。如果模型代码在中间做了任何 host 同步，捕获就会卡死或行为错误。这也解释了为什么 _register_ubatch_function 里的 dbo_yield 可以在图里「被记录为」空——它们本来就不发射任何 GPU 命令，图里只记录流上的内核和事件。

stream=compute_stream 指定了捕获的主流，pool=self.graph_pool 让 DBO 的图与常规图共享内存池，减少显存碎片。共享专家与 pre-capture 预取相关的 offloader 在捕获前后分别 sync_prev_onload() 与 join_after_forward()，防止拷贝流上的工作跨进图边界，造成未 join 的流错误。

回放

1
elif num_tokens in self.cudagraphs and cudagraph_runtime_mode is CUDAGraphMode.FULL:
2
    cudagraph_metadata = self.cudagraphs[num_tokens]
3
    get_offloader().sync_prev_onload()
4
    cudagraph_metadata.cudagraph.replay()
5
    return cudagraph_metadata.outputs

回放本身不需要任何多线程与 CPU 同步。图中已经录下了两条流上的全部内核与事件等待顺序——CUDA 图回放时会按捕获时的流拓扑重建依赖。这就是设计文档那句「一旦捕获过，就可以无需多线程与 CPU 同步地回放」。

只捕获 FULL，且只在均匀解码时捕获

1
allow_microbatching = (
2
    self.parallel_config.use_ubatching
3
    and cudagraph_runtime_mode == CUDAGraphMode.FULL
4
    and uniform_decode
5
    and check_ubatch_thresholds(config=..., num_tokens=batch_desc.num_tokens,
6
                                uniform_decode=uniform_decode)
7
)

三个限定：

1. 只有 FULL 图支持；PIECEWISE 图（把图切段、段间留出 CPU 可介入的边界）与多线程乒乓不兼容。
2. 只有均匀解码（每个请求恰好排一个 token）才捕获微批图。decode 是 DBO 收益最直接的场景：每步 token 少、通信占比高。
3. 超过阈值才捕获——阈值太低的形状图数量爆炸、收益又不明显。

其余形状仍然捕获非微批版本；到了运行时若形状对不上，UBatchWrapper 里那段 num_tokens not in self.cudagraphs 分支会退到线程执行（eager 或 piecewise）。这也是 vLLM 博客里所说「DBO 支持 prefill，但 prefill 侧主要靠 eager 线程路径」的原因。

为什么 DBO 只支持 FULL 图

如果只从实现细节去找，会看到「DBO 的图必须在 UBatchWrapper 里整体捕获」。从系统层面，原因更直接：PIECEWISE 图的核心价值是在图段之间让出 CPU，让调度器做下一批准备；而 DBO 本来就要在 CPU 上乒乓。两者叠加会让「谁在图内、谁在图外」变得难以推理；而且 PIECEWISE 的段边界是按单个线程设计的，无法容纳两个线程的交错。想更完整地理解两种图模式，可参考 vLLM 源码完全拆解（十二）：piecewise compilation 与 CUDA Graph——消灭 kernel 调度开销的两层机制 以及专门讨论跨框架桶与 pad 的 CUDA Graph 的桶与垫：五个推理框架如何把动态解码批塞进静态图。

1. 事件必须与流绑定

UBatchContext._signal_compute_done() 在 compute_stream 上 record 事件，_wait_compute_done() 让 comm_stream 等它。如果误把两个微批的同一类事件混在一起，就会出现「我在等对方的事件、对方在等我的事件」的活锁。代码里所有事件都是「本微批的通信完成/计算完成」，配对关系从 make_ubatch_contexts 里初始化时就固定了。

2. ROCm 下通信工作区必须同步

DeepEP HT 在 ROCm 上会复用 Buffer 内部的通信工作区：

1
def _sync_dbo_comm_if_needed(self):
2
    if self.sync_dbo_comm and dbo_enabled():
3
        torch.cuda.current_stream().synchronize()

如果不让上一个微批的 HT 内核结束，下一个微批就会复用一份还在被读的工作区。代价是整流的同步，但在 ROCm 上这是换取正确性的必要成本。

3. 模型输出的拼接顺序

_cat_ubatch_outputs() 以「按微批 id 排序后的结果」拼接，而不是按线程结束先后顺序。如果不排序，输出行序就会随机乱掉——采样拿到的 logits 行会对应错误的请求，结果表现为「答案串台」这种极其隐蔽的错误。

4. 当 ubatching 被否决时，避免捕获新图

abort/mis-vote 的路径（ubatch_slices is None 但形状已在 self.cudagraphs 里）必须手动把 runtime mode 降为 NONE。否则会进 cudagraph_wrapper，在运行时触发捕获——一个本该是 eager 的执行变成了一次缓慢的图捕获，甚至打断服务。

5. 只让「纯解码」走图捕获

前面提到 uniform_decode 的限制。混合批（含 prefill）走线程路径，positions 的二维处理、is_padding 的切片都在这条路径上验证过；而 FULL 图捕获只要均匀解码，形状单一、路径简单。