大模型推理加速:从 KV Cache 到 Continuous Batching 的实战复盘 大模型推理加速从 KV Cache 到 Continuous Batching 的实战复盘一、深夜告警GPU 没跑满请求却在排队某天凌晨监控面板突然报警——线上 LLM 推理服务的 P99 延迟从 800ms 飙到了 4.2s。排查下来发现并发量从 50 QPS 涨到 200 QPS 时GPU 利用率居然只有 35%。大部分时间不是花在计算上而是耗在请求调度和内存拷贝上。问题不在模型本身而是推理框架没把 GPU 喂饱。大模型推理的瓶颈往往不在计算密度而在于调度策略。请求调度、内存管理、批处理方式的粗放设计让 GPU 大量时间在“等数据”。这篇文章结合生产环境代码和压测数据聊聊 KV Cache 管理、Continuous Batching、Prefix Caching 这几个关键优化点。二、推理加速的三个关键点2.1 KV Cache避免重复计算注意力Transformer 自回归解码时每生成一个 token 都要重新计算前面所有 token 的注意力。KV Cache 把已经算好的 Key/Value 向量存下来下次直接用。不过 KV Cache 占用的显存会随着序列长度线性增长7B 模型在 seq_len4096 时光 KV Cache 就要占 2GB 左右。2.2 Continuous Batching别让短序列等长序列传统的 Static Batching 要求批内所有序列都跑完才能释放资源短序列只能干等着长序列。Continuous Batching 在每个迭代步动态插入新请求、移除已完成请求GPU 利用率能从 35% 提到 85% 以上。2.3 Prefix Caching复用公共前缀多轮对话里系统提示词和上下文前缀往往是一样的。Prefix Caching 把公共前缀的 KV Cache 跨请求复用后续请求直接命中缓存跳过 prefill 阶段。sequenceDiagram participant Client participant Scheduler participant KVCacheMgr participant GPU Client-Scheduler: 请求1 (prompt query) Scheduler-KVCacheMgr: 检查 prefix cache 命中 KVCacheMgr--Scheduler: 未命中分配新 slot Scheduler-GPU: prefill(prompt) decode(query) GPU--KVCacheMgr: 存储 KV Cache KVCacheMgr--Scheduler: 返回 token Client-Scheduler: 请求2 (相同 prompt 新 query) Scheduler-KVCacheMgr: 检查 prefix cache 命中 KVCacheMgr--Scheduler: 命中复用 prefix KV Scheduler-GPU: 仅 decode(query)跳过 prefill GPU--Scheduler: 返回 token延迟降低 60%三、代码实现与压测结果3.1 KV Cache 分页管理器KV Cache 最头疼的问题是显存碎片化。借鉴操作系统的虚拟内存分页机制把 KV Cache 切成固定大小的 Block按需分配。import torch from typing import Dict, List, Optional from dataclasses import dataclass, field dataclass class KVBlock: KV Cache 的一个物理块固定大小 block_id: int ref_count: int 0 # 引用计数支持 prefix cache 共享 device_tensor: Optional[torch.Tensor] None # 实际显存数据 class PagedKVCacheManager: 分页式 KV Cache 管理器 核心思路将 KV Cache 按固定 block_size 分页 逻辑序列通过 page table 映射到物理 block 避免显存预分配导致的碎片化问题 def __init__( self, num_blocks: int, block_size: int, num_kv_heads: int, head_dim: int, num_layers: int, dtype: torch.dtype torch.float16, ): self.block_size block_size self.num_layers num_layers # 预分配所有物理 block 的显存池 # 形状: [num_blocks, 2, num_kv_heads, block_size, head_dim] # 2 对应 K 和 V element_size torch.tensor([], dtypedtype).element_size() per_block_bytes 2 * num_kv_heads * block_size * head_dim * element_size total_bytes num_blocks * per_block_bytes * num_layers print(f[KVCache] 预分配显存池: {total_bytes / 1024**3:.2f} GB, f共 {num_blocks} 个 block) self.kv_pool torch.empty( (num_layers, num_blocks, 2, num_kv_heads, block_size, head_dim), dtypedtype, devicecuda ) # 空闲 block 链表 self.free_blocks: List[KVBlock] [ KVBlock(block_idi) for i in range(num_blocks) ] # 逻辑序列 - 物理 block 映射表 self.page_table: Dict[int, List[int]] {} # block_id - KVBlock 反向索引 self.block_map: Dict[int, KVBlock] { b.block_id: b for b in self.free_blocks } def allocate(self, seq_id: int, num_tokens: int) - List[int]: 为序列分配 KV Cache block 返回分配的物理 block_id 列表 num_needed (num_tokens self.block_size - 1) // self.block_size if len(self.free_blocks) num_needed: raise RuntimeError( f显存不足: 需要 {num_needed} 个 block, f仅剩 {len(self.free_blocks)} 个 ) allocated [] for _ in range(num_needed): block self.free_blocks.pop() block.ref_count 1 allocated.append(block.block_id) self.page_table[seq_id] allocated return allocated def free(self, seq_id: int) - None: 释放序列占用的所有 KV Cache block if seq_id not in self.page_table: return for block_id in self.page_table[seq_id]: block self.block_map[block_id] block.ref_count - 1 # 引用计数归零才真正回收支持 prefix cache 共享 if block.ref_count 0: block.ref_count 0 self.free_blocks.append(block) del self.page_table[seq_id] def copy_prefix( self, src_seq_id: int, dst_seq_id: int, prefix_len: int ) - List[int]: 复用 prefix 的 KV Cache零拷贝仅增加引用计数 用于多轮对话场景避免重复计算系统提示词 src_blocks self.page_table.get(src_seq_id, []) num_prefix_blocks prefix_len // self.block_size dst_blocks [] # 共享 prefix block增加引用计数零拷贝 for block_id in src_blocks[:num_prefix_blocks]: self.block_map[block_id].ref_count 1 dst_blocks.append(block_id) # 为新增 token 分配新 block remaining_tokens prefix_len % self.block_size if remaining_tokens 0: new_blocks self.allocate(dst_seq_id, remaining_tokens) dst_blocks.extend(new_blocks) self.page_table[dst_seq_id] dst_blocks return dst_blocks def get_physical_table(self, seq_id: int) - torch.Tensor: 返回序列的 page table用于 GPU kernel 中的地址映射 block_ids self.page_table.get(seq_id, []) return torch.tensor(block_ids, dtypetorch.int32, devicecuda)3.2 Continuous Batching 调度器import time from collections import deque from dataclasses import dataclass from typing import Deque, List, Set dataclass class Sequence: 推理序列状态机 seq_id: int prompt_token_ids: List[int] generated_tokens: List[int] field(default_factorylist) is_finished: bool False max_tokens: int 512 property def num_generated(self) - int: return len(self.generated_tokens) class ContinuousBatcher: 连续批处理调度器 核心逻辑每个 decode step 动态调整 batch 组成 已完成序列立即让出资源新请求即时填入 def __init__(self, max_batch_size: int 64): self.max_batch_size max_batch_size self.waiting_queue: Deque[Sequence] deque() self.running_batch: List[Sequence] [] self.finished_ids: Set[int] set() def add_request(self, seq: Sequence) - None: 新请求入队 self.waiting_queue.append(seq) def schedule(self) - List[Sequence]: 单步调度移除已完成序列填入新请求 返回当前 step 的活跃 batch # 移除已完成的序列 self.running_batch [ s for s in self.running_batch if not s.is_finished ] # 从等待队列填入新请求直到 batch 满 available_slots self.max_batch_size - len(self.running_batch) while available_slots 0 and self.waiting_queue: seq self.waiting_queue.popleft() self.running_batch.append(seq) available_slots - 1 return self.running_batch def step(self) - List[Sequence]: 执行一次 decode step 实际生产中此处调用 GPU kernel 执行推理 batch self.schedule() if not batch: return [] # 模拟 decode每个序列生成一个 token for seq in batch: # 实际场景调用模型 forward取 argmax token seq.generated_tokens.append(0) # placeholder if seq.num_generated seq.max_tokens: seq.is_finished True self.finished_ids.add(seq.seq_id) return batch def is_idle(self) - bool: return len(self.running_batch) 0 and len(self.waiting_queue) 03.3 压测数据加速效果对比在 A100 80GB 上部署 LLaMA-2-7B对比三种策略的吞吐与延迟策略QPSP50 延迟P99 延迟GPU 利用率Static Batching (batch32)451.2s4.1s38%Continuous Batching1200.6s1.8s82%Continuous Prefix Cache1650.35s1.1s88%数据很直观Continuous Batching 把吞吐提升了 2.7 倍加上 Prefix Cache 后达到 3.7 倍P99 延迟从 4.1s 降到了 1.1s。四、加速策略的代价显存、复杂度与一致性4.1 KV Cache 分页管理的显存开销分页管理解决了碎片化但也引入了 page table 的额外显存和查表开销。block_size 越小碎片越少但 page table 越大。实测 block_size16 是 7B 模型的甜点13B 模型建议 block_size32。4.2 Continuous Batching 的调度延迟每个 step 都要执行 schedule 逻辑在 batch_size64 时纯 Python 调度耗时约 0.3ms。对于 decode step 仅需 10ms 的场景调度占比 3%。如果 batch_size 超过 256得把调度逻辑下沉到 C/CUDA否则调度本身会成为瓶颈。4.3 Prefix Cache 的一致性风险共享 prefix block 用引用计数实现零拷贝但如果模型权重更新比如在线学习缓存的 KV 值和新权重不匹配输出质量会出问题。生产环境中模型权重更新时必须强制失效所有 prefix cache。4.4 不适合的场景显存极度紧张小于模型权重 1.2 倍时KV Cache 分页意义不大建议优先用 PagedAttention 的 swap 机制请求序列长度差异极大1 token vs 8192 token时Continuous Batching 的调度开销可能抵消收益单轮无前缀复用的场景Prefix Cache 完全没用五、总结大模型推理加速的核心是最大化 GPU 计算密度。KV Cache 分页管理消除显存碎片Continuous Batching 消除请求等待空洞Prefix Cache 消除重复计算——这三者分别从内存、调度、计算三个维度压缩浪费。压测数据表明三者叠加后 A100 上的推理吞吐提升了 3.7 倍P99 延迟降低了 73%。但每项优化都有代价分页引入查表开销连续批处理引入调度延迟前缀缓存引入一致性风险。性能优化从来不是免费午餐而是对具体场景的精确权衡。用代码说话用数据服人——这才是推理加速工程的正确打开方式。