大模型推理的 KV Cache 深度优化：PagedAttention 原理、vLLM 架构与显存管理实战
一、自回归推理的显存杀手——当 Key-Value 缓存放大了十亿参数模型的成本
Transformer 架构的推理过程是逐 Token 自回归生成的。以 LLaMA-2-70B 为例，生成一个长度为 2048 的序列时，每次解码都需要重新计算所有历史 Token 的注意力分数。如果不做任何优化，计算复杂度为 O(n²)，其中 n 是当前已生成的 Token 数量。
KV Cache 是解决这一问题的标准技术。它将每层 Attention 中已计算的 Key 和 Value 矩阵缓存下来，后续 Token 的注意力计算只需将新 Token 的 Query 与缓存中的历史 Key 进行点积。这意味着显存消耗与 Batch Size 和序列长度呈线性关系。
但传统的 KV Cache 管理方式存在显著的显存浪费。每个推理请求都会预留一段连续的显存区域用于 KV Cache，序列长度的不确定性导致大量的内部碎片（平均利用率仅 20%-40%）和外部碎片。在生产环境中，一块 80GB 的 A100 GPU，由于 KV Cache 的碎片化问题，实际可服务的并发请求数往往只有理论值的 50%。
此外，不同请求之间无法共享 KV Cache。即使两个请求共享相同的前缀 Prompt（如 System Prompt），它们的 Key-Value 缓存也无法复用，显存利用率进一步降低。
二、PagedAttention 的虚拟内存式设计
vLLM 项目提出的 PagedAttention 算法，借鉴了操作系统中虚拟内存和分页管理的思路。它将每个请求的 KV Cache 切分为大小固定的块（Block，默认 16 个 Token），通过块表（Block Table）映射逻辑位置到物理块，实现了 KV Cache 的非连续存储。
PagedAttention 带来的三个核心收益：
零显存碎片：所有 KV Cache 分配都以固定大小的块为单位，块用完即释放回空闲池，不会产生碎片化问题。实测表明，相比传统连续分配，显存利用率可从 30% 提升到 96% 以上。
前缀缓存共享：多个请求共享相同前缀时，前缀部分的 KV Cache 块只存储一份，引用计数机制保证在没有请求引用时才真正释放。这对于 AI 应用中常见的 System Prompt 共享场景效果显著，显存节省可达 40%-60%。
动态批处理：KV Cache 的块表机制允许不同请求独立管理各自的块映射，Scheduler 无需等待某个请求完成就可以将新请求加入 Batch，吞吐量大幅提升。
三、vLLM 的服务化架构与 Java 集成
vLLM 将 PagedAttention 封装为生产级推理服务，其核心组件包括：
/** * vLLM API 客户端，封装兼容 OpenAI 接口的推理请求。 * 生产环境中需处理重试、超时和连接池管理。 */ @Configuration public class VLLMClientConfig { /** * 配置连接池：vLLM 推理耗时较长（秒级），需调大超时和连接数。 */ @Bean public OkHttpClient vllmHttpClient() { // ConnectionPool 的 maxIdleConnections 控制空闲连接数上限 // 这里的 200 对应 200 个并发推理请求的连接复用 ConnectionPool pool = new ConnectionPool(200, 5, TimeUnit.MINUTES); return new OkHttpClient.Builder() .connectionPool(pool) .connectTimeout(10, TimeUnit.SECONDS) .readTimeout(120, TimeUnit.SECONDS) // 长文本生成可能耗时 > 60s .writeTimeout(30, TimeUnit.SECONDS) .retryOnConnectionFailure(true) .build(); } /** * vLLM 推理请求构建器，携带 KV Cache 优化参数。 */ public Map<String, Object> buildGenerationRequest(String prompt, int maxTokens) { Map<String, Object> request = new LinkedHashMap<>(); request.put("model", "llama-2-70b"); request.put("prompt", prompt); request.put("max_tokens", maxTokens); request.put("temperature", 0.7); // enable_prefix_caching: 开启前缀缓存共享，本质是 PagedAttention 的块复用 request.put("enable_prefix_caching", true); // use_beam_search: 关闭束搜索可以节省 KV Cache 分配 request.put("use_beam_search", false); return request; } } @Service public class VLLMInferenceService { @Autowired private OkHttpClient vllmHttpClient; @Value("${vllm.endpoint}") private String vllmEndpoint; /** * 发送推理请求到 vLLM 服务，带异常重试和降级。 */ public String generate(String prompt, int maxTokens) throws InferenceException { int maxRetries = 2; for (int attempt = 0; attempt <= maxRetries; attempt++) { try { Request request = new Request.Builder() .url(vllmEndpoint + "/v1/completions") .post(RequestBody.create( objectMapper.writeValueAsString(buildRequest(prompt, maxTokens)), MediaType.parse("application/json"))) .build(); try (Response response = vllmHttpClient.newCall(request).execute()) { if (!response.isSuccessful()) { // 服务端返回非 2xx，解析错误码 if (response.code() == 503 && attempt < maxRetries) { Thread.sleep((attempt + 1) * 1000L); // 退避重试 continue; } throw new InferenceException("vLLM service error: " + response.code()); } return parseResponse(response.body().string()); } } catch (IOException | InterruptedException e) { if (attempt >= maxRetries) { throw new InferenceException("vLLM request failed after retries", e); } } } throw new InferenceException("Unreachable"); } }关键设计考量：vLLM 推理是长耗时操作（生成 512 Token 约需 2-8 秒），客户端的readTimeout需要调大到 120 秒。连接池的maxIdleConnections不建议设置过大（200 左右即可），因为 vLLM 服务侧本身就通过动态批处理管理并发。
四、PagedAttention 的硬件适配与场景边界
PagedAttention 的效果受 GPU 架构影响较大。在 A100/H100 等支持大显存的 GPU 上，KV Cache 块大小为 16 时性能最优。但在 T4 等旧架构上，较小的块（如 8）可能更适合，因为 T4 的显存带宽限制了较大块内注意力计算的效率。
另一个限制来自 Prefix Caching 的命中率。当多个请求的系统提示词不完全一致（如根据用户角色动态拼接）时，前缀缓存的命中率会急剧下降。需要在上游做提示词规范化（将可变部分放到后缀），才能充分发挥共享优势。
vLLM 的 Continuous Batching 调度策略也有代价：当一个 Batch 中存在长度差异巨大的请求时，长请求会拖慢整体调度频率。建议在网关层按预估生成长度将请求分组，避免混入"长度炸弹"。
适合场景：高并发 LLM 推理服务（日均百万级请求）、共享相同 System Prompt 的 AI Agent 系统、需要最大化 GPU 利用率的成本敏感型部署。不适合场景：单请求低延迟要求的实时交互场景（vLLM 调度开销约 10-50ms）、生成序列极短的应用（KV Cache 块的固定开销占比过高）。
五、总结
KV Cache 的优化是 LLM 推理降本增效的核心战场。PagedAttention 通过操作系统的分页思想，将 KV Cache 管理从连续内存分配转变为离散块管理，解决了显存碎片和前缀无法共享的两大痛点。vLLM 在此基础上集成了 Continuous Batching 调度器，将推理吞吐提升了数倍。工程落地时，需根据 GPU 硬件调整 Block Size，在上游网关做提示词规范化以提高前缀缓存命中率，并对 Connection Pool 和超时参数做适配长耗时推理的调整。显存优化的收益最终会直接体现在每 GPU 并发请求数和单 Token 生成成本上。
