原理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