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