文章作者：徐家乐 张锐 赵军平。

系列文章回顾⬇️

大模型推理显存优化系列(1)：vTensor

大模型推理显存优化系列(2)：LayerKV

大模型推理显存优化系列(3)：FlowMLA——面向高吞吐的DP MLA零冗余显存优化

简介

显存管理是大模型推理高效、规模化部署的关键技术之一。蚂蚁集团ASystem团队联合上海交通大学研发的eLLM，挖掘激活与KV cache的动态管理机会，通过弹性显存管理技术实现可用显存的显著提升，在典型模型和显存受限场景下，相比基线方案（基于vLLM）实现3x batch增加和2.3x decode吞吐提升。方案精度无损，无需微调等额外操作。

以下简要介绍eLLM相关技术挑战、总体设计和初步性能评估，完整的技术报告见以下论文链接。

• eLLM: Elastic Memory Management Framework for Efficient LLM Serving 

  https://arxiv.org/abs/2506.15155 2025.6

显存管理挑战

显存需求

随着模型架构发展、长上下文包括思维链等新场景推广，显存需求构成悄然发生一些变化：

• 模型架构创新：DeepSeek MLA、GLA(Group Latent Attn，Tri Dao等，参考文献[2])、混合线性等通过更高效KV cache设计显著降低了KV用量。例如Jamba-52B仅需2块GPU支持200K上下文，其中激活与KV cache用量接近1:1。

• 上下文扩展：上下文长度从2K扩展至200K时，prefill阶段激活内存占比从0.3%跃升至30.8%；而如果是短输入、长思维链，prefill阶段的激活显存在decode阶段无法被使用。

• KV轻量化：针对KV的各种量化、压缩、稀疏探索层出不穷，而激活由于outlier的存在普遍难以轻量化。

  

图1：典型推理显存构成（A100部署，vLLM）

以上因素叠加促使我们重新关注显存的总体优化，特别是考虑当前主流的推理框架由于一些设计或实现的限制，导致了显存资源事实上的浪费（抛个观点“闲置即浪费”，浪费的单位是 GB * msec，即容量x时间，时间以毫秒计），轻易放弃了优化压榨的机会。

显存管理不足

当前vLLM、SGLang等主流框架经过了充分改进，取得了广泛普及。不过显存管理方面仍存在一些不足：

• 碎片化管理：模型权重、KV cache和激活三分天下，在框架启动时静态圈了3个小池子，互不通用。尤其是KV和激活，随不同模型、负载或SLO等存在更多动态性和潜在流通性。可惜，他们被碎片化的隔离了。

• 激活显存闲置：框架按prefill阶段所需的最大激活显存所预留的显存小池子，到了decode阶段由于激活需求降低了10^3~10^4倍基本在闲置（decode每次1个token，相比Prefill输入1K ~ 32K ctx，降低了1K ~ 32K倍），即使随生成变长，KV cache压力变大，进而触发了recompute/offload等兜底手段在拼命换气，一旁的激活显存仍在躺平。系统缺乏显存的全局、动态调度机制。

• chunk prefill的限制：chunk prefill可显著降低激活显存的占用（特别是prefill阶段），不过也存在一些限制：1）对一些模型不支持，例如至评测时不支持混合线性模型Jamba；2）影响TTFT指标。

图2：当前框架显存管理的局限和浪费优化机会

进一步，以vLLM为例，简单说明下主流框架的静态显存划分和管理策略：

• 总体逻辑：框架启动后先分配模型权重显存，再按启动参数估算激活内存（为prefill留够激活，prefill进行中不能OOM），剩余空间分配给KV cache。激活值由PyTorch的CUDACachingAllocator管理，KV cache则预分配为torch.Tensor。

• 隔离性弊端：静态划分的3类对象池不利于基于负载、SLO的动态显存，容易引起部分显存闲置。

  

图3：显存管理对比示意

eLLM设计

总体架构

eLLM采用三层架构设计：

1.虚拟张量抽象层：解耦逻辑地址与物理资源，为弹性管理奠定技术基础。

2.弹性内存机制：通过映射重分配实现激活/KV cache的动态空间调整，结合CPU内存作为弹性缓冲。

