https://github.com/deepseek-ai/FlashMLA

FlashMLA

FlashMLA 是针对Hopper架构GPU优化的高效MLA解码内核，专为变长序列服务场景设计。

当前已发布特性：

  • BF16支持
  • 分页式kvcache（块大小为64）

快速入门

安装

python setup.py install

性能测试

python tests/test_flash_mla.py

在H800 SXM5（CUDA 12.6环境）上实现内存受限配置下3000GB/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及以上版本
  • PyTorch 2.0及以上版本

致谢

FlashMLA的设计灵感来源于FlashAttention 2&3以及cutlass项目。