大模型推理显存优化：从碎片消除到按需分配的实战演进

导语

大模型参数量狂飙，单卡显存容量和访存带宽的增长却远远落后。在推理场景，尤其是小 Batch 场景下，访存带宽已成为绝对瓶颈。与此同时，Attention 算法（如 FlashAttn3、SparseAttn）快速演进，对显存管理机制的解耦性提出了更高要求。如果显存管理依然停留在粗放阶段，不仅会浪费昂贵的算力资源，更会直接拖垮推理服务的延迟和吞吐。

蚂蚁集团异构计算与推理团队基于 CUDA VMM（Virtual Memory Management）进行了一系列深度优化，形成了一套覆盖训练、单卡推理、多进程推理的显存优化全景图。本文将拆解这套方案的核心逻辑与实战经验。

核心问题与挑战

在 LLM 推理中，显存问题不仅是“不够用”，更是“用不好”，具体表现在以下四个核心痛点：

1. 显存碎片严重，且无法部分释放：CUDA 机制不支持直接释放部分空闲显存区域。LLM 推理时 KV Cache 的动态伸缩，极易在显存池中留下大量无法回收的碎片。
2. PagedAttn 适配成本高、性能有折损：业界代表方案 PagedAttn 虽然缓解了碎片，但对 Attention Kernel 的侵入性极强。适配新 Kernel 难度大、周期长（通常以月计），且早期版本不支持 TensorCore，存在性能开销，难以跟上 FA3 等算法的演进速度。
3. 高负载下 TTFT（首字延迟）激增：主流框架（如 vLLM）在 Prefill 阶段会按 Prompt 长度一次性 Reserve KV Cache。高并发下一旦显存不足，新请求只能排队等待 Decode 释放空间，导致 TTFT 出现数量级恶化。
4. 多进程推理缺乏显存共享：单机多卡或多进程部署同一模型时，模型权重、中间结果和 CUDA Context 大量重复加载，显存利用率极低。

方案与实践

底层基石：CUDA VMM 与 2 层指针

解决上述问题的底层钥匙是 CUDA VMM API。其核心思路是2 层指针与动态 Remapping：

• 虚拟地址：对用户可见，保证地址空间连续。
• 物理地址：以 Chunk（如 2MB）粒度按需分配，不要求物理连续。

通过这层解耦，当显存出现碎片时，只需在页表层面进行 Split 和 Remapping，即可按 Chunk 粒度回收物理显存，而无需搬移数据。

virtualTensor：解耦计算与存储，分钟级集成新 Kernel

针对 PagedAttn 的痛点，我们提出了 virtualTensor 方案，将 KV Cache 作为单一 Tensor 对象管理，实现显存管理与 Attn Kernel 的彻底解耦。

• 消除碎片：基于 VMM 的动态映射，KV Cache 显存碎片基本消除，效果媲美 PagedAttn，且支持按需弹性增长。
• 极简集成：Kernel 侧无需任何修改，只需 Import Package 并替换 2 个参数（3 行代码），即可在 1 分钟内完成新 Attn Kernel（如 FlashAttn3、PyTorch FlexAttn）到 vLLM 的适配。
• 性能提升：摆脱了 PagedAttn 的性能束缚，virtualTensor 结合 FA3 相比 FA2 可获得 2.1x 的加速。

LayerKV：分层按需分配，击穿 TTFT 瓶颈

针对高负载下显存不足导致的排队问题，LayerKV 的核心思想是“分层与按需”。

• 分层 Reserve 与 Alloc：Prefill 是逐层计算的，没必要按整个 Prompt 长度一次性 Reserve 全部 KV Cache。LayerKV 改为按层申请，极大降低了瞬态显存峰值。
• Offloading 与 Prefetch：将暂未参与计算的层卸载到 CPU 内存，并在计算前通过 PCIe 预取回 GPU。结合 Scheduler 的动态决策，计算与传输深度重叠。
• 实测效果：在 Llama2-7b 上，平均 TTFT 从 1024ms 降至 300ms；在 7b/34b/70b 多规格模型上，QPS 基本无损的前提下 TTFT 降低约 10 倍。结合 PD 分离架构，二字延迟也有显著优化。

全景扩展：训练与多进程推理的显存优化

