引言
部署大规模Mixture-of-Experts(MoE)模型如DeepSeek-R1,需要在延迟、吞吐量和成本之间实现精细平衡。这一挑战在性能不对称的硬件上尤为突出,例如H20 GPU,其内存带宽高但计算吞吐量较低。我们设计了一个服务栈,在满足高端GPU严格SLA的同时,利用H20的成本优势。
本文概述了实现这一目标的最佳实践,包括偏离传统方法的硬件感知部署策略,以及一系列系统和内核级优化:
- 硬件感知并行化:预填充阶段采用单节点TP-8,解码阶段采用小规模EP-16,满足延迟目标并缩小故障域。
- 内核级优化:FlashMLA-FP8和DeepGEMM swapAB,提升H20计算吞吐量。
- 调度与负载均衡:Single-Batch Overlap(SBO)提升小批量吞吐量,异步Expert Affinity Load Balancer减少跨节点通信。
- 轻量级可观测性:专为分布式MoE服务设计的诊断栈,快速定位瓶颈。
实验表明,使用我们的策略,每个节点在4096 token输入序列上实现16.5k输入token/s和5.7k输出token/s,这是H20上的SOTA性能,且是我们所知首次全面研究,涵盖部署、优化及大规模工业实践。
H20的挑战
H20的重要性
H20 GPU供应充足,使蚂蚁集团能构建超大规模集群。即使 modest 吞吐量提升,也能带来显著日常成本节省。
H20 vs. H800对比
| 规格 | H20-96G | H800-80G |
|---|---|---|
| FP8 Compute | 296 TFLOPS | 1979 TFLOPS |
| FP16/BF16 Compute | 148 TFLOPS | 989 TFLOPS |
| Memory Capacity | 96 GB | 80 GB |
| Memory Bandwidth | 4000 GB/s | 3352 GB/s |
| NVLink Bandwidth | 900 GB/s | 400 GB/s |
| RDMA NIC Bandwidth | 4 × 400 Gb/s | 8 × 400 Gb/s |
H20具备更大内存(96 GB)、更高内存带宽(4000 GB/s)和超2倍NVLink带宽(900 GB/s),但计算性能较弱,RDMA NIC带宽较低。推理尤其是解码阶段往往内存绑定,H20的高内存带宽和容量优势显著。我们据此设计优化,最大化推理吞吐量。
解决方案:H20上的优化与策略
部署策略
预填充(Prefill)
- SLA:预填充计算密集,多节点DP+EP会增加首token时间(TTFT),违反SLA。单节点TP保持TTFT在目标内。
- 弹性扩展:预填充需随KV缓存缩放,单节点TP简化资源与缓存管理。
解码(Decode)
- 硬件特性:H20牺牲计算换取更大内存和更高NVLink带宽(相较H800),高效利用KV缓存,将MoE通信置于高带宽NVLink。
- 故障域:小规模EP配置限制解码或GPU故障影响,EP高可用(HA)尚不成熟,小EP更可靠。
优化
预填充
观察
- 长序列下MLA比MHA成本更高。
- MoE延迟意外高,尽管计算量低。
embed/mlp all reduce + RMSNorm + fused_qkv_a_proj_with_mqa在TP中引入冗余通信和计算。
解决方案
- MHA/MLA:引入可调参数
se = extend × (extend + prefix),根据批次大小和序列长度选择MHA或MLA。 - MoE:优化
b_scale计算,用TMA重构down proj输入访问,并基于真实专家分布调优配置。 - TP优化:优化
embed/mlp reduce scatter + RMSNorm + fused_qkv_a_proj_with_mqa + all gather,减少计算和通信。
解码
负载均衡
Expert Affinity EPLB
标准EPLB仅均衡GPU内负载,忽略专家间相关性,常将频繁共激活专家分散到不同节点,增加跨节点通信。我们扩展EPLB,追踪top-k专家共激活构建专家亲和矩阵,后置调整确保高共激活专家同节点,较原版提升~5%性能。
异步动态负载调整
静态EPLB将负载均衡与推理紧耦,迁移阻塞推理导致延迟。我们解耦二者并行运行,采用分层传输策略最小化迁移影响,实现匹配或超静态EPLB性能,并保持>70%负载均衡率。
计算
FP8 MLA
BF16 FlashMLA性能良好但有优化空间。我们在Hopper(SM90)实现端到端FP8 attention,用TMA传输内存、WGMMA计算。两warp组流水QK^T和PV,减少共享内存压力并重叠计算内存。较BF16提升~70%速度,较前FP8版再增~5%。
SwapAB GEMM
Hopper WGMMA PTX约束N为8倍数、M固定64,导致粗粒度平铺浪费。我们引入swapAB,将M维度映射到N,实现更细粒度BLOCK_M(32),提升MoE中变M工作负载吞吐量。
SBO(Single-Batch-Overlap)
为何不用TBO
H20上TBO Decode收益有限:Hopper WGMMA block_m固定64,小批量引入冗余MLP GEMM;大批量(64/128)下低计算硬件无法满足TPOT SLA。
SBO工作原理
为提升Decode吞吐量不违SLA,采用SBO,修改DeepEP和DeepGEMM。设计基于通信-计算对齐粒度。观察到通信重叠中,token包因NIC等因素乱序到达接收端。