本文已于 2026.04.26 发表于公众号和知乎。

目录

• 1. Beam Search 是什么
• 2. 性能评测
• 3. Beam Search 工程的核心挑战
• 4. vLLM 的 Beam Search 实现方案
• 5. Transformers 的 Beam Search 实现方案简介
• 6. TensorRT-LLM 的 Beam Search 实现方案简介
• 7. SGLang 的实现方案
• 8. 设计思想总结——面向数据设计

1. Beam Search 是什么

Beam Search 是一种解码策略，与当前广泛使用的 Greedy Decoding（贪婪解码） 不同。Greedy Decoding 每一步只选择概率最高的一个 token，是一种局部最优的策略。而 Beam Search 每一步保留 beam_width 个累积对数概率最高的候选序列，通过在多条路径上并行探索，最终输出 beam_width 个全局累积对数概率最高的完整序列。Greedy Decoding 可以看作 beam_width = 1 的 Beam Search 特例。

 在当前以对话为主的 LLM 推理场景中，Beam Search 较少使用。如果以 Beam Search 为核心来组织推理流程，会显著增加代码复杂度，所以在当前的主流推理引擎中，都没有将 beam search 作为核心链路，SGLang 甚至相当长时间都不支持。在搜索推荐场景里，当大模型用作召回时，需要召回多条可能满足用户诉求的资源，这和 Beam Search 返回多结果的特性正好契合，因此部分公司会将 Beam Search 用于大模型召回，譬如快手的 OneRec。

本文接下来会先展示 SGLang 的性能评测数据（PR 中），然后介绍 Beam Search 的工程挑战、vLLM 和 TensorRT-LLM 的实现方案，最后重点讲解如何为 SGLang 添加 Beam Search 功能。

2. 性能评测

先给出各引擎在不同 beam width 下的 QPS（Queries Per Second）对比。SGLang 的实现在大 beam width 下可扩展性明显更好，整体领先 TensorRT-LLM，显著优于 vLLM。

评测环境

• • 模型：Qwen3-1.7B
• • 硬件：单张 NVIDIA L20 GPU
• • 数据集：ShareGPT（取 100 条样本，过滤 prompt 长度 < 100）
• • 采样参数：max_tokens = 10
• • 软件版本：
• • SGLang：https://github.com/sgl-project/sglang/pull/15645
• • TensorRT-LLM：v1.2.0rc4
• • vLLM：v0.13.0
• • Transformers：4.57.1

QPS 对比

Engine \ Beam Width	10	50	100	200	400
SGLang（本文方案）	11.53	10.23	8.73	6.83	4.27
TensorRT-LLM	12.23	10.70	8.79	6.21	3.50
Transformers	4.41	3.77	2.94	1.96	1.08
vLLM	5.54	2.01	0.80	0.23	0.06

可以看到：

• • 小 beam width（10、50）：SGLang 与 TensorRT-LLM 基本持平，略低 0.5~1 QPS；两者都把 vLLM 和 Transformers 甩开一个数量级。
• • 大 beam width（≥ 200）：SGLang 反超 TensorRT-LLM（BW=200 时 6.83 vs 6.21，BW=400 时 4.27 vs 3.50），且下降更缓；vLLM 在 BW=400 时已经跌到 0.06 QPS，基本不可用。

注：本文实现的 SGLang 功能完成于 2025.12 月份之前，因此评测的版本也较早。

3. Beam Search 工程的核心挑战

在介绍 Beam Search 的具体实现之前，我们先基于理论分析，初步了解 Beam Search 工程实践需要解决的几个核心挑战。

3.1 与推理引擎主流程的解耦：Beam Search 逻辑较复杂，如何避免 beam search 分支逻辑污染普通推理路径？

3.2 KV Cache 管理：Beam Search 的多条路径共享 prefix KV Cache，扩散、剪枝时又需要继承或释放各 beam 的 KV Cache，如何高效复用与回收？

3.3 请求调度与批处理：Beam Search 请求包含 beam_width 条独立路径，如何在 batch 中组织这些路径、是否与普通请求混合调度、请求完成后如何动态更新 batch？

3.4 Beam 的扩散、评分与剪枝：每步 decode 后需要累加 log prob、从候选集中选出 top beam_width 个存活路径，如何高效实现？

带着这些问题，我们来看看各家推理引擎是怎么实现 Beam Search 的。

4. vLLM 的 Beam Search 实现方案

vLLM 实现的 Beam Search 以 v0.6.2 为分界线，有两种实现，下面分开介绍。

 

4.1 v0.6.2 之前（含 v0.6.2）的实现

4.1.1 beam search 作为推理引擎主流程

在 v0.6.0 中，Beam Search 和 parallel sampling（best_of > 1）作为一等公民嵌入在 vLLM 的主流程中。一个 request 可能同时维护多条并行生成序列，这一设计决定了核心数据结构必须以"序列组"（SequenceGroup）而非单条序列为粒度：

• • SequenceGroup.seqs_dict: Dict[seq_id, Sequence]，允许一个 request 内动态地 fork/add/remove 多条 seq
• • SequenceGroupMetadata 中 seq_data 和 block_tables 也是 Dict[seq_id, ...]，为组内每条 seq 独立传递数据和 block 映射
• • 调度 budget 以 get_max_num_running_seqs() 预留位置，beam search 模式下始终预留 best_of 个 seq 位
• • 输出处理（SingleStepOutputProcessor）在每步 decode 后走不同分支：普通采样（best_of=1）直接 append token；parallel sampling（best_of > 1）fork 出多条独立序列并行生成；beam search 则执行 fork → score → prune → early stopping 的完整流程
• • SequenceOutput.parent_seq_id 专门为 beam search 的 fork 谱系追踪而设计

这种"组优先"的设计使得 beam search 的 fork/prune 操作可以在不改变调度器和 block manager 接口的前提下实现，但代价是：即使绝大多数请求只有 best_of=1（单条 seq），系统仍然在所有路径上承担了"组"抽象的开销。

注：best_of 参数在不同模式下含义不同：

