蓝迪哥玩转Ai(12)—FPGA本地算力研究KV 缓存原理、框架实现与性能优化实战在11中我们已经介绍了预填充和解码这两个 LLM 推理的核心阶段并初步认识到 KV 缓存作为关键加速机制在其中的作用。本章将在此基础上对 KV 缓存进行更为深入和全面的剖析。1 KV 缓存的原理与数学本质正如我们在第二章中介绍的KV 缓存是 Transformer 模型在推理过程中用于存储 Key (K) 和 Value (V) 向量的一种优化技术。理解其原理需要回顾自注意力机制的计算过程。注意力机制中 Key 和 Value 的作用与计算公式对于输入序列的第 i 个 token其 Query 向量 Qi 会与所有第 j 个 token 的 Key 向量 Kj 进行点积计算得到注意力分数。经过缩放和 Softmax 归一化后得到注意力权重。KV 缓存如何避免重复计算在自回归解码过程中每生成一个新的 token模型的输入序列长度就会增加。如果没有 KV 缓存为了计算下一个 token 的上下文表示模型需要对包括原始 Prompt 和所有已生成 token 在内的完整序列重新计算 Q、K 和 V 向量以及注意力权重。这会带来大量的重复计算尤其是在生成长文本时。KV 缓存的核心思想是将已经计算过的 Key 和 Value 向量存储起来并在后续的解码步骤中直接重用。GIF动画演示已分割为帧图片帧 1帧 2帧 3帧 4通过使用 KV 缓存在解码的每一步模型只需要计算当前生成 token 的 Query 向量并将其与缓存中已有的所有 Key 向量进行比较以计算注意力权重。这样就避免了对历史 token 的 Key 和 Value 向量的重复计算显著提升了解码效率。2 KV 缓存在主流推理框架中的实现不同的推理框架对 KV 缓存的实现和优化有所不同。vLLM 等框架中 PagedAttention 的原理与优势传统的 KV 缓存管理通常为每个推理请求在 GPU 显存中分配一块连续的内存空间来存储其 Key 和 Value 向量。这种做法在处理单个请求时可能没有太大问题但在高并发的场景下尤其是当不同请求的序列长度差异很大时就会暴露出一些明显的缺点• 内存碎片假设我们同时处理多个长度不一的序列为了容纳最长的序列可能会预先分配较大的连续内存块。当较短的序列结束后其占用的内存空间可能无法被其他请求充分利用导致内存碎片。• 内存浪费即使是长度接近的序列也需要各自独立的连续内存空间造成一定的内存浪费。• 难以支持动态长度变化在解码过程中序列长度是动态增长的。如果预分配的内存空间不足以容纳更长的生成序列就需要重新分配更大的内存块这会带来额外的开销。PagedAttention 的核心思想借鉴了操作系统中内存分页管理的机制。它将每个推理请求的 KV 缓存分割成多个固定大小的小块这些小块被称为页面Page。每个页面可以存储固定数量的 Key 和 Value 向量。PagedAttention 的优势• 极高的内存效率由于以页为单位进行分配可以更精细地管理内存显著减少内存碎片和浪费。• 灵活处理变长序列可以轻松地支持同一批次内不同长度的序列以及单个序列在解码过程中长度的动态增长只需分配新的页面即可。• 高效的 KV 缓存共享与拷贝对于具有相同 Prompt 前缀的多个请求vLLM 可以通过共享底层的内存页面来实现零拷贝的 KV 缓存共享。• 更好地支持连续批处理Continuous BatchingPagedAttention 的内存管理方式非常适合与 vLLM 的连续批处理技术结合使用。Transformers 框架中 KV 缓存的结构与操作代码分析在 Hugging Face Transformers 中KV 缓存通常作为模型 forward 函数的输出名为 past_key_values和输入进行传递。past_key_values 是一个包含每一层 Key 和 Value 缓存的元组。from transformersimportAutoModelForCausalLM, AutoTokenizerimporttorch model_namegpt2tokenizerAutoTokenizer.from_pretrained(model_name)modelAutoModelForCausalLM.from_pretrained(model_name)promptThe quick brown foxinput_idstokenizer.encode(prompt,return_tensorspt)past_key_valuesNoneforiinrange(10): outputsmodel(input_ids,past_key_valuespast_key_values,use_cacheTrue)next_token_logitsoutputs.logits[:, -1, :]next_token_idtorch.argmax(next_token_logits,dim-1).unsqueeze(-1)input_idstorch.cat([input_ids, next_token_id],dim-1)past_key_valuesoutputs.past_key_values print(fStep {i1}: Generated token - {tokenizer.decode(next_token_id[0])})generated_texttokenizer.decode(input_ids[0])print(f\nGenerated text: {generated_text})3 KV 缓存对不同模型架构的影响与适配KV 缓存是 Decoder-only 模型如 GPT、LLaMA实现高效自回归生成的核心。这类模型在每个解码步骤中都依赖于先前生成的 token 的信息来预测下一个 tokenKV 缓存完美地满足了这种需求。对于 Encoder-Decoder 模型如 T5、BART编码器负责处理输入序列并生成上下文表示这个上下文表示通常以某种形式传递给解码器。解码器在生成目标序列时也会进行自回归生成并且可以使用 KV 缓存来存储已生成 token 的 Key 和 Value。4 KV 缓存的优化策略与代码实践为了进一步提升推理性能和降低内存占用可以采用多种 KV 缓存的优化策略内存优化技巧• KV 缓存的量化可以将 KV 缓存中的 Key 和 Value 向量的数据类型从 FP16 降低到 INT8从而减少一半的内存占用。from vllmimportLLM model_nameQwen/Qwen2.5-7Bllm_int8LLM(modelmodel_name,quantizationint8)• KV 缓存的 CPU Offloading在 GPU 显存不足的情况下可以将部分 KV 缓存层卸载到 CPU 内存中。DeepSpeed 示例 (概念性)import deepspeedfrom transformers import AutoModelForCausalLM, AutoTokenizermodel_name “gpt2”tokenizer AutoTokenizer.from_pretrained(model_name)model AutoModelForCausalLM.from_pretrained(model_name)engine, _, _, _ deepspeed.initialize(modelmodel,model_parametersmodel.parameters(),config_params{“zero_optimization”: {“stage”: 0}, “fp16”: {“enabled”: True}})prompt “The quick brown fox”input_ids tokenizer.encode(prompt, return_tensors“pt”).to(engine.local_rank)outputs engine.generate(input_ids, max_length20)print(tokenizer.decode(outputs[0]))5 实操使用 KV 缓存加速长文本生成并分析性能提升我们将使用 vLLM 来演示 KV 缓存通过 PagedAttention 实现对长文本生成速度的影响。importtimeimporttorch from vllmimportLLM from transformersimportAutoTokenizer model_nameQwen/Qwen2.5-7BllmLLM(modelmodel_name)tokenizerAutoTokenizer.from_pretrained(model_name)prompts[The quick brown fox jumps over the lazy dog. *10]output_lengths[100,500,1000]print(Analyzing inference time with varying output lengths (same prompt):)forpromptinprompts: prompt_tokenslen(tokenizer.encode(prompt))print(fPrompt length (tokens): {prompt_tokens})formax_tokensinoutput_lengths: start_timetime.time()outputsllm.generate(prompt,max_tokensmax_tokens)end_timetime.time()inference_timeend_time - start_time tokens_per_secondmax_tokens / inference_time print(f Generating {max_tokens} tokens took {inference_time:.4f} seconds ({tokens_per_second:.2f} tokens/s).)print(\nAnalyzing inference time with varying prompt lengths (fixed output length):)prompts_long[The quick brown fox jumps over the lazy dog. *5,This is a longer starting sentence. *20,A very lengthy introduction to the topic at hand, spanning multiple paragraphs. *50]max_tokens200forpromptinprompts_long: prompt_tokenslen(tokenizer.encode(prompt))print(fPrompt length (tokens): {prompt_tokens})start_timetime.time()outputsllm.generate(prompt,max_tokensmax_tokens)end_timetime.time()inference_timeend_time - start_time tokens_per_secondmax_tokens / inference_time print(f Generating {max_tokens} tokens took {inference_time:.4f} seconds ({tokens_per_second:.2f} tokens/s).)内存监测forpromptinprompts:formax_tokensinoutput_lengths: torch.cuda.reset_peak_memory_stats()memory_beforetorch.cuda.memory_allocated()/(1024**2)# MBstart_timetime.time()outputsllm.generate(prompt,max_tokensmax_tokens)end_timetime.time()memory_aftertorch.cuda.memory_allocated()/(1024**2)# MBinference_timeend_time - start_time tokens_per_secondmax_tokens / inference_time print(f Generating {max_tokens} tokens: {inference_time:.4f}s ({tokens_per_second:.2f} tokens/s), Memory: {memory_after - memory_before:.2f} MB)总结本章我们深入探讨了 KV 缓存的原理、实现和优化策略。通过了解 KV 缓存的工作方式我们可以更好地理解现代 LLM 如何实现高效的推理。掌握 KV 缓存的相关知识对于我们未来学习更高级的推理优化技术至关重要。