执行摘要

• 一句话：修复 overlap 预构建行复用竞态，调度流等待前向流
• 推荐动作：值得精读。该 PR 修复的是调度核心路径上的竞态条件，涉及流同步和 FutureMap 的共享状态管理。理解其修复思路（通过显式流等待代替事件链）对处理类似并发问题有参考价值。此外，其测试编写方式（使用 mock 验证调用顺序）也值得学习。

功能与动机

PR body 明确指出这是一个 FutureMap 中的竞态条件 bug：publish() 可能在两种情况下被调用：(1) 前向流在 draft+verify 之后；(2) 调度流中，当 PD disagg 启用且 prefill 请求到达时。publish() 会修改 FutureMap 的共享状态，因此必须正确排序以避免竞态。publish_ready 事件不能用于同步，因为 stash() 在 record 之后也会修改 FutureMap 共享状态。Reviewer merrymercy 也确认这是实际存在的竞态：'an actual race condition when an already finished forward is still running, while PD-decode prebuilt requests come in and occupy the same lane.'

实现拆解

实现分三步：

1. 核心修复（decode.py）：在 get_new_prebuilt_batch 中，当 enable_overlap 为真时，在调用 new_batch.process_prebuilt(self.future_map) 之前，先调用 self.schedule_stream.wait_stream(self.forward_stream)，确保前向流中未完成的 publish() 操作已经完成，避免复用同一请求池行时出现数据竞争。

2. 简化 publish()（overlap_utils.py）：移除了 publish() 中对 publish_ready 事件进行链式等待的逻辑（device_module.current_stream().wait_event(self.publish_ready)），因为现在发布操作已经通过流同步保证顺序，不再需要此链式等待，避免了潜在的事件记录顺序问题。

3. 测试配套（test_priority_scheduling_disaggregation.py）：重构并新增测试。将原有 TestDecodePrebuiltPriority 测试类重构为 TestDecodePrebuilt，提取 _new_scheduler 辅助方法，并新增 test_overlap_waits_for_forward_before_processing_prebuilt 测试，验证 prepare_for_prebuilt、wait_stream、process_prebuilt 的执行顺序。

关键文件：

• python/sglang/srt/disaggregation/decode.py（模块 调度器；类别 source；类型 core-logic；符号 get_new_prebuilt_batch, process_decode_queue）: 核心修复位置：在 get_new_prebuilt_batch 中，当 enable_overlap 时，在 process_prebuilt 前添加 wait_stream，防止竞态。
• python/sglang/srt/managers/overlap_utils.py（模块 重叠调度；类别 source；类型 core-logic；符号 FutureMap.publish, FutureMap.stash, FutureMap.prepare）: 简化 publish() 的链式事件记录，避免潜在的问题。
• test/registered/unit/managers/test_priority_scheduling_disaggregation.py（模块 测试；类别 test；类型 test-coverage；符号 TestDecodePrebuilt, _new_scheduler, test_overlap_waits_for_forward_before_processing_prebuilt）: 新增测试覆盖竞态修复，验证执行顺序。

关键符号：get_new_prebuilt_batch, process_decode_queue, FutureMap.publish, FutureMap.stash, FutureMap.prepare

关键源码片段

python/sglang/srt/disaggregation/decode.py

核心修复位置：在 get_new_prebuilt_batch 中，当 enable_overlap 时，在 process_prebuilt 前添加 wait_stream，防止竞态。

# python/sglang/srt/disaggregation/decode.py
# 核心修复：在 prebuilt 批处理前等待前向流，防止竞态
def get_new_prebuilt_batch(self, running_batch):
    # ... 其他逻辑
    # 构造 fake completed prefill
    new_batch.prepare_for_prebuilt()
    if self.enable_overlap:
        # 已完成的请求可能仍有一个冗余前向在飞行中。
        # 在 prebuilt 请求复用可能被复用的请求池行之前，先 drain 它。
        self.schedule_stream.wait_stream(self.forward_stream)
    new_batch.process_prebuilt(self.future_map)
    return new_batch

python/sglang/srt/managers/overlap_utils.py

简化 publish() 的链式事件记录，避免潜在的问题。

