执行摘要

• 一句话：为推测解码各阶段添加 NVTX profiler span
• 推荐动作：建议合并。该 PR 规模小、逻辑清晰，属于低风险的可观测性改进。虽无专项测试，但工具函数设计合理（profiler 开启时生效，关闭时无开销），可直接合并。

功能与动机

为推测解码 pipeline 增加粗粒度的 profiling 可见性，使开发者能通过 nsys 等工具直观区分 draft / draft_extend / verify 各阶段的耗时，无需逐段手动打点。PR body 虽未详细展开，但功能性质明显。

实现拆解

1. 定义 spec_stage_span 工具函数（spec_utils.py）：新增函数，检查 torch.autograd._profiler_enabled() 后返回 torch.profiler.record_function(name) 或 nullcontext()，避免 profiler 关闭时引入额外开销。
2. 修改 frozen_kv_mtp_worker.py：在 forward_batch_generation 中的 draft_extend 和 draft 阶段入口处，将 spec_stage_span("draft_extend") 或 spec_stage_span("draft") 加入 with 上下文链。特别地，在 forward_draft_extend_after_decode 调用外也包裹了 draft_extend span，确保后续 decode 循环也被标记。
3. 修改 eagle_worker_v2.py：在 forward_batch_generation 中的 draft_extend（prefill 与 decode 两个分支）和 draft 阶段，同样插入 spec_stage_span("draft_extend") 和 spec_stage_span("draft")，保持与 frozen-kv 实现一致的 profiler 覆盖。
4. 清理无关变更（第二次提交）：移除了一个无调用者的大页 top-k 辅助函数，确保 PR 只包含 NVTX span 相关改动。

关键文件：

• python/sglang/srt/speculative/spec_utils.py（模块 推测解码；类别 source；类型 core-logic；符号 spec_stage_span）: 新增 spec_stage_span 函数，是本次注入 profiler span 的底层工具。
• python/sglang/srt/speculative/frozen_kv_mtp_worker.py（模块 推测解码；类别 source；类型 core-logic；符号 forward_batch_generation）: 在 forward_batch_generation 的 draft_extend 和 draft 阶段插入 span，并包裹 forward_draft_extend_after_decode 调用。
• python/sglang/srt/speculative/eagle_worker_v2.py（模块 推测解码；类别 source；类型 core-logic；符号 forward_batch_generation）: 在 Eagle v2 worker 的 draft_extend（prefill 和 decode 分支）和 draft 阶段插入 span，保持 profiler 覆盖一致。

关键符号：spec_stage_span

关键源码片段

python/sglang/srt/speculative/spec_utils.py

新增 spec_stage_span 函数，是本次注入 profiler span 的底层工具。

# python/sglang/srt/speculative/spec_utils.py
​
def spec_stage_span(name: str):
    """Profiler span for a coarse speculative-decoding stage (``draft`` /
    ``draft_extend`` / ``verify``).
    """
    # 仅在 torch profiler 启用时记录，否则返回无操作上下文以避免开销
    if torch.autograd._profiler_enabled():
        return torch.profiler.record_function(name)
    return nullcontext()

python/sglang/srt/speculative/frozen_kv_mtp_worker.py

在 forward_batch_generation 的 draft_extend 和 draft 阶段插入 span，并包裹 forward_draft_extend_after_decode 调用。

# python/sglang/srt/speculative/frozen_kv_mtp_worker.py
# 在 forward_batch_generation 方法中：
# 原有上下文管理器链中追加 spec_stage_span("draft_extend") 或 spec_stage_span("draft")
​
with (
    self.draft_tp_context(self.draft_model_runner.tp_group),
    speculative_moe_backend_context(),
    speculative_moe_a2a_backend_context(),
    spec_stage_span("draft_extend"), # <-- 新增：标记 draft_extend 阶段
):
    self.forward_draft_extend(...)

python/sglang/srt/speculative/eagle_worker_v2.py

在 Eagle v2 worker 的 draft_extend（prefill 和 decode 分支）和 draft 阶段插入 span，保持 profiler 覆盖一致。

# python/sglang/srt/speculative/eagle_worker_v2.py
# 在 forward_batch_generation 方法中：
# prefill 分支的 draft_extend 阶段
with (
    self.draft_worker.draft_tp_context(self.draft_worker.draft_runner.tp_group),
    speculative_moe_backend_context(),
    speculative_moe_a2a_backend_context(),
    spec_stage_span("draft_extend"), # <-- 新增：标记 draft_extend 阶段
):
    batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill(...)
​
# decode 分支的 draft 阶段
with (
    self.draft_worker.draft_tp_context(self.draft_worker.draft_runner.tp_group),
    speculative_moe_backend_context(),
    speculative_moe_a2a_backend_context(),
    spec_stage_span("draft"), # <-- 新增：标记 draft 阶段
):
    verify_input = self.draft_worker.draft(batch)

评论区精华

本 PR 无 review 评论。主要讨论来自提交历史：合并者在第三次提交中“drop verbose comments”精简了注释。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低：
  • spec_stage_span 在 profiler 关闭时退化为 nullcontext()，对性能无影响。
  • 仅在两个 speculative worker 的已有 with 上下文中追加 span，不会改变执行逻辑。
  • 无测试覆盖（profiler 行为难以在 CI 中稳定验证），但该功能属开发辅助性质，对生产无影响。
  • 影响：影响范围限定在 profiling 场景：开发者使用 nsys 或 torch.profiler 时可获得更清晰的 speculative 阶段边界。对非 profiling 运行无任何行为变化。
  • 风险标记：无专项测试, 仅 profiling 辅助

关联脉络

• 暂无明显关联 PR