Prhub

#25712 Pack scattered request logprob state into a dedicated container

原始 PR 作者 fzyzcjy 合并时间 2026-05-19 09:18 文件变更 9 提交数 1 评论 1 代码增减 +182 / -132

执行摘要

将 Req 的 logprob 字段封装为 ReqLogprob 数据类

根据 PR body 描述,目的是将原本散布在 Req 类各个地方的 logprob 相关状态(包括参数 top_logprobs_num、token_ids_logprob,以及输入输出返回值字段)统一装入一个专门的容器 ReqLogprob,从而降低 Req 的复杂度,使 logprob 处理逻辑更加内聚。

该 PR 是标准的代码重构,遵循“将相关状态集中管理”的设计原则,值得作为如何拆分大类/提取容器类的参考。但建议后续为 logprob 处理路径编写专项测试,以巩固重构的正确性。

讨论亮点

该 PR 是作者独立提交,review 评论数量为 0,未发现实质性技术讨论。唯一的 comment 来自 gemini-code-assist[bot] 的配额警告,与变更内容无关。

实现拆解

  1. 在 python/sglang/srt/managers/schedule_batch.py 中新增 @dataclass(slots=True, kw_only=True) 的 ReqLogprob 类,包含所有 logprob 属性(top_logprobs_num、token_ids_logprob、input_、output_ 等)。
  2. 修改 Req.init:移除原有的散落初始化代码,改为构造 self.logprob = ReqLogprob(...),并在 return_logprob=True 时条件设置 output_* 为空列表。
  3. 通过机械的正则替换,将所有引用 req.xxx(其中 xxx 属于迁移的字段)改为 req.logprob.xxx,涉及文件:disaggregation/utils.py、disaggregation/decode.py、disaggregation/decode_schedule_batch_mixin.py、layers/utils/logprob.py、managers/scheduler.py、managers/scheduler_components/logprob_result_processor.py、managers/scheduler_components/output_streamer.py、managers/scheduler_components/batch_result_processor.py。
  4. 保留 Req 上其他无关的 logprob 字段(如 input_logprob_sent、logprob_start_len、temp_* 等)。ScheduleBatch 中的批量数组保持原状(依然从 per-req scalars 构建)。
  5. 无测试文件配套修改,仅源码变动。
文件 模块 状态 重要度
python/sglang/srt/managers/schedule_batch.py 请求模型 modified 7.46
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py 对数概率处理 modified 7.26
python/sglang/srt/disaggregation/utils.py 解聚工具 modified 6.33
python/sglang/srt/managers/scheduler_components/output_streamer.py 输出流 modified 5.74
python/sglang/srt/managers/scheduler_components/batch_result_processor.py 批处理结果 modified 5.7

关键符号

ReqLogprob Req.__init__ SchedulerLogprobResultProcessor._process_input_token_logprobs SchedulerLogprobResultProcessor._process_input_top_logprobs SchedulerLogprobResultProcessor._process_input_token_ids_logprobs DisaggMetadataBuffer.set_buf prepare_abort OutputStream.accept BatchResultProcessor._apply_decode_logprobs LogprobComputer.add_output_logprobs_for_spec_v1

关键源码片段

python/sglang/srt/managers/schedule_batch.py core-logic

核心文件:新增 ReqLogprob dataclass,修改 Req 初始化逻辑,集中管理 logprob 状态。

@dataclasses.dataclass(slots=True, kw_only=True)
class ReqLogprob:
    """集中持有单个请求的 logprob 状态。"""
    top_logprobs_num: int
    token_ids_logprob: Optional[List[int]]
    input_token_logprobs_val: Optional[List[float]] = None
    input_token_logprobs_idx: Optional[List[int]] = None
    input_top_logprobs_val: Optional[List[List[float]]] = None
    input_top_logprobs_idx: Optional[List[List[int]]] = None
    input_token_ids_logprobs_val: Optional[List[List[float]]] = None
    input_token_ids_logprobs_idx: Optional[List[List[int]]] = None
    output_token_logprobs_val: Optional[list] = None
    output_token_logprobs_idx: Optional[list] = None
    output_top_logprobs_val: Optional[list] = None
    output_top_logprobs_idx: Optional[list] = None
    # 可以为 GPU 张量,针对 prefill-only 场景的延迟拷贝优化
    output_token_ids_logprobs_val: Optional[List[Union[List[float], torch.Tensor]]] = None
    output_token_ids_logprobs_idx: Optional[list] = None# 在 Req.__init__ 中使用:
self.logprob = ReqLogprob(
    top_logprobs_num=top_logprobs_num,
    token_ids_logprob=token_ids_logprob,
)
if return_logprob:
    self.logprob.output_token_logprobs_val = []
    self.logprob.output_token_logprobs_idx = []
    self.logprob.output_top_logprobs_val = []
    self.logprob.output_top_logprobs_idx = []
    self.logprob.output_token_ids_logprobs_val = []
    self.logprob.output_token_ids_logprobs_idx = []
# 移除了之前的分散字段定义,现在通过 self.logprob.xxx 访问。
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py core-logic

最大变动文件:大量属性访问从 req.xxx 改为 req.logprob.xxx,涉及输入输出 logprob 的批量处理。

# 以 _process_input_token_logprobs 为例,展示变更后的访问方式:
def _process_input_token_logprobs(
    self, req: Req, input_token_logprobs: List
) -> None:
    """处理输入 token logprobs 的值和索引。"""
    is_multi_item_scoring = self._is_multi_item_scoring(req)
​
    # 之前:req.input_token_logprobs_val = ...
    if is_multi_item_scoring:
        req.logprob.input_token_logprobs_val = input_token_logprobs
    else:
        req.logprob.input_token_logprobs_val = [None] + input_token_logprobs[:-1]
​
    # 索引处理
    if is_multi_item_scoring:
        delimiter_count = len(req.multi_item_delimiter_indices)
        input_token_logprobs_idx = [MIS_DELIMITER_TOKEN_ID] * delimiter_count
    else:
        input_token_logprobs_idx = req.origin_input_ids[req.logprob_start_len :]
​
    req.logprob.input_token_logprobs_idx = [
        x if x < self.model_config.vocab_size - 1 else 0
        for x in input_token_logprobs_idx
    ]

评论区精华

没有提炼出高价值讨论线程

当前评论区没有形成足够清晰的争议点或结论,后续有更多讨论时会体现在这里。

风险与影响

主要风险在于机械替换可能遗漏某些调用点,导致运行时 AttributeError。但由于替换通过正则覆盖了所有已知模式,且 PR 在 CI 中至少通过了基础测试(未提供具体 CI 状态),风险可控。另外,没有新增测试覆盖,如果未来重构引入新的 logprob 访问点可能未采用 req.logprob.xxx。对于用户透明,无性能影响。

影响范围:所有涉及 logprob 处理的代码路径,包括推理引擎的解码、分拆调度、打分器等。对用户无功能影响,仅内部 API 重构。开发者需要适应新的访问路径 req.logprob.xxx。对系统无性能变化。团队需确保所有分支同步此变更。

缺少测试覆盖 机械替换可能遗漏 影响多个调度组件

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论