分布式 LLM 系统中的 KV Cache 管理：放置、卸载与迁移策略

在自回归 Transformer 模型推理中，高并发 LLM 服务引擎瓶颈已逐渐从浮点计算能力（FLOPs）转向高带宽显存（HBM）容量与显存带宽的极限。为了高效运行诸如 DeepSeek-V3 或 Claude 3.5 Sonnet 等超大规模模型，现代 Serving 架构必须将 Key-Value (KV) Cache 视作跨越片上 SRAM、本地 GPU HBM、Host 系统 DRAM、本地 PCIe NVMe 存储以及远程 RDMA 网络共享内存池的弹性分布式内存层级结构。

在构建大规模推理集群或通过 n1n.ai 等 API 聚合平台路由高并发请求时，如何管理长上下文带来的显存开销是决定系统吞吐量的核心命题。深入理解 KV Cache 的生命周期、放置机制、异步卸载协议以及跨节点迁移算法，对于保障毫秒级的首 token 延迟（TTFT）与 Token 间生成延迟（ITL）至关重要。

KV Cache 显存占用与理论边界模型

自回归 LLM 生成的运行生命周期可划分为两个阶段：计算密集的 Prefill 阶段（并行处理 Prompt Token）与受显存带宽限制的 Decode 阶段（逐个串行生成 Output Token）。在 Decode 过程中，每个 Transformer 层都需要计算最新生成的 Query Token 与历史所有 Key Token 的 Attention 分数，导致 KV 值的显存占用随序列长度与并发数线性扩展。

+-----------------------------------------------------------------------------+
| Concurrent Requests (B) * Context Length (S) * KV Footprint                 |
|                                                                             |
|  [ Request 1: S=8192 tokens  ] ---> [ GPU 0 HBM: 80 GiB Capacity ]          |
|  [ Request 2: S=32768 tokens ] ---> [ Allocates Blocks In-Place  ]          |
|  [ Request 3: S=16384 tokens ] ---> [ Dynamic Paging Exhaustion  ]          |
|                                           |                                 |
|                                           v                                 |
|                     +-----------------------------------+                   |
|                     | Out-Of-Memory (OOM) / Head-of-Line|                   |
|                     | Blocking / Engine Stalls          |                   |
|                     +-----------------------------------+                   |
+-----------------------------------------------------------------------------+

为了精准建模显存占用，假设 Transformer 模型包含 $L$ 个网络层，KV 注意力头数为 $H_{\\text{KV}}$ ，单头维度为 $D_h$ ，总序列长度为 $S$ （包含 Prompt 长度 $S_{\\text{prompt}}$ 与生成长度 $S_{\\text{gen}}$ ），并发批处理大小为 $B$ ，数值精度标量字节数为 $P_{\\text{bytes}}$ （例如 FP16 = 2 字节，FP8 = 1 字节，INT4 = 0.5 字节）。

其未压缩的物理 KV Cache 基础占用 $M_{\\text{KV}}$ 字节数公式如下：

$M_{\\text{KV}} = 2 \\times L \\times H_{\\text{KV}} \\times D_h \\times S \\times B \\times P_{\\text{bytes}}$

常数系数 $2$ 分别代表 Key 与 Value 两个张量数组。不同的注意力机制对 $H_{\\text{KV}}$ 的影响显著：

Multi-Head Attention (MHA): $H_{\\text{KV}} = H_Q$ （等于 Query 头数）
Grouped-Query Attention (GQA): $H_{\\text{KV}} = \\frac{H_Q}{G}$ （ $G$ 为分组比例，通常 $G=8$ ）
Multi-Query Attention (MQA): $H_{\\text{KV}} = 1$

在现代分页内存管理架构（如 vLLM 或 TensorRT-LLM 的 PagedAttention）中，物理显存被划分为包含 $N_{\\text{block}}$ 个 Token 的固定大小物理块。考虑内部内存碎片后的实际分配显存 $M_{\\text{allocated}}$ 满足以下边界条件：

