FlashMLA：DeepSeek最新开源！MLA解码内核让NVIDIA Hopper开启性能狂暴模式，推理速度飙升至3000GB/s

蚝油菜花

已于 2025-02-24 22:15:57 修改

阅读量1.2k

点赞数 23

分类专栏：每日 AI 项目与应用实例文章标签：开源人工智能 deepseek

于 2025-02-24 21:46:22 首次发布

本文链接：https://blog.csdn.net/qq_19841021/article/details/145839340

版权

每日 AI 项目与应用实例专栏收录该内容

659 篇文章

订阅专栏

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！

🥦 AI 在线答疑 -> 智能检索历史文章和开源项目 -> 尽在微信公众号 -> 搜一搜：蚝油菜花 🥦

⚡ 「你的Hopper显卡还在摸鱼？DeepSeik放出MLA核弹：BF16精度+分块调度，实测推理速度破纪录！」
大家好，我是蚝油菜花。当你还在为LLM推理的这些问题头疼时——

▶️ 同一张H800显卡，别人的吞吐量是你的3倍
▶️ KV缓存管理吃掉30%显存，序列长度稍长就OOM
▶️ 自研Attention内核调试复杂，性能始终追不上FlashAttention…

FlashMLA 的横空出世彻底改变了游戏规则！这个由国产团队DeepSeek开源的高效MLA解码内核，专为Hopper架构打造：

✅ 页式KV缓存：64块精细管理，告别显存碎片
✅ 分块并行计算：580TFLOPS算力直接拉满
✅ 零适配成本：Python一键安装，兼容PyTorch生态

这个登上GitHub趋势榜的项目，正在引发LLM推理效率的二次革命——你的GPU准备好迎接性能狂暴模式了吗？

🚀 快速阅读

FlashMLA 是一款专为 NVIDIA Hopper 架构 GPU 优化的高效 MLA 解码内核，显著提升大语言模型推理性能。

核心功能：支持 BF16 精度、页式 KV 缓存，内存带宽可达 3000 GB/s，计算性能达 580 TFLOPS。
技术原理：通过分块调度与并行计算、优化内存访问模式，大幅提高计算效率和内存管理能力。

FlashMLA 是什么

FlashMLA 是由 DeepSeek 开源的一款针对 NVIDIA Hopper 架构 GPU 优化的高效 MLA（Multi-Head Linear Attention）解码内核。它专为处理可变长度序列设计，通过优化 KV 缓存机制和采用 BF16 数据格式，显著提升了内存和计算效率。在 H800 SXM5 GPU 上，FlashMLA 的内存带宽可达 3000 GB/s，计算性能可达 580 TFLOPS。

FlashMLA 的设计灵感来源于 FlashAttention 2&3 和 Cutlass 项目，支持分页缓存和低秩压缩等技术，进一步优化了内存管理和计算性能。它适用于大语言模型（LLM）的推理任务，在需要高效解码的自然语言处理（NLP）场景中表现出色。开发者可以通过简单的安装命令快速部署，并使用基准测试脚本验证性能。

FlashMLA 的主要功能

BF16 精度支持：支持 BF16 数据格式，兼顾性能与效率。
页式 KV 缓存：采用页式键值缓存机制，块大小为 64，实现更精细的内存管理。
极致性能表现：在 H800 SXM5 GPU 上，内存受限场景下可达到 3000 GB/s 的带宽，计算受限场景下可达到 580 TFLOPS 的算力。

FlashMLA 的技术原理

分块调度与并行计算：通过分块调度机制，将计算任务分解为多个小块并行处理，充分利用 GPU 的并行计算能力。
优化的内存访问模式：减少内存访问开销，在处理大规模数据时显著提升性能。

如何运行 FlashMLA

1. 环境准备

FlashMLA 需要以下硬件和软件环境：

硬件：NVIDIA Hopper 架构 GPU（例如 H800 SXM5）。
软件：CUDA 12.3 及以上版本；PyTorch 2.0 及以上版本。

2. 安装 FlashMLA

通过以下命令安装 FlashMLA：

python setup.py install

3. 运行基准测试

安装完成后，可以通过运行以下命令进行性能测试：

python tests/test_flash_mla.py

此脚本将验证 FlashMLA 在当前环境下的性能表现。例如，在 H800 SXM5 GPU 上，内存受限配置下可达到 3000 GB/s 的带宽，计算受限配置下可达到 580 TFLOPS 的算力。

4. 使用 FlashMLA

以下是 FlashMLA 的典型使用代码示例：

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