FlashMLA

FlashMLA 是由 DeepSeek 开发的开源高效 MLA（Multi-Head Linear Attention）解码内核，专为 NVIDIA Hopper 架构 GPU 优化，旨在提升处理可变长度序列的效率。通过采用 BF16 数据格式和优化的 KV 缓存机制，FlashMLA 在内存和计算性能方面均表现出色。在 H800 SXM5 GPU 上，其内存带宽可达 3000 GB/s，计算性能可达 580 TFLOPS。 FlashMLA 的设计借鉴了 FlashAttention 2 和 Cutlass 等项目，支持分页缓存和低秩压缩等技术，进一步优化了内存管理与计算效率。该工具适用于大语言模型（LLM）的推理任务，在自然语言处理（NLP）场景中具有显著优势。开发者可通过简单命令进行安装，并运行基准测试脚本验证性能表现。 FlashMLA 提供了高效的分块调度与并行计算能力，同时优化了内存访问模式，提升了大规模数据处理的性能。其使用需依赖 NVIDIA Hopper 架构 GPU 及 CUDA 12.3、PyTorch 2.0 等软件环境。完整代码与文档可在 GitHub 仓库中获取。