Sampler.forward()：采样分支与 logprob 回填#

执行模型前面的章节已经把 LogitsProcessorOutput、LogitsMetadata、speculative acceptance 和 finish 传播逐层讲开了，但仍然有一个最容易被一句“最后就是抽样”带过去的代码中枢：python/sglang/srt/layers/sampler.py。如果这里不单独讲透，执行模型在“真正选出下一个 token 的那一刻”仍然会缺一块关键解释。

这一章真正要回答的，不是“有哪些采样算法”，而是更贴近运行时的问题：logits 在进入采样前还会被怎样改写，不同请求为什么会走不同采样人格，return_logprob 为什么会显著改写后半段工作量，以及最终选出来的 token 又怎样沿着输出对象链继续往后传。

先把 Sampler.forward() 放回整条后半段执行链#

很多人第一次看到 Sampler.forward() 时，会自然把它想成某个 argmax 或 multinomial 的外壳。但把它放回整条执行链再看，就会发现它其实站在一个更敏感的位置：前面已经做完 LogitsProcessor、工作集裁剪和部分结构化约束准备，后面则要把采样结果、logprob 证据和 TP 同步现实一起收进输出对象。

下面这张图的作用，就是把它从“一个函数”重新放回“执行链中的最后一跳决策器”。

flowchart LR
    LP["LogitsProcessorOutput.next_token_logits"] --> Pre["_preprocess_logits(...)"]
    Pre --> Branch["greedy / probabilistic / deterministic / Ascend"]
    Branch --> IDs["batch_next_token_ids"]
    IDs --> Logprob["attach logprobs / top logprobs / token-id logprobs"]
    Logprob --> Sync["_sync_token_ids_across_tp(...)"]

这张图比纯文字多解释了一件事：采样不是一个原子动作，而是一条“预处理 -> 分支 -> 回填 -> 同步”的后半段链。

更稳的阅读顺序，不是直接看采样分支#

如果你一开始就跳进 greedy、top-k/top-p/min-p 和 Ascend 分支，很容易把 Sampler 读成“很多 if/else 的集合”。更稳的入口其实是 _preprocess_logits(...)，因为它先把两类横切逻辑统一收掉：

  • custom logit processor
  • NaN detection

这一步很关键，因为它提醒你：sampler 看到的 logits，不一定等于 LogitsProcessor 刚吐出来的那一份。也就是说，执行后半段在真正开始“选 token”之前，系统已经承认 logits 还可能被额外改写、清洗和裁剪。

从书稿结构上说，这一步也把本章和前面的 5.14 penalty、logit_bias 与 custom logit processor 的汇合 更紧地接了起来。前一章讲的是“哪些东西会改写 logits”，这一章则讲“这些改写后的 logits 怎样进入真正采样”。

采样人格不是算法分类，而是运行时分流#

Sampler.forward() 里最值得抓的不是“实现了几种采样”，而是系统怎样一步步缩小当前批次的复杂度。

第一层判断是：

  • sampling_info.is_all_greedy

这说明 greedy 在执行上并不只是“temperature 很小的近似情况”，而是真正的一条轻路径。一旦整批请求都属于 greedy 语义，系统就没有必要再走更重的概率抽样逻辑。换句话说，采样参数本身就会改写执行形态，而不是只改写数学分布。

第二层判断是：

  • simple_sampling_case

它对应的是当前批次没有更复杂的 top-k、top-p、min-p 等限制，于是系统又可以把一大批请求留在更轻的分支里。这种设计背后的工程判断非常典型：绝大多数简单路径应该尽早分出去，复杂路径只为真正需要它们的请求付费。

从好书的角度说，这比简单列出“支持 greedy / top-k / top-p”更重要。因为读者真正需要带走的，不只是算法名词，而是一个系统性判断：SGLang 在采样这里也在做按复杂度分层的运行时优化。

return_logprob 改变的不只是返回值，而是工作量#

如果只看接口，return_logprob 很容易被理解成“多返回一点信息”。但放回 Sampler.forward() 里看，它更像一条正式的执行分叉。因为一旦这个开关为真，sampler 就不再只选 token，还要继续准备：

  • next_token_logprobs
  • top logprobs
  • token-id-specific logprobs

