Beam Search 可视化文档
1. Beam Search 概述
Beam Search 是大语言模型推理阶段常用的解码策略。与贪心搜索每步只保留 1 个最优 token 不同，beam search 每步保留 beam width（本图 = 2）条得分最高的候选路径，在更大的搜索空间中寻找全局更优的输出序列。
- 每个 decode step，模型对所有活跃 beam 生成下一个 token 的概率分布
- 将所有 beam × vocab 的得分排序，选取全局 top-K 路径继续
- 这意味着两条 beam 可能选择同一个父节点（共享祖先）或各自独立扩展
2. 三种 Token 排列方法
当多条 beam 并行解码时，GPU 上需要将所有 token 组织成序列送入 attention 计算。如何排列这些 token，直接决定了 attention mask 的形状和 KV cache 的管理方式。
Method A Incremental Append（按生成顺序追加）
- 所有 beam 的 token 按累积概率降序交替追加到序列末尾
- 序列：[context, s0t0, s0t1, s1t0, s1t1, …]
- 优点：追加即可，无需搬移；缺点：不同 beam 的 token 交织在一起
Method B Beam-Grouped（按 beam 分组）
- 将属于同一 beam 的 token 聚合为连续块
- 序列：[context, beam0_tokens…, beam1_tokens…]
- 每步需根据 beam 的祖先路径重新组织各 beam 的 token 顺序
Method 0 Independent Sequences（独立序列 / Naive 方法）
- 每条 beam 持有完整的独立序列（context 被复制到每条 beam），batch_size 从 1 变为 B
- Beam 0: [context, s0t0, s1t0, …]；Beam 1: [context, s0t1, s1t1, …]
- 每条序列等价于 Method B 中对应 beam 的 token 序列加上相同的 context 前缀
- 优点：mask 始终是标准 causal（下三角），无需自定义 mask；缺点：context 的 KV 被复制 B 份，显存开销大
3. Attention Mask 的作用
并行解码多条 beam 时，必须通过 attention mask 隔离不同 beam 的 token，否则 beam 之间会发生信息泄漏，破坏搜索的独立性。
Method A 的 mask 特点：
- Context 部分（共享）可被所有 token attend
- 被剪枝的路径产生 dead token（灰色），其行列全部 masked
- Mask 矩阵呈不规则的斑块状，GPU 计算效率较低
Method B 的 mask 特点：
- 无 dead token — 每条 beam 只保留自己的活跃路径
- Mask 呈规整的块对角结构：context 块 + beam0 块 + beam1 块
- 结构规整利于 GPU 并行（可用 block-sparse attention 等加速）
Method 0 的 mask 特点：
- 每条 beam 拥有独立的 mask 矩阵（B 个 2D tensor）
- 每个 mask 都是标准的 causal mask（纯下三角矩阵），无任何自定义逻辑
- 最简单的 mask 形式，无需额外实现 beam 隔离
4. 共享祖先与 Beam 剪枝
Beam search 的 top-K 选择是全局的，两条 beam 可能在同一步选择了同一个父 token，此时产生共享祖先。
- Method A：共享祖先 token 的 beam membership 变为 shared（灰色），两条 beam 均可 attend；被抛弃路径上的 token 变为 dead
- Method B：共享祖先在两个 beam 组中各保留一份副本（加粗标记），物理上存两份，但逻辑上代表同一 token
- Method 0：与 Method B 类似，共享祖先在每条 beam 的独立序列中各存一份。beam 切换祖先时需将父 beam 的整条序列（含 context）拷贝到子 beam
5. KV Cache 管理
KV cache 存储了历史 token 的 key/value 张量，是推理的核心开销之一。两种方法对 KV cache 的管理有本质差异：
Method A：
- KV cache 是一整条追加序列，不同 beam 的 KV 交织存储
- 通过 attention mask 隔离不同 beam 的 KV
- Beam 切换祖先时，旧路径的 token 变 dead，其 KV 仍占用 cache 空间但不再被引用（内存浪费）
- 随着步数增加，dead token 的 KV 越积越多，cache 利用率逐步下降
Method B：
- KV cache 按 beam 分段管理，每条 beam 拥有独立的连续 KV 区间
- 无 dead token，cache 中不存在无用 KV，利用率始终 100%
- 但 beam 切换祖先时，需要将父 beam 的 KV 前缀拷贝到子 beam 的 cache 区间（拷贝开销）
- 实际系统（如 vLLM PagedAttention）通过页表映射避免物理拷贝，仅拷贝页表指针
Method 0：
- 每条 beam 独立持有完整的 KV cache（含 context 部分的副本）
- Context 的 KV 被复制了 B 份，显存开销最大（= B × 完整序列长度）
- Beam 切换祖先时同样需要拷贝，且拷贝量更大（包含 context 的 KV）
- 优点：无需自定义 attention kernel，标准 causal attention 即可直接使用
6. 性能对比与 Tradeoff
| Method A 追加式 | Method B 分组式 | Method 0 独立序列 |
|---|
| Token 排列 | 按生成时间交替追加 | 按 beam 分组聚合 | B 条独立完整序列 |
| KV Cache | 单条序列，有 dead KV 浪费 | 按 beam 分段，无浪费 | B 份完整序列（含 context 副本），显存最大 |
| Beam 切换 | 无需拷贝，标记 dead 即可 | 需拷贝/映射父 beam KV | 需拷贝完整序列（含 context） |
| Mask 形状 | 不规则斑块状 | 规整块对角结构 | B 个标准 causal 下三角 |
| GPU 效率 | mask 不规则，计算效率低 | 块结构利于并行加速 | 标准 attention，无额外开销 |
| 实现复杂度 | 简单 | 需要 KV 重组/页表管理 | 最简单，标准推理即可 |
| 显存利用率 | 中（有 dead KV） | 高（无浪费） | 低（context KV × B） |
| 适用场景 | beam width 小、序列短 | 长序列、大 beam width | 原型验证、小 beam width |
实际推理系统（vLLM、TensorRT-LLM、SGLang 等）普遍采用 Method B 变体 + PagedAttention，通过页表共享和 copy-on-write 机制实现高效管理。Method 0 虽然实现最简单，但因 context KV 冗余复制导致显存开销过大，仅适用于小规模场景或原型验证。