一眨眼,目前FlashMLA的github仓库代码已经来带了6000+的stars数量了,果然DeepSeek团队才是真正的OpenAI。
这次开源的技术,其实本质上就是在推理过程中优化了MLA。它其实总结下来就是两个关键的信息
适配DeepSeek中自研的MLA,能够有效加速MLA的推理速度。
其实MLA技术一开始就是DeepSeek-V2中训练的一种方法。在原始的mutil-head attention中,本质上在计算每一次attention的时候,都需要存储k、v、q三个变量,这会导致整体的显存过大。
之前也有论文曾经对multi-head attention进行优化,提出了共享KV的主要两种方式。一种是MQA,上图的右边。一层的所有Head,共享同一个𝑘,𝑣来计算Attention。另一种叫GQA,上图的中间这个图,是平衡了MQA和MHA的一种折中的方法,不是每个Head一个KV,也不是所有Head共享一个KV,而是对所有Head分组然后再进行存储。
而DeepSeek-V2提出的MLA(Multi-Head Latent Attention),在KV缓存数量上略大于GQA(2.25组),但具备恢复完整MHA效果的能力。其推理效果不仅优于GQA,更在部分场景中超越传统MHA,成为兼顾性能与效率的解决方案。
适配Hooper架构GPU
目前FlashMLA适配的是Hooper 架构GPU,更早的Ampere 架构以及RTX40 系用的 Ada 架构并不支持。
在H800 SXM5硬件环境下,搭配CUDA 12.6,在内存受限条件下可实现3000 GB/s的带宽与580 TFLOPS的浮点性能,展现了极致的硬件适配优化。
从部署使用上看,其用法相当简单,只需要安装对应的依赖,然后像下面一样调用即可
DeepSeek团队在首日就开源了FlashMLA,相信在接下来的几天内,社区会有更多实践和应用落地。DeepSeek正用行动证明,他们不仅仅是开源的践行者,更是AI技术普惠化的推动者。
领取专属 10元无门槛券
私享最新 技术干货