大模型 (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的数学本质与代码实现，以及解决显存碎片化的工业级标杆PagedAttention（vLLM 核心原理）。

1. 推理性能杀手：自回归解码与重复计算

LLM 的推理分为两个阶段：

1. Prefill（预填充阶段）：一次性处理用户的 Prompt，计算密集，高度并行化。
2. Decode（解码阶段）：逐字生成回复。生成第N+1N+1N+1个 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 Attention）KV Cache 机制的硬核手工实现：

importtorchimporttorch.nnasnnimporttorch.nn.functionalasFimportmathclassCausalSelfAttentionWithCache(nn.Module):def__init__(self,d_model):super().__init__()self.d_model=d_model# 融合投影层，加速计算self.c_attn=nn.Linear(d_model,3*d_model)self.c_proj=nn.Linear(d_model,d_model)defforward(self,x,layer_past=None):""" x: 当前输入，Prefill阶段形状为 [batch, seq_len, d_model] Decode阶段形状通常为 [batch, 1, d_model] layer_past: 历史缓存的 (K, V)，形状均为 [batch, past_seq_len, d_model] """B,T,C=x.size()# 1. 线性投影获取当前步的 q, k, vqkv=self.c_attn(x)q,k,v=qkv.split(self.d_model,dim=2)# 2. KV Cache 逻辑：如果存在历史缓存，则拼接iflayer_pastisnotNone:past_k,past_v=layer_past# 在序列长度维度（dim=1）进行拼接k=torch.cat((past_k,k),dim=1)v=torch.cat((past_v,v),dim=1)# 3. 更新当前的 KV Cache 供下一步使用present=(k,v)# 4. 计算 Attention Scores# q 的 seq_len=1 (Decode阶段)，k 的 seq_len=past_len+1# att shape: [B, 1, past_len+1]att=(q @ k.transpose(-2,-1))*(1.0/math.sqrt(k.size(-1)))# 因果掩码 (仅在 Prefill 阶段 T>1 时需要，Decode阶段 T=1 天然只看过去)ifT>1:causal_mask=torch.tril(torch.ones(T,T)).view(1,T,T).to(x.device)att=att.masked_fill(causal_mask==0,float('-inf'))att=F.softmax(att,dim=-1)# 5. 计算输出y=att @ v# shape: [B, T, C]y=self.c_proj(y)returny,present# ================= 测试验证 =================torch.manual_seed(42)d_model=64model=CausalSelfAttentionWithCache(d_model)# 模拟 Prefill 阶段：输入 "I love" (2 tokens)prompt=torch.randn(1,2,d_model)out_prefill,kv_cache=model(prompt)print(f"Prefill输出形状:{out_prefill.shape}")# [1, 2, 64]print(f"KV Cache形状(K):{kv_cache[0].shape}")# [1, 2, 64]# 模拟 Decode 阶段：基于前面生成的 token 继续生成下一个next_token_input=torch.randn(1,1,d_model)out_decode,new_kv_cache=model(next_token_input,layer_past=kv_cache)print(f"Decode单步输出形状:{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 框架）借鉴了操作系统的虚拟内存分页机制：

1. 将 KV Cache 切分成固定大小的 Block（例如每个 Block 容纳 16 个 Token 的 K 和 V）。
2. 在物理显存中，这些 Block 不需要连续存放。
3. 维护一个 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_id="TheBloke/Llama-2-7b-Chat-AWQ"# 实例化 LLM 引擎，底层自动初始化 PagedAttention 显存池# gpu_memory_utilization=0.9 表示将 90% 的可用显存全部分配给 KV Cache Blocksllm=LLM(model=model_id,quantization="awq",tensor_parallel_size=1,# 单卡推理gpu_memory_utilization=0.9,max_model_len=4096,# 最大上下文长度trust_remote_code=True)# 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_params=SamplingParams(temperature=0.7,top_p=0.95,max_tokens=512)# 4. 异步批量推理# vLLM 底层会利用 Continuous Batching 和 PagedAttention 最大化 GPU 吞吐量outputs=llm.generate(prompts,sampling_params)# 5. 解析输出foroutputinoutputs:prompt=output.prompt generated_text=output.outputs[0].textprint(f"Prompt:{prompt!r}")print(f"Generated:{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)切分策略。