3.轻量级调度策略：基于服务等级目标（SLO）实现资源分配与性能的动态权衡。

图4：eLLM总体架构

图5：vLLM和优化后eLLM的显存构成对比

eTensor对象

eTensor对象是扩展的PyTorch Tensor对象，保留现有接口，内部封装GPU虚拟内存管理（基于CUDA VMM API）弥合激活值与KV cache的抽象层差距：

• 双张量类型设计

  • KV eTensor：预分配与上下文等长的连续的虚拟地址空间，物理内存按需分配。
  • 激活eTensor：支持非均匀虚拟地址段，适配小粒度、高频访问特性。

• 地址对齐策略：虚拟地址段（张量槽）与物理内存块粒度严格对齐，平衡访问效率与碎片控制。

  

图6：eTensor对象和显存池

弹性显存分配

设计了tensor对象的动态扩、缩操作，打通KV cache和激活显存各自领地：

• 扩张（Inflation，KV cache扩容）：

1. 检测KV内存不足时向激活池发出借用请求；
2. 激活池通过轻量级GC回收非活跃内存块；
3. 逻辑转移内存所有权至KV池；
4. VMM动态映射至目标KV eTensor地址空间。

• 收缩（deflation）：逆向执行膨胀流程，采用惰性避免开销

图7：弹性显存扩缩示意

PD感知调度

针对推理Prefill、decode阶段特性优化资源分配：

• prefill阶段：KV cache卸载至CPU缓冲区，释放GPU空间处理激活数据和放入更多新请求（降低排队耗时）；

• decode阶段：激活内存需求骤降，将KV cache重新加载至GPU，利用膨胀机制扩大批量处理能力；

• 工程优化：采用层级KV cache传输机制，降低CPU-GPU数据迁移开销。

SLO感知弹性缓冲

通过逻辑缓冲区动态平衡TTFT与TPOT指标诉求：

• 动态缓冲区抽象：在固定物理容量内动态调整逻辑可用空间，避免数据逐出；

• 违规响应算法：

  • TPOT超标时缩小缓冲区，限制prefill请求；
  • TTFT超标时扩大缓冲区，优化首响应延迟；

• 自适应调节：通过5次迭代窗口内的违规次数触发调整，因子α=2控制调整速率。

eLLM结果对比

测试环境

• 基线方案：基于vLLM-0.5.5构建三组对比：

  • vLLM：标准实现；
  • vLLM-CP：开启chunked-prefill后的vLLM框架；
  • eLLM：本项目的实现方法。

• 评测模型：

  • 在线服务：Llama3-8B-262K（最大上下文262K）；
  • 离线推理：Llama3-8B-262K与Jamba-Mini。

• 数据集：

  • 在线：32K-2K、2K-2K合成负载，ShareGPT真实对话；
  • 离线：8K-2K、16K-2K、32K-4K、128K-4K合成负载。

在线推理性能

Llama3-8B在单A100（80GB）GPU的测试表明：

• TTFT优化：2K-2K负载下，eLLM相比vLLM和vLLM-CP分别实现295倍和140倍的首字加速。收益主要来自PD感知显存调度、基线由于显存不足，排队时间过长（更详细分析参考另一篇工作LayerKV[3]）。

• 弹性机制优势：vLLM-CP仅通过预留激活内存缓解KV压力，eLLM则结合GPU内弹性与GPU-CPU协同，在高并发场景下保持性能稳定；

• 批量处理增益：eLLM支持更大decode批次，输出性能提升显著。

  

图8：在线推理评测对比

离线推理优化

Jamba-Mini在不同负载下的吞吐表现：eLLM实现最大batch提高至3x，带来总吞吐与decode吞吐分别提升1.8x和2.3x。

• vLLM对混合线性模型不支持chunk prefill，导致Jamba模型因激活/KV占比更高，eLLM的弹性策略效果更显著。

• 输入序列增长时，eLLM通过内存借调维持大批次处理，避免vLLM的性能陡降。

图9：离线推理性能评测

参考

[1] eLLM: https://arxiv.org/abs/2506.15155

[2] GLA: https://arxiv.org/abs/2505.21487

[3] LayerKV: https://arxiv.org/abs/2410.00428