$M_{\\text{allocated}} = 2 \\times L \\times H_{\\text{KV}} \\times D_h \\times \\left( \\left\\lceil \\frac{S}{N_{\\text{block}}} \\right\\rceil \\times N_{\\text{block}} \\right) \\times B \\times P_{\\text{bytes}}$

$\\text{内部碎片开销比例} = \\frac{M_{\\text{allocated}} - M_{\\text{KV}}}{M_{\\text{KV}}} \\le \\frac{N_{\\text{block}} - 1}{S}$

+-----------------------------------------------------------------------------+
| Layer l: Paged Memory Block Layout (N_block tokens per physical slot)       |
|                                                                             |
| Block Index: 0x7F04                                                         |
| [ Key Head 0  | Token 0..N_block ] [ Key Head 1  | Token 0..N_block ] ...   |
| [ Val Head 0  | Token 0..N_block ] [ Val Head 1  | Token 0..N_block ] ...   |
|                                                                             |
| Block Index: 0x7F05 (Non-contiguous Physical HBM)                           |
| [ Key Head 0  | Token N_block..2N_block ] ...                               |
+-----------------------------------------------------------------------------+

多层级异构内存拓扑架构（Hierarchical Topology）

为避免高并发场景下的 Out-Of-Memory (OOM) 崩溃与队头阻塞（Head-of-Line Blocking），高性能 Serving 引擎将 KV Cache 映射到多级异构存储拓扑中：

                     +---------------------+ 
                     |       Request       |
                     +----------+----------+
                                |
                                v
                      +------------------+
                      | KV Cache Manager |
                      +--------+---------+
                               |
              +----------------+----------------+
              |                |                |
              v                v                v
          +---------+    +-----------+    +---------------+ 
          | GPU HBM |    | Host DRAM |    | Remote Memory |
          | Hot KV  |    |  Warm KV  |    | Cold/Shared KV|
          +----+----+    +-----+-----+    +-------+-------+
               |               |                  |
               +---------------+------------------+
                               |
                               v
                       +---------------+ 
                       | Decode Engine |
                       +---------------+

Tier 0 (SRAM / 寄存器堆): 流多处理器（SM）内部的片上高速缓存。在 FlashAttention 内核计算中用于暂存子块矩阵，访问延迟 < 1 ns。
Tier 1 (GPU HBM - 热数据层): 直接连通 Tensor Core 的高带宽显存。聚合带宽达 2.0–8.0 TB/s（如 NVIDIA H100/H200）。访问延迟 < 1 µs。
Tier 2 (Host DRAM - 温数据层): 通过 PCIe 总线连接的 CPU 系统内存。用于存放预缓存的 Prompt Prefix、暂停的会话块与溢出页。访问延迟 1–5 µs。
Tier 3 (本地 NVMe SSD - 冷数据层): PCIe NVMe 闪存，用于长会话持久化与低频 Prefix 存储。访问延迟 10–100 µs。
Tier 4 (远程网络共享内存池 - 分布式共享层): 通过 RoCEv2 或 InfiniBand 构建的分布式 DRAM 共享池。访问延迟 5–25 µs。

各存储层级性能指标与约束对比

存储层级	容量边界	单向传输带宽	典型访问延迟	主导系统瓶颈
GPU HBM (本地)	80–144 GiB / GPU	2.0–8.0 TB/s	`< 1 µs`	物理显存容量极限
Host DRAM (系统内存)	512–2048 GiB / 节点	30–64 GB/s (PCIe Gen5)	1–5 µs	PCIe 总线带宽争用
本地 PCIe NVMe SSD	3.8–30.7 TB / 盘	7–14 GB/s	10–100 µs	I/O 读写 IOPS 与闪存寿命
远程节点 DRAM	Multi-TB 共享池	25–50 GB/s (网卡链路)	5–25 µs	跨机架网络对分带宽
对等 GPU (远程 HBM)	80–144 GiB / GPU	450–900 GB/s (NVLink)	`< 1 µs`	NUMA NVLink 拓扑域限制

