剛剛，DeepSeek開源FlashMLA，推理加速核心技術，Star量飛漲中

机器之心發表於2025-02-24
原文網址 : https://www.jiqizhixin.com/articles/2025-02-24-2
上週五，DeepSeek 發推說本週將是開源周（OpenSourceWeek），並將連續開源五個軟體庫。

第一個專案，果然與推理加速有關。

北京時間週一上午 9 點，剛一上班（同時是矽谷即將下班的時候），DeepSeek 兌現了自己的諾言，開源了一款用於 Hopper GPU 的高效型 MLA 解碼核：FlashMLA。
image.png
該專案上線才 45 分鐘就已經收穫了超過 400 star！並且在我們截圖時，Star 數量正在瘋狂飆升。
image.png
專案地址：https://github.com/deepseek-ai/FlashMLA

眾所周知，MLA是DeepSeek大模型的重要技術創新點，主要就是減少推理過程的KV Cache，從而實現在更少的裝置上推理更長的Context，極大地降低推理成本。

此次 DeepSeek 直接開源了該核心技術的改進版本，可以說是誠意滿滿。

接下來，就讓我看下這個開源專案的核心內容。

據介紹，FlashMLA 是適用於 Hopper GPU 的高效 MLA 解碼核心，針對可變長度序列服務進行了最佳化。

目前已釋出的內容有：

  • BF16
  • 塊大小為 64 的分頁 kvcache

其速度非常快，在 H800 SXM5 GPU 上具有 3000 GB/s 的記憶體速度上限以及 580 TFLOPS 的計算上限。

在部署這個專案之前，你需要的有：

  • Hopper GPU
  • CUDA 12.3 及以上版本
  • PyTorch 2.0 及以上版本

快速啟動

  • 安裝
python setup.py install
  • 基準
python tests/test_flash_mla.py

使用 CUDA 12.6，在 H800 SXM5 上，在記憶體繫結配置下實現高達 3000 GB/s，在計算繫結配置下實現 580 TFLOPS。

  • 用法

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

tile_scheduler_metadata, num_splits = get_mla_metadata (cache_seqlens, s_q * h_q //h_kv, h_kv)

for i in range (num_layers):
...
o_i, lse_i = flash_mla_with_kvcache (
q_i, kvcache_i, block_table, cache_seqlens, dv,
tile_scheduler_metadata, num_splits, causal=True,
)
...

該專案釋出後也是好評如潮。
image.png
甚至有網友打趣地表示：「聽說第五天會是 AGI」。
image.png
最後，還是那句話：這才是真正的 OpenAI
image.png

相關文章