执行摘要

• 一句话：PD 分离下乐观预填充：提前开始 prefill 计算以降低 TTFT
• 推荐动作：该 PR 设计在性能与复杂性之间做了良好权衡（重试预算、fallback）。建议仔细审阅 pop_bootstrapped 中的循环控制与 metadata buffer 生命周期逻辑。对于使用 PD 的团队，建议在预发环境开启 --optimistic-prefill-retries=1 并监控 retry 指标，确认收益后逐步上调。同时关注后续与 HiCache 等组件的集成计划。

功能与动机

PD 分离中 prefill 必须等待 bootstrap（prefill 与 decode scheduler 到达时间偏差、decode 端分配延迟、bootstrap server/握手开销）完成后才能计算，直接增加 TTFT。乐观预填充旨在减少这一等待，降低首 token 延迟。

实现拆解

1. 配置入口与校验：在 server_args.py 添加 --optimistic-prefill-retries 参数，并在 _handle_other_validations 中校验：当 pipeline parallel size > 1、启用 hierarchical cache 或模型使用 mamba radix cache 时自动禁用并输出警告。
2. Bootstrap 队列扩展：在 prefill.py 中将 PrefillBootstrapQueue 的 add 方法拆解为 create_sender（创建 KV sender 并标记 pending_bootstrap）和 finalize_bootstrap（分配 metadata buffer、设置 bootstrap 完成时间、准备发送）。新增 ensure_metadata_buffer 确保索引可用。
3. 乐观调度循环：在 pop_bootstrapped 中，对 Bootstrapping 状态的请求不再等待，直接允许其进入 waiting 队列；在 get_num_optimistic_prefill_retries 中控制重试预算；在 scheduler 的 process_batch_result 和 handle_bootstrap_failure 中处理重试、回退与失败清理。
4. 请求时间统计更新：在 req_time_stats.py 新增 reset_prefill_retry_time 重置重试相关时间戳；新增 prefill_retry_count 指标；移除 alloc_wait_duration，调整 bootstrap_duration 语义。
5. 缓存命中率修正：在 metrics_reporter.py 中 cache hit 统计时减去 reprocessed_log_input_tokens 和 reprocessed_log_hit_tokens，避免重试导致虚高。
6. 失败注入与测试：在 disaggregation/utils.py 重构为 _poll_with_failure_injection 支持运行时失败概率注入；新增集成测试文件覆盖 gsm8k 准确率、logprob、故障注入等场景。

关键文件：

• python/sglang/srt/disaggregation/prefill.py（模块 预填充调度；类别 source；类型 core-logic；符号 should_force_retry, maybe_release_metadata_buffer, create_sender, ensure_metadata_buffer）: 核心实现文件，包含乐观预填充的全部流程：bootstrap 队列扩展、请求重试、metadata 缓冲区管理、失败处理。
• test/registered/disaggregation/test_disaggregation_optimistic_prefill.py（模块 乐观预填充测试；类别 test；类型 test-coverage；符号 rid_that_forces_retry, OptimisticPrefillRetryCounterMixin, _get_retry_counter, assert_retry_counter_increases）: 新增集成测试，覆盖乐观重试场景下的 gsm8k 准确率、logprob 正确性、故障注入等，确保功能正常且调度器健壮。
• python/sglang/srt/disaggregation/utils.py（模块 工具函数；类别 source；类型 dependency-wiring；符号 _get_failure_prob, _poll_with_failure_injection, is_aborted）: 重构失败注入机制，新增 _poll_with_failure_injection 和 is_aborted，支持运行时动态失败概率，用于测试和容错。
• python/sglang/srt/server_args.py（模块 服务配置；类别 source；类型 configuration）: 添加新参数和校验逻辑，确保乐观预填充仅在支持的配置下启用。
• python/sglang/srt/observability/req_time_stats.py（模块 请求时间统计；类别 source；类型 core-logic；符号 new_from_obj, reset_prefill_retry_time）: 新增 reset_prefill_retry_time 和更新时间统计语义，区分重试与正常时间。
• python/sglang/srt/managers/scheduler_components/metrics_reporter.py（模块 指标报告；类别 source；类型 core-logic）: 修正缓存命中率统计，减去 reprocessed 令牌避免虚高。
• python/sglang/srt/managers/schedule_policy.py（模块 调度策略；类别 source；类型 core-logic）: 调整调度策略以支持乐观重试请求的重新入队。
• python/sglang/srt/managers/scheduler.py（模块 调度器；类别 source；类型 core-logic）: 新增 handle_bootstrap_failure 等方法，串联重试与失败处理。
• python/sglang/srt/managers/schedule_batch.py（模块 批次管理；类别 source；类型 data-contract）: 添加少量属性支持重试标记。
• python/sglang/srt/environ.py（模块 环境变量；类别 source；类型 configuration）: 新增环境变量用于测试强制重试概率。
• test/registered/unit/managers/test_prefill_adder.py（模块 预填充加法器测试；类别 test；类型 test-coverage）: 兼容性修改，新增一行以适配新字段。

