执行摘要
- 一句话:将 logprob 组装逻辑从 Mixin 迁移到独立组件
- 推荐动作:建议快速合并。该 PR 是调度器解耦系列中简单且安全的一步,开发者可以重点关注其如何将静态方法迁移为实例方法并简化调用接口的设计模式。
功能与动机
作为调度器内部重构链的一环,本 PR 旨在将 logprob 结果处理逻辑从臃肿的 SchedulerOutputProcessorMixin 中分离,形成一个自包含的组件 SchedulerLogprobResultProcessor。这为后续的独立测试和未来优化奠定了基础。
实现拆解
- 填充组件类:在
logprob_result_processor.py 的 SchedulerLogprobResultProcessor 数据类中,新增了原本作为 @staticmethod 存在于混入类中的八个方法(如 _process_input_token_logprobs、add_logprob_return_values 等)。这些方法被转换为实例方法,并移除了显式的 self: "SchedulerLogprobResultProcessor" 类型注解。
- 从混入类中移除:在
scheduler_output_processor_mixin.py 中删除了对应的方法定义,并调整了导入语句(移除了 SchedulerLogprobResultProcessor 的类型导入以及 MIS_DELIMITER_TOKEN_ID,因为后者已被组件内部使用)。
- 更新调用者:
scheduler_output_processor_mixin.py 和 disaggregation/prefill.py 中原来的调用模式 self.<method>(self.logprob_result_processor, ...) 被替换为 self.logprob_result_processor.<method>(...),这是一次纯前缀变换。
关键文件:
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py(模块 调度器;类别 source;类型 core-logic;符号 _process_input_token_logprobs, _process_input_top_logprobs, _process_input_token_ids_logprobs, _calculate_relevant_tokens_len): 核心接收端,从空壳数据类变为包含完整 logprob 处理逻辑的组件。新增了八个方法,是本次迁移的目标文件。
python/sglang/srt/managers/scheduler_output_processor_mixin.py(模块 调度器;类别 source;类型 core-logic;符号 _process_input_token_logprobs, _process_input_top_logprobs, _process_input_token_ids_logprobs, _calculate_relevant_tokens_len): 原始混入类,删除了所有 logprob 处理方法,调用点改为委托给新组件。导入依赖也相应简化。
python/sglang/srt/disaggregation/prefill.py(模块 解聚合;类别 source;类型 core-logic): 包含两个调用点的更新,展示迁移后如何通过 self.logprob_result_processor 直接调用新组件的方法。
关键符号:_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
核心接收端,从空壳数据类变为包含完整 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
]
评论区精华
该 PR 没有产生任何 review 讨论或评论,说明变更直观且无争议。
风险与影响
- 风险:由于是纯粹的代码移动且方法体完全一致,回归风险极低。但需要注意所有调用路径是否正确更新,尤其是
disaggregation/prefill.py 中的两个调用点。现有的 CI 测试应能覆盖这些路径。
- 影响:对用户无任何功能或性能影响。对内部开发者而言,logprob 处理逻辑现在集中于一个组件,更易于理解和维护,也为未来添加新逻辑(如自定义 logprob 处理)提供了清晰的扩展点。
- 风险标记:回归风险低, 无行为变更
关联脉络
- PR #25635 Move output streaming to SchedulerOutputStreamer: 同属调度器组件化重构链,将输出流处理从混入类迁移到独立组件,模式相同。
- PR #25637 Move batch-result processing to SchedulerBatchResultProcessor and retire output_processor mixin: 类似的重构,将 batch 结果处理迁移到独立组件,与本 PR 共同推进 mixin 的拆分。
参与讨论