显存管理

大模型推理的KV Cache优化全解 KV Cache是大模型推理中最重要的优化技术，也是最大的内存瓶颈。理解KV Cache的工作原理和优化方法，是构建高效推理服务的基础。 KV Cache是什么 在Transformer的自回归生成中，每生成一个新token，需要计算它对所有之前token的注意力。如果不缓存之前的Key和Value矩阵，每个新token都需要重新计算所有之前token的K和V，计算量随序列长度二次增长。 KV Cache通过缓存之前计算过的K和V矩阵，将每步的计算复杂度从O(n²)降到O(n)。代价是内存占用线性增长——对于Llama-70B模型，生成4K token的KV Cache约占40GB显存。 内存瓶颈 KV Cache的内存占用可以用以下公式计算： KV Cache Size = 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size 以Llama-70B（80层, 64头, 128维, FP16）为例，单序列4096 token的KV Cache约40GB。这意味着一个80GB显存的A100只能服务两个并发请求——这是制约推理吞吐量的最大瓶颈。 PagedAttention：分页管理 vLLM团队提出的PagedAttention是KV Cache管理的革命性创新。灵感来自操作系统的虚拟内存分页机制。 传统分配的问题 传统方法为每个序列预分配一块连续的KV Cache空间，按最大序列长度分配。这导致严重的内存碎片——大多数序列不会用满预分配的空间，但多余的空间不能被其他序列使用。内存利用率通常只有20-40%。 分页方案 PagedAttention将KV Cache划分为固定大小的"页"（通常16个token），每个序列通过页表映射到物理页。页按需分配——序列增长时才分配新页。 效果是显著的：vLLM的内存利用率提升到90%以上，并发吞吐量提升2-4倍。碎片问题被彻底解决，因为不同序列的页可以散布在物理内存中的任意位置。 页的大小选择 页太小（如1 token）会增加页表开销，页太大（如256 token）则回到预分配的问题。16 token是在大多数场景下的最优选择——页表开销不到1%，内存浪费也不显著。 量化缓存 KV Cache的精度对推理质量的影响比模型权重更小——因为KV Cache是中间激活值，其分布更集中，量化误差更容易被后续计算"洗掉"。 FP8 KV Cache 将KV Cache从FP16量化到FP8，内存减半，几乎无损。现代GPU（H100及以后）原生支持FP8运算，所以推理速度也几乎不受影响。这可能是最简单且性价比最高的KV Cache优化。 INT4 KV Cache 更激进的方案是将KV Cache量化到INT4。内存减少到1/4，但精度损失开始显著——在长序列和需要精确注意力的任务上，INT4 KV Cache可能导致输出质量下降。 ...

大模型推理的KV Cache优化全解 KV Cache是大模型推理中最重要的优化技术，也是最大的内存瓶颈。理解KV Cache的工作原理和优化方法，是构建高效推理服务的基础。 KV Cache是什么 在Transformer的自回归生成中，每生成一个新token，需要计算它对所有之前token的注意力。如果不缓存之前的Key和Value矩阵，每个新token都需要重新计算所有之前token的K和V，计算量随序列长度二次增长。 KV Cache通过缓存之前计算过的K和V矩阵，将每步的计算复杂度从O(n²)降到O(n)。代价是内存占用线性增长——对于Llama-70B模型，生成4K token的KV Cache约占40GB显存。 内存瓶颈 KV Cache的内存占用可以用以下公式计算： KV Cache Size = 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size 以Llama-70B（80层, 64头, 128维, FP16）为例，单序列4096 token的KV Cache约40GB。这意味着一个80GB显存的A100只能服务两个并发请求——这是制约推理吞吐量的最大瓶颈。 PagedAttention：分页管理 vLLM团队提出的PagedAttention是KV Cache管理的革命性创新。灵感来自操作系统的虚拟内存分页机制。 传统分配的问题 传统方法为每个序列预分配一块连续的KV Cache空间，按最大序列长度分配。这导致严重的内存碎片——大多数序列不会用满预分配的空间，但多余的空间不能被其他序列使用。内存利用率通常只有20-40%。 分页方案 PagedAttention将KV Cache划分为固定大小的"页"（通常16个token），每个序列通过页表映射到物理页。页按需分配——序列增长时才分配新页。 效果是显著的：vLLM的内存利用率提升到90%以上，并发吞吐量提升2-4倍。碎片问题被彻底解决，因为不同序列的页可以散布在物理内存中的任意位置。 页的大小选择 页太小（如1 token）会增加页表开销，页太大（如256 token）则回到预分配的问题。16 token是在大多数场景下的最优选择——页表开销不到1%，内存浪费也不显著。 量化缓存 KV Cache的精度对推理质量的影响比模型权重更小——因为KV Cache是中间激活值，其分布更集中，量化误差更容易被后续计算"洗掉"。 FP8 KV Cache 将KV Cache从FP16量化到FP8，内存减半，几乎无损。现代GPU（H100及以后）原生支持FP8运算，所以推理速度也几乎不受影响。这可能是最简单且性价比最高的KV Cache优化。 INT4 KV Cache 更激进的方案是将KV Cache量化到INT4。内存减少到1/4，但精度损失开始显著——在长序列和需要精确注意力的任务上，INT4 KV Cache可能导致输出质量下降。 ...

