《deepseek FlashMLA ：高效的 MLA 解码内核》：此文为AI自动翻译

FlashMLA

GitHub - deepseek-ai/FlashMLA

FlashMLA 是适用于 Hopper GPU 的高效 MLA 解码内核，针对可变长度序列服务进行了优化。

当前发布：

• BF16、FP16
• 块大小为 64 的分页 kvcache

快速开始

安装

python setup.py install

基准

python tests/test_flash_mla.py

使用 CUDA 12.8，在 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,
    )
    ...

要求

• Hopper GPU
• CUDA 12.3 及更高版本
  • 但我们强烈建议使用 12.8 或更高版本以获得最佳性能
• PyTorch 2.0 及更高版本

确认

FlashMLA 的灵感来自 FlashAttention 2&3 和弯刀项目。

引文

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