执行摘要
- 一句话:为 logprob 状态创建独立处理器组件
- 推荐动作:建议合并后关注后续 PR(#25633 等)以了解完整迁移脉络。本次重构模式值得借鉴:先引入空数据类作为状态容器、再转换方法签名、最后通过提交链物理移动代码。
功能与动机
这是为后续将 logprob 处理逻辑物理迁移到独立组件所做的准备(PR body 明确指出 'Inplace prep for the introduce-logprob-result-processor mech move')。分离关注点、降低 scheduler.py 的复杂度,并让 logprob 状态以显式依赖的方式传递,是此次大规模调度器重构的目标。
实现拆解
-
创建新类骨架:在 scheduler_components/logprob_result_processor.py 中新建 SchedulerLogprobResultProcessor 数据类(@dataclass,frozen,slots),包含 server_args 和 model_config 两个只读字段。
-
注册到调度器:在 scheduler.py 的 Scheduler.__init__ 中导入并实例化 self.logprob_result_processor = SchedulerLogprobResultProcessor(server_args=..., model_config=...),插入位置紧挨 self.is_initializing = False 之前。
-
方法签名改造:在 scheduler_output_processor_mixin.py 中将 9 个 logprob 相关方法(_initialize_empty_logprob_containers、add_logprob_return_values、add_input_logprob_return_values、_is_multi_item_scoring、calculate_num_input_logprobs、_calculate_relevant_tokens_len、_process_input_token_ids_logprobs、_process_input_top_logprobs、_process_input_token_logprobs)加上 @staticmethod,并把第一个参数的类型标注改为 SchedulerLogprobResultProcessor;方法体保持不变。
-
调用点调整:在 scheduler_output_processor_mixin.py 自身以及 disaggregation/prefill.py 中的每次调用,都添加 self.logprob_result_processor 作为第一个实参;静态方法内部对同 Mixin 其他静态方法的调用改为 SchedulerOutputProcessorMixin.<method>(self, ...) 形式。
整个 PR 只做机械改造,不引入新逻辑。
关键文件:
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py(模块 调度器;类别 source;类型 core-logic;符号 SchedulerLogprobResultProcessor): 新建的 logprob 状态宿主类,本次重构的核心产出
python/sglang/srt/managers/scheduler_output_processor_mixin.py(模块 调度器;类别 source;类型 core-logic;符号 _process_input_top_logprobs, _process_input_token_ids_logprobs, _calculate_relevant_tokens_len, _calculate_num_input_logprobs): 核心改动文件:logprob 方法全部转为 @staticmethod,self 类型变更为新处理器
python/sglang/srt/managers/scheduler.py(模块 调度器;类别 source;类型 dependency-wiring): 导入并实例化 SchedulerLogprobResultProcessor,建立调度器与新组件的依赖关系
python/sglang/srt/disaggregation/prefill.py(模块 分离式预填充;类别 source;类型 core-logic): 两处调用 add_logprob_return_values / add_input_logprob_return_values 时补充新参数
关键符号:SchedulerLogprobResultProcessor, _initialize_empty_logprob_containers, add_logprob_return_values, add_input_logprob_return_values, _is_multi_item_scoring, calculate_num_input_logprobs, _calculate_relevant_tokens_len, _process_input_token_ids_logprobs, _process_input_top_logprobs, _process_input_token_logprobs
关键源码片段
python/sglang/srt/managers/scheduler_components/logprob_result_processor.py
新建的 logprob 状态宿主类,本次重构的核心产出
# python/sglang/srt/managers/scheduler_components/logprob_result_processor.py
# 新建文件:定义 logprob 结果处理器骨架,用于后续承载 logprob 状态与方法
from __future__ import annotations
from dataclasses import dataclass
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.server_args import ServerArgs
@dataclass(kw_only=True, slots=True, frozen=True)
class SchedulerLogprobResultProcessor:
"""
目前仅持有 server_args 与 model_config,
后续会将 logprob 处理方法物理迁移到此组件。
"""
server_args: ServerArgs
model_config: ModelConfig
python/sglang/srt/managers/scheduler_output_processor_mixin.py
核心改动文件:logprob 方法全部转为 @staticmethod,self 类型变更为新处理器
# python/sglang/srt/managers/scheduler_output_processor_mixin.py
# 展示方法签名的静态化改造模式
@staticmethod
def _process_input_top_logprobs(
# self 类型从 Scheduler 改为 SchedulerLogprobResultProcessor
self: "SchedulerLogprobResultProcessor",
req: Req,
) -> None:
"""Process input top logprobs."""
if req.top_logprobs_num <= 0:
return
# 旧调用 self._is_multi_item_scoring(req) 改为显式类名调用
is_multi_item_scoring = SchedulerOutputProcessorMixin._is_multi_item_scoring(
self, req
)
req.input_top_logprobs_val = [] if is_multi_item_scoring else [None]
req.input_top_logprobs_idx = [] if is_multi_item_scoring else [None]
评论区精华
本 PR 无实质 review 讨论,只有一个机器人发出的每日配额提醒,与变更无关。作者自行审查并合并。
风险与影响
- 风险:低风险。所有变更均为结构性,方法体字节完全一致,没有任何控制流或数值计算的改动。唯一潜在问题是:若未来
SchedulerOutputProcessorMixin 类被重命名或移除,静态方法内的显式引用(如 SchedulerOutputProcessorMixin._is_multi_item_scoring)会失效,但这是后续重构的正常代价。现有测试应能覆盖该路径。
- 影响:影响范围仅限调度器模块和分离式预填充模块。对用户无感知,不影响推理性能。对开发团队而言,这是调度器职责拆分的关键一步——logprob 状态从隐式
self 依赖变为显式参数传递,为后续物理迁移到独立文件打下基础。
- 风险标记:依赖后续 PR 完成迁移, 静态方法显式引用 Mixin 类名存在脆弱性
关联脉络
- PR #25633 Move logprob assembly to SchedulerLogprobResultProcessor: 紧随本 PR 的后续步骤,将 logprob 方法物理移动到新组件中,完成本 PR 启动的迁移过程。
- PR #25636 Carve out SchedulerBatchResultProcessor for batch-result state: 同一重构链中的另一组件提取 PR,体现调度器职责拆分的系统化推进。
- PR #25635 Move output streaming to SchedulerOutputStreamer: 类似手法将输出流逻辑析出,与本次 logprob 析出属同一系列。
参与讨论