蓝迪哥玩转Ai(11)---FPGA本地算力研究：推理加速核心-预填充（Prefill）与解码（Decode）的深度解析与实现

**推理加速核心-预填充Prefill与解码Decode的深度解析与实现**在第7次学习中我们了解了 LLM 推理的基本流程。本章将深入探讨推理过程中的两个核心阶段预填充Prefill和解码Decode它们是影响推理性能的关键因素。理解这两个阶段的原理和优化方法对于提升 LLM 推理效率至关重要。1 预填充Prefill阶段原理与优化预填充Prefill阶段发生在模型接收到完整输入 Prompt 之后但在开始生成第一个输出 token 之前。这个阶段的主要任务是处理输入的 Prompt计算出所有输入 token 的上下文表示并初始化后续解码阶段所需的关键数据结构——KV 缓存Key-Value Cache。计算流程的数学原理与图示说明假设我们有一个批次大小为 batch_size 的输入 Prompt每个 Prompt 的长度为 seq_len, 模型的隐藏层维度为 hidden_dim。输入 Embedding 层的输出形状为 (batch_size, seq_len, hidden_dim)。对于 Transformer 模型的每一层,自注意力机制是核心。对于输入 Hin形状为 (batch_size, seq_len, hidden_dim),我们通过线性变换得到 Query (Q)、Key (K) 和 Value (V) 矩阵。这些权重矩阵 WQ、WK、WV 的形状为 (hidden_dim, num_heads * head_dim)。因此Q、K、V 的形状分别为 (batch_size, seq_len, num_heads * head_dim)。通常会将 num_heads 维度分离出来得到形状为 (batch_size, num_heads, seq_len, head_dim) 的 Q、K、V。接下来计算注意力分数。对于每个注意力头我们将 Query 和 Key 的转置相乘。然后进行缩放和 Softmax 归一化得到注意力权重矩阵。最后将注意力权重与 Value 矩阵相乘得到自注意力层的输出。计算复杂度分析自注意力机制在预填充阶段的时间复杂度主要由注意力分数矩阵的计算和加权求和决定。这两个矩阵乘法的复杂度都是 O(batch_size · num_heads · seq_len² · head_dim)。由于 hidden_dim num_heads · head_dim所以整体复杂度可以表示为 O(batch_size · seq_len² · hidden_dim)。对于长序列的输入这个平方项会使得计算量显著增加成为预填充阶段的性能瓶颈之一。KV 缓存的初始化过程与代码实现在预填充阶段对于 Transformer 模型的每一层模型都会计算得到 Key (K) 和 Value (V) 矩阵。这些 K 和 V 矩阵会被缓存起来这就是所谓的 KV 缓存。对于每个输入 token 和每一层都会存储对应的 Key 和 Value 向量。KV 缓存的形状通常是 (batch_size, num_heads, seq_len, head_dim).在 vLLM 中当您调用 llm.generate(prompt) 时预填充阶段会自动进行 KV 缓存的初始化。vLLM 会在处理输入 Prompt 的过程中计算每一层的 Key 和 Value 向量并将它们存储在 GPU 内存中以供后续解码步骤使用.from vllmimportLLM model_nameQwen/Qwen2.5-7BllmLLM(modelmodel_name)promptThe capital of France is # 当调用 generate 时vLLM 会对 prompt 进行预填充并初始化 KV 缓存outputsllm.generate(prompt,max_tokens5)foroutputinoutputs: generated_textoutput.outputs[0].text print(fGenerated text: {generated_text})在上述代码中当 llm.generate(prompt, max_tokens5) 被调用时vLLM 会首先对 “The capital of France is 这个 Prompt 进行预填充。在此过程中对于 Prompt 中的每个 token“The”, capital”, of, France, isvLLM 会在模型的每一层计算其对应的 Key 和 Value 向量并将这些向量存储到 KV 缓存中。KV 缓存的结构可以想象成一个多维数组它按层、按注意力头、按 token 位置存储了 Key 和 Value 信息。预填充阶段的并行性与性能优化预填充阶段的一个关键优势在于其高度的并行性。由于整个输入 Prompt 在开始时是已知的模型可以同时计算所有 token 在每一层的表示。这使得预填充阶段能够充分利用 GPU 的并行计算能力显著提高处理速度.为了进一步优化预填充阶段的性能主流的技术包括• 高效注意力机制Efficient Attention Mechanisms如 FlashAttention通过重新组织注意力计算过程减少 GPU 内存读写次数显著加速计算尤其在处理长序列时表现出色.• 批处理Batching在实际应用中通常会同时处理多个独立的推理请求。预填充阶段可以将这些请求的 Prompt 组成一个批次进行处理更有效地利用 GPU 的计算资源.2 解码Decode阶段原理与实现解码Decode阶段在预填充阶段完成之后开始。在这个阶段模型以自回归的方式逐个生成输出 token。每生成一个 token该 token 就会被添加到已生成的序列中并作为下一步生成的输入.自回归生成机制详解与可视化解码阶段从预填充阶段处理的输入 Prompt 的最后一个 token 开始目标是生成后续的输出序列。假设预填充阶段处理了 n 个输入 token解码阶段的目标是生成接下来的 m 个输出 token (y₁, y₂, …, yₘ).在每一步 t从 1 到 m模型会基于已经生成的序列 (x₁, …, xₙ, y₁, …, yₜ₋₁) 来预测下一个 token yₜ。对于 Decoder-only 模型在解码的每一步通常只将上一步生成的 token 作为当前 Transformer 层的输入除了第一步输入是预填充的最后一个 token。KV 缓存的查询和更新机制KV 缓存在解码阶段是加速的关键。在预填充阶段我们已经为输入 Prompt 中的所有 token 计算了 Key 和 Value 向量并存储在 KV 缓存中。在解码的每一步假设模型生成了一个新的 token yₜ。为了预测下一个 token yₜ₊₁模型需要计算 yₜ 的 Key 和 Value 向量。然后这个新的 Key 和 Value 向量会被追加到 KV 缓存中扩展缓存的长度.当模型在某一步需要计算自注意力时对于当前要预测的 token yₜ₊₁需要计算其 Query 向量它会与 KV 缓存中所有历史的 Key 向量包括来自原始 Prompt 和之前已生成的 token进行比较计算注意力权重。然后使用这些权重对 KV 缓存中对应的 Value 向量进行加权求和得到上下文信息.以下代码示例展示了使用 vLLM 进行解码的过程。在 llm.generate 的调用中KV 缓存的查询和更新是自动处理的bash from vllmimportLLM model_nameQwen/Qwen2.5-7BllmLLM(modelmodel_name)promptThe weather today is # 指定生成最多 20 个 tokenoutputsllm.generate(prompt,max_tokens20)foroutputinoutputs: generated_textoutput.outputs[0].text print(fGenerated text: {generated_text})在上述代码中vLLM 在预填充 The weather today is 之后KV 缓存中已经包含了这 5 个 token 的 Key 和 Value 信息。在接下来的解码过程中每生成一个新的 token例如 “sunny”vLLM 会计算 “sunny”的 Key 和 Value 向量并将它们添加到 KV 缓存的末尾。当模型预测下一个 token时其 Query 向量会与 KV 缓存中所有 6 个 Key 向量进行注意力计算。这个过程会重复进行直到生成 20 个新的 token 或达到结束条件。 单步解码的实现与性能分析 在解码阶段的每一步模型主要进行以下操作 1. 接收上一步生成的 token 的 Embedding 2. 计算该 token 在所有 Transformer 层的 Query、Key 和 Value 向量 3. 将当前生成 token 的 Key 和 Value 向量更新到 KV 缓存中 4. 在自注意力计算中当前 token 的 Query 向量会与 KV 缓存中所有历史 token 的 Key 向量进行比较 5. 模型最后一层的输出会经过线性层和 Softmax 函数得到下一个 token 的概率分布 6. 根据解码策略例如采样从概率分布中选择下一个 token 单步解码的计算成本主要在于自注意力机制的计算其复杂度与当前 KV 缓存的长度等于原始 Prompt 长度加上已生成的 token 数量成正比。KV 缓存的关键作用在于它避免了在每一步都重新计算原始 Prompt 的 Key 和 Value 向量. 然而随着生成序列的长度增加KV 缓存的大小也会增长可能导致内存带宽成为瓶颈。 2.3 预填充与解码的协同工作与性能瓶颈 数据流分析与性能瓶颈识别 预填充和解码阶段共同完成了 LLM 的推理过程。预填充为解码准备了初始的 KV 缓存而解码阶段则迭代地利用和更新这个缓存来生成最终的输出. 数据流输入 Prompt → Tokenization → Embedding → 预填充 (生成 KV 缓存) → 解码 (利用并更新 KV 缓存逐个生成 token) → Detokenization → 输出文本 性能瓶颈识别 • 预填充阶段对于极长的输入 Prompt自注意力机制的计算量仍然很大可能成为计算瓶颈。同时加载模型权重和初始 KV 缓存的内存开销也需要考虑. • 解码阶段解码的串行自回归特性是主要的瓶颈。虽然 KV 缓存减少了重复计算但每一步仍然需要进行注意力计算并且随着生成序列的增长KV 缓存的大小也会增加可能导致内存带宽瓶颈. 实操对比分析不同 Prompt 长度下预填充和解码的耗时 bash import time from vllm import LLM from transformers import AutoTokenizer model_name Qwen/Qwen2.5-7B llm LLM(modelmodel_name) tokenizer AutoTokenizer.from_pretrained(model_name) output_tokens 100 # 固定生成 100 个 token prompt_lengths [10, 50, 100, 200] print(Analyzing inference time with varying prompt lengths:) results {} for prompt_len in prompt_lengths: prompt This is a test prompt. * prompt_len input_ids tokenizer.encode(prompt) prompt_token_count len(input_ids) start_time time.time() outputs llm.generate(prompt, max_tokensoutput_tokens) end_time time.time() total_time end_time - start_time results[prompt_token_count] total_time print(fPrompt length (tokens): {prompt_token_count}, Total inference time: {total_time:.4f} seconds) print(\nResults:) for length, time in results.items(): print(fPrompt Length: {length}, Inference Time: {time:.4f} seconds)分析观察不同 Prompt 长度下总推理时间的变化。当 Prompt 长度增加时预填充阶段的计算量会增加导致总推理时间变长。#而解码阶段生成的 token 数量固定其耗时相对稳定。可以尝试使用 matplotlib 等库将这些结果绘制成折线图横轴是 Prompt 长度纵轴是推理时间#这样可以更直观地看到 Prompt 长度对推理性能的影响。#随着 Prompt 长度的增加总的推理时间应该会呈现上升趋势这主要是由于预填充阶段的计算量增加所致。补充动态 Batching 的初步介绍原理与在预填充和解码阶段的应用动态 Batching 的初步介绍动态 Batching 是一种提高 LLM 推理吞吐量的关键技术尤其是在高并发的场景下.原理动态 Batching 允许推理系统在运行时动态地将多个独立的推理请求组合成一个批次进行处理。这个批次的组成可以根据请求的当前处理阶段预填充或解码、序列长度以及其他策略进行调整以最大化 GPU 的利用率.在预填充阶段的应用当多个新的推理请求到达时可以将它们的 Prompt 组成一个批次一起进行预填充计算。这可以显著提高 GPU 的并行计算效率.在解码阶段的应用动态 Batching 在解码阶段更为复杂但也更为关键。例如vLLM 采用了名为 PagedAttention 的技术它允许在解码过程中更灵活地管理 KV 缓存.