首页 / 资讯中心 / 文章详情

12G 显存跑 256K 上下文:我把 KV 缓存“赶“到了内存里

12G 显存跑 256K 上下文:我把 KV 缓存“赶“到了内存里 ★ FEATURED ARTICLE
12G 显存跑 256K 上下文我把 KV 缓存赶到了内存里12G 显存跑 256K 上下文我把 KV 缓存赶到了内存里系列文章记录一次真实的引擎改造。硬件RTX 3060 12GB / 32GB 内存 / Windows 10。目标让本地大模型吃下 256K 上下文且质量不掉。本篇是背景与选型——为什么明知山有虎偏向虎山行。一、起点一张 12G 的卡和一道过不去的墙我本地跑的是一个 27B 的三值ternary模型用的是 NInfer 引擎franken/v0.11 分支sm_86 本地编译。先给一份 dense常规配置的实测免得后面没有对比指标dense 160K上下文163,840 tokensprefill662 / 516 t/s 20K / 60Kdecode61.2 / 51.2 t/s 20K / 60K显存~11.1 GB / 12 GB质量PPL 5.639521261,167 tokens 语料很能打对吧但 160K 就是天花板——再往上加--kv-capacity直接被启动器拒绝12G 显存装不下 192K 的 KV 池。而我想要的是 256K。差了将近 100K tokens 的 KV 空间。二、三条路摆在面前的只有三条路路线 A认命留在 dense 160K。什么都不用改。代价是需求不满足。路线 B换 Linux。Linux 下没有 WDDM 那层显存记账KVMem 那套思路原样能跑。但我的生产环境、启动脚本、ZCode 的接入全在 Windows 上换系统等于把整套东西重做一遍。用户明确表态不换。路线 C继续改引擎。在 NInfer 里手工移植 KVMem 的核心思想设备 KV 池小于逻辑上下文溢出的部分无损存到主机内存需要时再取回来。我选了 C。理由很朴素——前两条路要么不满足需求要么动生产环境只有 C 是在现有地基上加一层改坏了也能靠环境变量关掉回到 dense。三、先说清楚我用了谁的东西写在正文前面免得后面像在冒功思想来源KVMem。设备池 主机溢出 按需召回这个架构不是我发明的是照着 KVMem 的思路移植的。它的原版检索策略mean-K 打分我也研究了——只是没照抄理由见第五节。先例一个第三方团队。据我了解已经有一个第三方团队走过在 NInfer 这条线上做环形 KV这条路而且据说做成了。但代码没有公开所以我是拿着 KVMem 的论文/思路 自己读引擎源码摸出来的没有复用他们的任何成果。如果他们哪天放出来我很想对一对实现差异。真正属于这次工作的部分把这套东西落到NInfer 引擎 Windows/WDDM这个具体组合上拆掉 8 层耦合检查、解决 pinned memory 计显存的问题、用 IDF 词法检索替代 mean-K、以及全套 PPL/召回验证。换句话说别人证成了能不能我这边解决的是在 Windows 12G 卡上怎么让它安全地跑起来。引用/转载时请照这个口径——idea 是 KVMem 的先例是那个团队的具体的坑和数据是我的。四、KVMem 的核心思想一句话版传统推理引擎的隐含假设是KV 缓存必须全部待在显存里。所以上下文长度 显存大小 ÷ 每 token 的 KV 开销是个死数。KVMem 打破的就是这个假设显存里只放一个较小的 KV 池我这里是 96K tokens算不下的 KV无损降级到主机内存32GB便宜且大每次算注意力之前判断哪些页真的需要把需要的从内存搬回显存不需要的永远不搬——省显存也省带宽。听起来像操作系统的虚拟内存 按需调页。没错本质就是把内存换页思想搬到 KV 缓存上。难的从来不是思想是引擎里到处都是KV 必须全在设备上的硬假设。五、第一道真正的墙我拿什么判断哪些页需要KVMem 原版用的是mean-K 打分每个 head 的 K 向量均值作为相关性代理需要专门写 CUDA kernel 在设备上算。我的问题我不想为此手写 kernel——那会引入一条全新的、没有测试覆盖的热路径风险太高。于是我换了个思路用词法相关性IDF 加权代替 mean-K。好处纯 CPU、纯文本直接从引擎已有的sequence.ledgertoken 流水账里算一行 CUDA 都不用写代价语义相关性不如 mean-K 精准——但对长文档里找回某段代码/某个编号这类检索型需求词法匹配反而更稳。这个取舍后面会被实测证明是对的第 4 篇有 A/B 对照。六、小结到这一步方案轮廓是代码已开源https://github.com/tancau/ninfer-kvmem-ringApache-2.0改造被拆成三个互相独立的开关任何一个关掉都能回到 dense环境变量作用NINFER_KV_RING1总开关允许 KV 池 逻辑上下文NINFER_KV_RETRIEVE开启主机页召回NINFER_HOST_PAGEABLE1主机 KV 用普通内存第 3 篇的血泪主角下一篇引擎里那 8 道KV 不许小于上下文的安全锁是怎么一道道拆掉的。系列目录1本文为什么折腾、三条路怎么选2拆掉引擎的 8 道安全锁3Windows 的坑cudaMallocHost 居然吃显存4实测PPL 一字不差256K 真能召回5收尾仓库、启动器以及还没吃完的性能红利
阅读完成 · 觉得有帮助?
咨询建站