Prhub

#25633 Move logprob assembly to SchedulerLogprobResultProcessor

原始 PR 作者 fzyzcjy 合并时间 2026-05-18 18:43 文件变更 3 提交数 1 评论 1 代码增减 +320 / -349

执行摘要

将 logprob 组装逻辑从 Mixin 迁移到独立组件

作为调度器内部重构链的一环,本 PR 旨在将 logprob 结果处理逻辑从臃肿的 SchedulerOutputProcessorMixin 中分离,形成一个自包含的组件 SchedulerLogprobResultProcessor。这为后续的独立测试和未来优化奠定了基础。

建议快速合并。该 PR 是调度器解耦系列中简单且安全的一步,开发者可以重点关注其如何将静态方法迁移为实例方法并简化调用接口的设计模式。

讨论亮点

该 PR 没有产生任何 review 讨论或评论,说明变更直观且无争议。

实现拆解

  1. 填充组件类:在 logprob_result_processor.pySchedulerLogprobResultProcessor 数据类中,新增了原本作为 @staticmethod 存在于混入类中的八个方法(如 _process_input_token_logprobsadd_logprob_return_values 等)。这些方法被转换为实例方法,并移除了显式的 self: "SchedulerLogprobResultProcessor" 类型注解。
  2. 从混入类中移除:在 scheduler_output_processor_mixin.py 中删除了对应的方法定义,并调整了导入语句(移除了 SchedulerLogprobResultProcessor 的类型导入以及 MIS_DELIMITER_TOKEN_ID,因为后者已被组件内部使用)。
  3. 更新调用者scheduler_output_processor_mixin.pydisaggregation/prefill.py 中原来的调用模式 self.<method>(self.logprob_result_processor, ...) 被替换为 self.logprob_result_processor.<method>(...),这是一次纯前缀变换。
文件 模块 状态 重要度
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py 调度器 modified 9.05
python/sglang/srt/managers/scheduler_output_processor_mixin.py 调度器 modified 8.86
python/sglang/srt/disaggregation/prefill.py 解聚合 modified 4.62

关键符号

_process_input_token_logprobs _process_input_top_logprobs _process_input_token_ids_logprobs _calculate_relevant_tokens_len calculate_num_input_logprobs _is_multi_item_scoring add_input_logprob_return_values add_logprob_return_values

关键源码片段

python/sglang/srt/managers/scheduler_components/logprob_result_processor.py core-logic

核心接收端,从空壳数据类变为包含完整 logprob 处理逻辑的组件。新增了八个方法,是本次迁移的目标文件。

from __future__ import annotations
from dataclasses import dataclass
from typing import List, Tuple
import torch
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.server_args import MIS_DELIMITER_TOKEN_ID, ServerArgs
​
​
@dataclass(kw_only=True, slots=True, frozen=True)
class SchedulerLogprobResultProcessor:
    server_args: ServerArgs
    model_config: ModelConfig
​
    # 以下方法是从 SchedulerOutputProcessorMixin 迁移而来,
    # 原本是 @staticmethod,现改为普通实例方法,
    # self 隐式指向本类实例,不再需要显式类型注解。
​
    def _process_input_token_logprobs(
        self, req: Req, input_token_logprobs: List
    ) -> None:
        """Process input token logprobs values and indices."""
        is_multi_item_scoring = self._is_multi_item_scoring(req)
​
        if is_multi_item_scoring:
            req.input_token_logprobs_val = input_token_logprobs
        else:
            # 为普通请求在开头插入 None,并移除最后一位(采样 token)
            req.input_token_logprobs_val = [None] + input_token_logprobs[:-1]
​
        if is_multi_item_scoring:
            # 多条目评分时,分数来自 input_token_ids_logprobs,
            # 但 pipeline 要求 input_token_logprobs_idx 具有相同长度,
            # 因此用分隔符 token ID 填充占位
            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 :]
​
        # 裁剪图像 token 的填充 hash 值,防止反标记化错误
        req.input_token_logprobs_idx = [
            x if x < self.model_config.vocab_size - 1 else 0
            for x in input_token_logprobs_idx
        ]

评论区精华

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

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

风险与影响

由于是纯粹的代码移动且方法体完全一致,回归风险极低。但需要注意所有调用路径是否正确更新,尤其是 disaggregation/prefill.py 中的两个调用点。现有的 CI 测试应能覆盖这些路径。

对用户无任何功能或性能影响。对内部开发者而言,logprob 处理逻辑现在集中于一个组件,更易于理解和维护,也为未来添加新逻辑(如自定义 logprob 处理)提供了清晰的扩展点。

回归风险低 无行为变更

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论