Prhub

#26779 [core] Compute dimensions/return_pooled_hidden_states in ForwardBatch.init_new

原始 PR 作者 hnyls2002 合并时间 2026-05-31 08:38 文件变更 2 提交数 7 评论 0 代码增减 +34 / -44

执行摘要

将 dimensions 等计算移至 ForwardBatch.init_new

根据 PR 描述,dimensionsreturn_pooled_hidden_states 两个字段只在 ForwardBatch 中被读取,调度器内部(ScheduleBatch 或 downstream 模块)并不使用。因此将它们的计算从 prepare_for_extend 移到 ForwardBatch.init_new,直接基于 batch.reqs 计算,从而删除 ScheduleBatch 上的冗余状态。

值得阅读该 PR 以了解如何将只被下游消费者使用的计算从调度器迁移至前向初始化阶段,这是一种典型的分层清理方法。

讨论亮点

无实质讨论,审阅者 gemini-code-assist[bot] 评论表示变更干净正确,没有进一步建议。

实现拆解

  1. forward_batch_info.py 中新增 _maybe_init_prefill_only 方法,从 batch.reqs 聚合 dimensionsreturn_pooled_hidden_statesmulti_item_delimiter_indices;通过 is_prefill_only 门控仅对 prefill-only 批次设置。
  2. ForwardBatchInfo.init_new 中,在构造实例后调用 ret._maybe_init_prefill_only(batch),替换原来直接从 batch 复制值的代码。
  3. schedule_batch.py 中删除 dimensionsreturn_pooled_hidden_statesmulti_item_delimiter_indices 字段定义,并从 prepare_for_extend 方法中移除对应的计算逻辑。
  4. 调整参数传递:从 init_new 的构造函数调用中移除 dimensionsreturn_pooled_hidden_states,改为通过 helper 方法注入。
文件 模块 状态 重要度
python/sglang/srt/model_executor/forward_batch_info.py 前向批处理 modified 7.24
python/sglang/srt/managers/schedule_batch.py 调度器 modified 6.38

关键符号

ForwardBatchInfo.init_new ForwardBatchInfo._maybe_init_prefill_only ScheduleBatch.prepare_for_extend

关键源码片段

python/sglang/srt/model_executor/forward_batch_info.py data-contract

新增 `_maybe_init_prefill_only` 方法并在 `init_new` 中调用,核心重构发生于此。

def _maybe_init_prefill_only(self, batch):
    """从 batch.reqs 派生每个请求的字段,用于非生成(embedding/reward/scoring)前向传播;否则为空操作。"""
    # 仅在 prefill-only 批次中设置
    if not self.is_prefill_only:
        return
​
    # 只有 Matryoshka 模型且至少有一个请求指定了 dimensions 时,才构建 dimensions 列表
    if batch.model_config.is_matryoshka and any(
        r.dimensions is not None for r in batch.reqs
    ):
        self.dimensions = [
            # 若请求未指定则使用模型的 hidden_size 作为默认值
            r.dimensions if r.dimensions else batch.model_config.hidden_size
            for r in batch.reqs
        ]
​
    # 对 pooled hidden states,只要任意一个请求要求就返回
    self.return_pooled_hidden_states = any(
        r.return_pooled_hidden_states for r in batch.reqs
    )
​
    # 多项目评分(MIS)所需的 delimiter indices,仅在 --enable-mis 启用时才计算
    if get_global_server_args().enable_mis and any(
        r.multi_item_delimiter_indices is not None for r in batch.reqs
    ):
        # score 端点保证所有请求都携带 delimiter indices
        assert all(
            r.multi_item_delimiter_indices is not None for r in batch.reqs
        ), "MIS batch 必须每个请求都有 delimiter indices"
        self.multi_item_delimiter_indices = [
            torch.tensor(r.multi_item_delimiter_indices, dtype=torch.int64)
            for r in batch.reqs
        ]# 在 ForwardBatchInfo.init_new 尾部的调用:
ret = ForwardBatchInfo(...)
ret._maybe_init_prefill_only(batch)

评论区精华

Code Review by gemini-code-assist other

No review comments, changes look clean and correct.

结论:No further feedback. · 已解决

风险与影响

主要风险是 multi_item_delimiter_indices 原本在 prepare_for_extend 中无条件生成,现在仅在 is_prefill_only 为 True 时生成。如果 decode 路径中需要此字段,将导致 None 访问。但根据上下文,该字段仅被 ForwardBatch 在 prefill-only 时使用,因此行为保持。另外,is_prefill_only 门控条件可能与上游调度器行为耦合,需确保调度器在所有 prefill 场景下正确设置此标志。缺少针对迁移后逻辑的测试覆盖。

影响范围小,仅涉及内部两个源文件。外部 API 和用户行为无变化。团队内代码维护性提升,ScheduleBatch 状态减少。

核心路径变更 缺少测试覆盖 条件分支变更

关联 Issue

未识别关联 Issue

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

完整报告

参与讨论