执行摘要

• 一句话：decode 阶段 H2D 拷贝改为异步
• 推荐动作：值得精读：这是一个小而精的优化案例，展示了理解 CUDA stream 语义和 overlap scheduler 后，如何通过一次 non-blocking 拷贝消除不必要的同步。对理解 SGLang 的 decode 调度路径也有帮助。

功能与动机

prepare_for_decode 中 latest_output_ids 的构造使用了阻塞式的 torch.tensor(..., device=...)，而同一路径下的其他 per-step H2D 拷贝均使用 non_blocking=True。在 overlap scheduler 下，这会导致每 decode 步都触发 cudaStreamSynchronize，使 CPU 被阻塞在正在执行的前向之后，在小 batch / 低延迟场景下尤为明显。对于不需要读取此张量的 penalty（如 min_new_tokens），这是纯开销。

实现拆解

1. 在 python/sglang/srt/managers/schedule_batch.py 的 prepare_for_decode 方法中，将原来嵌套的 list comprehension 直接传给 torch.tensor(..., device=self.device) 的方式拆分为两步：先构建 Python 列表 last_tokens，再通过 torch.tensor(last_tokens, dtype=torch.int64).to(self.device, non_blocking=True) 进行异步拷贝。
2. 精简了内联注释，移除已过时的描述。
3. 由于 forward_stream 已经等待 schedule_stream，异步拷贝能保证在前向消费之前完成排序。
4. 此改动仅影响 penalty 路径；当 penalizer_orchestrator.is_required 为 False 时，整个代码块被跳过，路径不变。

关键文件：

• python/sglang/srt/managers/schedule_batch.py（模块 调度器；类别 source；类型 core-logic）: 核心调度文件，prepare_for_decode 方法是 decode 阶段准备 batch 数据的关键路径，本次变更直接修改了该方法的 latest_output_ids 构造逻辑。

关键符号：prepare_for_decode

关键源码片段

python/sglang/srt/managers/schedule_batch.py

核心调度文件，prepare_for_decode 方法是 decode 阶段准备 batch 数据的关键路径，本次变更直接修改了该方法的 latest_output_ids 构造逻辑。

# python/sglang/srt/managers/schedule_batch.py
​
if self.sampling_info.penalizer_orchestrator.is_required:
    # Under overlap batch.input_ids is just a placeholder here -- the
    # real token is relayed via future_map and resolved at forward
    # entry. So take the last output token from Req directly
    # (origin_input_ids[-1] on the first decode, before any output).
    last_tokens = [
        req.output_ids[-1] if len(req.output_ids) else req.origin_input_ids[-1]
        for req in self.reqs
    ]
    # Non-blocking H2D so this per-step copy doesn't sync behind the forward.
    latest_output_ids = torch.tensor(last_tokens, dtype=torch.int64).to(
        self.device, non_blocking=True
    )
    self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
        latest_output_ids
    )

评论区精华

该 PR 没有 review 评论，仅作者自己合并。但 PR body 中包含了详细的 benchmark 数据和异步安全性的论证：forward_stream 已经等待 schedule_stream，所以异步拷贝的排序得以保证。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低。改动仅涉及 latest_output_ids 的构建方式，且 non_blocking=True 在 CUDA 语义下是安全的，因为后续的 cumulate_output_tokens 在同一 stream 上执行，且该 stream 在 forward 之前已同步。唯一的风险是如果未来在 prepare_for_decode 中增加了对 latest_output_ids 的同步读取，但当前代码路径中没有。
• 影响：影响范围：仅作用于 decode 阶段且启用 penalty 的场景（如 min_tokens、repetition_penalty 等）。对于不启用 penalty 的 decode，无变化。对于非 overlap 模式，也可能因减少一次同步而略有收益。影响程度：在 benchmark 中 decode 吞吐提升约 29%，前向占用率提升 15.5 个百分点，属于显著性能优化。
• 风险标记：暂无

关联脉络

• PR #28465 Batch EAGLE draft/draft-extend replay memcpys via grouped foreach copy: 同样针对 decode 路径中的内存拷贝进行优化，体现了近期对 decode 性能的关注。