蚂蚁百灵公开 Ling-2.5-1T-base 超长上下文训练提速方法
蚂蚁百灵针对万亿参数预训练模型 Ling-2.5-1T-base 给出超长上下文训练的工程优化方案，在 256K 上下文训练中实现 token 吞吐 2.5 倍提升。方案包括 AllGather 与 All2All 两种 CP 并行实现，以及 linghe（灵核）高性能算子库的算子融合，后者在显存仅增加 2.8G 的情况下额外带来 43.7% 性能收益。团队还记录了 Ring Attention 调用 cuDNN 在 causal=False 下的性能问题、Linear Attention 朴素实现的 kernel launch 开销，以及 Triton 精度与竞态条件等踩坑经验。
g Attention 调用的 cuDNN 性能问题。 在初版实现中，我们用了Ring Attention。在超长上下文训练的实践中却发现当 MLA 结合特定的算子时，Ring Attention 会暴露严重的性能问题。Ring Attention 的核心机制是将长序列切分成多个块，GPU 在计算完本地的 Q 和本地的 K、V 后（即对角线块），需要接收来自其他 GPU 的历史 KV 块进行计算。 在处理历史 KV 块时，历史 token 对当前 Q 是可见的，不需要掩码，因此需要调用 causal=False 的 fused attention 算子。在实际测试中，我们发现底层 cuDNN 实现的 MLA 算子在 causal=False 这种配置下 backward 存在性能问题，算子的性能问题导致 Ring Attention 并不是一个合适的选择。 踩坑二，Linear Attention 部分如果使用简单的实现，会遭遇严重的 kernel launch cpu overhead 问题。 我们以一段描述变长序列（varlen）情形下的伪代码来描述我们的第一版实现： 在初始版本的
匹配字段正文