DeepSeek 开源记-第1天-FlashMLA

DeepSeek 开源工作第1天

🚀 Day 1 of #OpenSourceWeek: FlashMLA

Honored to share FlashMLA - our efficient MLA decoding kernel for Hopper GPUs, optimized for variable-length sequences and now in production.

✅ BF16 support
✅ Paged KV cache (block size 64)
⚡ 3000 GB/s memory-bound & 580 TFLOPS compute-bound on H800

🔗 Explore on GitHub: https://github.com/deepseek-ai/FlashMLA

---

🚀 第1天 #OpenSourceWeek：FlashMLA

很荣幸分享FlashMLA-我们用于Hopper GPU的高效MLA解码内核，针对可变长度序列进行了优化，现已投入生产。

✅ BF16支持
✅ 分页KV缓存（块大小64）
⚡ H800上3000 GB/s内存绑定和580 TFLOPS计算绑定

🔗 在GitHub上探索：https://github.com/deepseek-ai/FlashMLA

FlashMLA 开源项目介绍

1. 项目概述

FlashMLA 是一个由 DeepSeek AI 开发的高性能 CUDA 库，主要用于实现 Multi-head Linear Attention (MLA) 的计算加速。该项目采用 MIT 许可证开源。

2. 核心功能

项目主要实现了以下功能：

1. MLA 计算加速：

   • 支持 FP16 和 BF16 数据类型
   • 针对 NVIDIA SM90 架构(如 H100)优化
   • 支持 KV Cache 机制
   • 实现了高效的分块计算策略

2. 关键特性：

   • 固定 head dimension = 576
   • Value dimension = 512
   • 支持 Shared KV 模式
   • 实现了高效的 softmax 计算

3. 项目结构

FlashMLA/
├── csrc/                    # C++/CUDA 核心实现
│   ├── cutlass/            # NVIDIA CUTLASS 库(子模块)
│   ├── flash_api.cpp       # C++ API 接口
│   ├── flash_fwd_*.cu      # CUDA 核心实现
│   └── *.h                 # 头文件
├── flash_mla/              # Python 包
├── benchmark/              # 性能测试
├── tests/                  # 单元测试
└── setup.py               # 构建脚本

4. 核心实现

1. 计算流程：

• 输入处理和数据加载
• 分块计算 attention scores
• Softmax 计算
• 输出结果合并

2. 优化策略：

• 使用 CUDA GMMA (Tensor Core) 加速矩阵运算
• 实现了高效的共享内存访问模式
• 使用双缓冲技术优化数据传输
• 支持 batch 处理和序列长度的动态适配

5. 性能对比

项目提供了多个基准测试实现：

• Torch 原生实现
• FlashMLA 实现
• FlashInfer 实现
• Triton 实现

可以通过 benchmark/bench_flash_mla.py 进行性能测试和对比。

6. 使用方法

1. 安装依赖：

git clone https://github.com/deepseek-ai/FlashMLA.git
cd FlashMLA
pip install -e .

2. 基本使用：

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)

# 执行 MLA 计算
output, lse = flash_mla_with_kvcache(
    q, blocked_k, block_table, cache_seqlens, dv,
    tile_scheduler_metadata, num_splits, causal=causal,
)

7. 限制条件

1. 仅支持 NVIDIA SM90 架构
2. head_dim 固定为 576
3. value_dim 固定为 512
4. 需要 CUDA 工具链支持

这个项目主要面向大规模语言模型的推理加速，特别适合需要高性能 MLA 计算的场景。通过深度优化和精心的工程实现，在特定硬件上能够获得显著的性能提升。

DeepSeek 开源组织及X账号

• https://github.com/deepseek-ai
• https://x.com/deepseek_ai