• • Beam search 模式（use_beam_search=True）：best_of = beam width，控制每步保留的束宽
• • 普通采样模式：best_of = 从同一 prompt fork 出的并行生成序列数，最终从中按 cumulative logprob 选出 top-n 返回；默认 best_of = n，即生成几条返回几条
• • 在当前最新的架构里 best_of 参数仅用于 Beam Search 模式。

4.1.2 KV Cache 的复用和 Copy on Write

在 KV Cache 管理上，采用了引用计数 + Copy-on-Write 策略。每一步 decode，sampler 产生 2×beam_width 个候选，output processor 先将这些候选扩散为 child 序列（Python 对象的 deepcopy），然后经过评分与剪枝，选出最多 beam_width 个 running child 继续推理。对于这些选中的 child，调用 scheduler.fork_seq(parent, child)，新建一个 BlockTable 容器，共享 parent 的所有 PhysicalTokenBlock 引用，并增加每个物理 block 的引用计数，不拷贝 GPU 上的 KV Cache 数据本身。当下一步 decode 需要往最后一个 block 追加 token 且该 block 被多个 seq 共享（ref_count > 1）时，触发 CoW：分配新的物理 block，在 GPU 上执行真正的 KV 数据拷贝，然后将该 child 的 block table 指向新 block。如果 block_size=1，由于每个 block 始终是满的，新 token 必然落在新分配的 block 上，CoW 不会触发。

4.1.3 请求调度和批处理

在请求调度与批处理上，vLLM 将一个 beam search 请求视为一个 SequenceGroup（即一条请求），与普通请求混合调度、放在同一个 batch 中执行。区别在于资源预算的计算：调度器通过 get_max_num_running_seqs() 获取每个请求所需的 seq slot 数量，普通请求返回 1，beam search 请求始终返回 beam_width。这意味着一个 beam_width=4 的 beam search 请求会在 max_num_seqs 预算中占用 4 个名额。模型前向推理时，同一 batch 内会同时包含普通请求的 seq 和 beam search 请求的多条 beam，GPU 不区分它们。推理完成后，引擎逐个 SequenceGroup 处理输出结果：普通请求走快速路径直接返回，beam search 请求走独立的扩散-评分-剪枝分支，各请求的处理互不影响。请求完成后，调度器调用 free_finished_seq_groups() 释放资源，腾出的 budget 立即可供后续请求使用。

4.1.4 Beam 的扩散、评分与剪枝

主要分散在 Sampler 层（候选生成）和 Output Processor 层（选择和剪枝）。

Sampler 层：_beam_search_sample()

• • Prompt 阶段：只有 1 条 seq，直接从 [1, vocab_size] 的 logprobs 中选出 Top 2 * beam_width 个 token 作为初始候选
• • Generation 阶段：将 [beam_width, vocab_size] 的当前步 logprobs 与各 beam 的 cumulative logprob 相加（利用 GPU 广播），然后在 beam_width * vocab_size 的展平空间中选出 Top 2 * beam_width 个候选，通过整除和取模还原出 parent_id（来自哪个 beam）和 token_id（选了哪个词）

Output Processor 层：_process_sequence_group_outputs()

• • Fork 构造候选：根据 Sampler 返回的 2 * beam_width 个 (parent_id, token_id) 对，从 parent seq fork 出新的 child seq，每个 child 追加对应的 token。没有被任何候选引用的 parent seq 被直接标记为 FINISHED_ABORTED 并释放
• • Stop 检查：对所有 child seq 执行 detokenize 和 stop condition 检查（EOS、stop token、max_tokens），将触发 stop 的 seq 标记为 finished
• • Finished seq 管理：将历史已完成的 finished seq 与本轮新 finished 的 seq 合并，按 get_beam_search_score（含 length_penalty）降序排序，只保留 Top beam_width 个 finished seq，淘汰低分的
• • Running seq 选择：将未完成的 running child seq 也按 beam search score 降序排序，选出 Top beam_width 个继续下一步生成，其余被淘汰，当 Top beam_width 不足 beam_width 时，不做特殊处理
• • Early stopping 判定：判断是否可以提前终止整个 beam search，有三种策略：
  • • early_stopping=True：只要凑齐 beam_width 个 finished seq 就立即停止
  • • early_stopping=False：当最好的 running seq 即使继续生成也不可能超过当前最差的 finished seq 时停止
  • • early_stopping="never"：考虑理论上最长序列的最高可达分数，只有确定无法超越时才停止

4.2 v0.6.2 之前的 Beam Search 实现的部分缺点

4.2.1 未采用面向数据的设计

Beam search 请求会扩散出 beam_width 条路径，vLLM 采用面向对象的做法来管理这些路径，带来多方面的开销：

• • GPU → CPU 的非批次同步拷贝：_beam_search_sample() 在 Python for 循环中逐 seq_group 对 logprobs 做 torch.topk 再 .tolist()，每次 .tolist() 都触发一次 GPU-CPU 同步。作为对比，普通采样是先批量 .cpu() 再切片，只需一次同步。
• • 大量临时对象构造与浪费：Sampler 层选出 2 * beam_width 个候选 token，在 _build_sampler_output() 中包装成 SequenceOutput 对象；再到 Output Processor 层，通过 parent.fork() 对 Sequence 做 copy.deepcopy（每个都包含完整的历史 token 序列和 logprobs）。每个 parent 的最后一个 child 会复用 parent 对象，但其余全部需要深拷贝。经过评分和剪枝后，只保留 beam_width 个继续运行，大量 fork 和深拷贝被浪费。
• • 逐对象操作无法向量化：fork（copy.deepcopy）、append_token_id 以及 KV Cache 复用的 block table 拷贝等操作，全部在 CPU/Python 层逐个 Sequence 对象执行，无法利用 GPU 并行或批量张量操作。

逐对象操作示例：

for child_sample in child_samples[:-1]:
    # 对每个候选做一次 copy.deepcopy
    new_child_seq_id: int = next(self.seq_counter)
    child = parent.fork(new_child_seq_id) # copy.deepcopy(self)
    child.append_token_id(child_sample.output_token,
                          child_sample.logprobs)
    child_seqs.append((child, parent))

整体来说，在 beam width 比较大时，面向对象的设计，使得 Sequence 临时对象的构造、KV Cache 复用时的 block table 列表拷贝都造成比较大的浪费。

4.2.2 较多的重复计算

