大模型推理优化：KVCache、PageAttention、Continuous Batching 与 Streaming

问题

大模型推理为什么慢？显存都去哪了？为什么 7B 模型跑 128K 上下文时显存不够用？vLLM 凭什么比 HuggingFace 原生推理快 2-4 倍？

分析

大模型推理的瓶颈不在计算，在显存带宽和KVCache 显存占用。Decoder-only 模型每生成一个 token，都需要重新计算之前所有 token 的 attention。如果每次从头算，7B 模型生成 100 个 token，复杂度是 O(n²)，n=100 时还好，但 n=4096 时单次 attention 的计算量都是 100 的 1600 倍。所以 KVCache 是必然的优化——把之前算好的 key/value 缓存下来，避免重复计算。

但 KVCache 本身也成了新的瓶颈。7B 模型，4K 上下文，KVCache 占约 1GB；128K 上下文时飙升到 32GB。而模型权重本身（FP16）才 14GB。显存大头变成了 KVCache，这直接限制了 LLM 的部署密度和上下文长度。

深入每一层优化

1. KVCache 原理

KVCache 的核心思路：在 Decoder-only 的 self-attention 中，计算第 t 个 token 的 attention 时，Q 是当前 token 的 query，K 和 V 是前 t-1 个 token 的 key/value 矩阵。解码过程中，每步只生成一个 token，但 attention 需要看到所有历史 token。如果每次重新计算历史的 K/V，复杂度是 O(t²×d)，缓存后降为 O(t×d)。

注意力计算流程（文字时序图）：

Step 1: 输入 token "我"
         Q1 = W_q * "我"     ← 计算当前 token 的 Q
         K1 = W_k * "我"     ← 计算当前 token 的 K
         V1 = W_v * "我"     ← 计算当前 token 的 V
         缓存: K=[K1], V=[V1]
         输出: Attention(Q1, [K1], [V1])

Step 2: 输入 token "是"
         Q2 = W_q * "是"
         K2 = W_k * "是"     → 追加到缓存
         V2 = W_v * "是"     → 追加到缓存
         缓存: K=[K1, K2], V=[V1, V2]
         输出: Attention(Q2, [K1,K2], [V1,V2])  ← 复用历史 K/V

Step 3: 输入 token "谁"
         缓存: K=[K1,K2,K3], V=[V1,V2,V3]
         输出: Attention(Q3, [K1,K2,K3], [V1,V2,V3])

如果没有 KVCache：Step 3 时重新计算 Step 1 和 Step 2 的 K/V
有 KVCache：Step 3 只算 Q3, K3, V3，K1,K2 和 V1,V2 从缓存读

import torch

class KVCache:
    """简化的 KVCache 实现，展示核心逻辑"""
    def __init__(self, max_batch_size, max_seq_len, n_layers, n_heads, head_dim, dtype=torch.float16):
        self.max_batch_size = max_batch_size
        self.max_seq_len = max_seq_len
        self.n_layers = n_layers
        # 预分配连续显存
        self.k_cache = [torch.zeros(max_batch_size, max_seq_len, n_heads, head_dim, dtype=dtype) for _ in range(n_layers)]
        self.v_cache = [torch.zeros(max_batch_size, max_seq_len, n_heads, head_dim, dtype=dtype) for _ in range(n_layers)]
        self.seq_len = 0

    def append(self, layer_idx, k, v):
        """追加当前 token 的 K/V 到缓存"""
        batch_size, seq_len = k.shape[0], k.shape[1]  # seq_len 通常是 1（逐 token 生成）
        start = self.seq_len
        end = start + seq_len
        self.k_cache[layer_idx][:batch_size, start:end] = k
        self.v_cache[layer_idx][:batch_size, start:end] = v
        self.seq_len = end

    def get(self, layer_idx):
        """返回当前所有缓存的 K/V"""
        return (self.k_cache[layer_idx][:, :self.seq_len],
                self.v_cache[layer_idx][:, :self.seq_len])

# 推理时使用
cache = KVCache(max_batch_size=1, max_seq_len=4096, n_layers=32, n_heads=32, head_dim=128)
for step in range(100):
    # 模型只计算当前 token 的 K/V
    k, v = model_layers[0].compute_kv(current_token)
    cache.append(0, k, v)
    # attention 时使用完整缓存的 K/V
    full_k, full_v = cache.get(0)
    attn_output = attention(q, full_k, full_v)

