https://beaiera.top/tags/%E6%8E%A8%E7%90%86%E5%BC%95%E6%93%8E/Recent content in 推理引擎 on Richelieu's BlogHugozh-cnMon, 01 Jun 2026 11:10:00 +0800https://beaiera.top/posts/vllm-ubatch-barrier/Mon, 01 Jun 2026 11:10:00 +0800https://beaiera.top/posts/vllm-ubatch-barrier/<h2 id="背景什么是-ubatchwrapper">背景：什么是 UBatchWrapper</h2> <p>在 vLLM 的 GPU 推理引擎中，<code>UBatchWrapper</code>（位于 <code>vllm/v1/worker/gpu_ubatch_wrapper.py</code>）是一个<strong>模型包装器</strong>。它拦截对原始模型的调用，在内部将一次大 batch 的 forward 拆分成多个 <strong>micro-batch（ubatch）</strong>，并利用多线程 + CUDA stream 实现计算与通信的重叠。</p> <p>核心目标：<strong>解决 MoE 模型中 all2all 通信开销大的问题</strong>——让一个 ubatch 在做计算时，另一个 ubatch 的通信已经在背后并行执行。</p> <hr> <h2 id="线程模型与同步原语">线程模型与同步原语</h2> <h3 id="threadingbarrier集结号">threading.Barrier：集结号</h3> <p><code>UBatchWrapper</code> 在初始化时创建<strong>一个</strong> <code>threading.Barrier</code>：</p> <div class="highlight"><pre tabindex="0" class="chroma"><code class="language-python" data-lang="python"><span class="line"><span class="cl"><span class="bp">self</span><span class="o">.</span><span class="n">ready_barrier</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Barrier</span><span class="p">(</span><span class="n">num_ubatches</span> <span class="o">+</span> <span class="mi">1</span><span class="p">)</span> </span></span></code></pre></div><p><strong><code>threading.Barrier</code> 不关心具体有哪些线程，只计数</strong>。内部维护一个计数器，每当有线程调用 <code>barrier.wait()</code>，计数器就 +1。当总数达到构造时指定的 <code>parties</code> 值（这里是 <code>num_ubatches + 1</code>），所有正在 <code>wait()</code> 的线程<strong>同时被释放</strong>，计数器归零。</p> <p>这里的 <code>+1</code> 包含了 <strong>N 个 ubatch 线程 + 1 个主线程</strong>。</p> <p><code>ready_barrier.wait()</code> 在三处被调用：</p> <table> <thead> <tr> <th>文件</th> <th>行号</th> <th>调用方</th> </tr> </thead> <tbody> <tr> <td><code>gpu_ubatch_wrapper.py</code></td> <td>261</td> <td><strong>主线程</strong> 在 <code>_capture_ubatches</code> 中</td> </tr> <tr> <td><code>gpu_ubatch_wrapper.py</code></td> <td>325</td> <td><strong>主线程</strong> 在 <code>_run_ubatches</code> 中</td> </tr> <tr> <td><code>ubatching.py</code></td> <td>56</td> <td><strong>每个 ubatch 线程</strong> 在 <code>UBatchContext.__enter__</code> 中</td> </tr> </tbody> </table> <p><strong>少一个线程到达，所有人全部卡住。</strong></p>