动态放置与重算 vs 传输决策数学模型

放置引擎需动态评估：某一请求的 KV 块应当保留在 HBM、卸载至 Host DRAM、通过 RDMA 远程拉取，还是直接丢弃并在本地重新计算？

设 $T_{\\text{recompute}}(S)$ 为对 $S$ 个 Token 执行 Prefill 重算所需的时间， $T_{\\text{transfer}}(\\text{Tier}_k, S)$ 为从第 $k$ 层存储拉取数据的时间（传输带宽为 $B_k$ ，基础建连延迟为 $\\alpha_k$ ）：

$T_{\\text{transfer}}(\\text{Tier}_k, S) = \\alpha_k + \\frac{M_{\\text{KV}}(S)}{B_k}$

$T_{\\text{recompute}}(S) \\approx \\frac{\\text{FLOPs}_{\\text{prefill}}(S)}{\\text{Attainable FLOPS}_{\\text{GPU}}} = \\frac{2 \\times N_{\\text{params}} \\times S + 4 \\times L \\times H_Q \\times D_h \\times S^2}{\\text{Attainable FLOPS}_{\\text{GPU}}}$

只有当跨层传输与内存调度开销严格小于本地前向重算成本时，从异构层级 $\\text{Tier}_k$ 迁移数据才是正向收益的：

$\\mathbb{E}[\\text{Benefit}] = P(\\text{Reuse}) \\times \\left( T_{\\text{recompute}}(S) - T_{\\text{transfer}}(\\text{Tier}_k, S) \\right) > 0$

其中 $P(\\text{Reuse})$ 表示该 Prefix 或序列状态在被淘汰前再次被请求调用的经验概率。

+-----------------------------------------------------------------------------+
| Placement Boundary Decision Logic                                           |
|                                                                             |
|      +---------------------------------------------------------+            |
|      | T_transfer(Tier_k, S) &lt; T_recompute(S) ?                |            |
|      +----------------------------+----------------------------+            |
|                                   |                                         |
|                 +-----------------+-----------------+                       |
|                 | YES                               | NO                    |
|                 v                                   v                       |
|   +---------------------------+       +---------------------------------+   |
|   | Fetch from Tier_k via     |       | Recompute KV states via         |   |
|   | DMA / RDMA Stream Pipeline|       | Local Prefill Kernel Execution  |   |
|   +---------------------------+       +---------------------------------+   |
+-----------------------------------------------------------------------------+

异步卸载与预取流水线（Asynchronous Offloading）

当 GPU HBM 占用率触及警戒线时，主动卸载机制可将非活跃的物理页换出至 Host DRAM，在不破坏计算上下文的前提下释放 HBM。

       EVICTION PIPELINE                           PREFETCH PIPELINE
     +-------------------+                       +-------------------+
     |   GPU HBM Block   |                       |  Decode Request   |
     +---------+---------+                       +---------+---------+
               |                                           |
               | (Async D2H Copy)                          v
               v                                 +-------------------+
     +-------------------+                       | Block Location?   |
     |   Host DRAM       |                       +----+----+----+----+
     +---------+---------+                            |    |    |
               |                                      |    |    +---> Remote: Network Fetch
               | (Page Write)                         |    +--------> Host: PCIe Prefetch
               v                                      +-------------> HBM: Direct Access
     +-------------------+
     | Local NVMe Flash  |
     +-------------------+

关键架构原语：

锁页内存分配（Pinned Host Memory）: 为实现最高效的直接内存访问（DMA），Host 端接收缓冲区必须通过 cudaHostRegister 锁页，避免操作系统内部二次拷贝，从而吃满 PCIe 总线带宽。
异步 CUDA 流流水线: 内存换入换出操作运行在独立的 CUDA Copy Stream (cudaStream_t) 上，与 Decode 算子重叠执行。Decode 引擎在处理第 $N$ 步生成时，Copy 流同步预取第 $N+1$ 步所需的 KV 块。
双缓冲区滑动窗口: 对于超长上下文，系统在 HBM 中保留当前活跃的滑动窗口，后台线程则通过 PCIe 从 Host 内存异步拉取后续窗口块。

