原理Beam Search束搜索是一种在序列生成任务中常用的搜索算法用于在每一步生成时平衡计算效率和结果质量文字。主要作用减少搜索空间相对于贪心搜索而言在序列生成如机器翻译、文本生成、语音识别时如果每一步都保留所有可能的候选词穷举搜索计算量将指数增长。Beam Search 在每一步只保留 top-kk 称为 beam size个最有可能的候选序列大幅降低计算负担。近似最优解它并不是穷举所有路径那样保证全局最优但计算不可行而是用局部贪心扩展结合一定宽度搜索得到一个接近最优的序列。平衡生成质量与速度Beam Size 越小速度越快但可能错过更好的序列。Beam Size 越大质量通常更好更接近穷举搜索但计算越慢。工作流程假设 beam_size 2生成句子从起始符开始第一步模型输出每个词的概率保留概率最高的 2 个词如 A, B。对于这 2 个候选词分别预测下一个词的概率这时会得到 2 × V 种可能V 是词表大小。从这 2 × V 个候选序列长度为 2中选择总概率或对数概率之和最高的 2 个继续扩展。重复直到遇到结束符最后从 beam 中选择分数最高的序列输出。vLLM 实现beamsearch 可视化https://huggingface.co/spaces/m-ric/beam_search_visualizerif__name____main__:llmLLM(modelfacebook/opt-125m)paramsBeamSearchParams(beam_width2,max_tokens50)prompts[TextPrompt(promptThe future of artificial intelligence)]outputsllm.beam_search(prompts,params)foroutputinoutputs:generated_textoutput.sequences[0].textprint(fGenerated text:{generated_text!r})参数含义max_tokens: 迭代次数for range(max_token)beam_width: 束宽度。有下面两个作用1. 针对每个输入限制模型推理保留 top 2*beam_width 的 token2. 控制每次迭代需要参与推理的 token 数量自回归阶段 promot [“The future of artificial intelligence”]itemprefilldecodeinput_ids (token_id)[2, 133, 499, 9, 7350, 2316, 0, 0][13]positions[0, 1, 2, 3, 4, 5, 0, 0]logits_indices50hidden_states model_output tensor(8, 768) model_output tensor(1, 768)sample_hidden_states hidden_states[logits_indices] tensor(1, 768) hidden_states[logits_indices] tensor(1, 768)logitstensor(1, 50257)tensor(1, 50257)GPUModelRunner._update_states()更新后的状态input_batches会被 _prepare_inputs 函数使用更新 input_ids以创建模型所需的GPU张量输入。这是GPU模型运行器中状态管理和输入准备的关键连接点。未使用 beam_search : sampling_metadata.max_num_logprobs None使用 beam_search: sampling_metadata.max_num_logprobs 10