NInfer 分页 KV 缓存实现原理:页、副本与地址空间如何高效管理长上下文
【免费下载链接】ninferHigh-performance single-GPU inference for selected model checkpoints and GPUs.项目地址: https://gitcode.com/gh_mirrors/ni/ninfer
NInfer 是一个面向单 GPU 的高性能推理引擎,其分页 KV 缓存用固定 64 token 的页面、Device/Host 双副本和独立地址空间来管理长上下文显存。这篇文章带你完整理解 NInfer 分页 KV 缓存是如何做到"长上下文不爆显存、前缀可以共享、推理不搬数据"的。
1️⃣ 为什么长上下文绕不开 KV 缓存管理
Transformer 推理时,每生成一个新 token,Attention 都要回看前面所有 token 的 Key/Value(即 KV 缓存)。上下文越长,KV 缓存越大——一个 27B 级别模型跑 128K 上下文,KV 缓存可能达到数十 GB。
如果按"每条请求分配一整块连续显存"的传统方式管理,会遇到三个经典问题:
| 传统连续分配的问题 | 后果 |
|---|---|
| 显存内部碎片 | 申请 100GB 实际只用了 60GB |
| 无法多请求共享前缀 | 相同系统提示词重复 prefill,浪费算力 |
| 扩容需要搬迁 | 上下文增长时整块 KV 要拷贝,卡顿明显 |
NInfer 的答案借鉴了操作系统虚拟内存的"分页"思想:把 KV 切成小页面,用"逻辑页 → 物理页"的映射表间接寻址。核心实现位于 src/core/paged_kv_cache.h 与 src/core/paged_kv_storage.h,设计合同见官方文档 docs/maintainer/paged-kv-cache.md。
上图是 NInfer CLI 的多模态测试样本:模型要"看懂"山间小屋并回答图中数字(24)。这类多模态长上下文请求,正是分页 KV 缓存需要伺候好的典型负载。
2️⃣ 核心概念一页读懂:页、池、地址空间
NInfer 把"管理 KV"拆成了三层独立概念,各管各的粒度:
| 概念 | 通俗理解 | NInfer 中的设定 |
|---|---|---|
| 页(Page) | KV 的最小分配单位 | 固定P=64个 token 一页 |
| 池(Pool) | 互相隔离的显存"仓库" | 启动时固定,Main/MTP/Draft 各用一个池 |
| 地址空间(Address Space) | 每条序列自己的"门牌号序列" | 逻辑连续,物理可不连续 |
页的大小为什么是 64?文档 paged-kv-cache.md 解释得很清楚:64 恰好对齐 Attention 的 32/64-key 计算 tile 与 128-token 的 prefill 分块,同时让 block table 元数据足够小。页面边界不等于语义边界——有效 frontier 可以落在页内任意 token 位置,这一点保证了按 token 精确的截断与回滚。
池为什么按类型隔离?不同生命周期的 KV 绝不能混仓:
- Main Text 池:目标模型全注意力层的 K/V;
- MTP 池 / Draft Full 池:投机解码(MTP、DFlash)后端的辅助 KV,各自独立 frontier、独立提交、独立裁剪。
池内只收"同 frontier、同生命周期、同页大小"的数据面(K、V、量化 code、scale 面共享同一个页组 ID)。固定大小的页组意味着没有变长碎片,也无需显存压缩整理。
3️⃣ 物理页怎么分配:租约、预留与 block table
每个池内部区分三种容量的状态,这也是"active 请求不被挤占"的保证:
已分配页(lease) + 已预留页(reservation) + 全局可用页 = 池总容量- 预留(Reservation):为某条 active 序列锁定的"未来额度",不会被其他请求抢走;
- 租约(Lease):预留落地为真实物理页后的持有凭证,带 generation 代号防止释放后被旧句柄误用;
- Block table:每条 active 序列租用一行,把"逻辑块号 → 物理页号"批量发布到 GPU 上的固定矩阵
block_tables[N_logical, C]。
decode 每步通常只在跨越页边界时才申请一个新页,尾页还有空位时分配器完全不用干活——这是长上下文 decode 不掉速的关键之一。批量预留接口见 reserve_device_kv_page_bundle。
4️⃣ 副本机制:显存装不下就沉到主机内存
长上下文的另一个杀手锏是Device/Host 双副本:
- Device 副本用"消费端原生"的页平面布局,供 Attention kernel 直接读;
- Host 副本是紧凑打包布局,存放在启动时固定的 pinned 内存 arena(见 src/core/host_kv_arena.h)。
迁移遵循严格的发布顺序:预留目标位置 → 拷贝整页 → 校验内容 epoch 与覆盖范围 → 发布新副本 → 才释放旧副本。拷贝期间两端都被 pin 住,任何时刻都只有一个有效副本对外可见。Inactive 的上下文(如会话缓存)可以整体沉到 Host 内存,把显存腾给 active 请求,激活时再按需搬回——这就是"上下文缓存"能同时容纳更多并发会话的底气。
5️⃣ 前缀共享与写时复制:一份 KV 服务多条请求
多个请求共享同一系统提示词时,重复计算 prefill 是最浪费的。NInfer 的共享策略分两档:
- Move(私有延续):源 checkpoint 没有其他引用时,地址空间直接"过户",零拷贝;
- Fork + COW(共享前缀):frontier 之前的完整页只增加不可变引用,大家共享同一份物理页;只有当要写入的页被多方共享时,才对页内的部分尾页做一份私有拷贝(Copy-on-Write),后续追加只写私有尾页。
以P=64、checkpoint frontier = 1000 为例:前 15 个完整页被共享,第 16 页里的 40 个已提交 token 复制进私有尾页。frontier不向下取整到 960——分配粒度与 token 级有效性彻底解耦。这套逻辑落在 src/models/qwen3_5/program/storage/kv_store.h 的 LogicalKVPage 与 KVAddressSpace 存储中。
6️⃣ 地址翻译:kernel 如何不搬数据地读页
Attention kernel 从不接触分配器,它拿到的只是一个非拥有视图PagedKVLayerView(平面张量 + 一行 block table,定义见 src/core/paged_kv_cache.h)。地址翻译就三步:
逻辑块号 = position >> 6 页内偏移 = position & 63 物理页 = block_table[逻辑块号]随后按固定的"页内平面顺序"算出元素地址即可。由于页大小固定为 64,翻译可以退化成移位和掩码,在 page/tile 粒度算一次、内层循环复用。整个 growing-cache 路径不经过任何"攒成连续 KV"的 gather 拷贝——分页不引入任何与上下文长度成正比的中间搬运,这就是"直接分页执行"的性能红利。消费端契约见 include/ninfer/ops/softmax_attention.h 与 include/ninfer/ops/kv_cache_append.h。
7️⃣ 一页省多少:量化存储对容量的杠杆
分页管的是"怎么放",量化管的是"每页放多少"。NInfer 的 KV 池支持多种封闭存储 profile(布局解析见 paged_kv_storage.h),以 D256 每 token/每头 的物理占用为例:
| KV Profile | 每 token/head 的 K+V 占用 | 相对 BF16 |
|---|---|---|
| BF16 | 1024 B | 1.0x |
| INT8-G64 | 528 B | ≈0.52x |
| FP8-E4M3FN-row256 | 516 B | ≈0.50x |
| K8V4(非对称) | 402 B | ≈0.39x |
| NVFP4-G16 | 288 B | ≈0.28x |
同样的显存预算下,NVFP4 能让"可用页"数量翻倍以上,长上下文容量随之放大——分页机制保证了这些不同字节宽度的页组在同一个框架里被同等地分配、共享与迁移。
8️⃣ 一张表总结:长上下文高效管理的五个支点
| 支点 | 做法 | 效果 |
|---|---|---|
| 分页分配 | 固定 64-token 页组 + block table | 无碎片、无需整理、页边界不碰语义 |
| 容量三态 | 已分配 / 已预留 / 全局可用 分账 | active 请求额度不被缓存策略偷走 |
| 双副本 | Device/Host 整页迁移 + epoch 校验 | 上下文缓存下沉,显存留给活跃请求 |
| 共享 + COW | Move / Fork / 部分尾页私有化 | 相同前缀零重复 prefill |
| 直接分页执行 | 移位翻译、无 gather 拷贝 | kernel 不搬数据,decode 稳定高速 |
写在最后
NInfer 的分页 KV 缓存本质上把操作系统的虚拟内存哲学搬进了 GPU 推理:逻辑上连续、物理上离散、共享靠引用、写入靠租约。理解"页、副本、地址空间"这三层分离之后,你会发现长上下文推理引擎里最复杂的显存账目,其实是被这套严格的粒度划分管得井井有条。
📚 想继续深入,可以阅读:
- 分页 KV 设计合同:docs/maintainer/paged-kv-cache.md
- 资源调度与上下文缓存:docs/maintainer/resource-scheduling-and-context-cache.md
- 设备端池与租约实现:src/core/paged_kv_cache.cpp
- Host 副本 arena:src/core/host_kv_arena.cpp
- 程序级逻辑页与地址空间:src/models/qwen3_5/program/storage/kv_store.h
【免费下载链接】ninferHigh-performance single-GPU inference for selected model checkpoints and GPUs.项目地址: https://gitcode.com/gh_mirrors/ni/ninfer
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考