• GMLake（训练场景）：将 VMM 思路引入大模型训练，重点解决映射元数据管理与策略控制，全程无数据 Copy，在 76 个微调工作负载下实现了平均 15%、最高 33% 的吞吐提升。
• 跨进程细粒度透明共享（多进程推理）：针对模型权重只读的特点，基于 Hash 和字节比对实现类似存储“重删”的细粒度透明共享。6 个进程的总显存从 13.2GB 降至 4.6GB（降幅 65%），对用户完全透明，不影响精度与性能。

原则与方法论沉淀

在显存优化的实战中，我们沉淀出四条核心工程原则：

1. 显存虚拟化与按需映射：打破物理显存连续的执念，用虚拟地址连续 + 物理按需 Chunk 映射从根上解决碎片。
2. 计算与存储解耦：显存管理机制绝不能绑定计算 Kernel，只有解耦才能让算法迭代不受基础设施拖累。
3. 分层与按需分配：尊重计算生命周期（如 Prefill 逐层特征），改一次性占用为按需申请与卸载，用时间换空间。
4. 透明共享与重删：跨进程显存共享应做到细粒度、自动发现、对业务透明，将存储领域的成熟思路降维打击到显存管理。

总结与行动建议

大模型推理的显存优化不是单点突破，而是从底层 API 到上层调度、从单卡到多进程的体系化工程。建议工程团队在落地时遵循以下路径：

1. 优先解决碎片与解耦：接入 virtualTensor 思路，消除 KV Cache 碎片，同时解放 Attention Kernel 的迭代速度，这是 ROI 最高的一步。
2. 攻坚延迟瓶颈：在服务高负载场景下，引入 LayerKV 的分层申请与 Offloading 策略，打破 TTFT 排队魔咒。
3. 挖掘存量红利：在多进程部署场景，开启跨进程透明共享，直接砍掉冗余的权重与激活显存占用。

开放问题与延伸方向

1. virtualTensor 的 2 层指针与动态 remapping 在极端高并发请求下，VMM API 的调用频次与内核态切换开销是否会成为新的性能瓶颈？（白帽-技术挖掘：需在极限并发下压测 VMM 映射开销底线。）
2. LayerKV 在分层卸载至 CPU 时，对 Host 内存带宽与 PCIe 带宽的占用率如何量化评估，是否会引发与数据加载等其他操作的带宽争抢？（白帽-可行性：需建立带宽争抢的量化模型与隔离机制。）
3. 跨进程透明共享依赖 hash 和字节比对机制，在复杂部署环境下是否容易引发难以调试的隐性数据竞争或内存一致性问题？（红帽-隐性担忧：极简接口背后的状态一致性是线上排障的深水区。）
4. virtualTensor 宣称 3 行代码实现分钟级集成，这种高度封装的极简接口是否会掩盖底层显存真实状态从而导致线上排障极度困难？（红帽-隐性担忧：需配套建设显存状态的可观测性工具。）
5. LayerKV 依赖 Offloading 到 CPU 内存，在超长上下文（如 100k+ tokens）场景下，是否会导致 Prefetch 严重跟不上 Decode 速度而产生停顿？（黑帽-局限性：长上下文可能打破计算与传输的重叠假设。）
6. 跨进程共享机制在模型权重量化（如 INT8/INT4）或叠加 LoRA 微调的场景下，hash 命中率是否会大幅下降甚至导致共享完全失效？（黑帽-反驳：动态量化与微调会破坏数据一致性前提。）
7. virtualTensor 的“计算与存储解耦”设计，能否顺带解决 MoE 模型中 Expert 动态激活带来的显存碎片与调度难题？（黄帽-扩展：虚拟化思路天然契合 MoE 离散激活特征。）
8. 将 LayerKV 的分层按需分配思想逆向应用于大模型训练场景的 Activation Checkpointing，能否大幅降低训练峰值显存？（黄帽-迁移：分层卸载思路在训练重计算场景有想象空间。）
9. 面对 PagedAttn 的改造成本，是否考虑过基于 CXL 内存扩展协议来直接扩展物理显存池，而非走 CPU Offload 的迂回路径？（绿帽-替代路径：硬件架构演进可能提供新解法。）
10. 结合 virtualTensor 的虚拟化与跨进程共享能力，能否构建一种显存池化的 Serverless 推理架构，实现请求级别的显存弹性伸缩？（绿帽-新组合：向显存云原生演进。）
11. 从 VMM 基础到 LayerKV 再到跨进程共享，这套显存优化体系在实际落地时，应依据怎样的优先级与实施路径来规划才能最快收敛 ROI？（蓝帽-下一步：需结合业务痛点给出标准落地路径。）