大模型推理的 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 生成成本上。