但这个实现有一个问题：预分配显存。max_seq_len=4096 时，批处理 batch=4 的情况下，32 层每层 32 个 head 共 128 维，KVCache 显存 = 4 × 4096 × 32 × 32 × 128 × 2 bytes ≈ 4GB。如果 max_seq_len 设成 128K，显存直接变 128GB——离了个大谱。而且不同请求的序列长度不同，预分配造成大量浪费。

KVCache 显存占用对比表（FP16, 32 层, 32 head, 128 dim）：

上下文长度	batch=1	batch=4	batch=16	7B 模型权重
4K	1 GB	4 GB	16 GB	14 GB
8K	2 GB	8 GB	32 GB	14 GB
32K	8 GB	32 GB	128 GB	14 GB
128K	32 GB	128 GB	512 GB	14 GB

结论： 上下文到 128K 时，batch=1 的 KVCache 就已经超过模型权重本身。这就是为什么 128K 上下文部署时单卡几乎跑不动——A100 80GB 上一张卡只能跑 1 个请求（KVCache 32GB + 权重 14GB + 中间激活 ≈ 50GB）。

2. PageAttention：vLLM 的杀手锏

vLLM 的 PageAttention 借鉴了操作系统的分页内存管理思想。传统 KVCache 是连续分配的，但连续内存有两大问题：

• 内部碎片：预分配了 max_seq_len 但实际只用了 10%，剩下 90% 浪费
• 外部碎片：不同请求的序列长度不同，释放的空间不连续，无法被新请求复用

分页 vs 连续分配对比（文字示意）：

连续分配（传统方式）：
  ┌──────────────────────────────────────────────┐
  │ 请求A: 预分配 4096 个 token 空间              │
  │ ｜ 实际用 512 个 ｜▓▓▓▓▓▓▓▓▓▓▓▓ 空白 ｜       │
  │ 浪费 3584 个 token 的显存                     │
  ├──────────────────────────────────────────────┤
  │ 请求B: 预分配 4096 个 token 空间              │
  │ ｜ 实际用 128 个 ｜▓▓▓▓▓▓▓▓▓▓▓▓ 空白 ｜       │
  └──────────────────────────────────────────────┘
  → 内部碎片率: 50-90%

PageAttention（分页方式）：
  物理 page 池: [page0][page1][page2][page3]...[pageN]
  请求A 页表: 逻辑页0→page3, 逻辑页1→page7, 逻辑页2→page1
  请求B 页表: 逻辑页0→page5, 逻辑页1→page2
  → 按需分配，用完释放，碎片率接近 0%

class PageAttention:
    """PageAttention 的简化示意，展示分页管理核心逻辑"""
    def __init__(self, page_size=16, total_pages=1024, head_dim=128, n_heads=32, dtype=torch.float16):
        self.page_size = page_size
        self.total_pages = total_pages
        # 物理 page 池，一次分配所有页
        self.k_page_pool = torch.zeros(total_pages, page_size, n_heads, head_dim, dtype=dtype)
        self.v_page_pool = torch.zeros(total_pages, page_size, n_heads, head_dim, dtype=dtype)
        # 空闲页列表
        self.free_pages = list(range(total_pages))
        # 页表：seq_id -> [page_id, ...]
        self.page_tables = {}

    def alloc(self, seq_id, num_pages):
        """为序列分配物理页"""
        pages = self.free_pages[:num_pages]
        self.free_pages = self.free_pages[num_pages:]
        self.page_tables[seq_id] = pages

    def write(self, seq_id, token_idx, k, v):
        """将 K/V 写入对应 token 的物理页"""
        page_idx = token_idx // self.page_size
        offset = token_idx % self.page_size
        phy_page = self.page_tables[seq_id][page_idx]
        self.k_page_pool[phy_page, offset] = k
        self.v_page_pool[phy_page, offset] = v

    def read(self, seq_id, token_idx):
        """读取对应 token 的 K/V"""
        page_idx = token_idx // self.page_size
        offset = token_idx % self.page_size
        phy_page = self.page_tables[seq_id][page_idx]
        return (self.k_page_pool[phy_page, offset],
                self.v_page_pool[phy_page, offset])

    def free(self, seq_id):
        """释放序列占用的页，回收到空闲池"""
        self.free_pages.extend(self.page_tables.pop(seq_id))

这个设计直接解决了显存碎片问题。当一个请求结束，释放的 page 立即就可以被新请求复用，不需要 GC。而且多个请求共享相同 prefix 时（如系统 prompt 相同），可以共享 KVCache 的 page，进一步节省显存。

实测效果：vLLM 在相同硬件上比 HuggingFace 原生推理吞吐量高 2-4 倍，核心原因就是 PageAttention 消除了 KVCache 的显存碎片化，让 GPU 显存利用率接近 100%。

