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

简介: 本文将介绍蚂蚁集团ASystem团队在推理显存优化上的新工作FlowMLA

文章作者:陶思睿 张锐 赵军平

系列文章回顾⬇️

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

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

本期我们将介绍蚂蚁集团ASystem团队在推理显存优化上的新工作FlowMLA,以提高DeepSeek等大模型推理性能(batch增加2.8x),适合带有DP并行(包括DP+TP并行)下显存受限场景(需要带NVLink)。方案对精度无影响,且无需额外微调操作,目前正在提交到开源社区(SGLang等)。

背景与问题

DeepSeek大模型推理

以DeepSeek为代表的大模型得到广泛普及,为了提升规模部署的性能并降低成本,DeepSeek官方开源了一系列优化技术,开源社区包括SGLang、vLLM等做了很多持续努力,一时间仿佛全民上阵,使得推理性能包括延迟和吞吐得到了快速优化。我们注意到,不同的业务场景对延迟和吞吐指标存在一定偏好,例如强化学习(RL)倾向长序列、高吞吐。本文侧重优化吞吐指标,对延迟指标无优化或略微牺牲。

显存压力问题

DeepSeek的MLA架构(及其他MQA)在典型attention并行策略下会面临显存容量方面的技术挑战,特别是国内普遍使用的H800机型。

  • TP并行下的KV cache显存冗余:MLA在decode阶段通过矩阵吸收的等价计算,将实际KV头数压缩至1,显著降低了KV cache的显存占用。然而,这种设计的一个副作用是,TP并行处理Attention时,每个rank需要冗余复制所有的KV cache,导致巨大的显存压力(特别是大batch、长序列下),基本不推荐单独使用。特别是H800 80GB的显存容量已经捉襟见肘。
  • DP并行下的attention权重冗余:为了解决了KV cache的冗余问题,Deepseek官方报告中采用了DP与TP(或SP)混合的attention并行策略。这个手段不小心又引入了一个副作用,每个DP rank不得不冗余存储attention权重。两者其害取其轻,大家尽力了。

注:以上的技术分析主要围绕attention module,和MoE部分无关,MoE可以采用常见的EP或TP并行。

