张刚AI拆解

FlashMLA：DeepSeek开源的大模型推理加速核心引擎

作者：

张刚

在

FlashMLA（Multi-head Latent Attention Kernels）是由DeepSeek AI开发的高性能注意力计算内核库，专为英伟达Hopper架构GPU（H100/H800/H200/B200等）优化，旨在解决大模型推理中变长序列处理效率低的核心痛点，已在DeepSeek-V3和V3.2系列模型中大规模部署验证。

一、核心定位与价值

FlashMLA是大模型推理的”性能加速器”，专注于优化Transformer架构中的多头潜注意力（MLA） 计算，通过硬件级优化与算法创新，将高端GPU的算力与带宽利用率提升至极限，同时大幅降低显存占用，为大模型生产环境部署提供关键技术支撑。

二、核心功能

FlashMLA提供完整的注意力计算解决方案，覆盖大模型训练与推理全流程：

功能模块具体实现特点
稀疏注意力内核1. 预填充阶段token级稀疏注意力2. 解码阶段token级稀疏注意力（支持FP8 KV缓存）适配DeepSeek Sparse Attention (DSA)，显存占用降低93.3%
稠密注意力内核1. 预填充阶段稠密注意力2. 解码阶段稠密注意力支持MQA（Multi-Query Attention）与MHA（Multi-Head Attention）两种模式
分页KV缓存64字节块大小的分页管理机制动态分配显存，消除静态填充导致的70%空间浪费
混合精度计算BF16主计算 + FP8 KV缓存（稀疏模式）兼顾速度与精度，显存占用进一步降低

三、技术亮点与创新

1.  Hopper架构深度优化：充分利用SM90/SM100架构特性（如Tensor Core、异步拷贝），实现内存带宽与计算性能双突破

2.  变长序列并行优化：动态调度计算资源，解决序列长度差异导致的资源浪费问题

3.  极致带宽利用率：H800上内存绑定场景可达3000 GB/s，逼近硬件理论极限

4.  计算性能飙升：计算绑定场景最高达660 TFLOPS（H800 SXM5，CUDA 12.8），较传统实现提升50%+

5.  FP8稀疏解码内核：创新地将KV缓存量化为FP8，主计算保持BF16，实现性能与精度平衡

四、性能表现

场景硬件性能指标
稠密MLA解码（计算绑定）H800 SXM5660 TFLOPS
稠密MLA解码（内存绑定）H800 SXM53000 GB/s
稀疏MLA解码（FP8 KV）H800 SXM5410 TFLOPS
稀疏MLA预填充H800 SXM5640 TFLOPS
稀疏MLA预填充B2001450 TFLOPS
稠密MHA预填充（前向）B2001460 TFLOPS

注：2025年4月更新版本为计算密集型工作负载带来5%~15%额外性能提升，接口完全兼容旧版本

五、适用场景

FlashMLA特别适合以下大模型应用场景：

1.  长文本处理：文档摘要、法律/医疗文本分析、代码理解等上下文窗口大且长度不一的任务

2.  对话系统：多轮对话中历史上下文动态增长的场景，KV缓存优化效果显著

3.  多模态推理：图像+文本混合输入的动态序列处理，如视觉问答、图文生成

4.  低成本部署：在有限硬件资源下实现更大模型或更长上下文的推理服务

5.  大模型训练：稠密MHA预填充内核支持前向/反向计算，提升训练效率

六、技术栈与部署要求

类别具体要求
硬件SM90/SM100架构GPU（H100/H800/H200/B200等）
软件CUDA 12.8+（SM100需12.9+），PyTorch 2.0+
语言C++（65.7%）、CUDA（25.7%）、Python（8.2%）
许可证MIT开源许可，商业使用友好

七、安装与使用

安装步骤

git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla
git submodule update --init --recursive
pip install -v .

快速使用示例

1. 稠密MLA解码

from flash_mla import get_mla_metadata, flash_mla_with_kvcache
 
# 初始化元数据（解码循环前调用一次）
metadata = get_mla_metadata(
    s_q=1, h_kv=1, h_q=32, d_qk=576, d_v=512, 
    page_block_size=64, is_fp8_kvcache=False
)
 
# 解码循环中调用
out, lse = flash_mla_with_kvcache(
    q, kv_cache, block_table, metadata, sm_scale
)

2. 稀疏MLA解码（FP8 KV缓存）

# 启用FP8 KV缓存
metadata = get_mla_metadata(..., is_fp8_kvcache=True)
out, lse = flash_mla_with_kvcache(
    q, kv_cache, block_table, metadata, sm_scale,
    indices=sparse_indices  # 稀疏索引张量
)

3. 稀疏MLA预填充

from flash_mla import flash_mla_sparse_fwd
 
out, max_logits, lse = flash_mla_sparse_fwd(
    q, kv, indices, sm_scale
)

八、生态支持与社区发展

FlashMLA已形成跨硬件平台的适配生态，支持多种国产与国际GPU：

● MetaX GPU：MetaX-MACA/FlashMLA

● 摩尔线程GPU：MooreThreads/MT-flashMLA

● 海光DCU：OpenDAS/MLAttention

● 燧原NNP：Intellifusion/tyllm

● 天数智芯GPU：Deep-Spark/FlashMLA

● AMD Instinct：AITER/MLA

项目GitHub已获得11.8k stars、905 forks，社区活跃，持续迭代优化。

九、总结

FlashMLA是大模型推理加速领域的”性能标杆”，通过硬件感知优化与算法创新，将Hopper架构GPU的潜力发挥到极致，为大模型落地提供了关键的”软加速”方案。对于需要处理变长序列、追求高吞吐量与低延迟的大模型应用，FlashMLA是理想的推理优化选择，尤其适合资源受限但需高性能推理的中小团队与企业。

更多文章