执行摘要

• 一句话：提取 InputLogprobProcessor 并迁移 logprob 处理逻辑
• 推荐动作：推荐阅读该 PR，特别是：
  • 如何通过依赖注入（get_logits_fn）解耦 lm_head 的调用。
  • 如何在不改变行为的前提下安全地从大类别中提取职责。
  • 机械验证方法（AST 等价性检查）可作为重构的参考实践。

功能与动机

当前 LogitsProcessor 承担了过多职责（logits 计算、input logprob 处理、采样等），且 logprob 相关函数分散在 layers/utils/logprob.py 中。PR 旨在为 logprob 处理建立一个独立的模块，先从 input（prefill）侧开始提取 InputLogprobProcessor，为后续提取 OutputLogprobProcessor 和统一 logprob 结果数据结构做准备。PR body 明确说明 'Introduce python/sglang/srt/layers/logprob_processor.py as the dedicated home for logprob processing, starting with the input (prefill) side.'

实现拆解

1. 文件搬迁与重命名：将 layers/utils/logprob.py 的全部内容（包括 InputLogprobsResult、get_top_logprobs_prefill、get_token_ids_logprobs_prefill、compute_spec_v2_logprobs 等工具函数）移至 layers/logprob_processor.py，并在此文件中新建 InputLogprobProcessor 类。
2. 提取核心类：从 LogitsProcessor 中剥离 process_input_logprobs 和 process_input_logprobs_by_chunk 两个方法，贴上新的类 InputLogprobProcessor。新类通过构造函数读取环境变量（SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK 和 SGLANG_LOGITS_PROCESSER_CHUNK_SIZE），并通过 forward 入口向下分发。
3. 重构 LogitsProcessor：移除 LogitsProcessor 中与 input logprob 处理直接相关的代码（如环境变量读取、should_skip_chunking 判断、process_input_logprobs 等方法），改为在 __init__ 中创建 InputLogprobProcessor 实例，并在 forward 中委托调用。解耦点是通过 get_logits_fn 回调（原 self._get_logits）注入 lm_head 计算，将 dp_attention 的跳过标志作为参数传入。
4. 更新导入路径：将所有引用 sglang.srt.layers.utils.logprob 的文件改为 sglang.srt.layers.logprob_processor，涉及 sampler.py、eagle_worker_common.py、ngram_worker.py、lora/layers.py、lora/utils.py。其中 lora/layers.py 和 lora/utils.py 仅改动文档字符串中的方法引用。
5. 机械验证：通过 AST 等价性检查确认提取后的两个方法与原版字节等同（除两处重命名外，self._get_logits → get_logits_fn，self.do_tensor_parallel_all_gather_dp_attn → skip_chunking_for_dp_attn），并在 CI 上运行 test_srt_endpoint.py、test_lora_hf_sgl_logprob_diff.py、test_spec_ngram.py、test_original_logprobs.py 均通过。

关键文件：

• python/sglang/srt/layers/logprob_processor.py（模块 logprob处理层；类别 source；类型 rename-or-move；符号 InputLogprobProcessor, init, forward, process_input_logprobs）: 新建模块，包含从 LogitsProcessor 提取的 InputLogprobProcessor 类以及原 logprob.py 的所有工具函数，是本次重构的核心文件。
• python/sglang/srt/layers/logits_processor.py（模块 logits处理；类别 source；类型 core-logic；符号 LogitsProcessor.init, LogitsProcessor.forward, process_input_logprobs, process_input_logprobs_by_chunk）: 主要修改文件之一，删除了约 230 行与 input logprob 处理相关的代码，改为委托给 InputLogprobProcessor。
• python/sglang/srt/layers/sampler.py（模块 采样器；类别 source；类型 dependency-wiring）: 改动了导入路径，从 layers.utils.logprob 改为 layers.logprob_processor。
• python/sglang/srt/speculative/eagle_worker_common.py（模块 推测解码；类别 source；类型 dependency-wiring）: 更新 compute_spec_v2_logprobs 的导入路径。
• python/sglang/srt/speculative/ngram_worker.py（模块 推测解码；类别 source；类型 dependency-wiring）: 更新 compute_spec_v2_logprobs 的导入路径。
• python/sglang/srt/lora/layers.py（模块 LoRA；类别 source；类型 dependency-wiring）: 更新文档字符串中的方法引用。
• python/sglang/srt/lora/utils.py（模块 LoRA；类别 source；类型 dependency-wiring）: 更新文档字符串中的方法引用（与 layers.py 相同）。

关键符号：InputLogprobProcessor.init, InputLogprobProcessor.forward, InputLogprobProcessor.process_input_logprobs, InputLogprobProcessor.process_input_logprobs_by_chunk, LogitsProcessor.init, LogitsProcessor.forward

关键源码片段

python/sglang/srt/layers/logprob_processor.py

新建模块，包含从 LogitsProcessor 提取的 InputLogprobProcessor 类以及原 logprob.py 的所有工具函数，是本次重构的核心文件。

class InputLogprobProcessor:
    """Input (prefill) logprob processing: single-pass or chunked.
​
    Logits are computed through the injected ``get_logits_fn(hidden_states,
    lm_head, logits_metadata)`` callable, so this class stays decoupled from
    the lm_head / TP-gather machinery in LogitsProcessor.
    """
​
    def __init__(self):
        # 从环境变量读取是否启用分块处理
        self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get()
        # 分块大小
        self.logprobs_chunk_size = envs.SGLANG_LOGITS_PROCESSER_CHUNK_SIZE.get()