图1:DP attention混合并行部署方案 (图源https://lmsys.org/blog/2025-05-05-large-scale-ep/)


现在,MLA权重冗余开销成为一些典型场景下制约吞吐的首要瓶颈。特别是,只有数台H800机器部署满血版Deepseek R1 BF16模型(例如用于RL场景下的样本生成),单卡attention权重占用近22GB,严重挤压了KV cache空间,而KV cache的大小基本直接决定了batch size和吞吐上限。

为此,我们提出显存优化方案FlowMLA来重点优化attention权重冗余问题。

注:显存容量因素可能直接影响了一些条件下AMD MI300 vs. H100/H200的推理性能(见AMD的评测链接中文版英文版)。

FlowMLA思路:以通信换显存,释放更多KV cache

我们借鉴训练场景ZeRO的思想,将attention权重分片至多卡,计算时通过all gather通信取回,以释放显存空间来支持大batch推理。这样带来的额外开销是每层引入了一次通信操作。进一步,我们通过设计layer间的计算-通信overlap的设计来掩盖通信开销。因此,FlowMLA的核心设计包括:

  • double buffer:两块共享的显存分别用于当前层的attention计算及下一层的weight prefetch。
  • overlapping:  当前层的计算与下一层attention权重的prefetch并发起来,隐藏prefetch开销。
  • layer-wise scheduling: 细粒度分层的计算和通信调度,异步流水执行。

图2:FlowMLA模块示意(以o_proj为例)


FlowMLA设计

以下重点介绍分层的显存复用和流水通信overlap的设计。

1.显存分区

  • 私有显存区(每卡固定)存储本地权重分片(图3实线块),总大小为weight / dp_size。
  • 共享显存区(跨层复用)all gather后完整权重的暂存区(图3虚线块),分奇/偶层复用,总大小为 2 * (weight / n_layers) * (dp_size - 1) / dp_size。

图3:FlowMLA显存复用机制,以4层模型 & DP4为例


我们扩展了此前介绍的kv cache显存优化工作vTensor(详见大模型推理显存优化系列(1):vTensor)的使用场景,同样也是借助cuda VMM(Virtual Memory Management) 来实现权重的sharding和re-sharding(类似ZeRO-3)。

2. 异步预取

接下来我们考虑计算-通信流水线。当第i层attention计算完成时,我们立即发起非阻塞all gather操作,跨GPU收集第i+2层所需的权重分片。该预取操作与当前层的后续计算及第i+1层的计算并行执行,实现通信时延的有效隐藏。为保障计算正确性,奇数层与偶数层分别复用各自共享显存区,comp_stream和comm_stream间通过cuda event同步状态,如图4所示。在具体实现中,我们的统一管理两类cuda stream:comp_stream上完成所有计算,奇数层、偶数层的权重预取在各自的comm_stream上执行。我们通过适当的stream管理使得上述操作完美兼容cuda graph。

图4:计算-通信流水线,以4层模型为例


FlowMLA实现

基于SGLang,针对占MLA权重超过60%的o_proj进行了开发实现,并向SGlang社区提交了PR (https://github.com/sgl-project/sglang/pull/7005),用户可以通过--vtensor_enable启动命令一键开启优化。

FlowMLA使得8xH20机可开启DP attention部署满血版FP8 Deepseek R1,这是SGlang开源版此前难以达到的。

FlowMLA结果对比

硬件环境:H20 SXM *8(单机, NVLink 900GB/s);H800 SXM *32(四机,NVLink 400GB/s),cuda 12.4,GPU驱动版本535

软件版本:sglang v0.4.7

负载情况:input 1024,output 2048,占用满KV cache空间的batch size

关注指标:output tok/s

8xH20 FP8 FlowMLA启动和测试命令示意:

python3 -m sglang.launch_server --model-path /path-to-DeepSeek-R1-FP8 --tp 8 --trust-remote-code --attention-backend flashinfer --port 8008 --speculative-num-steps 0 --speculative-eagle-topk 1 --speculative-num-draft-tokens 0 --disable-radix-cache --enable-dp-attention --mem-fraction-static 0.965 --cuda-graph-bs 136 --chunked-prefill-size 1024 --max-prefill-tokens 1024 --dp 8 --vtensor-enable

python -m sglang.bench_serving --backend sglang --model /path-to-DeepSeek-R1-FP8 --port 8008 --dataset-name random --random-input-len 1024 --random-output-len 2048 --random-range-ratio 1 --dataset-path /path-to-ShareGPT-dataset

image.gif

实验结果:

8xH20,DeepSeek FP8模型

32xH800 DeepSeek BF16模型

DP w/o FlowMLA

TP w/o FlowMLA

ours (DP)

DP w/o FlowMLA

 ours (DP)

每卡吞吐 (tokens/s)

(显存不足,无法运行)

128

181

280

400

batch size

52

136

30

83

实验结果表明,


  • 8卡H20场景下,FlowMLA的DP attention相比于TP attention基线提升batch 2.6x吞吐41%。且基线DP无法运行。
  • 32卡H800场景下,FlowMLA可增加DP attention batch 2.8x,吞吐提升43%

注:由于batch的提升以及异步通信对SM的占用,TPOT有所下降,可通过NCCL_MAX_CTAS控制NCCL的SM占用(推荐用最新版的NCCL 2.27.*)。


profile结果展示,对于8xH20 FP8的部署场景,batch size为1时就能做到通信时延的完美隐藏。

图5:通过torch profiler抓取的单请求下推理过程

小结

FlowMLA通过权重分片+显存复用+异步流水,释放了DP MLA中的冗余显存,极大地提升了系统吞吐。不止MLA,其他无法充分TP的GQA、MQA模型理论上均可基于该方法进一步提升。

进一步工作:

  • 针对TP MLA场景下的KV cache冗余进行优化。
  • 对qkv_proj的sharding实现。
  • 支持更多的模型。

号外

蚂蚁集团ASystem团队新开社招岗位(北京、杭州),对大模型训+推工程优化(特别是强化学习场景)感兴趣的同学欢迎关注、推荐,详见招聘链接:⬇️

https://talent.antgroup.com/off-campus-position?positionId=25052904956438&qrCode=A6PX3Z1H%2F8pHC%2FwtNmu3yvBgV1kTQgocnQfA2IvhOeqz8WkKFQAF8aawUf0wSZ5vJA%2FTLGrTgOxvs_mC10AuQQ%3D%3D

详细信息可加wx:18301691605

邮箱:george.zr@antgroup.com

相关文章
|
3月前
|
消息中间件 人工智能 资源调度
云上AI推理平台全掌握 (5):大模型异步推理服务
针对大模型推理服务中“高计算量、长时延”场景下同步推理的弊端,阿里云人工智能平台 PAI 推出了一套基于独立的队列服务异步推理框架,解决了异步推理的负载均衡、实例异常时任务重分配等问题,确保请求不丢失、实例不过载。
|
3月前
|
存储 机器学习/深度学习 缓存
阿里云AirCache技术实现多模态大模型高效推理加速,入选国际顶会ICCV2025
阿里云研发的AirCache技术被计算机视觉顶会ICCV2025收录,该技术通过激活跨模态关联、优化KV缓存压缩策略,显著提升视觉语言模型(VLMs)的推理效率与存储性能。实验表明,在保留仅10%视觉缓存的情况下,模型性能下降小于1%,解码延迟最高降低66%,吞吐量提升达192%。AirCache无需修改模型结构,兼容主流VLMs,已在教育、医疗、政务等多个行业落地应用,助力多模态大模型高效赋能产业智能化升级。
325 1
|
3月前
|
人工智能 缓存 资源调度
云上AI推理平台全掌握 (4):大模型分发加速
为应对大模型服务突发流量场景,阿里云人工智能平台 PAI 推理服务 PAI-EAS 提供本地目录内存缓存(Memory Cache)的大模型分发加速功能,有效解决大量请求接入情况下的推理延迟。PAI-EAS 大模型分发加速功能,零代码即可轻松完成配置。
|
1月前
|
机器学习/深度学习 算法 数据可视化
从零开始训练推理模型:GRPO+Unsloth改造Qwen实战指南
推理型大语言模型兴起,通过先思考再作答提升性能。本文介绍GRPO等强化学习算法,详解其原理并动手用Qwen2.5-3B训练推理模型,展示训练前后效果对比,揭示思维链生成的实现路径。
229 2
从零开始训练推理模型:GRPO+Unsloth改造Qwen实战指南
|
1月前
|
机器学习/深度学习 人工智能 JSON
微软rStar2-Agent:新的GRPO-RoC算法让14B模型在复杂推理时超越了前沿大模型
Microsoft Research最新推出的rStar2-Agent在AIME24数学基准测试中以80.6%的准确率超越超大规模模型DeepSeek-R1,展现“思考更聪明”而非“更长”的AI推理新方向。
128 8
微软rStar2-Agent:新的GRPO-RoC算法让14B模型在复杂推理时超越了前沿大模型
|
1月前
|
存储 缓存 负载均衡
LLM推理成本直降60%:PD分离在大模型商业化中的关键价值
在LLM推理中,Prefill(计算密集)与Decode(访存密集)阶段特性不同,分离计算可提升资源利用率。本文详解vLLM框架中的PD分离实现及局限,并分析Dynamo、Mooncake、SGLang等主流方案,探讨KV缓存、传输机制与调度策略,助力LLM推理优化。建议点赞收藏,便于后续查阅。
719 1
|
3月前
|
机器学习/深度学习 负载均衡 C++
MoR vs MoE架构对比:更少参数、更快推理的大模型新选择
本文将深入分析递归混合(MoR)与专家混合(MoE)两种架构在大语言模型中的技术特性差异,探讨各自的适用场景和实现机制,并从架构设计、参数效率、推理性能等多个维度进行全面对比。
234 0
MoR vs MoE架构对比:更少参数、更快推理的大模型新选择
|
3月前
|
数据采集 编解码 人工智能
Gemma 3n正式版开源:谷歌全新端侧多模态大模型,2GB 内存就能跑,重点提升编码和推理能力!
6月底,Google正式开源发布了全新端侧多模态大模型 Gemma 3n!相较此前的预览版,最新的 Gemma 3n 完整版进一步提升性能表现,支持在 2GB 内存的硬件上本地运行,重点提升了编码和推理方面的能力。
473 1

热门文章

最新文章