beam search 的扩散和剪枝逻辑在 CPU 侧也不够精细，譬如：selected_child_seqs 将完成的和未完成的序列混在同一个容器里，导致后续遍历时需要对每个序列都做 is_finished() 判断；SequenceGroup 的 get_finished_seqs() 每次调用都要遍历整个 seqs_dict 做过滤。在整理已完成请求时，beam search 没有缓存，每一次都需要重新调用 get_beam_search_score 计算。

def get_finished_seqs(self) -> List[Sequence]:
    return [seq for seq in self.seqs_dict.values() if seq.is_finished()]

# Sort the finished sequences by their scores.
all_finished_seqs.sort(key=lambda x: x[0].get_beam_search_score(
    length_penalty=length_penalty, eos_token_id=x[0].eos_token_id),
                       reverse=True)

尽管代码性能上有诸多考虑不周的细节点，但因推理耗时主要在 GPU 上，CPU 侧的粗放实现对端到端性能影响较小。后来 Beam Search 从 vLLM 主流程中迁出，性能缺陷并非主因，而是架构方面的原因：beam search 的逻辑相比普通推理更复杂（涉及序列的 fork、free、评分、剪枝），作为一个使用率较低的功能，放在主流程中会增加代码的维护负担，拖慢普通推理路径的迭代效率。

4.3 v0.6.2 之后的实现

 v0.6.2 之后将 Beam Search 从 vLLM 的推理主流程中剥离，变成了一个在入口层通过多轮普通推理请求编排实现的上层功能。代码涉及两条路径——离线 API（LLM.beam_search，位于 vllm/entrypoints/llm.py）和在线 API（OpenAIServing.beam_search，位于 vllm/entrypoints/openai/engine/serving.py），以及共享的数据结构层（vllm/beam_search.py）。下面沿用与 4.1 节相同的四个维度展开。

4.3.1 与推理引擎主流程的解耦

Beam search 的编排逻辑被提升到入口层（entrypoint），引擎内核中没有任何 beam search 特有代码。入口层通过请求中的 use_beam_search 字段识别 beam search 请求，将其分流到 beam_search() 编排循环，而非直接提交给引擎内核。在编排循环内部，每条 beam 被包装为一个普通的独立推理请求提交给引擎内核：

• • 每次只生成 1 个 token，但要求引擎内核返回 top 2 * beam_width 个 logprobs。
• • 引擎内核将这些请求与普通推理请求一视同仁地处理，无需任何特殊路径。
• • 引擎内核返回 logprobs 后，入口层在 CPU 侧完成候选累加、路径筛选，构造新的 beam 作为下一步的请求再次提交给引擎，如此循环直至达到最大步数或所有 beam 完成。

Beam search 的状态管理涉及三个轻量 dataclass（BeamSearchSequence、BeamSearchOutput、BeamSearchInstance），完全在入口层维护，引擎内核对其无感知。

4.3.2 KV Cache 管理

每条 beam 对引擎内核而言是完全独立的请求，引擎内核无法感知它们之间的 KV Cache 共享关系。每条 beam 在入口层将原始 prompt 与当前已生成的 token 序列重新组装为完整输入，作为一条新请求提交，引擎内核不知道这些请求共享同一个 prompt 前缀。因此 KV Cache 复用依赖 prefix caching：不同 beam 请求的公共 prompt 前缀在哈希匹配后可以复用已缓存的 KV 块。

被剪枝掉的 beam 只是在入口层从 beam search 管理的活跃列表中移除，其对应的引擎内核请求在返回结果后已因 max_tokens=1 自然结束，KV Cache 由引擎内核按正常请求完成流程释放。

4.3.3 请求调度与批处理

Beam search 请求对调度器（Scheduler）完全透明。入口层在每一步将所有活跃 beam 展平为独立请求并发提交，调度器看到的只是若干条 max_tokens=1 的普通请求，无需为 beam search 预留特殊的调度预算。这些请求被引擎内核按到达顺序正常调度，与其他用户请求混合处理。

4.3.4 Beam 的扩散、评分与剪枝

扩散、评分与剪枝逻辑全部集中在入口层的 CPU 侧，没有利用 GPU 批量计算，仅在 EOS 检测和 top-k 选择时使用 numpy 做 CPU 侧向量化加速。下面以在线 API 为例重点介绍 logprob 得分累积。

每条 beam 返回 2 * beam_width 个候选 token 及其 logprob。入口层遍历所有 beam 的返回结果，将每个候选的 logprob 与其父 beam 的累积 logprob 相加，展平到一个列表中：

all_beams_token_id = []
all_beams_logprob = []
for i, result in enumerate(output):   # 遍历 beam_width 条 beam
    current_beam = all_beams[i]
    logprobs = result.outputs[0].logprobs[0]  # 该 beam 的 top 2*beam_width 候选
    all_beams_token_id.extend(list(logprobs.keys()))
    all_beams_logprob.extend([
        current_beam.cum_logprob + obj.logprob  # 父 beam 累积 logprob + 候选 logprob
        for obj in logprobs.values()
    ])
# 此时 all_beams_logprob 共有 beam_width * 2 * beam_width 个候选

相比 v0.6.2 之前在引擎内核实现的 beam search，搜索空间大幅缩小——从 beam_width × vocab_size 降为 beam_width × (2 × beam_width)，但 logprob 累加部分使用 Python for 循环逐候选处理，计算效率较低。

此外，当前实现也没有实现 early_stopping 策略，仅靠 max_tokens 步数耗尽停止。

4.4 v0.6.2 之后的 Beam Search 实现的部分缺点

当前实现的核心优点是与推理主流程彻底解耦，引擎内核保持简洁。但在性能方面缺点较为明显：

• • 每条 beam 每步都是独立请求：每条 beam 的每一步都需要走完整的请求生命周期（调度、KV Cache 分配、前向推理、输出处理），max_tokens=1 的短请求与长请求承担相同的调度开销，开销被放大了 beam_width × max_tokens 倍。
• • KV Cache 复用依赖 prefix caching：每条 beam 每步都需要通过哈希表查找可复用的前缀块，而非像老版本那样在 fork 时直接通过引用计数共享父 beam 的 block table，查找和匹配的开销更大。
• • logprob 累加在 CPU 侧用 Python 循环完成：搜索空间虽然从 beam_width × vocab_size 缩小为 beam_width × (2 × beam_width)，但累加和排序均在 CPU 侧以 Python 循环或 numpy 完成，无法利用 GPU 并行。
• • 每步大量对象创建和 prompt 重建：每步每个候选都需要新建 BeamSearchSequence 对象，并通过 get_prompt() 将完整 token 序列重新拼接为引擎输入，涉及列表拷贝和内存分配。