3. Continuous Batching：动态调度取代静态批处理

传统推理的 batching 是静态的：一个 batch 的所有请求同时开始，同时结束，中间不能加新请求。如果 batch 里有一个长序列，所有短序列都得等它。这在 LLM 场景下非常低效——不同请求的序列长度差异巨大，等待时间差异也大。

静态批次 vs 动态批次对比（文字时序）：

静态 Batching（传统方式）：
  时间轴
  ────▶
  Batch 1: [A(50t) | B(50t) | C(200t)]  ← C 是长序列
  Batch 2: [D(30t) | E(30t) | F(30t)]  ← 必须等 Batch 1 全部结束
  Batch 3: ...
  → GPU 利用率 30-50%，大量时间在等待

Continuous Batching（vLLM 方式）：
  时间轴
  ────▶
  Step 1: [A | B | C]  ← 三人同时生成第 1 个 token
  Step 2: [A | B | C]  ← 三人同时生成第 2 个 token
  ...
  Step 50: [A(完成) | B(完成) | C]  ← A 和 B 结束，释放显存
  Step 51: [C | D | E]  ← D 和 E 立即加入，不等待
  ...
  → GPU 利用率 85-95%，几乎没有空闲期

class ContinuousBatchingScheduler:
    """
    Continuous Batching 调度器简化示意
    每步都重新调度，没有固定 batch 边界
    """
    def __init__(self, max_batch_size=64):
        self.running = []      # 正在生成的请求
        self.waiting = []      # 等待的请求
        self.completed = []    # 完成的请求
        self.max_batch_size = max_batch_size

    def step(self):
        """每一轮 token 生成结束后调度"""
        # 移除已完成的请求
        still_running = []
        for req in self.running:
            if req.eos or len(req.tokens) >= req.max_new_tokens:
                self.completed.append(req)
            else:
                still_running.append(req)
        self.running = still_running

        # 从等待队列中调度新请求，填满 batch
        slots = self.max_batch_size - len(self.running)
        if slots > 0 and self.waiting:
            # 按优先级调度（FCFS / Shortest Job First 等策略）
            new_reqs = self.waiting[:slots]
            self.waiting = self.waiting[slots:]
            self.running.extend(new_reqs)

        # 当前 batch 的所有请求生成下一个 token
        batch = self.running
        if not batch:
            return

        # 生成每个请求的下一个 token
        for req in batch:
            next_token = model.generate_next_token(req.tokens, kvcache=req.kvcache)
            req.tokens.append(next_token)

Continuous Batching 让 GPU 利用率从静态批处理的 30-50% 提升到 85-95%。核心收益：短序列不会被长序列卡住，长序列也不会因为等待短序列而浪费 GPU 时间。

4. Streaming：首 token 延迟

Streaming 是降低用户感知延迟的关键技术。传统方式：等模型生成完所有 token 再一次性返回，用户要等 N 秒才能看到第一个字。Streaming 方式：模型生成第一个 token 后就开始输出，用户几乎立即看到第一个字，后面的字逐 token 流式到达。

Streaming 的工程实现不复杂，但有一个容易被忽略的问题：Token 的逐字输出频率。模型生成 token 的速度取决于 GPU 算力和模型大小，7B 模型在 A100 上大约 30-50 tokens/s，13B 约 20-30 tokens/s。如果每个 token 都单独发一次 SSE 事件，前端渲染压力大；如果攒一批再发，用户感知延迟变差。实践中用 token buffer 策略：每 3-5 个 token 或每 50ms flush 一次，平衡感知延迟和渲染效率。

import asyncio
from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

async def generate_stream(prompt: str):
    """流式生成，使用 token buffer 控制输出频率"""
    tokens = tokenize(prompt)
    kvcache = KVCache(...)
    buffer = []
    last_flush = time.time()

    for step in range(256):
        logits = model.forward(tokens[-1:], kvcache=kvcache)
        next_token = sample(logits)
        tokens.append(next_token)
        token_text = detokenize(next_token)
        buffer.append(token_text)

        # 每 50ms 或 buffer 满 5 个 token 时 flush
        if len(buffer) >= 5 or (time.time() - last_flush) > 0.05:
            yield "".join(buffer)
            buffer.clear()
            last_flush = time.time()

        if next_token == EOS_TOKEN:
            break

    if buffer:
        yield "".join(buffer)

@app.post("/chat/stream")
async def chat_stream(prompt: str):
    return StreamingResponse(generate_stream(prompt), media_type="text/event-stream")

生产环境踩坑

坑 1：KVCache 量化导致精度下降