关键符号：should_force_retry, maybe_release_metadata_buffer, create_sender, ensure_metadata_buffer, finalize_bootstrap, advance_logprob_pt, reset_prefill_retry_time, _poll_with_failure_injection, is_aborted, get_num_optimistic_prefill_retries

关键源码片段

python/sglang/srt/disaggregation/prefill.py

核心实现文件，包含乐观预填充的全部流程：bootstrap 队列扩展、请求重试、metadata 缓冲区管理、失败处理。

# 文件 : python/sglang/srt/disaggregation/prefill.py
​
import hashlib
from sglang.srt.environ import envs
from sglang.srt.managers.schedule_batch import Req
​
​
def should_force_retry(req: Req) -> bool:
    """测试钩子：基于环境变量概率强制请求进入乐观重试路径。
    用于集成测试中模拟重试场景。
    """
    retry_prob = envs.SGLANG_TEST_FORCE_OPTIMISTIC_PREFILL_RETRY_PROB.get()
    # 仅针对尚未重试且未被撤回的请求采样
    if retry_prob <= 0 or req.time_stats.prefill_retry_count > 0 or req.is_retracted:
        return False
    digest = hashlib.sha256(str(req.rid).encode()).digest()
    return int.from_bytes(digest[:8], "big") < retry_prob * 2 ** 64
​
​
def maybe_release_metadata_buffer(
    req: Req, allocator: ReqToMetadataIdxAllocator
) -> None:
    """安全释放请求关联的 metadata buffer 索引。
    在重试、失败、完成等路径中确保索引被归还。
    """
    if req.metadata_buffer_index >= 0:
        allocator.free(req.metadata_buffer_index)
        req.metadata_buffer_index = -1
​
​
class PrefillBootstrapQueue:
    # ... 其他方法省略 ...
​
    def create_sender(self, req: Req, num_kv_heads: int) -> bool:
        """创建 KV sender 但不入队，返回 False 表示容量超限。"""
        if self._check_if_req_exceed_kv_capacity(req):
            return False
        backend = (
            TransferBackend.FAKE
            if self.server_args.disaggregation_transfer_backend == "fake"
            else TransferBackend.NCCL
        )
        req.disagg_kv_sender = self.kv_manager.create_sender(
            req,
            kv_class=self.kv_class,
            tp_rank=self.tp_rank,
            tp_size=self.tp_size,
            pp_rank=self.pp_rank,
        )
        self._process_req(req)
        req.pending_bootstrap = True
        return True
​
    def ensure_metadata_buffer(self, req: Req) -> bool:
        """确保请求已分配 metadata buffer 索引，若无法分配返回 False。"""
        if req.metadata_buffer_index >= 0:
            return True
        if self.req_to_metadata_buffer_idx_allocator.available_size() == 0:
            return False
        req.metadata_buffer_index = self.req_to_metadata_buffer_idx_allocator.alloc()
        assert req.metadata_buffer_index is not None
        return True