当通过 n1n.ai 统一路由跨模型请求时，这种 Host 端卸载缓冲能够极大提升 Serving 节点抵御突发并发流量的能力，确保连接不断连。

分布式 Prefill-Decode 解耦与 RDMA 迁移

在现代 Prefill-Decode 解耦架构中，Prefill（计算密集型）与 Decode（带宽密集型）被拆分到不同的 GPU 节点群中。这要求 KV Cache 能够跨物理节点进行高效迁移。

+-------------------+                       +-------------------+
|   Prefill Pool    |                       |    Decode Pool    |
| (Compute-Dense)   |                       | (Bandwidth-Dense) |
|   [ GPU Worker ]  |                       |   [ GPU Worker ]  |
+---------+---------+                       +---------+---------+
          |                                           ^
          | 1. Export Paged Blocks                    | 4. Ingest Blocks
          v                                           |
+-------------------+                       +---------+---------+
| Local RDMA Subsys |                       | Local RDMA Subsys |
+---------+---------+                       +---------+---------+
          |                                           ^
          |        2. Kernel-Bypassing Transfer       |
          +===========================================+
                    3. RoCEv2 / InfiniBand Fabric

迁移执行步骤：

GPUDirect RDMA 传输: 传输过程彻底绕过 CPU Host 内存。利用 PCIe P2P 与网卡（如 NVIDIA ConnectX HCA），物理 KV 页通过零拷贝的 IBV_WR_RDMA_WRITE 直接从源 GPU HBM 写入目标 GPU HBM。
布局打包（Layout Swizzling）: PagedAttention 分配的非连续物理页会在发送前打包为连续缓冲区，或通过 RDMA 离散/聚合列表（Scatter/Gather Lists）进行一次性传输，极大提升网络 MTU 传输效率。
两阶段握手协议:
- 阶段 1（块分配与注册）: Decode 节点预分配物理块，通过 RPC 控制信道将内存 RKEY 与虚拟地址偏移发送至 Prefill 节点。
- 阶段 2（RDMA 触发与 Barrier 校验）: Prefill 节点发起 RDMA Write 并附带完成信号（IBV_SEND_WITH_IMM），Decode 节点收到通知后方可将块绑定至活跃页表。

基于树状 Radix Tree 的感知重算成本淘汰算法

传统的 LRU 算法在 LLM 推理场景中表现不佳，因为它忽视了共享 Prefix 结构与不同长度 Prefix 的重算成本差异。

+-----------------------------------------------------------------------------+
| Radix Prefix Tree Layout with Reference Counts                              |
|                                                                             |
| [ Root / Shared System Prompt: 4096 tokens ] (Ref Count = 48)  &lt;-- PINNED   |
|         |                                                                   |
|         +---&gt; [ Task Sub-Prompt A ] (Ref Count = 12)          &lt;-- PROTECTED |
|         |           |                                                       |
|         |           +---&gt; [ Request Context 1 ] (Ref = 1)     &lt;-- EVICTABLE |
|         |                                                                   |
|         +---&gt; [ Task Sub-Prompt B ] (Ref Count = 1)           &lt;-- EVICTABLE |
+-----------------------------------------------------------------------------+

淘汰与压缩策略：

