Compute logits and past_key_values once for the initial context rather than `num_beams` times?

快速结论：这个问题不是一个报错，而是一个性能优化特性请求（feature request）。在 Hugging Face Transformers 库的 `generate()` 方法中，对于 batch size 为 `num_beams` 或 `num_return_sequences` 的初始上下文，建议先计算一次 logits 和 past_key_values（即 prefill 阶段），再广播到各个 beam，以消除冗余计算。目前（截至 Issue 关闭时）尚未合入主分支，属于待实现优化。

问题场景

用户在使用 Hugging Face transformers 库的 generate() 方法进行文本生成时，调用了以下解码策略之一：sample、beam search、beam sample、group beam search。当初始上下文（prefix/prompt）很长（例如数千 tokens）时，模型会在 prefill 阶段将输入 batch 立即扩展到 num_beams 或 num_return_sequences，导致同一个初始上下文被重复推理多次，产生显著开销。

报错原文

这不是一个报错（error），而是性能优化相关的功能请求。核心问题描述如下：

It looks like (at least) the following methods:
- sample
- beam search
- beam sample
- group beam search

all immediately expand their input batch size to either `num_beams` or `num_return_sequences`. This means that the initial context, which can be very long, will be run through the model `num_beams`/`num_return_sequences` times, which seems to be unnecessary and causes a large overhead, especially when the initial context is long.

原因分析

该问题的核心原因是 generate() 方法内部的实现顺序：在预填充（prefill）阶段之前，输入就已经根据 num_beams 或 num_return_sequences 进行了 batch 维度扩展。这导致模型对 N 份完全相同的初始上下文重复计算了 N 次 logits 和 past_key_values。理论上，只需计算一次，然后将输出广播（broadcast）到各个 beam，即可省去冗余计算。

Hugging Face 核心贡献者 @gante 确认这是一个长期期望的优化点，并指出可以在 GenerationMixin 中创建一个 prefill 函数，在进入不同的生成模式之前进行处理，从而在不破坏公共 API 的前提下实现加速。

环境排查

该问题与特定环境版本无关，但建议在尝试复现性能差异时确认以下环境信息：

Hugging Face transformers 库版本（建议 >= 4.30，且确认生成逻辑未发生破坏性变更）
torch 版本（建议 >= 2.0）
CUDA 版本及 GPU 型号（影响推理速度对比）
num_beams 和 num_return_sequences 的取值（数值越大，冗余计算越明显）
初始上下文的 token 长度（越长，prefill 开销占比越大）

解决步骤

注意：截至 Issue 关闭时，这是一个待实现的 feature，没有官方已合并的补丁。以下步骤基于 Issue 讨论中的推测性解决方案，请酌情测试。

在 GenerationMixin 中实现 _prefill() 函数：将初始上下文的推理单独提取出来，在 batch 扩展之前先计算一次 logits 和 past_key_values。
调整生成入口逻辑：在 generate() 方法中，在调用 # 10. go into different generation modes 步骤之前，先调用 _prefill()，然后将预填充的输出根据 num_beams 或 num_return_sequences 进行广播。
注意缓存兼容性：由于缓存（cache）现在支持 batch 扩展，预填充的输出需要能正确地按 beam 数展开到后续的自回归生成中。
可选：在本地临时 hack：如果希望测试效果，可以参考 Issue 原作者提供的 hack 示例（位于 Issue 评论中，非正式 PR），在本地对 generation_utils.py 进行修改，将 batch 扩展操作后移到 prefill 之后。

验证方法

可使用性能测试脚本来对比优化前后的推理速度。参照 Issue 中的测试脚本，关键指标为多次生成的平均推理时间（elapsed time）：

start = time.time()
for _ in range(10):
    outputs = model.generate(**inputs, max_new_tokens=15, do_sample=False, num_beams=6, num_return_sequences=6)
print("elapsed", (time.time()-start)/10)

如果优化生效，在长前缀 + 多 beam 场景下，应观察到明显的延时下降。例如 Issue 中原作者的测试结果显示优化前为 4.59 秒，优化后有望显著降低。