12,617
// 项目简介
FlashMLA 是一个高性能 Attention Kernels 库，专为驱动 DeepSeek-V3 和 DeepSeek-V3.2 模型而设计。它为预填充和解码阶段的稀疏和密集 Attention 机制提供了优化实现。该库支持 FP8 KV cache 等高级功能，并兼容包括 SM90 和 SM100 在内的多种 GPU 架构。
// 技术分析
FlashMLA 是一个专门的高性能注意力算子库，旨在为 DeepSeek 的 Multi-Head Latent Attention (MLA) 模型提供支持。通过为稠密和稀疏注意力提供高度优化的实现，该项目解决了大规模 transformer 推理（特别是在预填充和解码阶段）中固有的计算瓶颈。其设计优先考虑 NVIDIA 架构上的硬件级效率，利用 FP8 KV cache 量化等技术在保持模型精度的同时最大化吞吐量。
// 核心亮点
01
提供专门的稀疏注意力算子，支持 DeepSeek Sparse Attention (DSA) 的预填充和解码。
02
支持 FP8 KV cache 量化，在保持 bfloat16 计算精度的同时显著降低解码过程中的内存开销。
03
在 NVIDIA H800 GPU 上实现高性能，在计算密集型工作负载中达到高达 660 TFlops 的性能。
04
包含针对标准 Multi-Head Attention (MHA) 和 Multi-Query Attention (MQA) 模式的优化稠密注意力算子。
05
提供灵活的 API 以集成自定义稀疏索引，允许开发者在 token 级别控制注意力计算。
06
保持与现代 GPU 架构（包括 SM90 和 SM100）的兼容性，并支持各种硬件特定的优化。
// 典型使用场景
01
用于预填充和解码阶段的 Token-level 稀疏 Attention
02
用于高性能预填充和解码的密集 Attention Kernels
03
支持 FP8 KV cache 以优化内存和计算效率
// 快速开始
要开始使用 FlashMLA，请克隆存储库，初始化子模块，并使用 'pip install -v .' 安装该包。安装完成后，您可以通过使用 'get_mla_metadata' 准备 tile scheduler 元数据，然后在解码循环中调用 'flash_mla_with_kvcache'，将这些算子集成到您的推理流水线中。有关具体的实现示例，请参考 'tests/' 目录中提供的测试脚本。