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

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的主要特点

• BF16支持：FlashMLA支持BF16（Bfloat16）数据类型，这使得它在计算和内存使用上更加高效。
• 分页KV缓存：通过分页机制管理键值（KV）缓存，块大小为64，这使得它能够高效处理大规模序列。
• 高性能：在搭载CUDA 12.6的H800 SXM5上，在受内存限制的配置下可实现高达3000 GB/秒的速度，在受计算能力限制的配置下可达到580万亿次浮点运算每秒（TFLOPS）。

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

1. 克隆代码仓库：git clone https://github.com/deepseek-ai/FlashMLA.git cd FlashMLA
2. 安装依赖：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。

注意事项

• FlashMLA 专为 Hopper 架构 GPU 设计，不兼容其他架构的 GPU。
• 确保 CUDA 和 PyTorch 版本满足要求，否则可能会出现兼容性问题。

通过以上步骤，即可快速上手使用 FlashMLA，体验其在高效解码和推理加速方面的优势。