剛剛,DeepSeek開(kāi)源FlashMLA,推理加速核心技術(shù),Star量飛漲中
上周五,DeepSeek 發(fā)推說(shuō)本周將是開(kāi)源周(OpenSourceWeek),并將連續(xù)開(kāi)源五個(gè)軟件庫(kù)。
第一個(gè)項(xiàng)目,果然與推理加速有關(guān)。
北京時(shí)間周一上午 9 點(diǎn),剛一上班(同時(shí)是硅谷即將下班的時(shí)候),DeepSeek 兌現(xiàn)了自己的諾言,開(kāi)源了一款用于 Hopper GPU 的高效型 MLA 解碼核:FlashMLA。
該項(xiàng)目上線才 45 分鐘就已經(jīng)收獲了超過(guò) 400 star!并且在我們截圖時(shí),Star 數(shù)量正在瘋狂飆升。
項(xiàng)目地址:https://github.com/deepseek-ai/FlashMLA
眾所周知,MLA是DeepSeek大模型的重要技術(shù)創(chuàng)新點(diǎn),主要就是減少推理過(guò)程的KV Cache,從而實(shí)現(xiàn)在更少的設(shè)備上推理更長(zhǎng)的Context,極大地降低推理成本。
此次 DeepSeek 直接開(kāi)源了該核心技術(shù)的改進(jìn)版本,可以說(shuō)是誠(chéng)意滿滿。
接下來(lái),就讓我看下這個(gè)開(kāi)源項(xiàng)目的核心內(nèi)容。
據(jù)介紹,F(xiàn)lashMLA 是適用于 Hopper GPU 的高效 MLA 解碼內(nèi)核,針對(duì)可變長(zhǎng)度序列服務(wù)進(jìn)行了優(yōu)化。
目前已發(fā)布的內(nèi)容有:
- BF16
- 塊大小為 64 的分頁(yè) kvcache
其速度非常快,在 H800 SXM5 GPU 上具有 3000 GB/s 的內(nèi)存速度上限以及 580 TFLOPS 的計(jì)算上限。
在部署這個(gè)項(xiàng)目之前,你需要的有:
- Hopper GPU
- CUDA 12.3 及以上版本
- PyTorch 2.0 及以上版本
快速啟動(dòng)
- 安裝
python setup.py install
- 基準(zhǔn)
python tests/test_flash_mla.py
使用 CUDA 12.6,在 H800 SXM5 上,在內(nèi)存綁定配置下實(shí)現(xiàn)高達(dá) 3000 GB/s,在計(jì)算綁定配置下實(shí)現(xiàn) 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,
)
...
該項(xiàng)目發(fā)布后也是好評(píng)如潮。
甚至有網(wǎng)友打趣地表示:「聽(tīng)說(shuō)第五天會(huì)是 AGI」。
最后,還是那句話:這才是真正的 OpenAI