​
    def finalize_bootstrap(self, req: Req) -> bool:
        """bootstrap 完成后初始化 sender，分配 metadata buffer，设置完成时间。
        返回 False 表示 metadata buffer 分配失败（非终结性失败）。
        """
        assert req.pending_bootstrap, "finalize_bootstrap is not idempotent"
        if not self.ensure_metadata_buffer(req):
            return False
        req.time_stats.set_bootstrap_done_time()
        num_kv_indices = len(req.origin_input_ids)
        decode_prefix_len = req.disagg_kv_sender.pop_decode_prefix_len()
        req.start_send_idx = decode_prefix_len
        num_kv_indices_to_send = num_kv_indices - decode_prefix_len
        # ... 续接页面分配与发送准备
        return True

评论区精华

Review 中聚焦以下问题：

• 死锁风险（head-of-line blocking）（correctness）：gemini-code-assist[bot] 指出 pop_bootstrapped 中使用 break 导致 metadata buffer 耗尽时后续失败请求无法回退，可能死锁。作者接受并改为 continue。
• 测试成功计数器（testing）：ShangmingCai 询问 test_survive_requests 是否应增加成功断言。作者说明所有请求预期失败，不需要计数器；ShangmingCai 接受。
• KV cache 释放（correctness）：ShangmingCai 担心 requeue 前未释放 KV cache 导致泄漏。作者确认 finalize_bootstrap 失败时已通过 assertion 保障释放。
• 文档建议（documentation）：ShangmingCai 建议将新特性写入文档。作者同意计划后续默认启用时更新。

• HiCache 兼容性（design）：weibingo 认为 HiCache 前缀匹配在前可能不影响，作者回应 abort 和取消操作复杂，需后续重构。

• pop_bootstrapped 中 break 引起死锁风险 (correctness): 作者接受并修改为 continue，确保所有请求都能被处理。

• 测试中 success 计数器缺失 (testing): ShangmingCai 接受，不添加计数器。
• KV cache 在 requeue 前释放问题 (correctness): 问题澄清，已通过断言保障。
• 文档建议添加新特性说明 (documentation): 推迟到特性默认启用时更新文档。
• HiCache 兼容性讨论 (design): 当前暂不集成，留待后续重构。

风险与影响

• 风险：
  1. 重试风暴：高负载下大量请求重试可能加剧资源竞争；重试预算（最多 N 次）和 fallback 路径可缓解，需告警监控 num_prefill_retries 指标。
  2. Metadata buffer 耗尽：若分配失败则请求可能长时间等待或重复重试；代码已添加 fallback 回 bootstrap 队列，死锁风险通过 continue 修复。
  3. 兼容性限制：HiCache、Mamba、PP > 1 时自动禁用，但用户可能误配置导致静默降级；需文档明确说明。
  4. 性能影响：重试增加 scheduler 开销（释放/重新排队），但测试显示 TPOT/吞吐基本不变；极端场景下可能影响 ITL。
  5. 可观测性缺口：新增 retry 指标，但运维流程（告警阈值、面板）可能需要补充。
     - 影响：对用户端，启用后 TTFT 可降低 40%-50%，延迟敏感场景收益明显，默认向后兼容。对系统，PD 部署的 scheduler 逻辑复杂度增加，metadata buffer 占用小。对团队，核心调度、可观测性、配置模块均受影响，后续需规划与 HiCache/Mamba/PP 的集成重构。
     - 风险标记：死锁风险已修复, 与 HiCache 不兼容, 与 Mamba 不兼容, 与 PP>1 不兼容, metadata buffer 耗尽 fallback, 重试风暴需监控

关联脉络

• PR #28450 [AMD] Fuse shared-expert append + DeepEP remap into one Triton kernel: 同属 disaggregation 功能线（修改了相同的 moe/topk.py），但技术领域不同（MOE 内核融合 vs 预填充优化）。
• PR #28237 [AMD] fix(moe): correct fused shared-expert scaling on aiter/DeepEP path (mori all-to-all): 同为 disaggregation 下 moe 相关修复，与本 PR 无直接依赖，但共享部分上下文。