DeepSeek 启动 “开源周”，发布高效 MLA 解码内核 FlashMLA

FlashMLA 开源发布

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

需求与安装

FlashMLA 需要以下环境：
  • Hopper GPU
  • CUDA 12.3 及以上版本
  • PyTorch 2.0 及以上版本
安装步骤：
python setup.py install

性能测试

使用以下命令进行基准测试：
python tests/test_flash_mla.py
在使用 CUDA 12.6 和 H800 SXM5 的内存受限配置下，FlashMLA 可达到 3000 GB/s 的带宽；在计算受限配置下，可达到 580 TFLOPS 的算力。

使用示例

以下是 FlashMLA 的使用示例代码：
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,
    )

引用方式

如果您在研究中使用了 FlashMLA，请引用以下文献：
@misc{flashmla2025,
      title={FlashMLA: Efficient MLA decoding kernel}, 
      author={Jiashi Li},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}

结论

DeepSeek 启动 “开源周”，并发布了针对 Hopper GPU 优化的高效 MLA 解码内核 FlashMLA。该内核专为处理可变长度序列而设计，灵感来自 FlashAttention 2&3 和 cutlass 项目。FlashMLA 在内存受限配置下可达到 3000 GB/s 的带宽，在计算受限配置下可达到 580 TFLOPS 的算力。用户可以通过简单的安装和使用步骤，快速集成并测试该内核。此次开源发布不仅为开发者提供了强大的工具，也促进了社区的技术交流与合作。