操作系统启发：虚拟内存 PagedAttention的核心灵感来自操作系统的虚拟内存管理。在OS中，进程的虚拟内存被分成固定大小的页（Page），物理内存不必连续分配。 PagedAttention将这一思想应用于KV Cache管理： 传统KV Cache: [token0, token1, ..., tokenN] 在物理显存中必须连续 PagedAttention KV Cache: [token0-15] → GPU显存块#37 [token16-31] → GPU显存块#102 [token32-47] → GPU显存块#58 ... （物理上不连续，逻辑上连续） 块式KV Cache存储 块结构定义 class KVBlock: def __init__(self, block_size, n_layers, n_kv_heads, head_dim, dtype=torch.float16): self.block_size = block_size # 每块存储的token数（如16） self.n_layers = n_layers self.n_kv_heads = n_kv_heads self.head_dim = head_dim # 每块的形状: [n_layers, block_size, n_kv_heads, head_dim] self.k = torch.zeros( n_layers, block_size, n_kv_heads, head_dim, dtype=dtype ) self.v = torch.zeros_like(self.k) self.ref_count = 0 # 引用计数（用于共享） self.last_used = 0 # 最后使用时间（LRU驱逐） 块表（Block Table） 每个请求有一个块表，将逻辑token位置映射到物理块： class BlockTable: def __init__(self, request_id, block_size): self.request_id = request_id self.block_size = block_size self.physical_blocks = [] # 物理块索引列表 self.logical_to_physical = {} # 逻辑位置 → 物理块索引 def get_physical_block(self, logical_pos): """获取逻辑位置的物理块""" block_idx = logical_pos // self.block_size offset = logical_pos % self.block_size if block_idx >= len(self.physical_blocks): return None, None # 尚未分配 return self.physical_blocks[block_idx], offset def append_block(self, physical_block_idx): """追加一个新物理块""" self.physical_blocks.append(physical_block_idx) 物理块管理 块池（Block Pool） 所有物理块由中央块池管理： ...