并把这些内容回填回 LogitsProcessorOutput。这说明 logprob 返回从来不是纯附加信息，而是采样函数自己承担的一项主要职责。也正因为这样，前面讲输出后半段和 finish 语义的章节，最后都必须再次回到 sampler。

一个特别值得强调的细节是 SGLANG_RETURN_ORIGINAL_LOGPROB。这个开关说明“logprob”这个词本身也有口径差异：你要的是接近原始 logits 的 logprob，还是采样链最终使用过的 logprob。很多系统文章会把两者混成一回事，但从工程实践看，这种区分非常重要，因为它会直接影响后续评估、可解释性和调试口径。

backend 也会改写采样路径#

这一章另一个最容易被低估的点，是硬件 backend 不只影响前向，也会影响采样本身。_forward_ascend_backend(...) 和 _sample_from_logits(...) 的存在说明，某些后端并不一定要先显式 softmax 成概率再抽样。换句话说，采样分支不仅受算法参数影响，也受 backend 能力约束。

这和前面 graph runner、attention backend 那几章形成了非常自然的回扣：执行模型里的很多“人格”从来都不是纯算法层的抽象，而是算法、硬件和工程路径一起决定的。

同样，rl_on_policy_target 和 enable_deterministic_inference 这些分支也提醒读者：采样并不总是概率抽样。在 on-policy 或更强调可复现的场景下，系统会转向 deterministic sampling。这再次说明，采样真正解决的是“如何做出运行时决策”，而不只是“如何按分布取样”。

另一半关键逻辑在“回填”和“同步”#

如果 _preprocess_logits(...) 是入口，那么 _attach_logprobs_to_output(...) 就是这棵树真正的另一半。它把：

  • next_token_logprobs
  • top logprobs
  • token-id-specific logprobs

正式写回 LogitsProcessorOutput，也让 sampler 成为“采样决策”和“返回证据”同时汇合的地方。只看前半段分支，而不看回填，就很容易误把 sampler 读成纯决策函数。

同样，_sync_token_ids_across_tp(...) 也不能被当作一个可忽略的小尾巴。它揭示的是另一层更现实的运行时问题：采样的最终 token id 并不总是天然在所有 TP rank 上完全一致。系统平时会尽量保持轻路径，但在 grammar 或某些不稳定场景里，又必须显式同步。这是一种很典型的“默认优化，特殊场景保守回退”的策略。

这章和前几章不是重复，而是补最后一跳#

从位置上看，这一章容易和前面的几章看起来重叠：

但三者的职责其实不同。前面这些章节分别回答的是：

  • 采样与 speculative 的整体执行位置
  • 从 logits 到 finish 的语义传播
  • 哪些位置会继续进入 logits / logprobs / hidden-state 处理链

这一章补的是最后一跳：当工作集已经确定、logits 已经准备好之后，Sampler.forward() 这棵代码树怎样真正做出 token 选择，并把 logprob 与同步现实一起带回输出对象。也就是说，它承担的是执行模型里的“代码级决策解释”。

最容易出现的三种误判#

第一，误以为采样就是一次 argmax 或 multinomial。
实际上前后还夹着 logits 预处理、分支选择、logprob 回填和 TP 同步。

第二，误以为 logprobs 是 output processor 后面才补上的。
真正的回填点就在 sampler 内部。

第三，误以为 backend 只改写前向，不改写采样。
Ascend 路径已经说明采样本身也会出现 backend-specific 分支。

真正顺着源码读时，推荐怎么走#

如果你准备把这棵树真正打开读，最稳的顺序仍然是：

  1. _preprocess_logits(...)
  2. forward(...)
  3. _sample_from_probs(...) / _sample_from_logits(...)
  4. _attach_logprobs_to_output(...)
  5. _sync_token_ids_across_tp(...)

这样读，你先理解“输入 logits 还会被怎样改写”，再理解“分支怎样做出选择”，最后再看“结果怎样带回输出对象”。这比一上来盯某个 top-p 分支要稳得多，也更符合技术书应有的阅读引导方式。

小结#

Sampler.forward() 的真正价值，不在于它支持很多采样算法，而在于它把执行后半段的最后一次决策真正落成了代码：logits 怎样被预处理，哪条采样人格被启用，哪些 logprob 证据会被带回，以及在什么场景下必须对最终 token 做更保守的同步。

到这里，执行模型里的“真正选出 token 的那一刻”才算被正式讲透。