DeepSeek 開源周day1 推出 FlashMLA！ 專為Hopper 架構(gòu) GPU打造的 超高效 MLA (Multi-Layer Attention) 解碼內(nèi)核，現(xiàn)在已經(jīng)正式開源啦！??

什么是 FlashMLA？??

簡單來說，F(xiàn)lashMLA就是一個為了讓 英偉達Hopper架構(gòu) GPU跑得更快更溜的 MLA 解碼加速器！ 它主要針對 變長序列 (variable-length sequences) 的場景進行了深度優(yōu)化，尤其是在大模型推理服務(wù)中，效率提升簡直肉眼可見！??

根據(jù)官方介紹，F(xiàn)lashMLA 目前已經(jīng) 在生產(chǎn)環(huán)境中使用，穩(wěn)定性杠杠的！

FlashMLA 有啥亮點？ ?

• ? ? BF16 精度支持： 緊跟潮流，支持最新的 BF16 精度，性能與效率兼得！
• ? ? Paged KV Cache (頁式鍵值緩存)： 采用頁式 KV 緩存，塊大小為 64，更精細的內(nèi)存管理，效率更高！
• ? ?? 極致性能： 在 H800 SXM5 GPU 上，內(nèi)存受限場景下可達 3000 GB/s 的驚人速度，計算受限場景也能達到 580 TFLOPS 的算力！ 而且，這一切都基于 CUDA 12.6 實現(xiàn)的！

快速上手，體驗飛一般的感覺！ ??

DeepSeek AI 也非常貼心地給出了快速上手指南，只需簡單幾步，就能體驗 FlashMLA 的強大！

安裝：

python setup.py install

復(fù)制

跑個 Benchmark 試試水：

python tests/test_flash_mla.py

?? GitHub 倉庫傳送門： https://github.com/deepseek-ai/FlashMLA

技術(shù)細節(jié) & 引用

FlashMLA 的背后，離不開對 FlashAttention 2&3 以及 cutlass 等優(yōu)秀項目的學習和借鑒。DeepSeek AI 在這些基礎(chǔ)上進行了創(chuàng)新和優(yōu)化，才有了今天的 FlashMLA。

寫在最后

溫馨提示： FlashMLA 需要 Hopper 架構(gòu) GPU、CUDA 12.3 及以上 以及 PyTorch 2.0 及以上 版本支持哦！ 使用前請確保你的環(huán)境滿足要求！