DeepSeek 代码库开源进度 1/5：为 Hopper GPU 优化的高效 MLA 解码内核 FlashMLA

时间2025-03-13 09:08:08发布胡小编分类AI浏览655

导读：IT之家 2 月 24 日消息，DeepSeek 今日启动“开源周”，首个开源的代码库为 FlashMLA—— 针对 Hopper GPU 优化的高效 MLA 解码内核，专为处理可变长度序列而设计。据介绍，FlashMLA 的灵感来自 FlashAttention 2&3 和 cutlass 项目。IT之家附开源地址：https://github.c...

2 月 24 日消息，DeepSeek 今日启动“开源周”，首个开源的代码库为 FlashMLA—— 针对 Hopper GPU 优化的高效 MLA 解码内核，专为处理可变长度序列而设计。据介绍，FlashMLA 的灵感来自 FlashAttention 2&3 和 cutlass 项目。

附开源地址：https://github.com/deepseek-ai/FlashMLA

官方对其的介绍大意如下：

需求：

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

安装：

python setup.py install

benchmark：

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,     )

引用：

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