LLM 推理底层：KV Cache 显存几何、GQA/MLA 压缩与 PagedAttention 调度

自回归解码与 KV Cache 的物理来源

设因果（Causal）Transformer 包含 $L$ 层，隐藏层维度为 $d_{\text{model}} = n_h d_h$ （ $n_h$ 为查询头数， $d_h$ 为单头维度）。在第 $t$ 步生成时，输入序列 $x_{1:t}$ 的最后位置隐状态 $h_t \in \mathbb{R}^{1 \times d_{\text{model}}}$ 与历史位置 $1 \dots t$ 的键、值矩阵做缩放点积注意力：

q_t = h_t W_Q, \qquad k_i = h_i W_K, \qquad v_i = h_i W_V

o_t = \operatorname{softmax}\!\left(\frac{q_t K_{1:t}^{\mathsf T}}{\sqrt{d_h}}\right) V_{1:t}, \qquad K_{1:t} = \begin{bmatrix} k_1 \\ \vdots \\ k_t \end{bmatrix} \in \mathbb{R}^{t \times d_h}

由于因果掩码保证了先前的隐状态 $h_i$ （ $i < t$ ）不依赖后续 Token，矩阵 $K_{1:t-1}$ 与 $V_{1:t-1}$ 在步间严格不变。将其缓存在高带宽显存（HBM）中即为 KV Cache：单步解码的计算量从重新计算完整前缀的 $O(t d_{\text{model}}^2 + t^2 d_{\text{model}})$ 降为 $O(d_{\text{model}}^2 + t d_{\text{model}})$ ，代价是显存占用与 HBM 读取带宽随序列长度 $t$ 线性增长。

显存占用与 Roofline 算术强度分析

KV Cache 容量公式

设并发批次大小为 $B$ ，序列上下文长度为 $S$ ，每层包含 $n_{kv}$ 个键值头，单头维度为 $d_h$ ，数值精度字节数为 $b_{\text{dtype}}$ （FP16/BF16 为 $2$ ，FP8/INT8 为 $1$ ）。全部 $L$ 层的 KV Cache 显存占用为：

M_{\text{KV}} = 2 \cdot L \cdot n_{kv} \cdot d_h \cdot S \cdot B \cdot b_{\text{dtype}}

以 Llama-3-70B（ $L=80, n_h=64, n_{kv}=8, d_h=128$ ，BF16）在 $B=16, S=8192$ 下为例：

若采用传统多头注意力（MHA， $n_{kv}=64$ ）：单 Token 占用 $2 \times 80 \times 64 \times 128 \times 2 = 2.5\ \text{MiB}$ ，总计 $320\ \text{GiB}$ ，远超单台 8 卡 80G 节点的权重剩余显存；
采用分组查询注意力（GQA， $n_{kv}=8$ ）：单 Token 降至 $320\ \text{KiB}$ ，同批次占用降至 $40\ \text{GiB}$ 。

Prefill 与 Decode 的算术强度对比

定义硬件峰值算力与显存带宽之比为机器平衡点（Ridge Point） $\mathcal{I}_{\text{ridge}} = \frac{\text{Peak FLOP/s}}{\text{Peak HBM BW}}$ 。对于 NVIDIA H100 SXM（BF16 峰值 $989\ \text{TFLOPS}$ ，HBM3 带宽 $3.35\ \text{TB/s}$ ）， $\mathcal{I}_{\text{ridge}} \approx 295\ \text{FLOPs/Byte}$ 。

考察单层注意力核心运算 $q_t K_{1:t}^{\mathsf T}$ 与权重矩阵乘法的算术强度 $\mathcal{I} = \frac{\text{FLOPs}}{\text{Bytes}}$ ：

Prefill 阶段（输入长度 $S \gg 1$ ）：对于线性投影 $Y = XW$ （ $X \in \mathbb{R}^{B S \times d}$ ， $W \in \mathbb{R}^{d \times d}$ ），算术强度为：

\mathcal{I}_{\text{GEMM}} \approx \frac{2 B S d^2}{(2 B S d + 2 d^2 + 2 B S d)} \xrightarrow{BS \ll d} BS