# python/sglang/srt/managers/overlap_utils.py
# publish() 简化：移除链式事件记录，因为已通过流同步保证顺序
def publish(
    self,
    future_indices: torch.Tensor,
    new_seq_lens: torch.Tensor,
    confidence: Optional[torch.Tensor] = None,
) -> None:
    indices = future_indices
    if indices.shape[0] == 0:
        return # DP idle
    self.new_seq_lens_buf[indices] = new_seq_lens.to(self.new_seq_lens_buf.dtype)
    publish_confidence = self.needs_confidence_relay and confidence is not None
    if publish_confidence:
        self.confidence_relay.scatter(indices, confidence)
    # 只有 spec_v2 需要事件；它门控 seq_lens D2H 在私有流上。
    if self.spec_algo.is_some():
        if self.publish_ready is None:
            self.publish_ready = torch.get_device_module(self.device).Event()
        self.publish_ready.record()
        self._publish_fresh = True
    if publish_confidence:
        self.confidence_relay.issue_ring_copy(
            stream=self.fwd_prepare_d2h_stream,
            publish_ready=self.publish_ready,
        )

test/registered/unit/managers/test_priority_scheduling_disaggregation.py

新增测试覆盖竞态修复，验证执行顺序。

# test/registered/unit/managers/test_priority_scheduling_disaggregation.py
# 新增测试：验证 overlap 时等待前向流
def test_overlap_waits_for_forward_before_processing_prebuilt(self):
    scheduler = self._new_scheduler(enable_overlap=True)
    scheduler.waiting_queue = [MagicMock(rid="request")]
​
    call_order = []
    new_batch = MagicMock()
    new_batch.prepare_for_prebuilt.side_effect = lambda: call_order.append("prepare")
    scheduler.schedule_stream.wait_stream.side_effect = lambda _: call_order.append("wait")
    new_batch.process_prebuilt.side_effect = lambda *_: call_order.append("process")
​
    with patch(
        "sglang.srt.disaggregation.decode.ScheduleBatch.init_new", return_value=new_batch,
    ), get_context().override_server_args(disaggregation_decode_enable_radix_cache=False):
        ret = SchedulerDisaggregationDecodeMixin.get_new_prebuilt_batch(
            scheduler, scheduler.running_batch
        )
​
    self.assertIs(ret, new_batch)
    scheduler.schedule_stream.wait_stream.assert_called_once_with(scheduler.forward_stream)
    self.assertEqual(call_order, ["prepare", "wait", "process"])

评论区精华

审查者 hnyls2002 指出：LGTM, thanks for the fix. my previous PR #30435 was trying to fix the similar PR, but I didn't consider the stashing data races between prebuilt batch and normal batches. 这表明此修复解决了之前未覆盖的 stash 数据竞争问题，与历史 PR #30435 相关。merrymercy 确认这是实际竞态，Jialin 的评论强调了竞态发生的场景。

• 竞态条件的确认 (correctness): 确认修复必要。
• 与历史修复的关联 (design): 本修复补充了 stash 竞态修复。

风险与影响

• 风险：风险较低，但需注意：
  • 仅影响启用 overlap 的 PD 分离场景：修复在 enable_overlap 条件下新增了流同步，可能带来轻微的性能开销，但这是必要开销，且只在 PD 分离时发生。
  • 依赖流等待的正确性：wait_stream 的调用假定 schedule_stream 和 forward_stream 均存在且有效，如果这些属性未正确初始化，可能会引发异常。但测试中已覆盖。
  • 移除链式记录逻辑：若未来有其他路径在前向流之外调用 publish()，且不经过 get_new_prebuilt_batch 的同步，可能会导致事件顺序问题回归。需要确保所有调用路径都受保护。
  • 影响：影响范围限于启用 overlap 的 PD 分离部署（enable_overlap 且 disaggregation_decode_enable_radix_cache 等配置）。修复后，请求池行的复用不会出现竞态，提高了稳定性和正确性。对性能影响极小，仅增加一次流等待调用。测试覆盖了优先级调度和 overlap 场景，降低了回归风险。
  • 风险标记：核心路径变更, 新增测试覆盖

关联脉络

• PR #30435 Unknown: 评审者提及之前的修复尝试，与本 PR 相关。