为什么KV Cache如此重要？ 在自回归生成中，每生成一个新token都需要计算它与之前所有token的注意力。如果不缓存，生成第N个token需要重新计算前N-1个token的K和V，复杂度为O(N²)。KV Cache将已计算的K/V存储起来，使每步生成只需O(N)的计算。 但缓存本身带来了巨大的显存压力。以Llama-3-70B为例，FP16精度下单个token的KV Cache占用约320KB。在32K上下文、batch_size=32的场景下，KV Cache总量达到约320GB——远超模型本身的140GB。 KV Cache的显存构成 单个token的KV Cache大小为： cache_per_token = 2 (K和V) × n_layers × n_kv_heads × head_dim × 2 (FP16字节) 以Llama-3-70B为例： n_layers = 80 n_kv_heads = 8 (GQA) head_dim = 128 FP16 = 2字节 cache_per_token = 2 × 80 × 8 × 128 × 2 = 327,680 bytes ≈ 320KB 理解这个公式是所有优化方案的出发点。 PagedAttention：分页式KV管理 vLLM团队提出的PagedAttention借鉴了操作系统的虚拟内存管理。核心思想： 块式管理 将KV Cache分为固定大小的块（Block），每个块存储若干token的KV。块不需要在物理上连续，通过块表（Block Table）映射逻辑位置到物理位置。 class PagedKVCache: def __init__(self, n_blocks, block_size, n_layers, n_kv_heads, head_dim): self.block_size = block_size # 每块的token数 self.n_blocks = n_blocks # 预分配所有块的存储 self.k_blocks = torch.zeros(n_blocks, n_layers, block_size, n_kv_heads, head_dim) self.v_blocks = torch.zeros(n_blocks, n_layers, block_size, n_kv_heads, head_dim) self.free_blocks = list(range(n_blocks)) self.block_tables = {} # seq_id -> [block_indices] def allocate(self, seq_id, n_tokens): n_blocks_needed = (n_tokens + self.block_size - 1) // self.block_size blocks = [] for _ in range(n_blocks_needed): if not self.free_blocks: raise OOMError("No free KV cache blocks") blocks.append(self.free_blocks.pop()) self.block_tables[seq_id] = blocks def append(self, seq_id, token_idx, k, v): block_idx = token_idx // self.block_size offset = token_idx % self.block_size physical_block = self.block_tables[seq_id][block_idx] self.k_blocks[physical_block, :, offset] = k self.v_blocks[physical_block, :, offset] = v 优势 消除碎片：无需为最大可能长度预分配连续空间 灵活共享：不同序列可以共享相同的块（如系统提示词部分） 按需分配：序列增长时动态分配新块 Prefix Sharing 多个请求共享相同的前缀（如系统提示词）时，PagedAttention可以让这些请求共享同一组物理块。这在对话AI场景中可节省30-50%的显存。 ...

