vLLM GPU UBatchWrapper 中的 Barrier 机制解析
背景：什么是 UBatchWrapper 在 vLLM 的 GPU 推理引擎中，UBatchWrapper（位于 vllm/v1/worker/gpu_ubatch_wrapper.py）是一个模型包装器。它拦截对原始模型的调用，在内部将一次大 batch 的 forward 拆分成多个 micro-batch（ubatch），并利用多线程 + CUDA stream 实现计算与通信的重叠。 核心目标：解决 MoE 模型中 all2all 通信开销大的问题——让一个 ubatch 在做计算时，另一个 ubatch 的通信已经在背后并行执行。 线程模型与同步原语 threading.Barrier：集结号 UBatchWrapper 在初始化时创建一个 threading.Barrier： self.ready_barrier = threading.Barrier(num_ubatches + 1) threading.Barrier 不关心具体有哪些线程，只计数。内部维护一个计数器，每当有线程调用 barrier.wait()，计数器就 +1。当总数达到构造时指定的 parties 值（这里是 num_ubatches + 1），所有正在 wait() 的线程同时被释放，计数器归零。 这里的 +1 包含了 N 个 ubatch 线程 + 1 个主线程。 ready_barrier.wait() 在三处被调用： 文件 行号 调用方 gpu_ubatch_wrapper.py 261 主线程 在 _capture_ubatches 中 gpu_ubatch_wrapper.py 325 主线程 在 _run_ubatches 中 ubatching.py 56 每个 ubatch 线程 在 UBatchContext.__enter__ 中 少一个线程到达，所有人全部卡住。 ...