​
    def forward(
        self,
        pruned_states: torch.Tensor,
        sample_indices: Optional[torch.Tensor],
        input_logprob_indices: torch.Tensor,
        token_to_seq_idx: list[int],
        lm_head: VocabParallelEmbedding,
        get_logits_fn: Callable,
        logits_metadata: LogitsMetadata,
        skip_chunking_for_dp_attn: bool = False,
    ) -> Tuple[InputLogprobsResult, torch.Tensor]:
        # 判断是否需要跳过分块（三种条件之一成立即跳过）
        should_skip_chunking = (
            not self.enable_logprobs_chunk
            or pruned_states.shape[0] <= self.logprobs_chunk_size
            or skip_chunking_for_dp_attn
        )
​
        if should_skip_chunking:
            # 单次计算所有 logits
            logits = get_logits_fn(pruned_states, lm_head, logits_metadata)
            sampled_logits = (
                logits[sample_indices] if sample_indices is not None else logits
            )
            input_logits = logits[input_logprob_indices]
            del logits
​
            logprobs_result = self.process_input_logprobs(input_logits, logits_metadata)
        else:
            logprobs_result, sampled_logits = self.process_input_logprobs_by_chunk(
                pruned_states,
                sample_indices,
                input_logprob_indices,
                token_to_seq_idx,
                lm_head,
                get_logits_fn,
                logits_metadata,
            )
​
        return logprobs_result, sampled_logits
​
    def process_input_logprobs(self, input_logits, logits_metadata: LogitsMetadata):
        # 计算 log_softmax
        input_logprobs = torch.nn.functional.log_softmax(input_logits, dim=-1)
​
        # 如果请求 top-k logprob，则计算
        if logits_metadata.extend_return_top_logprob:
            (
                input_top_logprobs_val,
                input_top_logprobs_idx,
            ) = get_top_logprobs_prefill(input_logprobs, logits_metadata)
        else:
            input_top_logprobs_val = input_top_logprobs_idx = None
​
        # 收集 input 位置上的 token logprob
        input_token_logprobs = torch.gather(
            input_logprobs, dim=-1, index=logits_metadata.input_token_ids
        ).squeeze(-1)
​
        # 如果请求 token_ids_logprobs，则计算
        if logits_metadata.extend_return_token_ids_logprobs:
            (
                input_token_ids_logprobs_val,
                input_token_ids_logprobs_idx,
            ) = get_token_ids_logprobs_prefill(input_logprobs, logits_metadata)
        else:
            input_token_ids_logprobs_val = input_token_ids_logprobs_idx = None
​
        return InputLogprobsResult(
            input_token_logprobs=input_token_logprobs,
            input_top_logprobs_val=input_top_logprobs_val,
            input_top_logprobs_idx=input_top_logprobs_idx,
            input_token_ids_logprobs_val=input_token_ids_logprobs_val,
            input_token_ids_logprobs_idx=input_token_ids_logprobs_idx,
        )

python/sglang/srt/layers/logits_processor.py

主要修改文件之一，删除了约 230 行与 input logprob 处理相关的代码，改为委托给 InputLogprobProcessor。

# 在 LogitsProcessor.__init__ 中，删除了原有的 chunk 环境变量读取，改为创建处理器
self.input_logprob_processor = InputLogprobProcessor()
​
# 在 LogitsProcessor.forward 中，原本大量的 chunk 判断与计算代码被替换为一行委托
logprobs_result, sampled_logits = self.input_logprob_processor.forward(
    pruned_states=pruned_states,
    sample_indices=sample_indices,
    input_logprob_indices=input_logprob_indices,
    token_to_seq_idx=token_to_seq_idx,
    lm_head=lm_head,
    get_logits_fn=self._get_logits,
    logits_metadata=logits_metadata,
    skip_chunking_for_dp_attn=self.do_tensor_parallel_all_gather_dp_attn,
)

评论区精华

Review 没有产生讨论分歧。作者通过 PR body 附上了机械验证脚本的结果（Gist 链接）和 CI 绿色报告，合入者 hnyls2002 在评论中 /rerun-test 指定测试集并全部通过。唯一一个 gemini-code-assist[bot] 的消息是每日配额提示，与技术内容无关。

• 暂无高价值评论线程

风险与影响

• 风险：本次重构严格保持行为不变，通过 AST 等价性验证确保提取的方法签名和逻辑一致。风险点主要在：
  • 文件重命名：若外部项目或脚本直接导入 layers.utils.logprob，会断裂。本仓库内已全部更新（共 5 个导入文件），无剩余占用。
  • 属性移除：LogitsProcessor 中删除了 enable_logprobs_chunk 和 logprobs_chunk_size 属性，若其他代码直接访问这些属性（当前无），则可能报错。
  • 无新增测试：依赖现有测试覆盖，但重构本身不引入新功能，现有测试的通过提供了足够信心。
• 影响：对用户无功能影响。对开发者：
  • 所有 logprob 相关函数的导入路径从 sglang.srt.layers.utils.logprob 变为 sglang.srt.layers.logprob_processor，需要更新外部引用。
  • LogitsProcessor 的职责减轻，新增 input_logprob_processor 属性，未来扩展 OutputLogprobProcessor 时可以在同一模块下对称实现。
  • 降低 LogitsProcessor 的复杂度，有利于独立单元测试。
  • 风险标记：核心路径变更, 依赖注入解耦, 无测试回归

关联脉络

• PR #31498 [Fix] Enable chunked input-logprob processing by default to cap peak memory: 该 PR 修改了同一个 logits_processor.py 文件并涉及 chunked logprob 处理逻辑，是本次重构的上游上下文。本次重构提取的 InputLogprobProcessor 正是为了承接这类逻辑。