推理优化的三个维度 显存优化 ↑ | 计算优化 ←──→ 通信优化 LLM 推理的瓶颈不是计算，而是显存带宽。7B 模型生成 1 个 Token 需要读取全部 7B 参数，但只做 1 次矩阵乘法。 显存分析 7B 模型显存分布 模型权重 (FP16): 14 GB KV Cache (2K ctx): 1 GB KV Cache (8K ctx): 4 GB KV Cache (32K ctx): 16 GB ← 比模型还大！ KV Cache (128K ctx): 64 GB 总显存 (8K ctx): 18 GB → A100 40G 可以跑 总显存 (32K ctx): 30 GB → A100 40G 勉强 总显存 (128K ctx): 78 GB → 需要 A100 80G KV Cache 计算 def kv_cache_size( num_layers: int, # 层数 num_heads: int, # 注意力头数 head_dim: int, # 每头维度 seq_len: int, # 序列长度 batch_size: int, # 批大小 dtype_size: int = 2, # FP16 = 2 bytes ): """计算 KV Cache 显存""" # K 和 V 各一份 return 2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size # Qwen3-7B 示例 size = kv_cache_size( num_layers=32, num_heads=32, head_dim=128, seq_len=8192, batch_size=1, ) print(f"{size / 1024**3:.2f} GB") # 4.00 GB 优化技术一：量化 INT8 量化 # 权重 INT8：14GB → 7GB # KV Cache INT8：显存减半 # 性能损失：1-3% from transformers import AutoModelForCausalLM import torch model = AutoModelForCausalLM.from_pretrained( "Qwen3-7B-Instruct", load_in_8bit=True, # 权重 INT8 device_map="auto", ) INT4 量化 # 权重 INT4：14GB → 3.5GB # 性能损失：3-8% # 适合消费级 GPU model = AutoModelForCausalLM.from_pretrained( "Qwen3-7B-Instruct", load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4", ) AWQ 量化 # AWQ：比 INT4 更好的量化方案 # 保护重要权重，精度损失更小 # vLLM 原生支持 # 量化模型 python -m awq.quantize \ --model-path /models/Qwen3-7B \ --output-path /models/Qwen3-7B-AWQ \ --w-bits 4 # vLLM 加载 llm = LLM(model="/models/Qwen3-7B-AWQ", quantization="awq") 优化技术二：KV Cache 优化 PagedAttention（vLLM） # 传统：预分配最大长度的连续显存 # 浪费率：60%（大量预分配空间未使用） # PagedAttention：按需分配，像操作系统的分页机制 # 浪费率：<5% # vLLM 默认启用 llm = LLM( model="...", gpu_memory_utilization=0.9, # 95% 显存用于 KV Cache ) KV Cache 量化 # KV Cache 用 INT8 存储，显存减半 llm = LLM( model="...", kv_cache_dtype="int8", # 或 "fp8" ) KV Cache 共享 # 多个请求共享相同的 system prompt KV Cache # vLLM 的 prefix caching llm = LLM( model="...", enable_prefix_caching=True, ) # 效果： # Request 1: "You are a helpful assistant. [Q1]" → 生成 KV Cache # Request 2: "You are a helpful assistant. [Q2]" → 复用 KV Cache # 节省：system prompt 部分的计算和显存 优化技术三：计算优化 Flash Attention # Flash Attention：减少 GPU 显存读写 # 标准 Attention：O(n²) 显存 # Flash Attention：O(n) 显存，速度提升 2-4x # vLLM 默认启用 llm = LLM( model="...", enforce_eager=False, # 启用 CUDA Graph + Flash Attention ) CUDA Graph # CUDA Graph：预编译计算图，减少 kernel launch 开销 # 适合：固定形状的推理（批处理） llm = LLM( model="...", enforce_eager=False, # 启用 CUDA Graph ) # 延迟降低 10-20% Continuous Batching # 传统批处理：等所有请求完成才处理下一批 # Continuous Batching：每个 step 动态加入/移除请求 # vLLM 默认使用 continuous batching # 效果：吞吐量提升 3-10x # 原理： # Time 0: [A, B, C] 开始生成 # Time 5: A 完成，D 加入 → [B, C, D] # Time 8: B 完成，E 加入 → [C, D, E] # GPU 始终满载 优化技术四：Speculative Decoding # 原理：用小模型猜，大模型验 # 小模型生成 5 个 token → 大模型一次验证 5 个 # 如果全部正确：5x 加速 # 如果 3 个正确：3x 加速 llm = LLM( model="/models/Qwen3-7B-Instruct", speculative_model="/models/Qwen3-0.5B-Instruct", num_speculative_tokens=5, ) # 实测效果： # 简单任务（翻译、摘要）：2.0x 加速 # 复杂任务（推理、创作）：1.3x 加速 # 代码生成：1.8x 加速 优化技术五：模型架构优化 GQA（Grouped Query Attention） # MHA：每个 head 有独立的 K, V → KV Cache 大 # MQA：所有 head 共享 K, V → KV Cache 小，但质量降 # GQA：折中，g 个 head 共享 K, V # Qwen3 使用 GQA，KV Cache 减少为 MHA 的 1/4 # 显存节省 75%，质量损失 <1% 模型蒸馏 # 大模型蒸馏到小模型 # Teacher: Qwen3-72B → Student: Qwen3-7B # 保留 90% 能力，推理快 10x，显存省 10x # 适用场景：固定任务，不需要通用能力 基准测试 测试方法 import time import asyncio async def benchmark(llm, prompts, max_tokens=512): start = time.time() results = await asyncio.gather(*[ llm.generate(p, SamplingParams(max_tokens=max_tokens)) for p in prompts ]) elapsed = time.time() - start total_tokens = sum(len(r.outputs[0].token_ids) for r in results) return { "total_time": elapsed, "total_tokens": total_tokens, "throughput": total_tokens / elapsed, # tokens/s "avg_latency": elapsed / len(prompts), } 参考数据 配置 模型 吞吐量 延迟 A100 80G Qwen3-7B 5000 tok/s 50ms A100 80G Qwen3-72B 800 tok/s 200ms 4090 24G Qwen3-7B (AWQ) 3000 tok/s 80ms 4090 24G Qwen3-7B (FP16) OOM - 2x 4090 Qwen3-72B (AWQ) 400 tok/s 250ms 生产环境配置建议 # 配置决策树 def recommend_config(gpu, model_size, concurrency): if gpu == "A100 80G": if model_size <= 7: return {"tp": 1, "quant": "fp16", "max_seqs": 256} elif model_size <= 72: return {"tp": 1, "quant": "awq", "max_seqs": 64} elif gpu == "4090 24G": if model_size <= 7: return {"tp": 1, "quant": "awq", "max_seqs": 32} elif model_size <= 72: return {"tp": 2, "quant": "awq", "max_seqs": 16} return {"error": "configuration not supported"} 结论 LLM 推理优化的核心原则： ...