2025年02月24日，DeepSeek在其OpenSourceWeek第一天发布了FlashMLA，那么FlashMLA是什么？有什么特点，如何使用呢？一起来看看。
目录
- FlashMLA是什么？
- FlashMLA的主要特点
- FlashMLA的应用场景
- 如何使用FlashMLA？
- 环境准备
- 安装 FlashMLA
- 运行 Benchmark 测试
- 使用 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
运行 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，体验其在高效解码和推理加速方面的优势。