4.5 附：vLLM 社区新旧版本 Beam Search 的讨论

• • vllm GitHub Issue #6226: https://github.com/vllm-project/vllm/issues/6226
• • vllm GitHub Issue #8306: https://github.com/vllm-project/vllm/issues/8306

5. Transformers 的 Beam Search 实现方案简介

Transformers 的 Beam Search 实现集中在 src/transformers/generation/utils.py 的 GenerationMixin 类中，核心方法是 _beam_search()。与 vLLM 不同，Transformers 的 Beam Search 是纯张量操作的实现——所有 beam 的扩散、评分、剪枝都通过 PyTorch 张量运算完成，没有面向对象的序列管理，也没有 HTTP 层的编排开销。
和 vLLM 老版本的实现一样，支持 early_stop 策略，另外在实现细节上采用面向数据设计，大量的使用 GPU 加速。譬如：

• • 使用 tensor 一次性计算整个词表的 log_probs 累加：

log_probs = log_probs + running_beam_scores[:, :, None]

• • 一次性构造下一轮的子序列

topk_current_beam_indices = topk_indices // vocab_size
topk_running_beam_indices = self._gather_beams(running_beam_indices, topk_current_beam_indices)
topk_running_sequences = self._gather_beams(running_sequences, topk_current_beam_indices)
topk_ids = topk_indices % vocab_size

# Update sequences for the K top-k new sequences.
topk_running_sequences[:, :, cur_len] = topk_ids

在 beam search 的扩散、评分、剪枝等算法逻辑上，Transformers 的纯张量实现比 vLLM 的两个版本都更高效。但纯张量操作的实现也付出了可读性的代价——一大批平铺的 tensor 没有层次感，相比 vLLM 新版本基于 Python 对象的直白编排，阅读门槛更高。另外，在 prefill 阶段 Transformers 为了代码规整，会一次性推理 beam_width 个重复的请求，这有性能浪费。

6. TensorRT-LLM 的 Beam Search 实现方案简介

6.1 整体思路

与 vLLM 新版本将 Beam Search 剥离到入口层不同，TensorRT-LLM 将 Beam Search 作为一等公民嵌入在推理主流程中，提供了 C++ 和 Python 两套完整实现，默认使用 C++ 实现。两套实现在核心思路（如 cache_indirection 映射表）上是一致的，差异主要体现在与主流程的解耦方式、以及扩散、评分、剪枝的执行形式上。

C++ 版采用策略模式：DecodingLayer 在构造时根据 DecodingMode 一次性选择 BeamSearchLayer 或其他采样、投机 layer 实例，beam search 与普通采样在调用栈上是平行的两条路径，互不交织；扩散、打分、剪枝逻辑由专门的 CUDA kernel 实现。而 Python 版没有这一层抽象，beam search 通过 if self._use_beam_search 分支散布在 sampler 的多个方法中，与普通采样流程存在一定程度的交织；扩散、打分、剪枝则通过 PyTorch 算子（如 torch.topk）拼装实现。

C++ 版将核心计算（logprob 累加、TopK、early stopping 判断）拆成连续的多个 CUDA kernel 顺序发射完成，核心计算流水内部零 GPU-CPU 同步，并针对不同 beam width 有定制的性能优化，代码较为复杂，这里不展开介绍。下面仅基于 Python 代码介绍 KV Cache 的特别实现（该实现思路与 C++ 版一致）。

6.2 Python 版 KV Cache 管理实现

TensorRT-LLM 的 beam search 在 KV Cache 上采用"按 beam 独占 slot + cache_indirection 两级间接寻址"的策略：

• • 索引槽位预留、KV Cache 按需分配：索引表（block offset 表）在 Runtime 启动时就按 [maxBatch, maxBeam, 2, maxBlocksPerSeq] 一次性为每条 beam 预留槽位，generation 阶段（即 decode 阶段）每生成一个新 token，每条 beam 都会将自己的 KV 写入到 block offset 记录的 block 里。另外，block offset 表里记录的 block ID 是按需申请的，不会提前预留真实的显存。
• • 扩散、剪枝不搬运 KV 数据，而是通过 cache_indirection 映射表记录每条 beam 在每个历史时间步继承自哪条 src beam（值是 0~beam_width-1 的整数），读取时再用这个 src beam 去 block offset 表里查到对应的物理 block。beam 扩散时，按"每条新 beam 继承自哪条旧 beam"的关系更新这张映射表。
• • 读取时两级间接寻址：attention kernel 先查 cache_indirection[beam][step] → src_beam，再用 src_beam 去 block offset 表里查到具体的 block_id 和块内偏移，最后从物理 block 里取到 KV。

整套机制其实是 Block（物理显存）、Block Offset 表（beam→block 的索引）、Cache Indirection 表（beam→beam 的重映射） 三件东西的分工与联动——天然把"写、扩散、读"三种操作解耦到了不同的数据结构上：写只关心 block offset 表、扩散只关心 cache_indirection 表、读则同时穿过两张表完成两级寻址。

总的来说，这种设计通过增加一层映射，避免了早期 vLLM 在 block size > 1 时、被复用 KV Cache 的末尾未填满 block 的 CoW。代价主要有两点：

• • 显存占用偏高：block offset 里已申请的 block，即便对应 beam 的历史已被 cache_indirection 改写为继承自其他 beam（其自身存的 KV 已"过时"无用），也无法提前释放，必须等整个请求结束才统一回收。
• • 多一张表的索引开销：相比 vLLM 只有一张 block table，这里在 block table 之上还多了一张 cache_indirection 表，每个 generation step 都要重写一次（对应 vLLM 在 beam fork 时才重构 block table，TRT-LLM 重写频率更高）；attention kernel 读 KV 时也要多走一级寻址。

