DeepSeek 開源周首日:開源 FlashMLA,AI 推理速度再進化!
DeepSeek 開源周day1 推出 FlashMLA! 專為Hopper 架構 GPU打造的 超高效 MLA (Multi-Layer Attention) 解碼內核,現在已經正式開源啦!??
什么是 FlashMLA???
簡單來說,FlashMLA就是一個為了讓 英偉達Hopper架構 GPU跑得更快更溜的 MLA 解碼加速器! 它主要針對 變長序列 (variable-length sequences) 的場景進行了深度優化,尤其是在大模型推理服務中,效率提升簡直肉眼可見!??
根據官方介紹,FlashMLA 目前已經 在生產環境中使用,穩定性杠杠的!
FlashMLA 有啥亮點? ?
- ? ? BF16 精度支持: 緊跟潮流,支持最新的 BF16 精度,性能與效率兼得!
- ? ? Paged KV Cache (頁式鍵值緩存): 采用頁式 KV 緩存,塊大小為 64,更精細的內存管理,效率更高!
- ? ?? 極致性能: 在 H800 SXM5 GPU 上,內存受限場景下可達 3000 GB/s 的驚人速度,計算受限場景也能達到 580 TFLOPS 的算力! 而且,這一切都基于 CUDA 12.6 實現的!
快速上手,體驗飛一般的感覺! ??
DeepSeek AI 也非常貼心地給出了快速上手指南,只需簡單幾步,就能體驗 FlashMLA 的強大!
安裝:
python setup.py install
復制
跑個 Benchmark 試試水:
python tests/test_flash_mla.py
?? GitHub 倉庫傳送門: https://github.com/deepseek-ai/FlashMLA
技術細節 & 引用
FlashMLA 的背后,離不開對 FlashAttention 2&3 以及 cutlass 等優秀項目的學習和借鑒。DeepSeek AI 在這些基礎上進行了創新和優化,才有了今天的 FlashMLA。
寫在最后
溫馨提示: FlashMLA 需要 Hopper 架構 GPU、CUDA 12.3 及以上 以及 PyTorch 2.0 及以上 版本支持哦! 使用前請確保你的環境滿足要求!