DeepSeek代码开源第一弹Flash MLA,揭秘大模型训练低成本关键

360影视 2025-02-25 01:19 1

摘要:随着DeepSeek大模型开源引发全球热潮后,2月21日DeepSeek在社交平台X发文称,这周起会陆续开源5个代码库。

【TechWeb】2月24日消息,随着DeepSeek大模型开源引发全球热潮后,2月21日DeepSeek在社交平台X发文称,这周起会陆续开源5个代码库。

今天DeepSeek开源首个代码库Flash MLA,引发极大关注,截至目前github Star星数已经超过4.5k。

Flash MLA是DeepSeek针对英伟达Hopper GPU优化的高效MLA解码内核,其特别针对可变长度序列作了优化,现已投入生产。

FlashMLA的使用基准为:Hopper GPU、CUDA 12.3及以上版本、PyTorch 2.0及以上版本。

经实测,FlashMLA在H800 SXM5平台上(CUDA 12.6),在内存受限配置下可达最高3000GB/s,在计算受限配置下可达峰值580 TFLOPS。

这种优化可以确保FlashMLA在高性能硬件上有效地处理大语言模型和其他AI应用程序的密集计算需求。

目前已经发布的内容包括:采用BF16,块大小为64的分页kvcache(键值缓存)。

团队在致谢部分表示,FlashMLA的设计参考了FlashAttention 2&3以及CUTLASS的技术实现。

资料显示,FlashAttention 是一种针对Transformer模型注意力计算的高效优化算法,由斯坦福团队于2022年提出,核心目标是通过硬件感知的内存管理和计算流程重构,显著降低长序列处理时的显存占用与计算延迟。

CUTLASS(CUDA Templates for Linear Algebra Subroutines)是NVIDIA推出的开源高性能计算库,专为GPU加速的线性代数计算(尤其是矩阵乘法和卷积)设计。其核心目标是通过模块化模板和硬件级优化,为开发者提供灵活、高效的底层计算内核,支撑AI训练、科学计算与图形渲染等领域。

根据DeepSeek过往提交的论文,DeepSeek大模型训练成本大幅下降有两项关键技术,一是MoE,另一个就是今天开源的MLA(多头潜注意力)。

DeepSeek的成本涉及两项关键的技术:一个是MoE,一个就是MLA(Multi-head Latent Attention,多头潜注意力)。

MLA旨在优化传统Transformer架构的效率与性能,其核心原理包括:

KV压缩与潜在变量:将键(Key)和值(Value)联合压缩为低维潜在向量,显著减少推理时的KV缓存,降低内存占用。计算时通过升维恢复原始信息,平衡压缩效率与计算精度。

低秩降维技术:对查询(Queries)进行低秩压缩(降维后再升维),减少训练中的激活内存(activation memory),但需注意此操作不影响KV缓存。

动态序列处理:针对可变长度输入序列优化,支持高效处理不同长度的句子(如长文本对话场景)。

MLA可将每个查询KV缓存量减少93.3%,显著减少了大模型训练和推理过程中的内存占用

今天开源的MLA是DeepSeek在注意力机制上的重要创新,通过KV压缩、低秩降维等技术实现高效长序列处理与资源优化,成为其模型性能领先的关键技术之一。

本周后续,DeepSeek还将陆续开源4个代码库,期待一下!(宜月)

来源:TechWeb

相关推荐