7. SGLang 的实现方案

SGLang 原生不支持 Beam Search。本章介绍如何在 SGLang 的核心推理流程上添加 Beam Search 能力，同时尽量避免污染普通推理路径。整体思路是：继承 vLLM 老版本的"一等公民"设计，但把面向对象的 Sequence 管理替换为面向数据的 tensor + 轻量 dataclass 组合，并且从调度器到 KV Cache 分配、batch 过滤、回包四条链路上，都用 mixin 把 beam search 分支隔离到独立文件。

7.1 SGLang 的整体架构

SGLang 的请求从接入到返回经过四个进程或线程：

• • TokenizerManager：在 server 进程内，负责把 HTTP 请求的文本 tokenize 成 token id，组装 GenerateReqInput 通过 ZMQ 送给 Scheduler；并从 DetokenizerManager 接收回包，按 rid 匹配到对应的 async generator 返回给 HTTP handler。
• • Scheduler：独立进程，维护 waiting_queue 和 running_batch，每一步从 waiting queue 调度请求进入 ScheduleBatch，驱动 TpWorker 跑 prefill 或 decode，然后在 scheduler_output_processor_mixin 里根据 sampler 的输出更新请求状态，通过 stream_output 把 BatchTokenIDOutput 送给 DetokenizerManager。
• • TpWorker、ModelRunner、Sampler：执行前向 + 采样，返回 logits 或采样后的 token id。
• • DetokenizerManager：独立进程，把 token id 转成文本，再转给 TokenizerManager 回到 HTTP 层。

几个关键数据结构与 beam search 直接相关：

• • ScheduleBatch：本轮要推理的一批请求，持有 req_pool_indices（每条 seq 在 req_to_token_pool 中占一行，行里记录该 seq 的每个 token 在 token_to_kv_pool 中的物理槽位）、seq_lens、out_cache_loc 等一组平铺 tensor。SGLang 要求同一个 batch 内所有请求在这些 tensor 里是连续的，因此删除请求时必须对 batch 做紧凑整理。
• • Req：单条请求的完整状态对象，持有 origin_input_ids、output_ids、req_pool_idx、采样参数、stop conditions 等。
• • req_to_token_pool、token_to_kv_pool：前者是 [max_num_reqs, max_context_len] 的行式索引表，每行属于一个 seq；后者是底层的物理 KV Cache 池。Beam search 的多条 beam 各占 req_to_token_pool 的一行，但这些行会通过复制的方式共享同一段 prefix 物理 KV。
• • Sampler：普通路径下直接返回 next_token_ids；beam search 模式下绕过采样，只把 log_softmax(logits) 塞进 logits_output.logprobs 供上层做 top-k 与剪枝。

7.2 Beam Search 需要考虑的关键问题

对照第 3 章提出的四个通用挑战，结合 SGLang 的架构特点，我们需要回答以下四个具体问题：

1. 1. 与推理引擎主流程的解耦：SGLang 的 prepare_for_decode、filter_batch、process_batch_result，以及 TokenizerManager、DetokenizerManager、OpenAI 兼容层，都是所有请求的共同路径。beam search 涉及"1 条请求 → beam_width 条 seq"的结构突变，以及"1 条请求 → beam_width 条结果 + beam_score"的回包差异，如何避免这些分支逻辑污染普通推理路径？
2. 2. KV Cache 管理：prefill 后，1 条请求要"原地"扩散成 beam_width 条 seq，prefix 部分的物理 KV block 应当被所有 beam 共享；decode 阶段每步剪枝后，又需要把幸存 beam 的完整 KV 继承到新 beam 位置，并释放被剪掉的 beam 独占的 block。如何在 req_to_token_pool + token_to_kv_pool 这套体系上高效完成"共享 prefix、继承幸存 beam 的 KV、回收被剪枝 beam 独占的槽位"？
3. 3. 请求调度与批处理：一个 beam_width 的请求会在 batch.req_pool_indices 里占 beam_width 个连续槽位，且 SGLang 要求同 batch 内请求在平铺 tensor 里连续。如何在 waiting → running、batch 过滤等调度动作中以"请求粒度"维护这些连续槽位？
4. 4. Beam 的扩散、评分与剪枝：每步前向得到 [beam_width, vocab_size] 的 logprobs，需要累加父 beam 的 cum_logprob 选出 top beam_width 存活路径，同时处理 EOS、stop string、stop regex、max_new_tokens 多种终止条件。如何尽量做成 GPU 向量化？

下面 7.3～7.6 四节依次对应上述四个挑战，给出 SGLang 的具体实现。

使用边界：同一批次内允许不同 beam_width 的 beam search 请求共存；但在服务级别上，结合之前和社区开发者的讨论，beam search 通过启动参数开启，开启后该实例只接受 beam search 请求。

7.3 与主流程的解耦

SGLang 对"解耦"的回答由三部分组成：代码层面用 mixin 注入、对外协议复用 OpenAI 的 n、跨进程协议在 BatchTokenIDOutput 上扩一个字段——共同保证普通请求路径几乎不被影响。

7.3.1 mixin 文件分工与 if 分支入口

最终的实现以"mixin + 分支 if"的形式落地，所有新增代码都集中在独立文件中，普通推理路径几乎不受影响。新增的核心文件如下：

文件	职责
managers/beam_search_type.py	定义 BeamSearchSequence、BeamSearchList 两个数据结构
managers/schedule_batch_beam_search_mixin.py	给 ScheduleBatch 和 Req 注入 beam 相关的 batch 整理、请求初始化逻辑
managers/scheduler_beam_search_processor_mixin.py	给 Scheduler 注入 prefill、decode 后处理，beam 扩散、剪枝，以及 KV Cache 复制与释放
managers/beam_search_tokenizer_manager_mixin.py	在 TokenizerManager 侧把 beam 结果打平回标准回包结构
managers/beam_search_detokenizer_mixin.py	在 DetokenizerManager 侧把每条 beam 的 token 序列批量 decode
entrypoints/openai/openai_beam_search_mixin.py	在 OpenAI 兼容层把 beam_results 展开为多条 choices

Scheduler 主流程里仅新增了少量一行的 if is_beam_search: 分支作为入口，例如：

