大模型 (LLM) 推理加速核心技术解析：从 KV Cache 到 PagedAttention 实战

大模型 (LLM) 推理加速核心技术解析从 KV Cache 到 PagedAttention 实战在 AI Infrastructure (AI Infra) 领域大语言模型LLM的部署与推理优化是当前最具挑战性的方向。与传统的 CV 模型不同LLM 的推理过程尤其是生成阶段由于自回归Autoregressive的特性面临着极其严重的**内存墙Memory Wall**问题呈现出明显的 Memory-bound访存瓶颈而非 Compute-bound计算瓶颈特征。本文将深入底层拆解 LLM 推理加速的两大基石KV Cache的数学本质与代码实现以及解决显存碎片化的工业级标杆PagedAttentionvLLM 核心原理。1. 推理性能杀手自回归解码与重复计算LLM 的推理分为两个阶段Prefill预填充阶段一次性处理用户的 Prompt计算密集高度并行化。Decode解码阶段逐字生成回复。生成第N1N1N1个 Token 时需要计算前NNN个 Token 的 Attention。如果不做优化每次生成新词都会把前面所有词的 Q、K、V 矩阵重新计算一遍导致极大的算力浪费O(N2)\mathcal{O}(N^2)O(N2)。2. 核心技术一KV Cache 机制与手写实现为了避免 Decode 阶段的重复计算我们将之前时间步生成的 Key 和 Value 张量缓存Cache在显存中。当生成新 Token 时只需计算当前单一 Token 的 Query并与缓存的 K、V 拼接计算 Attention。以下为基于 PyTorch 的单头注意力Single-Head AttentionKV Cache 机制的硬核手工实现importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimportmathclassCausalSelfAttentionWithCache(nn.Module):def__init__(self,d_model):super().__init__()self.d_modeld_model# 融合投影层加速计算self.c_attnnn.Linear(d_model,3*d_model)self.c_projnn.Linear(d_model,d_model)defforward(self,x,layer_pastNone): x: 当前输入Prefill阶段形状为 [batch, seq_len, d_model] Decode阶段形状通常为 [batch, 1, d_model] layer_past: 历史缓存的 (K, V)形状均为 [batch, past_seq_len, d_model] B,T,Cx.size()# 1. 线性投影获取当前步的 q, k, vqkvself.c_attn(x)q,k,vqkv.split(self.d_model,dim2)# 2. KV Cache 逻辑如果存在历史缓存则拼接iflayer_pastisnotNone:past_k,past_vlayer_past# 在序列长度维度dim1进行拼接ktorch.cat((past_k,k),dim1)vtorch.cat((past_v,v),dim1)# 3. 更新当前的 KV Cache 供下一步使用present(k,v)# 4. 计算 Attention Scores# q 的 seq_len1 (Decode阶段)k 的 seq_lenpast_len1# att shape: [B, 1, past_len1]att(q k.transpose(-2,-1))*(1.0/math.sqrt(k.size(-1)))# 因果掩码 (仅在 Prefill 阶段 T1 时需要Decode阶段 T1 天然只看过去)ifT1:causal_masktorch.tril(torch.ones(T,T)).view(1,T,T).to(x.device)attatt.masked_fill(causal_mask0,float(-inf))attF.softmax(att,dim-1)# 5. 计算输出yatt v# shape: [B, T, C]yself.c_proj(y)returny,present# 测试验证 torch.manual_seed(42)d_model64modelCausalSelfAttentionWithCache(d_model)# 模拟 Prefill 阶段输入 I love (2 tokens)prompttorch.randn(1,2,d_model)out_prefill,kv_cachemodel(prompt)print(fPrefill输出形状:{out_prefill.shape})# [1, 2, 64]print(fKV Cache形状(K):{kv_cache[0].shape})# [1, 2, 64]# 模拟 Decode 阶段基于前面生成的 token 继续生成下一个next_token_inputtorch.randn(1,1,d_model)out_decode,new_kv_cachemodel(next_token_input,layer_pastkv_cache)print(fDecode单步输出形状:{out_decode.shape})# [1, 1, 64]print(f更新后的KV Cache形状(K):{new_kv_cache[0].shape})# [1, 3, 64]3. 核心技术二从显存碎片到 PagedAttention普通的 KV Cache 虽然解决了计算墙但带来了更严重的内存墙问题。在实际的高并发推理服务中不同请求的生成长度是不可预知的。如果预先为每个请求分配最大长度的显存如 2048 token会造成极其严重的内部碎片如果动态分配又会引发内存对齐和外部碎片。实测表明传统方案中显存的实际利用率不到 30%。UC Berkeley 提出的PagedAttention已落地于 vLLM 框架借鉴了操作系统的虚拟内存分页机制将 KV Cache 切分成固定大小的 Block例如每个 Block 容纳 16 个 Token 的 K 和 V。在物理显存中这些 Block 不需要连续存放。维护一个 Block Table页表记录逻辑 Token 到物理 Block 的映射。这种机制彻底消除了外部碎片将显存利用率提升至 90% 以上使得单卡能同时处理的并发请求数Batch Size飙升。4. 工业级落地vLLM 引擎极简集成在真实的 AI Infra 建设中我们不会手动管理庞大的物理页表而是直接接入 vLLM 或 TensorRT-LLM。以下展示如何使用 Python 快速拉起一个支持 PagedAttention 且具备高并发处理能力的推理引擎并开启 AWQ (Activation-aware Weight Quantization) INT4 量化。# pip install vllmfromvllmimportLLM,SamplingParams# 1. 引擎初始化参数设置# 采用 AWQ 量化版本模型大幅降低显存占用如 7B 模型只需 ~5GB 显存model_idTheBloke/Llama-2-7b-Chat-AWQ# 实例化 LLM 引擎底层自动初始化 PagedAttention 显存池# gpu_memory_utilization0.9 表示将 90% 的可用显存全部分配给 KV Cache BlocksllmLLM(modelmodel_id,quantizationawq,tensor_parallel_size1,# 单卡推理gpu_memory_utilization0.9,max_model_len4096,# 最大上下文长度trust_remote_codeTrue)# 2. 定义多并发请求prompts[Write a detailed C multi-threading tutorial.,Explain the concept of AI Infrastructure in simple terms.,Write a python script to monitor Linux system metrics.]# 3. 设置采样参数 (Temperature, Top-p 等)sampling_paramsSamplingParams(temperature0.7,top_p0.95,max_tokens512)# 4. 异步批量推理# vLLM 底层会利用 Continuous Batching 和 PagedAttention 最大化 GPU 吞吐量outputsllm.generate(prompts,sampling_params)# 5. 解析输出foroutputinoutputs:promptoutput.prompt generated_textoutput.outputs[0].textprint(fPrompt:{prompt!r})print(fGenerated:{generated_text[:100]}...\n)5. 总结与进阶展望从朴素的 KV Cache 到操作系统的分页思想PagedAttention再到工程落地的连续批处理Continuous Batching和量化技术LLM 推理优化是一场对 GPU 架构、显存带宽和操作系统原理的深度压榨。对于想要深耕 AI Infra 或大模型后端的开发者而言下一步的进阶路线应当是研究FlashAttention-2的算子融合Operator Fusion与 SRAM 读写优化以及在多机多卡场景下的Tensor Parallelism (TP)与Pipeline Parallelism (PP)切分策略。