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 SXM5 | 660 TFLOPS |
| 稠密MLA解码（内存绑定） | H800 SXM5 | 3000 GB/s |
| 稀疏MLA解码（FP8 KV） | H800 SXM5 | 410 TFLOPS |
| 稀疏MLA预填充 | H800 SXM5 | 640 TFLOPS |
| 稀疏MLA预填充 | B200 | 1450 TFLOPS |
| 稠密MHA预填充（前向） | B200 | 1460 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是理想的推理优化选择，尤其适合资源受限但需高性能推理的中小团队与企业。