# managers/scheduler_output_processor_mixin.py
if batch.reqs and batch.reqs[0].is_beam_search:
    self.process_beam_search_prefill_result(batch, logits_output)
    return

# managers/scheduler.py  process_batch_result
if batch.forward_mode.is_decode():
    if batch.reqs and batch.reqs[0].is_beam_search:
        self.process_beam_search_decode_result(batch, result)
    else:
        self.process_batch_result_decode(batch, result)

这样主流程的代码改动被压缩到几行 if 分支，beam search 的复杂逻辑全部位于 mixin 中。

另外，为了避免功能组合带来的维护负担，server 启动时若开启 --enable-beam-search，会自动关掉 PD separation、pipeline parallelism、overlap schedule 等还未实现的功能（见 server_args._handle_beam_search）。

7.3.2 对外协议：复用 n 与 beam_score

对外接口复用 OpenAI 兼容的 n 参数表示 beam width，避免引入新字段：

• • 启动：python -m sglang.launch --model-path <model> --enable-beam-search
• • 请求：n > 1 的请求会被当作 beam search 请求（受全局 enable_beam_search 控制）
• • GenerateReqInput._handle_beam_search_parallel_sampling 会在 beam search 模式下把 parallel_sample_num 强制改为 1，这样 beam search 请求不会再被入口层像 best_of 那样展开为 n 条独立请求；n 的语义在这里重定义为 beam_width。
• • 返回：每条请求返回 n 条 choices，按 beam_score 从高到低排序；每条 choice 带一个 sglext.sequence_score 字段（就是 beam_score）。
• • 限制：beam search 模式下 return_logprob、speculative decoding、grammar 约束被自动关闭或拒绝。

7.3.3 模块间协议：BatchTokenIDOutput 扩展

为了让 beam search 沿用 SGLang 现有的跨进程通信管道（ZMQ + pickle），我们在 BatchTokenIDOutput、BatchStrOutput 上各加了一个字段：

# managers/io_struct.py
beam_search_output: List[BeamSearchOutput] = None

BeamSearchOutput 里只装一个 sequences: List[BeamSearchSequence]。这条路径和普通请求的回包路径完全并行，上下游只要看到 beam_search_output 非空，就走多结果分支，普通路径的字段（text、output_ids、meta_info）仍然用第一条 beam 的内容填充，保持后向兼容。

7.4 KV Cache 管理

KV Cache 管理思路，整体总计为三句话：扩散时多条 beam 共享同一批物理 block 索引；decode 时每步新 token 写在新分配的 block 上；剪枝时把"上一轮索引集合"和"本轮幸存索引集合"做差集，就得到可以释放的物理 block。

7.4.1 扩散时的 prefix KV 共享

扩散发生在从 prefill 过渡到第一轮 decode 时的 _prepare_for_new_beam_search：

beam_req_pool_indices = self.req_to_token_pool.alloc_by_count(total_slots)   # 一次性申请 beam_width 行
...
for req in new_reqs:
    normal_idx = self.req_pool_indices[skip_idx + i : skip_idx + i + 1]
    seq_len    = self.seq_lens[skip_idx + i : skip_idx + i + 1].squeeze()
    beam_indices   = beam_req_pool_indices[beam_start:beam_end]
    normal_kvcache = req_pool[normal_idx, :seq_len].squeeze(0)
    req_pool[beam_indices, :seq_len] = normal_kvcache   # 并行广播复制
    new_req_pool_indices_list.append(beam_indices)
    ...
    req.beam_list.batch_slot_start_idx = current_idx

几个关键点：

• • 复制 block 索引实现 KV Cache 复用。物理 KV block 仍然由最初 prefill 的那行所拥有，这里只是让 beam_width 条 beam 的 req_to_token_pool 行指向同一批物理 block 索引。
• • 原 prefill 槽位 req.req_pool_idx 不变，只有 batch.req_pool_indices 被替换。这样请求最终完成时走同一条 release_kv_cache(req, tree_cache) 路径把 prefill KV 释放给 radix cache，和普通请求保持一致。

7.4.2 decode 时的 beam 继承与释放

每一步 decode 之后，先把上一轮所有 beam 在 decode 段（[prompt_len, seq_len)）用过的 KV 索引集合记下来，再把幸存 beam 的同一段 KV 复制到本轮新分配的槽位上、并记下本轮真正读到的索引集合；这两个集合合起来做一次去重，只出现在上一轮里的索引就是被剪枝 beam 留下的"孤儿" block，直接释放即可。整个判断在 GPU 上一次张量运算完成，路径里没有 per-seq 的引用计数维护，也没有 Python 循环里逐 block 的 free。

# last_beam_kv_indices：上一轮 beam_width 条 beam decode 段用过的 KV 索引（合并去重）
# keep_kv_indices   ：本轮幸存 beam 在复制时真正读到的源索引（合并去重）
uniques, counts = torch.unique(
    torch.cat([last_beam_kv_indices, keep_kv_indices]), return_counts=True
)
free_kv_indices = uniques[counts == 1]           # 只出现在上一轮 = 被剪枝 beam 的孤儿 block
self.token_to_kv_pool_allocator.free(free_kv_indices)

7.4.3 请求完成时的 KV 清理

_cache_finished_beam_search 负责把已完成 beam search 请求的 KV 一次性释放：

beam_decode_kv_indices, beam_pool_indices = self._collect_beam_req_decode_kv_indices(batch, finished_reqs)
self.token_to_kv_pool_allocator.free(beam_decode_kv_indices)     # 释放 decode 段物理 block
self.req_to_token_pool.free_by_indices(beam_pool_indices.tolist())  # 释放 req_to_token 行
for req in finished_reqs:
    release_kv_cache(req, self.tree_cache)                        # 通过 req.req_pool_idx 把 prefill 段交给 radix cache

注意只释放 decode 段（[prompt_len, seq_len)）的物理 block，prefill 段交给 tree_cache，由 radix cache 根据 prefix 共享策略决定是否保留——这就天然让下一个带相同 prompt 的 beam search 请求复用到同一段 prefix KV。

7.5 请求调度与批处理

请求的调度和批处理主要是维护：beam search 请求在 batch 里是"每请求 beam_width 个连续槽位"的块状结构，这要求在"从 prefill 过渡到 decode"和"过滤完成请求"这两个时机做请求粒度的槽位管理。