当 $BS \ge 295$ 时运算进入 Compute-Bound 区域，GPU Tensor Core 利用率饱和。

Decode 阶段（单步生成 1 个 Token）：注意力部分对每个 KV 头共享 $g = n_h / n_{kv}$ 个查询头。计算 $g$ 个查询向量与长度 $S$ 的一条 KV 序列交互，浮点运算次数为 $4 g S d_h$ ，而从 HBM 读取 KV Cache 的字节数为 $4 S d_h$ （BF16）：

\mathcal{I}_{\text{Attn-Decode}} = \frac{4 g S d_h}{4 S d_h + 2 g d_h} \approx g = \frac{n_h}{n_{kv}}

注意到分母与分子中的序列长度 $S$ 相消：无论 Batch Size $B$ 增大到多少，由于不同请求的上下文各自独立（不共享前缀时），注意力核的算术强度恒被锁死在 $g = n_h / n_{kv}$ （MHA 为 $1$ ，GQA 通常为 $4 \sim 8$ ），远低于 $295\ \text{FLOPs/Byte}$ 。因此长序列解码阶段的 Attention 严格受限于 HBM 带宽（Memory-Bandwidth Bound）。

注意力头架构演进：MHA、GQA 与 MLA

为同时压缩 KV Cache 容量并提升解码阶段的算术强度，主流模型经历了从 MHA、MQA/GQA 到多头潜在注意力（Multi-Head Latent Attention, MLA）的演进：

架构	单层每 Token 缓存元素数	Attention 解码算术强度（BF16）	核心机制
MHA	$2 n_h d_h$	$\approx 1$	每个查询头独占一组 $(K, V)$
GQA	$2 n_{kv} d_h$	$\approx n_h / n_{kv}$	每组 $g = n_h/n_{kv}$ 个查询头共享 $(K, V)$
MLA	$d_c + d_r$	$\approx \dfrac{2 n_h (d_c + d_r)}{d_c + d_r} = 2 n_h$	低秩联合压缩隐向量 $c_t^{KV}$ 与解耦 RoPE 键 $k_t^R$

MLA 的矩阵吸收（Matrix Absorption）推导

DeepSeek-V2/V3 提出的 MLA 将隐状态 $h_t$ 投影为低维潜在向量 $c_t^{KV} \in \mathbb{R}^{d_c}$ （ $d_c \ll n_h d_h$ ）：

c_t^{KV} = h_t W^{DKV}, \qquad k_t^C = c_t^{KV} W^{UK}, \qquad v_t^C = c_t^{KV} W^{UV}

如果在推理时将 $c_t^{KV}$ 显式上投影恢复成 $k_t^C$ ，计算量与带宽将退化。利用矩阵乘法结合律，对于第 $i$ 个头的上投影块 $W_i^{UK} \in \mathbb{R}^{d_c \times d_h}$ ：

q_{t,i} (k_{j,i}^C)^{\mathsf T} = (h_t W_i^Q) (c_j^{KV} W_i^{UK})^{\mathsf T} = \bigl(h_t (W_i^Q (W_i^{UK})^{\mathsf T})\bigr) (c_j^{KV})^{\mathsf T}

因此可以将 $(W_i^{UK})^{\mathsf T}$ 吸收到查询侧投影中，将 $W_i^{UV}$ 吸收到输出投影 $W_i^O$ 中。但若直接在 $k_t^C$ 上施加旋转位置编码 $\operatorname{RoPE}(t)$ ，中间夹入与位置差 $t - j$ 相关的正交旋转矩阵 $\mathcal{R}_{t-j}$ ， $W_i^{UK}$ 将无法与 $W_i^Q$ 静态合并。MLA 为此引入解耦 RoPE（Decoupled RoPE）：单独切分出一个小维度共享键 $k_t^R = \operatorname{RoPE}(h_t W^{KR}) \in \mathbb{R}^{d_r}$ 承载位置信息，从而只需缓存 $(c_t^{KV}, k_t^R)$ ，在 $n_h = 128$ 时将注意力解码的算术强度从 $O(1)$ 提升至 $O(n_h)$ ，使长上下文解码可利用 Tensor Core 矩阵乘算力（即 FlashMLA 所优化的形态）。

