FlashMLA – DeepSeek发布的开源项目，提升大语言模型的解码效率和性能

2025年02月24日，DeepSeek在其OpenSourceWeek第一天发布了FlashMLA，那么FlashMLA是什么？有什么特点，如何使用呢？一起来看看。

FlashMLA是什么？

FlashMLA是由DeepSeek最新推出的开源项目，一个专为Hopper GPU（特别是H800）优化的高效多层注意力（MLA）解码内核，针对可变长度序列服务进行了优化，旨在提升大语言模型（LLM）的解码效率和性能，尤其适用于高性能AI推理任务。要求是Hopper GPU、CUDA 12.3 及更高版本及PyTorch 2.0 及更高版本。

FlashMLA - DeepSeek发布的开源项目，提升大语言模型的解码效率和性能

FlashMLA的主要特点

针对Hopper GPU优化：专为H800 GPU设计，充分利用其硬件特性。
高性能表现：
- 实现3000 GB/s的内存带宽。
- 提供580 TFLOPS的计算性能。
支持BF16浮点格式：提升计算效率和精度。
分页KV缓存：块大小为64，优化内存管理。
灵感来源：借鉴了FlashAttention 2&3和Cutlass项目的技术优势。

FlashMLA的应用场景

FlashMLA主要用于需要高效解码和快速响应的AI任务，具体包括：

实时AI推理：适用于对性能要求极高的实时生成任务，如智能客服、实时翻译等。
聊天机器人：加速对话生成，提升交互流畅度。
文本生成：提高文本生成效率，适用于内容创作、文案生成等场景。

如何使用FlashMLA？

FlashMLA开源地址：https://github.com/deepseek-ai/FlashMLA

以下是使用 FlashMLA 的详细指南：

环境准备

FlashMLA 需要以下环境支持：

硬件要求：NVIDIA Hopper 架构 GPU（如 H800）。
软件要求：
- CUDA 12.3 及以上版本。
- PyTorch 2.0 及以上版本。

安装 FlashMLA

克隆代码仓库：git clone https://github.com/deepseek-ai/FlashMLA.git cd FlashMLA
安装依赖：python setup.py install这将完成 FlashMLA 的安装。

运行 Benchmark 测试

安装完成后，可以运行官方提供的 Benchmark 测试脚本，验证 FlashMLA 的性能：

python tests/test_flash_mla.py

在 H800 SXM5 GPU 上，FlashMLA 在内存受限场景下可达到 3000 GB/s 的带宽，在计算受限场景下可实现 580 TFLOPS 的算力。

使用 FlashMLA 的代码示例

以下是 FlashMLA 的基本使用示例：Python

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
)

# 在多层解码中使用 FlashMLA
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。