7.5.1 prepare_for_beam_search_decode：从 prefill 到 decode 的 batch 形状突变

Prefill 阶段一条 beam search 请求在 batch 里仍然只占 1 个槽位，和普通请求没有区别；真正的"物理扩散"发生在进入第一轮 decode 之前的 prepare_for_beam_search_decode——这一步把请求在 batch 里的占位从 1 个扩成 beam_width 个，并把 prefill 阶段算好的 KV 一次性广播复制给这 beam_width 条 beam。这是整个 beam search 流程里 batch 形状唯一一次发生突变的时机：扩散完成后，下游的 ModelRunner、Sampler、attention kernel 都把它当成普通 batch 去跑，不需要任何特殊分支。每条请求在这时记下自己在 batch 中的起始偏移，后续的过滤和剪枝都靠这个偏移定位。

7.5.2 filter_beam_search_batch：请求粒度的槽位重排

请求完成、被抢占或切换分块 prefill 时，batch 需要把不再参与 decode 的请求从各个 tensor 里剔除。普通请求一条 seq 只占一个槽位，按请求下标 slice 就行；beam search 下一条请求占 beam_width 个连续槽位，必须按"请求粒度"整块保留或整块剔除，并在重排后同步更新每条存活请求记录的批内起始偏移——否则下一轮剪枝就会按错位的偏移去定位 beam，整个 KV 继承逻辑都会崩。另外 beam search 模式不走 sampling_info / spec_info，所以这里也没有与之对应的 filter 负担。

7.6 Beam 扩散、评分与剪枝：如何高效挑 top beam_width

每一步 decode 之后都要回答同一个问题：给定 beam_width 条存活 beam，以及它们在最新位置上的 logits，如何最快选出下一步的 beam_width 条新存活 beam？朴素做法是对每个请求的 beam_width × vocab_size 打分矩阵单独做一次 topk——vocab_size 动辄几万，而且 N 个请求就有 N 次 kernel 启动，GPU 规模和 launch 粒度都吃亏。SGLang 围绕两个点把这件事压下去：把搜索空间在源头就砍到很小，把 kernel 启动粒度从每请求拉到整个 batch。

候选空间的主动收窄。不在 vocab_size 维度做 topk，而是每条父 beam 只保留自己的 top-2*beam_width 个候选 token，把搜索空间从 beam_width × vocab_size 一步砍到 beam_width × 2*beam_width，后续所有打分、排序、分支判断都跑在这个极小矩阵上。2*beam_width 的冗余是刻意留的缓冲——哪怕一半候选被 EOS / stop token 剪掉，也够凑齐 beam_width 条存活 beam。代价是候选池始终比"用完整 vocab_size logprob 加历史 cum_logprob 再做全局 topk"（Transformers 的做法）小一些：极端长尾分布下有概率漏掉本应胜出的父子对，但在主流生成任务上几乎观察不到质量差异。

一次 GPU kernel 服务整 batch。整个 batch 所有请求的 top-2*beam_width 在同一次 kernel 里算完，再到 CPU 侧按请求切片。同一 batch 里不同请求的 beam_width 可以不一样，所以 topk 取全 batch 的 max_k、并且必须 sorted=True——这样每个请求在自己那行的前 2*beam_width 个位置就是它真正的 top，直接切片即可。进入逐请求循环后，还有一次把父 beam 累计 logprob 广播到候选上、再做 flat topk 选 2*beam_width 个 (parent, token) 父子对——这次张量规模极小，代价可以忽略。整个 decode 步骤里真正值得关注的 GPU 开销就是这两次 topk：

# ① 整个 batch 一次 topk：logprobs [N_total_beams, vocab_size] → [N_total_beams, max_k]
max_k = max(req.beam_candidates for req in batch.reqs)   # beam_candidates = 2 * beam_width
top_tokens_all, top_logprobs_all = logprobs.topk(max_k, dim=1, sorted=True)

# ② 逐请求：父 beam cum_logprob 广播相加 + flat topk，拿到 2*beam_width 个 (parent, token)
for req in batch.reqs:
    top_tokens    = top_tokens_all   [start_idx:end_idx, :req.beam_candidates]  # [bw, 2*bw]
    top_logprobs  = top_logprobs_all [start_idx:end_idx, :req.beam_candidates]

    all_cum_logprobs = req.beam_list.cum_logprobs.unsqueeze(1) + top_logprobs   # [bw, 2*bw]
    topk_values, topk_indices = torch.topk(all_cum_logprobs.flatten(),
                                           k=req.beam_candidates, largest=True)
    parent_idx = topk_indices // req.beam_candidates       # 候选来自哪条父 beam
    token_ids  = top_tokens.flatten()[topk_indices]        # 候选 token
    # topk_values（新 cum_logprob）随候选一起交给下游 —— 更新存活 beam 的累计打分
    # EOS / stop 分流：一次 torch.isin 批量判定 + CPU list 分组，不做 per-beam 循环

7.7 实现细节中的精益求精：tensor 化 per-beam 循环

这一节专门讲"怎么做得更省"。核心思想：把 per-beam 的 Python 循环替换成一次 GPU 批量操作——不论是 topk、集合差集，还是变长范围收集。

7.7.1 一次大 topk 服务整个 batch（_extract_beam_topk_data）

同 batch 里不同请求的 beam_width 可能不等，对应的 beam_candidates = 2 * beam_width 也不同。取 batch 内的最大值一次性算完，每条请求再按自己的 beam_candidates 切片即可：

max_k = max(req.beam_candidates for req in batch.reqs)
top_logprobs, top_tokens = torch.topk(logprobs, k=max_k, dim=1, sorted=True)

sorted=True 让切片后天然保持降序，不再二次排序。对比 vLLM 老版本"每个 SequenceGroup 单独发一次 topk"，kernel 启动开销从 O(len(batch.reqs)) 降到 O(1)。

7.7.2 用集合差集替代引用计数（_handle_beam_kv_cache）

vLLM 老版本给每个 block 维护 ref_count：fork 时 ++、释放时 --、归零才 free。这里换成把"上一轮用过"和"本轮要保留"两个集合并起来做一次 unique：