PagedAttention：虚拟内存与非连续块表调度

连续预分配的三类显存碎片

早期推理框架为每个请求预先分配形状为 $[B, n_{kv}, S_{\max}, d_h]$ 的连续显存张量，引发三类严重浪费：

预留浪费（Reserved Fragmentation）：为未来可能生成的 Token 预留但当前未使用的槽位；
内部碎片（Internal Fragmentation）：请求提前遇到 <eos> 终止，剩余预留槽位直至请求释放前无法被其他请求复用；
外部碎片（External Fragmentation）：不同长短请求交错分配与释放后，HBM 中剩余大量不连续空洞，无法满足新到达长上下文请求的连续分配需求。

逻辑块与物理块映射

vLLM 引入的 PagedAttention 借鉴操作系统虚拟内存分页机制，将连续的逻辑 KV 序列切分为固定大小 $B_{\text{block}}$ （通常取 $16$ 或 $32$ 个 Token）的逻辑块（Logical Blocks），通过每个请求维护的**块表（Block Table）**映射到 HBM 池中的非连续物理块（Physical Blocks）。

对于序列的第 $j$ 个 Token（ $0$ -indexed），其所在的逻辑块号与块内偏移分别为：

\text{block\_idx} = \left\lfloor \frac{j}{B_{\text{block}}} \right\rfloor, \qquad \text{block\_offset} = j \bmod B_{\text{block}}

物理显存地址由块表查表给出： $\text{phys\_block} = \text{BlockTable}[\text{seq\_id}][\text{block\_idx}]$ 。如此一来，内部碎片严格被限制在每个请求的最后一个未填满物理块内，至多浪费 $B_{\text{block}} - 1$ 个 Token，显存有效利用率可超过 $96\%$ 。

Prefix Caching 与 Copy-on-Write

当多个请求共享相同的系统提示词（System Prompt）、Few-shot 示例或多轮对话历史时，可以将填满的物理块以链式哈希标识：

H_k = \operatorname{Hash}\!\left(H_{k-1},\; \text{Tokens}_{k B_{\text{block}} : (k+1) B_{\text{block}}}\right), \qquad H_{-1} = 0

基数树 / 哈希表前缀复用（Automatic Prefix Caching）：空闲物理块在引用计数归零后不立即抹除，而是放入按 $H_k$ 索引的 LRU 驱逐池；新请求在 Prefill 前沿前缀匹配 $H_0, H_1, \dots$ ，命中块只需递增 ref_count 并把跳过部分从 Prefill 计算中裁除。
写时复制（Copy-on-Write）：在并行采样（Parallel Sampling）或 Beam Search 中，多条候选分支共享同一个未填满尾块（ref_count > 1）。当某条分支向该块追加新 Token 时，分配器才从空闲池申请一个新物理块、拷贝已有槽位并递减原块引用计数。

块分配器与 Prefix Caching 参考实现

下面给出支持引用计数、链式哈希前缀复用与 Copy-on-Write 的块分配器核心实现：

from collections import OrderedDict
from dataclasses import dataclass, field
import hashlib


def compute_block_hash(parent_hash: int, token_ids: tuple[int, ...]) -> int:
    h = hashlib.blake2b(digest_size=8)
    h.update(parent_hash.to_bytes(8, "little", signed=False))
    for tid in token_ids:
        h.update(tid.to_bytes(4, "little", signed=False))
    return int.from_bytes(h.digest(), "little", signed=False)


@dataclass
class PhysicalBlock:
    block_id: int
    block_size: int
    ref_count: int = 0
    block_hash: int | None = None
    tokens: list[int] = field(default_factory=list)

    @property
    def is_full(self) -> bool:
        return len(self.tokens) == self.block_size