Radix Tree 引用计数淘汰: 显存块在基数树（Radix Tree）中进行组织。共享根节点（如公共 System Prompt）保持高引用计数（RefCount > 1）并被锁定（Pinned）；淘汰管理器严格从引用计数为 1 的叶子节点开始清理。
感知重算成本的淘汰权重: Prefill 计算复杂度与 Token 长度呈二次方关系 $\\mathcal{O}(S^2)$ 。淘汰一个长度 $S=16384$ 的块会导致极其高昂的重算代价；因此淘汰引擎倾向于优先淘汰短期 Prefix。
Attention 权重稀疏化剪枝: 类似 Heavy-Hitter Oracle ( $H_2O$ ) 与 StreamingLLM 的算法会在生成过程中动态分析 Attention 权重，将长期低贡献的 Token 块剪枝或压缩为低精度格式（如 FP8/INT4）。

综合淘汰得分 $\\Phi_i$ 的计算公式如下：

$\\Phi_i = \\omega_1 \\cdot \\text{Recency}_i + \\omega_2 \\cdot \\text{RefCount}_i + \\omega_3 \\cdot \\frac{T_{\\text{recompute}}(\\text{Block}_i)}{M_{\\text{KV}}(\\text{Block}_i)} - \\omega_4 \\cdot \\text{SLO\\_Deficit}_i$

系统将优先淘汰得分 $\\Phi_i$ 最低的 Candidate 块。

生产环境常见故障模式与规避方案

故障模式	根本原因机制	系统外部表现	生产级规避策略
HBM 显存抖动 (Thrashing)	工作集在 HBM 边界频繁换入换出	PCIe 总线饱满，Decode 延迟剧增	设置阈值预留 Margin (`> 15%`)；主动降级投机采样长度
内存注册停顿 (Reg MR Stall)	在推理主线程中动态调用 `ibv_reg_mr`	执行线程被阻塞，出现 ITL 尖刺	节点启动时预先分配并注册静态内存 Block Pool
迁移竞争条件 (Race Condition)	RDMA 传输未落盘即触发 Decode 算子	内存访问非法，生成文本乱码	严格执行 `IBV_SEND_WITH_IMM` 完成队列轮询同步
故障转移重算雪崩	主节点宕机导致热 Cache 丢失，备节点迎接收束	TTFT 严重超标，队列超时崩溃	在独立故障域节点上异步冗余复制 Root 级 Prefix

实践指南：异步 KV Cache 卸载器 Python 实现

以下 Python 示例展示了一个基于 PyTorch 的多层级 KV Cache 管理器，实现了在 GPU HBM 与锁页 Host DRAM 之间利用非阻塞 CUDA 流异步换出/预取 KV 块的功能：

import torch
import typing
import time

class TieredKVCacheManager:
    def __init__(
        self, 
        num_layers: int, 
        num_heads: int, 
        head_dim: int, 
        block_size: int, 
        hbm_block_capacity: int, 
        host_block_capacity: int,
        dtype: torch.dtype = torch.float16
    ):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.block_size = block_size
        self.dtype = dtype

        # 单个 Block 的 Shape ( Key 和 Value 两个数组)
        self.block_shape = (2, num_layers, num_heads, block_size, head_dim)
        
        # 1. 初始化 Tier 1: 本地 GPU HBM 显存池
        self.hbm_pool = torch.empty(
            (hbm_block_capacity, *self.block_shape),
            dtype=dtype,
            device="cuda:0"
        )
        
        # 2. 初始化 Tier 2: Host 锁页内存池 (Pinned Memory) 以实现 DMA 高速传输
        self.host_pool = torch.empty(
            (host_block_capacity, *self.block_shape),
            dtype=dtype,
            device="cpu"
        ).pin_memory()

        # 跟踪 Block 状态
        self.free_hbm_blocks = list(range(hbm_block_capacity))
        self.free_host_blocks = list(range(host_block_capacity))
        self.block_mapping: typing.Dict[int, typing.Dict[str, typing.Any]] = \{\}

        # 创建用于异步传输的专用 CUDA Stream
        self.offload_stream = torch.cuda.Stream(device="cuda:0")

    def allocate_block(self, logical_block_id: int) -> int:

参考来源：https://dev.to/wantsvibes/kv-cache-management-in-distributed-llm-systems-placement-offloading-and-migration-p1n