四路24GiB显卡跑32K上下文：显存账本与KV Cache优化

手上这张 24 GiB 的显卡已经成了很多本地大模型玩家的默认上限。模型量化完权重刚好塞进 24 GiB这种成就感我很懂——毕竟这年头谁没为那几百 MB 的剩余空间折腾过。可一旦你想再上四路并行、把上下文拉到 32K问题就来了权重占掉的只是静态空间上下文却要吃掉一份动态内存而且这个动态内存的增长速度比你想象中夸张得多。这篇文章就围绕“四路 32K 上下文”这个组合把显存账本摊开算一算到底还装不装得下顺便把里面最容易被忽略的几个坑也一并踩一遍。如果你正准备用四张 24 GiB 卡做本地推理或者已经在单卡 24 GiB 上跑过量化模型、想再往长上下文方向挤一挤这篇内容应该能帮你省下不少试错时间。下面不会只给结论而是会把每一步估算方式、关键公式、框架预留规则都讲明白让你换一个模型也能照着自己算。1. 静态空间和动态空间24 GiB 的分配逻辑很多人对显存规划的理解停留在“模型权重多大推理时就要占多大”。这句话在加载阶段是对的但不是推理阶段的全貌。24 GiB 卡能不能装下权重和你跑 32K 上下文时会不会 OOM其实是两本账。1.1 权重占的是“买房面积”KV Cache 是“装修”我习惯把显存分成两部分来看权重是买了多大面积的房子上下文则是往里添了多少家具。房子面积是固定的你买的时候是多少之后不会变家具却会随着你招待的客人增多而越摆越多。权重的“面积”很好算。一个 70B 级模型如果以 BF16 精度加载每个参数占 2 字节总大小就是 140 GB 左右。放到四张 24 GiB 的卡上做张量并行不考虑任何开销每张卡也要吃掉约 35 GB显然放不下。这也是为什么本地跑大模型通常要先量化INT4 量化后每个参数大约只占 0.50.6 字节同样 70B 模型能压到 40 GB 上下四卡分片后每卡 10 GB 左右才算是有了讨论余地。上下文这个“家具”则复杂一些。Transformer 做自注意力时需要把已经生成的 token 对应的 Key 和 Value 缓存下来方便后续 token 做注意力计算这部分缓存就叫 KV Cache。它的大小和你输入了多少 token 直接挂钩32K 上下文的含义就是最多缓存 32768 个 token 的 K/V。KV Cache 单 token 大小的估算公式是这样的单 token KV 大小 层数 × 2 × KV头数 × 头维度 × 每元素字节数为什么要乘 2因为 Key 和 Value 各有一套。以常见的 70B 级 GQA 模型为例假设 80 层、8 个 KV 头、每个头维度 128、BF16 精度下每元素 2 字节那么每个 token 需要80 × 2 × 8 × 128 × 2 327680 字节约 320 KB单看 320 KB 觉得不大但乘以 32K 上下文之后320 KB × 32768 ≈ 9.77 GB也就是说光是这一份完整 KV Cache就差不多要吃掉 10 GB。放到单张 24 GiB 卡上如果权重已经占了 22 GB那 32K 上下文基本没有容身之地。这个数字才是标题里那句“还装得下吗”的真正来源——权重能不能装进 24 GiB 是一回事权重和上下文一起装又是另一回事。1.2 除了 KV还有哪些看不见的住户权重和 KV Cache 是显存里的两大主力但它们不是全部。真正部署时还有三四个“隐形住户”平时看不出存在感一到极限就出来捣乱。第一个是 CUDA context。只要 GPU 上跑起 CUDA 程序驱动就会预留一块内存来做 CUDA context 和深度学习框架的运行时基础。这块空间在小卡上可能只占几百 MB但在 24 GiB 的卡上某些框架预留给 allocator 的冗余、cuBLAS/cuDNN 的 workspace加起来轻松超过 1 GB。如果用的是 PyTorch 的默认缓存分配器即使你只申请了很小一段显存它也可能为了后续分配预留一大块。第二个是激活值。前向推理时中间层输出、Attention 的计算结果、FFN 的中间矩阵都要临时占显存。Prompt 特别长时激活值可能暴涨尤其在 prefill 阶段因为要同时处理所有输入 token中间张量会比 decode 阶段大得多。后文我会专门讲这个坑这里先记住KV Cache 只是稳态值临时峰值是另一回事。第三个是并行通信的临时 buffer。四卡张量并行时每一层都要做 all-reduce 之类的集合通信通信结果要放缓冲区。虽然可以复用但峰值时也会占掉几百 MB 到 1 GB。很多人把显存预算算得刚刚好结果实际部署时发现“少了 12 GB”多半就是这些开销没算进去。所以一份更稳妥的单卡显存预算表应该是这样的权重分片占大头量化后按每卡分片计算KV Cache按目标上下文长度计算并预留至少 20% 余量CUDA context 与框架 allocator预留 11.5 GB激活值长上下文下预留 24 GB视 token 批量大小而定通信 buffer并行卡数越多越明显预留 0.51 GB把这个框架放在脑子里再去想“四路 32K”就不会只盯着权重量化的那一点空间算死账了。2. 四路并行的两种理解显存账完全是两本“四路”这个词在不同场景里有不同含义。一种是四卡张量并行把一个大模型拆到四张卡上协同推理另一种是四张卡各自跑独立实例每人各处理各的请求。两者的显存逻辑差别很大很多人混着混着就把预算算错了。2.1 四路张量并行把权重和 KV 都切成四份张量并行Tensor Parallelism的做法是把每一层 Transformer 的权重拆成 4 份分别放到 4 张卡上计算时通过通信汇总部分结果。这是“四路”最常见的技术含义。在这种模式下权重被摊薄成原来的 1/4KV Cache 也会按注意力头切分摊到四张卡上。还是用上面那个 70B 级模型的例子。INT4 量化后总权重约 40 GBTP4 时每卡只放约 10 GB。KV Cache 总量约 9.77 GB按注意力头切分后每卡约 2.44 GB。也就是说单卡上真正被权重和 KV 占掉的部分大概是 10 2.44 12.44 GB距离 24 GiB 还有很大余量激活值和 CUDA context 就算再吃掉 34 GB也仍然能装得下。这就是四路张量并行的好处它不只是把权重分摊还把长上下文的 KV 压力也一起分摊了。单卡 24 GiB 跑 32K 上下文可能很吃力四卡 TP 后反而可能很宽裕。所以“权重装进了 24 GiB”和“四路 32K 上下文装不装得下”在 TP 模式下其实要分开看权重本来就只占一部分剩下的是不是够用取决于模型规模、量化精度和框架预留。但要注意TP 不是没有代价。张量并行要求所有显卡都在同一台机器上通过 NVLink 或 PCIe 互联通信量很大。四张卡中间结果频繁 all-reduce如果互联带宽不行推理速度可能比单卡还慢显存是够了性能却可能崩。这也是为什么不少本地玩家宁可用四卡独立实例也不愿意随便开 TP。另外还得提醒一点TP 分 KV 时并不是简单地把 10 GB 除以四就完事。实际切分还要考虑注意力头数能不能被 4 整除。如果模型 KV 头数少于并行度某些头会被完整复制KV 反而不会严格变成 1/4。所以算单卡 KV 时最好按“每卡分配到的最多注意力头数 × 单头缓存大小”去推而不是直接除。2.2 四卡独立部署上下文并不会分摊另一种“四路”是四张 24 GiB 卡互不通信各自加载一个完整模型各自处理不同请求。这种模式看起来总显存也是 96 GiB但每个请求的 KV Cache 并不会被分到多张卡上而是每一张卡都要完整吃下属于自己的那部分上下文。同样一个 70B INT4 模型四卡独立部署时每张卡要先装约 40 GB 权重但 24 GiB 卡根本装不下所以这种模式下 70B 级模型首先就被排除了。换成 32B 级模型INT4 权重约 20 GB每卡还剩 4 GB 左右32K 上下文的 KV Cache 如果总量是 5 GB单卡就要全吃 5 GB直接超了。你看同样叫“四路”一个能跑一个不能跑差异就在“上下文是否分摊”上。所以当你看到“四路 32K 上下文”这种说法时得先确认是哪种“四路”。如果只是把四个独立推理进程塞进四张卡那“32K”的实际开销是每张卡各付一次全款和四路并行一分钱关系都没有。我见过不少人在不同框架里反复折腾 OOM后来发现是框架默认按“每卡独立模型”部署根本没有把上下文分摊出去。这种问题查代码不如先查部署拓扑拓扑清楚了显存账自然就清楚。3. 32K 上下文放在 24 GiB 卡上预算表与量化选择光讲原理还不够还是得要一套能直接抄的数字表。下面我用几个常见的模型量级分别算一遍 32K 上下文下每卡要占多少显存以及怎么靠量化把“装不下”变成“装得下”。3.1 先算清楚 KV Cache 的真实体积KV Cache 的关键参数是层数、KV 头数、头维度和精度。下面这张表给出了三种典型模型的单 token KV 大小模型量级示例层数KV 头数头维度精度单 token KV轻量级小模型324128BF16约 64 KB中型模型648128BF16约 256 KB70B 级大模型808128BF16约 320 KB实际模型配置会有出入但这个量级基本可信。再往上乘小模型 32K 上下文 KV ≈ 64 KB × 32768 ≈ 2 GB中型模型 32K 上下文 KV ≈ 256 KB × 32768 ≈ 8 GB70B 级模型 32K 上下文 KV ≈ 320 KB × 32768 ≈ 9.77 GB这几个数字是单份完整 KV Cache 裸大小尚不含任何量化压缩。如果还想再精确一点可以打开模型配置文件看num_hidden_layers、num_key_value_heads、head_dim这几个参数套上面的公式手算。模型官方文档通常都会写不写也可以通过transformers加载后打印 config 拿到。手算的好处是换任何模型你都能独立判断不用等推理框架报 OOM 再回头查。3.2 一份可直接参考的四路 TP 显存预算表假设四张 24 GiB 卡做张量并行每卡可用显存按 22 GiB 算预留 2 GiB 给 CUDA context 等开销。下面这张表是不同权重方案下的估算结果模型与量化权重总量四路后每卡权重32K KV 总量四路后每卡 KV每卡合计24 GiB 下是否够7B BF1614 GB3.5 GB2 GB0.5 GB约 4 GB非常够30B INT418 GB4.5 GB5 GB1.25 GB约 5.75 GB非常够70B INT440 GB10 GB9.77 GB2.44 GB约 12.44 GB够70B INT870 GB17.5 GB9.77 GB2.44 GB约 19.94 GB比较紧70B BF16140 GB35 GB9.77 GB2.44 GB约 37.44 GB不够这张表里最值得注意的是 70B INT4 和 70B BF16 的差距同样一个模型、同样 32K 上下文权重精度改变了 24 GiB 卡上“装得下”还是“装不下”的结局。这就是为什么量化在长上下文场景下几乎成了必选项它省下来的不是一点点内存而是整个部署方案能否成立的差别。如果只跑单卡不用 TP那 32K 上下文 KV Cache 要把“总量”那一列全部吃下。比如 70B INT4 单卡光 KV 就要 9.77 GB权重 40 GB 根本装不进 24 GiB所以单卡 24 GiB 玩 70B 级 32K基本可以直接放弃。真正能在单卡 24 GiB 上跑 32K 的通常是量化后的中型模型或者 KV 优化特别激进的方案。3.3 权重量化和 KV 量化的取舍量化是显存规划里最立竿见影的手段但凡事有代价。INT4 权重虽然省空间但推理速度可能受反量化开销影响模型输出质量也可能轻微下降。如果你对生成质量敏感建议先试 INT8 或 FP8确认质量可以接受后再往 INT4 走。还有个常被忽略的维度是 KV Cache 量化。推理框架里部分实现支持把 KV Cache 从 2 字节压到 1 字节甚至更低这能直接减小长上下文的内存压力。比如 70B 级模型的 9.77 GB KV如果 KV 量化到 INT8就能压到约 4.9 GB四路 TP 后每卡只剩 1.2 GB 左右。KV 量化的精度损失一般比权重量化更可控因为注意力得分对绝对数值没那么敏感不过有一点要提醒KV 量化可能会让长文本生成时的重复度变高输出质量是否可接受需要实测。从经验上看权重决定你能不能启动KV 量化决定你能把上下文撑到多长。如果卡里权重已经占了 20 GB上下文还想要 32K先看 KV 能不能量化、能不能换更省 KV 的注意力实现比继续压权重更划算。反过来如果权重占了 22 GBKV 怎么压缩都救不回来那就得考虑换更小模型或更多卡。4. 实操中的显存陷阱与排查实录算清楚了预算表部署时仍然可能翻车。这一节不是我按文档念的而是实际踩过坑之后总结的几个高频场景。4.1 预填充prefill阶段才是 OOM 高发区不少人有个错觉上下文越长越到后面显存越满OOM 应该发生在生成最后一两个 token 的时候。实际上最危险的往往是首批 token 的 prefill 阶段。因为 prefill 要把整段上下文一次性拿进注意力计算中间激活值的规模很大它和 KV Cache 是叠加出现的不是错开的。举个例子API 请求带着 30K 字符的 prompt 进来prefill 期间框架可能要为这个长 prompt 开辟远超 KV Cache 本身的临时张量空间。KV Cache 稳态估算 9.77 GB但 prefill 峰值有可能临时冲到 1415 GB。如果四路 TP 后每卡原本只留了 5 GB 余量这一下就把显存顶穿了。应对办法是把长 prompt 分段处理也就是 chunked prefill。主流推理框架里开启后框架会按 chunk 大小分批 prefill避免一次性把 30K 字符全部塞进显存。代价是 prefill 的并行度被削低一些但换来的是显存峰值可控。我一直建议做长上下文部署测试时不要只看最终的 KV 稳态还要专门拿一个最长 prompt 请求压一遍 prefill。只有 prefill 和 decode 两个阶段都跑通了这个配置才算真稳。4.2 显存“说没就说没”的三种典型情形我遇到过的 OOM 基本能归成三类症状不一样处理方式也不一样。第一类启动就 OOM。模型加载进显存后立刻报 CUDA OOM但按显存计算明明够。这通常不是权重的问题而是框架默认预留了过多显存。很多框架在启动时会让 CUDA allocator 吃掉大量显存比如预留给后续 KV Cache 的池子。解决办法是调低gpu_memory_utilization给它一个比较保守的值比如 0.850.9。第二类跑一段时间后 OOM。这类最烦人刚开始是好的随着上下文增长KV Cache 池一点点耗尽最终爆掉。这时优先检查是不是上下文长度设置超过了卡的实际能力或者框架的 KV Cache 池没有根据剩余显存自动伸缩。可以考虑缩短最大序列长度、开启 KV 量化或者换更小的模型。第三类显存看着还有却一直 OOM。这类通常和碎片有关。频繁申请释放大小不一的显存块会让 allocator 没法把碎片凑成一个大块。现象是nvidia-smi显示还剩下好几个 GB但程序报错申请不到连续显存。最直接的解决办法是重启进程或者减少并发请求数也可以在框架配置里调小并发批量降低单次内存申请的峰值。4.3 几个可以马上用的排查动作上面这些坑与其等它爆出来不如在部署前就排查一轮。我的习惯是 nvidia-smi 只看一部分更关键的是看框架自己的日志和统计信息。首先服务启动后先打一个短请求比如 1K 上下文确认模型推理正常、显存申请稳定记下空闲显存数值。然后逐步加大到 4K、8K、16K、32K每档都观察显存峰值变化。如果 16K 正常、32K 爆了基本可以确认是 KV Cache 池不够而不是权重问题这时优先考虑 KV 量化和降低预留比例。其次用nvidia-smi dmon看每张卡的实时显存占用注意 TP 模式下四张卡数值应当保持基本一致。如果某张卡明显偏高很可能是注意力头切分不均匀或者通信 buffer 在个别卡上堆积。这时候检查一下模型配置里的头数和 TP 并行度是否匹配有时候把并行度从头数整除关系上再核对一遍就能找出问题。最后不要忽略框架日志。比如 vLLM 会在日志里打印当前 KV Cache 可用块数和最大模型长度这些信息比nvidia-smi准确得多。看到Free blocks不断下降但Max model len没变时就该调整配置了。经常有人问我要一个“万能显存公式”但实话讲推理框架、CUDA 版本、GPU 互联方式都会影响实际占用。固定不变的反而是权重和 KV 的裸大小所以上面那些公式和表格才是判断配置能不能跑通的地基。地基算好了剩下的无非是在框架参数里做微调。最后分享一个我自己的习惯算完 32K 上下文的显存后我会故意给 KV Cache 的估算值再乘 1.2。这个 20% 的余量不是拍脑袋而是为了把 CUDA context、临时激活、通信 buffer 和一点点碎片都兜进去。只要乘完还在卡的可控范围内实际部署基本稳了如果乘完已经逼近 24 GiB 上限我会果断砍上下文长度或者换个更省显存的量化方案。这个“预算不乘系数不上线”的习惯帮我避开了绝大部分推理过程中的意外 OOM。