DeepSeek开源第一弹！突破H800性能上限，FlashMLA重磅开源

FlashMLA 是一个针对 Hopper GPU 优化的高效 MLA（Multi-Head Latent Attention）解码内核，支持变长序列处理，现在已经投入生产使用。

FlashMLA 通过优化 MLA 解码和分页 KV 缓存，能够提高 LLM（大语言模型）推理效率，尤其是在 H100 / H800 这样的高端 GPU 上发挥出极致性能。

说人话就是，FlashMLA 是一种专门为 Hopper 高性能 AI 芯片设计的先进技术——一种「多层注意力解码内核」。简单来说，它就像是一个超级高效的「翻译器」，能让计算机更快地处理语言信息。它能让计算机处理各种长度的语言信息，而且速度特别快。

---

---

Stars 数	9957
Forks 数	609

已发布内容

•  BF16

•  块大小为 64 的分页 kvcache

其速度非常快，在 H800 SXM5 GPU 上具有 3000 GB/s 的内存速度上限以及 580 TFLOPS 的计算上限。

项目环境

•  Hopper GPU

•  CUDA 12.3 及以上版本

•  PyTorch 2.0 及以上版本

快速启动

• 安装

python setup.py install

• 基准

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

GitHub：GitHub的 - deepseek-ai/FlashMLA