class PagedKVAllocator:
    def __init__(self, num_blocks: int, block_size: int = 16) -> None:
        self.block_size = block_size
        self.blocks = [PhysicalBlock(i, block_size) for i in range(num_blocks)]
        # evictable_pool 保存 ref_count == 0 的块，按 LRU 顺序组织
        self.evictable_pool: OrderedDict[int, PhysicalBlock] = OrderedDict(
            (i, self.blocks[i]) for i in range(num_blocks)
        )
        self.hash_to_block: dict[int, PhysicalBlock] = {}

    def _acquire_raw_block(self) -> PhysicalBlock:
        if not self.evictable_pool:
            raise MemoryError("Out of KV cache physical blocks")
        _, blk = self.evictable_pool.popitem(last=False)
        if blk.block_hash is not None:
            self.hash_to_block.pop(blk.block_hash, None)
            blk.block_hash = None
        blk.tokens.clear()
        blk.ref_count = 1
        return blk

    def allocate_for_prompt(self, token_ids: list[int]) -> tuple[list[int], int]:
        """为新请求分配块表，返回 (block_table, 命中的缓存 token 数)。"""
        block_table: list[int] = []
        cached_tokens = 0
        parent_hash = 0
        prefix_Caching_active = True

        for start in range(0, len(token_ids), self.block_size):
            chunk = token_ids[start : start + self.block_size]
            if len(chunk) == self.block_size and prefix_Caching_active:
                b_hash = compute_block_hash(parent_hash, tuple(chunk))
                cached = self.hash_to_block.get(b_hash)
                if cached is not None:
                    if cached.ref_count == 0:
                        self.evictable_pool.pop(cached.block_id, None)
                    cached.ref_count += 1
                    block_table.append(cached.block_id)
                    cached_tokens += self.block_size
                    parent_hash = b_hash
                    continue
                prefix_Caching_active = False

            blk = self._acquire_raw_block()
            blk.tokens.extend(chunk)
            if blk.is_full:
                b_hash = compute_block_hash(parent_hash, tuple(chunk))
                blk.block_hash = b_hash
                self.hash_to_block[b_hash] = blk
                parent_hash = b_hash
            block_table.append(blk.block_id)

        return block_table, cached_tokens

    def append_slot(self, block_table: list[int], new_token: int) -> None:
        """自回归解码追加单 Token，处理尾块满或 Copy-on-Write 分裂。"""
        last_blk = self.blocks[block_table[-1]]
        if last_blk.is_full:
            new_blk = self._acquire_raw_block()
            new_blk.tokens.append(new_token)
            block_table.append(new_blk.block_id)
            return

        if last_blk.ref_count > 1:
            # Copy-on-Write：分裂共享的未满尾块
            last_blk.ref_count -= 1
            cow_blk = self._acquire_raw_block()
            cow_blk.tokens.extend(last_blk.tokens)
            cow_blk.tokens.append(new_token)
            block_table[-1] = cow_blk.block_id
            return

        last_blk.tokens.append(new_token)

    def free_sequence(self, block_table: list[int]) -> None:
        for bid in block_table:
            blk = self.blocks[bid]
            blk.ref_count -= 1
            if blk.ref_count == 0:
                self.evictable_pool[bid] = blk

Continuous Batching、Chunked Prefill 与抢占恢复

在迭代级调度（Iteration-level Scheduling / Continuous Batching）中，所有活跃请求随解码步数增长不断消耗物理块。当空闲块耗尽、无法为当前批次所有序列分配下一个槽位时，调度器按后进先出（LIFO）或优先级策略触发抢占（Preemption）：

重计算（Recompute）：直接释放被抢占请求的物理块，将其生成的输出 Token 拼接回原始输入提示词，放入等待队列首部；恢复时作为一次新的 Prefill 重新计算 KV Cache。由于单 Token 重计算在 Prefill 下处于高算术强度区间，当 PCIe 带宽受限时往往快于卸载。
显存置换（Swap to Host RAM）：通过异步 cudaMemcpyAsync 将被抢占序列的物理块拷贝到主机锁页内存（Pinned Memory）的 CPU 块池中，唤醒时再拷回 HBM。
Chunked Prefill（分块预填充）：为避免长提示词请求（如 $S = 32\text{K}$ ）独占单个迭代步导致同批次正在 Decode 的请求出现数百毫秒的尾延迟（ITL Spike），调度器将长 Prefill 切分为固定 Token 预算（如每步 $2048$ 个 Token）的分块，与当前正在 Decode 的请求拼入同一个混合批次（Hybrid Batch）协同执行，既用 Decode 的闲置算力摊薄权重读取开销，又保障了流式输出的时间均匀性。