DeepSeek 代碼庫(kù)開(kāi)源進(jìn)度 1/5：為 Hopper GPU 優(yōu)化的高效 MLA 解碼內(nèi)核 FlashMLA

IT之家 2 月 24 日消息，DeepSeek 今日啟動(dòng)“開(kāi)源周”，首個(gè)開(kāi)源的代碼庫(kù)為 FlashMLA—— 針對(duì) Hopper GPU 優(yōu)化的高效 MLA 解碼內(nèi)核，專(zhuān)為處理可變長(zhǎng)度序列而設(shè)計(jì)。據(jù)介紹，F(xiàn)lashMLA 的靈感來(lái)自 FlashAttention 2&3 和 cutlass 項(xiàng)目。

IT之家附開(kāi)源地址：https://github.com/deepseek-ai/FlashMLA

官方對(duì)其的介紹大意如下：

需求：

• Hopper GPU

• CUDA 12.3 及以上版本

• PyTorch 2.0 及以上版本

安裝：

python setup.py install

benchmark：

python tests/test_flash_mla.py

使用 CUDA 12.6，H800 SXM5 在內(nèi)存受限配置下可達(dá) 3000 GB/s 帶寬，在計(jì)算受限配置下可達(dá) 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,
    )

引用：

@misc{flashmla2025,
      title={FlashMLA: Efficient MLA decoding kernel}, 
      author={Jiashi Li},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}

廣告聲明：文內(nèi)含有的對(duì)外跳轉(zhuǎn)鏈接（包括不限于超鏈接、二維碼、口令等形式），用于傳遞更多信息，節(jié)省甄選時(shí)間，結(jié)果僅供參考，IT之家所有文章均包含本聲明。