用 INT8 KVCache 时，显存减半但长序列场景下精度下降明显。实测：4K 上下文时 INT8 和 FP16 的 PPL（困惑度）差异 < 0.1，但 32K 上下文时差异扩大到 0.5-1.0。解决：混合精度方案——前 4K 用 FP16，后面用 INT8，平衡精度和显存。

坑 2：Batch Size 调太大反而变慢

Continuous Batching 并不是 batch 越大越好。当 batch size 超过某个阈值后，单次 token 生成时间变长，导致每个请求的 TTFT（Time to First Token）飙升。经验值：A100 80GB 上 7B 模型，batch 64 是甜点值，超过 64 后吞吐量不升反降。

坑 3：Prefix Cache 对长 prompt 反而没用

系统 prompt 超过 512 token 时，prefix cache 的命中率会显著下降。因为每个请求的 prompt 虽然前缀相同，但后面的个性化部分差异巨大，导致 GPU 的 page 表查找开销超过了计算节省。实践：只对 128 token 以内的系统 prompt 开启 prefix cache，长的直接走全量计算。

坑 4：Streaming 的 SSE 连接超时

前端 SSE 连接如果长时间没有 token 回来（比如 prefill 阶段耗时 > 30s），浏览器会断开连接。解决：prefill 阶段发心跳事件，或者用 WebSocket 替代 SSE。

生产环境的优化组合

面试官真正想听的是：这些优化怎么组合使用。以下是生产环境的最佳实践组合：

1. KVCache + INT8 量化：KVCache 的显存占用直接减半。vLLM 从 0.4.0 版本开始支持 FP8 KVCache（H100 硬件），这是 2025 年大规模部署的标配。
2. Prefix Caching + PageAttention：公共 prompt 前缀（如系统 prompt）的 KVCache page 可以在多个请求间共享，减少重复计算。这对 Chat 场景特别有效——系统 prompt 通常是固定的，缓存后可以节省 30-50% 的 prefill 时间。
3. Speculative Decoding：用小模型（如 1.3B）先快速预测 5-10 个候选 token，大模型（7B）一次验证，速度提升 2-3 倍。2025 年 vLLM 和 TensorRT-LLM 都原生支持。
4. Chunked Prefill：长 prompt 的 prefill 阶段按 chunk 分次处理，避免单次显存峰值过高。vLLM 的默认策略，对 128K 上下文特别有效。
5. Disaggregated Prefill/Decode (DPO)：2025 年大厂新实践，将 prefill 阶段和 decode 阶段分离到不同 GPU 上。Prefill GPU 用高端 GPU 做高算力，一次处理 batch 的 prefill；decode GPU 用大显存 GPU 做低延迟的逐 token 生成。实现 RPS 提升 2-3 倍。

FAQ（面试高频题）

Q1：KVCache 为什么不在训练时用？ A：训练时是 teacher forcing，每个 token 的 K/V 同时计算，不存在"复用"的问题。而且训练时 batch 内所有序列 padding 到相同长度，也不需要分页管理。

Q2：MHA 和 MQA/GQA 的 KVCache 有什么区别？ A：MHA（Multi-Head Attention）有 32 组 Q/K/V，KVCache 存 32 组。MQA（Multi-Query Attention）所有 head 共享 K/V，KVCache 只有 1 组，显存降到 1/32。GQA（Grouped Query Attention）折中，8 组 Q 共享 1 组 K/V，显存降到 1/4。LLaMA 2 70B 用 GQA，LLaMA 3 全系列用 GQA。

Q3：为什么 vLLM 的 PageAttention 在 prefill 阶段反而比原生慢？ A：prefill 阶段需要构建页表，存在额外的元数据开销。但 decode 阶段因为页表已经建好，可以高效复用。所以 vLLM 的加速主要在 decode 阶段，prefill 阶段差异不大。

Q4：FlashAttention 和 PageAttention 的关系？ A：FlashAttention 是 attention 计算层面的优化（IO-aware，减少显存读写），PageAttention 是 KVCache 管理层面的优化（分页，减少碎片）。两者互补，vLLM 同时使用了 FlashAttention 和 PageAttention。

总结

大模型推理优化的核心就一句话：KVCache 是显存瓶颈，所以一切优化都围绕怎么更高效地管理 KVCache。PageAttention 解决了显存碎片化问题，Continuous Batching 解决了 GPU 利用率问题，Streaming 解决了用户感知延迟问题。这三者不是孤立的，vLLM 把它们整合在一起，才实现了 2-4 倍的吞吐量提升。

参考：vLLM 论文（Kwon et al.）、FlashAttention 论文、Continuous Batching 论文、TensorRT-LLM 文档