TensorRT LLM--Beam Search

Beam Search思想介绍：如何通俗的理解beam search？ - 知乎 (zhihu.com)

LLM里的Beam Search

在模型解码过程中，模型是根据前一个结果继续预测后边的，依次推理，此时为了生成完整的句子，需要融合多个step的输出，目标就是使得输出序列的每一步的条件概率相乘最大。

最直接的方法就是贪心算法（greedy search），每步取概率最大的输出，然后将从开始到当前步的输出作为输入，取预测下一步，直到句子结束。如下图所示，第一步先去概率最大的A,依次是B,C。这种的缺点就是如果中间某一步取值不对，将影响后续的结果。

beam search对贪心算法做了优化，在每个step取beam num个最优的tokens。下图beam num=2，第一步先选取最优的2个值{A,C}；在第二步，A和C作为输入，继续预测，则有10个可能得结果，然后再取其中最大的2个，作为输入，进行第三步，取结果中概率最大的2个,得到2个结果，然后再在2个结果中取最优的一个作为输出。

TensorRT LLM中对Beam Search的支持

GPTAttention算子支持波束搜索，在上下文阶段，每个输入序列计算单个波束。在生成阶段，MHA/MQA/GQA内核使用额外的张量来重建每个波束的正确路径。