last_beam_kv_indices = self._batch_collect_range_kv_indices(
    beam_req_pool_indices, beam_req_seq_lens, batch.device, prompt_lens
)  # 上一轮所有 beam 在 [prompt_len, seq_len) 范围内用过的 block
keep_kv_indices = self._copy_kvcache_for_beams(...)  # 本轮幸存 beam 真正要读的 block

uniques, counts = torch.unique(
    torch.cat([last_beam_kv_indices, keep_kv_indices]), return_counts=True
)
free_kv_indices = uniques[counts == 1]   # 只在 last 里出现 = 已被剪枝，可释放
self.token_to_kv_pool_allocator.free(free_kv_indices)

counts == 1 把 per-block 引用计数压缩成一行 GPU 运算。外层 free_group_begin/end 再把整批请求的 free 合并成一次 allocator 调用。

7.7.3 变长范围的向量化收集（_batch_collect_range_kv_indices）

每条 beam 要收集 [prompt_len_i, seq_len_i) 范围内的 KV 索引，各自长度不同。一般写法是 Python for 循环逐条 slice 再 cat，这里改成纯 tensor：

max_range_len = (seq_lens - prefix_lens).max().item()
position_indices = torch.arange(max_range_len, device=device).unsqueeze(0) \
                 + prefix_lens.unsqueeze(1)        # [num_reqs, max_range_len]
mask = position_indices < seq_lens.unsqueeze(1)    # 越界位置 mask 掉
batch_kv_indices = self.req_to_token_pool.req_to_token[
    pool_indices.unsqueeze(1), position_indices    # 二维高级索引一次取齐
]
return batch_kv_indices[mask].unique()

思路是"对齐到最长 + mask"：所有范围对齐到 max_range_len，一次二维高级索引把整张表读出来，越界位置用 mask 筛掉。末尾 .unique() 合并跨 beam 的重复索引（刚扩散出来的兄弟 beam 通常共享一大段 prefix），顺带缩小下游运算的输入规模。这段同时被 _handle_beam_kv_cache（剪枝释放）和 _collect_beam_req_decode_kv_indices（请求完成清理）复用。

7.7.4 快慢双路径（_copy_kvcache_for_beams）

每步 decode 后要把幸存父 beam 的 [prompt_len, seq_len) 拷到子 beam 行。同一请求内所有 beam 长度一致，batch 里常见的情况就是"所有 beam 同长度"，所以分了快慢两条路径：

if len(prompt_lens.unique()) == 1 and len(seq_lens_batch.unique()) == 1:
    # 快路径：所有 beam 同长度，直接一次 _copy_kvcache_group 完事
    return self._copy_kvcache_group(src, dst, prompt_len, seq_len)
else:
    # 慢路径：按 (prompt_len, seq_len) 分组，组内还是批量
    copy_groups = {}
    for beam_idx, key in enumerate(zip(prompt_lens_cpu, seq_lens_cpu)):
        copy_groups.setdefault(key, []).append(beam_idx)
    ...

快路径的判断全程在 GPU 上（unique() + len()），直到 prompt_lens[0].item() 才同步一次。慢路径按 (prompt_len, seq_len) 分组，每组内仍是单次 GPU copy，不退化到 per-beam 循环。这里还故意允许 src == dst 的冗余读——少数 beam 的 src 其实就是 dst，本可以跳过，但让它们参与批量 copy 比加分支判断更快，是 throughput 优先的取舍。

7.7.5 用 return_inverse 去重 src（_copy_kvcache_group）

beam_width * 2 个父子候选里，常有多条来自同一条 parent beam（一条强势 parent 贡献好几个 token），它们的 src_indices 重复。直接按原序读会把同一行读好几遍，这里先去重再广播回去：

unique_src_indices, inverse_indices = torch.unique(
    src_indices, return_inverse=True
)
kvcache_batch_unique = self.req_to_token_pool.req_to_token[
    unique_src_indices, prefix_len:seq_len
].clone()
kvcache_batch = kvcache_batch_unique[inverse_indices]   # 用 inverse 还原成原始顺序
self.req_to_token_pool.req_to_token[dst_indices, prefix_len:seq_len] = kvcache_batch
return kvcache_batch_unique.flatten().unique()

这里使用 return_inverse=True：去重后读一次，再用 inverse_indices 广播回每个原始位置——beam_width 越大收益越明显。返回的 kvcache_batch_unique.flatten().unique() 就是本步真正读到的物理 block 集合，作为 keep_kv_indices 交给 §7.7.2 参与释放判断。对比 vLLM 老版本 for seq in beams: copy_block_table(...) 一条 beam 一次调用，这里整组合并成一次 kernel。

8. 设计思想总结——面向数据设计

对比前面介绍的几个方案，SGLang 的 beam search 有一条贯穿始终的设计原则：面向数据设计（Data-Oriented Design）——围绕"数据怎么被访问"来组织数据布局，而不是围绕"什么是一个 beam"来构造对象抽象。这一原则在两个点上落地：

1. 状态按访问模式分层。
参与每步打分与剪枝的数值字段（累计 logprob、上一轮产出的 token、prompt 长度等）集中以平铺 tensor 存放，打分与剪枝退化成一次广播加法 + 一次 topk；历史 tokens、finish_reason、文本这类不参与数值计算的字段，才落回 Python 容器里按需处理。既避免了"全部塞进对象、每步都要遍历"，也避免了"全部塞进 tensor、失去可读性"。

2. 关系也是数据，而不是指针。
beam 之间的 KV 复用、幸存 beam 的继承、被剪枝 beam 的回收，全部转化为对索引表的批量运算：集合差识别不再使用的槽位，批量索引运算完成从幸存 beam 继承到新 beam。不需要维护每个 block 的引用计数，也不需要在 Python 循环里逐 block 释放。

这两条共同带来的效果是：整条 decode 路径的每一步都能走批量路径——大批量 GPU 打分与剪枝、批量回包、批量 KV 复制与释放，只有剪枝后的少量候选才回落到 CPU 侧处理。没有"对每条 beam 逐个 deepcopy、逐个 append_token、逐个复制 block_table"这类逐对象密集操作。

本文所在：https://www.cnblogs.com/cswuyg/p/19934085