执行摘要

• 一句话：修复 draft 模型共享 logits buffer 形状不匹配
• 推荐动作：建议精读该 PR，特别是 _copy_logits_to_buffer 的最终实现，它展示了一种在性能优化与正确性之间平衡的实用模式：只在形状完全匹配时复用缓冲区，否则回退到安全路径。review 中关于条件切片的建议也值得学习。

功能与动机

PR #29779 引入的共享 logits 缓冲区在 reduced-vocab draft 模型（如 EAGLE/FR-Spec 使用 --speculative-token-map）和独立投机解码时，缓冲区形状与计算出的 logits 不匹配，导致 CI 测试失败。例如，draft 的 lm_head 可能被切片为 32768 大小，但共享缓冲区尺寸为 128256；或者 logits 有 8 行但缓冲区只有 1 行。

实现拆解

1. 在 _copy_logits_to_buffer 中增加条件切片：现在先检查 logits.shape[-1] > self.vocab_size，仅在超过时截取 [:, :self.vocab_size]，避免在不必要时创建 view 增加开销。
2. 将匹配条件从仅检查 vocab 宽度改为检查完整形状：原逻辑只比较 logits_buffer.shape[-1] == self.vocab_size，现在改为比较 tuple(logits_buffer.shape) == tuple(logits.shape)，确保批大小和 vocab 维度都匹配后才复用缓冲区。
3. 调整 fallback 行为：如果不匹配，直接 logits.float() 返回正确的 float 拷贝，而不再硬编码截取 [:, :self.vocab_size]，因为截取操作已在条件切片中完成。
4. 仅修改了 python/sglang/srt/layers/logits_processor.py 文件，没有新增测试，因为现有 CI 测试已经覆盖了这些回归场景。

关键文件：

• python/sglang/srt/layers/logits_processor.py（模块 logits 处理；类别 source；类型 core-logic；符号 _copy_logits_to_buffer）: 核心修改文件：调整了 _copy_logits_to_buffer 方法中的缓冲复用条件，增加了条件切片和形状完全匹配检查。

关键符号：_copy_logits_to_buffer

关键源码片段

python/sglang/srt/layers/logits_processor.py

核心修改文件：调整了 _copy_logits_to_buffer 方法中的缓冲复用条件，增加了条件切片和形状完全匹配检查。

    def _copy_logits_to_buffer(
        self, logits: torch.Tensor, logits_metadata: LogitsMetadata
    ) -> torch.Tensor:
        logits_buffer = logits_metadata.next_token_logits_buffer
        # 仅在 logits 宽度超过 vocab 大小时切片，避免不必要的 view 创建
        if logits.shape[-1] > self.vocab_size:
            logits = logits[:, : self.vocab_size]
        # 检查共享缓冲区的形状是否与 logits 完全匹配（包括批大小和 vocab 宽度）
        if logits_buffer is not None and tuple(logits_buffer.shape) == tuple(
            logits.shape
        ):
            assert logits_buffer.dtype == torch.float
            logits_buffer.copy_(logits)
            logits = logits_buffer
        else:
            # 形状不匹配时，直接返回 float 拷贝，不尝试复用缓冲区
            logits = logits.float()
        return logits

评论区精华

Reviewer gemini-code-assist[bot] 提出了一个性能优化建议：避免在 hot path 中无条件 slice 张量，仅在 logits.shape[-1] > self.vocab_size 时切片。该建议被作者采纳并应用到了最终代码中。此外，cctry 在 review 时表示 LGTM，但提出了疑问：为什么 #29779 的 CI 测试全通过了？这可能意味着 #29779 的测试覆盖不够完整，没有覆盖到 reduced-vocab draft 等场景。

• 避免不必要的张量切片 (performance): 作者采纳了建议，修改为条件切片。
• CI 测试在 #29779 全部通过的原因 (question): 未在评论中明确回复，但本 PR 的测试覆盖了这些场景。

风险与影响

• 风险：本 PR 修改了 logits 处理的 hot path，但修改范围窄（仅一个函数），且通过形状完全匹配来决策，逻辑相对安全。主要风险在于：若 logits_buffer 与 logits 形状匹配但语义不对应（例如 batched 场景下 buffer 被错误复用），可能导致静默错误。不过现有测试覆盖了基本的形状不匹配场景，且 assert 会捕获 dtype 异常。

• 影响：对用户：修复了 reduced-vocab draft 模型（EAGLE/FR-Spec with --speculative-token-map）和独立投机解码的运行时崩溃，使这些功能恢复正常。
  对系统：在形状不匹配时放弃了共享缓冲区优化，会额外进行一次 float() 拷贝，但在不可用场景下这是必要的，性能影响可以忽略。
  对团队：提供了一个清晰的模式来处理共享缓冲区形状不兼容问题，可推广到类似场景。

• 风险标记：hot path 变更, 缺少新增测试覆盖

关联脉络

• PR #29779 Share one logits output buffer across prefill/decode/draft cuda-graph runners: 本 PR 修复了 #29779 引入的共享缓冲区形状不兼容问题。
• PR #29458 Unknown: Oasis-Git 在 Issue 评论中怀疑 #29779 和 #29458 的混合导致了问题，但本 PR 并未明确引用 #29458。