摘要:DeepSeek开源周正式拉开帷幕,为AI社区带来了一系列开源代码库,北京时间周一上午九点,DeepSeek公布了开源周的第一个项目:FlashMLA。
DeepSeek开源周正式拉开帷幕,为AI社区带来了一系列开源代码库,北京时间周一上午九点,DeepSeek公布了开源周的第一个项目:FlashMLA。
FlashMLA是DeepSeek专为英伟达Hopper GPU设计的高效 MLA(Multi-Head Latent Attention)解码内核,用于优化可变长度序列的推理服务,其目标是在H100等Hopper GPU上实现更快的推理速度,且所有代码均经过实际业务场景验证,目前正处于持续发布中。
发布后,FlashMLA迅速成为全球开发者关注的焦点,在GitHub上的Star数已突破5000。
按照官方说法,FlashMLA在H800 SXM5平台上(CUDA 12.6),在内存受限配置下可达最高3000GB/s的内存带宽,以及580TFLOPS的计算性能,代表了当前AI计算领域的顶级水准。
推理性能突破在大语言模型的推理过程中,高效的序列解码对于降低延迟和提高吞吐量至关重要。FlashMLA针对可变长度序列和分页KV缓存进行了优化,使其非常适用于这一需求。
那么FlashMLA如何实现的呢?这里不得不提到MLA,DeepSeek-V2模型中引入的MLA是一种注意力机制的变体,主要解决推理时KV Cache占用过多内存导致的性能瓶颈。MLA通过引入低秩KV压缩技术,有效减少了KV Cache的大小。而DeepSeek-V3则进一步采用了MLA和DeepSeekMoE,以实现高效推理和节约训练成本。
FlashMLA可以看做是针对Hopper GPU的进一步优化。它是DeepSeek为MLA开发的GPU内核,并针对H800 SXM5等Hopper架构GPU进行了专门优化,主要用于大语言模型的解码阶段(即按顺序生成新token的过程)。该项目借鉴了Flash Attention 2和3以及英伟达cutlass库的一些设计思路。
要运行 FlashMLA,需要配备 Hopper 架构的 GPU,例如 H800 SXM5,并确保安装了 CUDA 12.3 或更高版本,以及 PyTorch 2.0 或更高版本。FlashMLA 目前支持 BF16 精度,并具备分页 KV 缓存功能,它以 64 为块大小进行内存管理,实现了高效的内存利用。
在性能方面,FlashMLA 展现出了卓越的表现。在内存带宽方面,FlashMLA 能够达到 3000 GB/s,这已经非常接近 H800 SXM5 的硬件理论极限 3350 GB/s。这种优异的带宽表现很可能得益于类似 Flash Attention 所采用的分块和缓存技术。在计算性能方面,FlashMLA能够达到580TFLOPS,这是理论峰值的87%
有开发者推测,尽管代码库中注明使用的是BF16,但如此高的性能可能暗示了更高级的优化手段,如FP8等低精度格式。但更可能的情况是,DeepSeek在Tensor Core上进行了特殊设计,或针对MLA的特点定制了高效的计算核心,从而在特定场景下实现了性能突破。
FlashMLA的设计思路借鉴了Flash Attention 3(后者在H100上能达到740 TFLOPS的常规注意力计算)。虽然FlashMLA在H800上的580 TFLOPS看似较低,但考虑到它是专门为MLA优化的,这一性能实则非常亮眼。尤其值得注意的是,FlashMLA在性能稍弱的H800 GPU上,实现了接近H100上Flash Attention 3的表现。
总的来说,FlashAttention在注意力计算方面表现出色,但在解码任务上的优化有限,更适用于固定长度序列模型的训练与推理。相比之下,FlashMLA凭借高达3000 GB/s的内存带宽和580 TFLOPS的计算性能,更适合智能助手、翻译和文本生成等需要实时处理可变长度序列的应用场景。
场景适配:从推理加速到成本优化FlashMLA的高效性能使其在多个应用领域具有显著优势,例如:
低延迟应用: 聊天机器人、翻译服务、语音助手等应用需要快速响应。FlashMLA的高内存带宽和计算吞吐量确保了这些应用能够快速、高效地提供结果。
批处理: 在需要同时处理多个序列的场景(如批量推理)中,FlashMLA处理可变长度序列及高效管理内存的能力确保了最佳性能。
研发加速: 从事新型AI模型或算法研究的开发者可以利用FlashMLA加速实验和原型设计,尤其是在处理大规模模型和数据集时。
DeepSeek-V3采用了MLA和DeepSeekMoE以实现高效推理和节约训练成本。
成本控制: FlashMLA是DeepSeek模型中的一个关键部分。例如DeepSeek-V2总参数量2360亿,每个token激活210亿参数,支持128K上下文长度。MLA将KV缓存减少93.3%,节省42.5%的训练成本,并提升生成吞吐量5.76倍。 DeepSeek-V3参数量增加到6710亿,展现了MLA的可扩展性。并且在两个月内以558万美元的成本高效训练完成。DeepSeek自身的应用案例充分证明了FlashMLA在支持大规模、高效率语言模型方面的重要作用,这也顺应了当前AI发展追求经济效益的趋势。
“这才是真正的OpenAI”相比于发表技术论文,开源代码库更注重实际应用和社区协作。通过共享可运行的源代码,DeepSeek推动了技术普及和开发者之间的合作。FlashMLA的开源,已经让全球开发者赞叹DeepSeek的开放精神。
大模型行业从业者tomsheep认为,FlashMLA的主要贡献在于思路层面,虽然其本身的兼容性有限(例如,可能无法直接用于国产GPU),但这种思路可以启发其他项目开发类似功能。
tomsheep表示,他最期待DeepSeek能够开放训练部分的代码和数据流程。不过,根据DeepSeek此前在X上发布的信息,本周开源的5个代码库都是DeepSeek在线服务中的“基础模块”,并且都经过了生产环境的实战检验。DeepSeek官方还在相关repo中发布了一篇关于AI基础设施的文章:《Fire-Flyer AI-HPC: A Cost-Effective Software-Hardware Co-Design for Deep Learning》。
“参考OpenAI,先发布的往往是应用范围较窄的技术。FlashMLA更多地针对后端服务器运维的开发者。不过,看起来这一周的开源内容可能都围绕AI基础设施。”
绝大多数企业和用户是模型的使用者而非训练者。高效的推理框架能让AI基础设施提供商和终端用户同时受益:用户体验到更快的响应速度,企业也能以更低的成本服务更多客户。
通过开源FlashMLA,DeepSeek将成本优势共享给整个AI社区,让更多企业和开发者能够以更低的硬件投入,获得更高的推理性能,真正实现了“DS非常省硬件成本,开源出来大家一起省”。
FlashMLA的开源仅是DeepSeek开源周的第一天,不仅展示了DeepSeek的技术实力,更体现了其推动行业生态构建的开放态度,期待接下来会有更多惊喜。
来源:硅星人