DeepSeek AI 团队开源了 FlashMLA，这是一个针对 Hopper GPU 优化的 MLA 解码内核，适用于可变长度序列服务。FlashMLA 实现了高达 3000 GB/s 的内存带宽和 580 TFLOPS 的计算性能，支持 BF16 和分页 KV 缓存。这是 DeepSeek AI 团队 #OpenSourceWeek 活动的一部分，展示了小团队在 AGI 领域的创新能力和开源精神。
DeepSeek AI 开源周震撼来袭
“我们是一个探索 AGI 的小团队。” DeepSeek AI 在其官方声明中这样写道。但这支小团队却有着大大的野心。从下周开始，他们将启动 #OpenSourceWeek 活动，连续开源 5 个项目，以完全透明的方式分享他们在 AGI 领域的进展。
DeepSeek AI 是一家专注于通用人工智能（AGI）研究的初创公司。尽管团队规模不大，但他们在 AGI 领域已经取得了一系列令人瞩目的成果。
这可不是闹着玩的。DeepSeek AI 团队表示，这些开源项目都是经过生产环境验证、文档齐全的构建模块。这意味着，这些代码不仅仅是理论上的探索，而是已经在实际应用中发挥作用的、可靠的工具。
在“内卷”严重的AI届，开源或许才是真正的未来！
核心技术：FlashMLA 究竟是什么？
MLA 简介
多层注意力机制（Multi-Layer Attention, MLA）是现代深度学习模型，特别是 Transformer 模型中的一个重要组成部分。它允许模型在处理序列数据时，同时关注来自不同层次的信息。
MLA 的核心思想是将输入序列分成多个层次，每个层次提取不同尺度的特征。然后，通过注意力机制，将这些不同层次的特征进行融合，从而得到更全面的信息表示。这种多层次的特征提取和融合方式，使得模型能够更好地捕捉序列中的长距离依赖关系和复杂模式。
FlashMLA 的特性与优势
FlashMLA 是 DeepSeek AI 团队为 Hopper GPU 专门优化的 MLA 解码内核。它特别针对可变长度序列的服务进行了优化，这意味着它在处理不同长度的文本输入时，能够更加高效地利用计算资源。
那么，FlashMLA 到底有哪些亮点呢？
-
1. BF16 支持： FlashMLA 支持 BF16 数据格式，这是一种在深度学习中越来越流行的低精度格式，可以在减少内存占用的同时，保持较高的计算精度。 -
2. 分页 KV 缓存： FlashMLA 采用了分页 KV 缓存技术，块大小为 64。这种技术可以更有效地管理内存，特别是在处理长序列时，可以避免不必要的内存浪费。
与 FlashAttention 等类似技术相比，FlashMLA 最大的优势在于对 Hopper GPU 的深度优化。通过充分利用 Hopper GPU 的硬件特性，FlashMLA 实现了更高的内存带宽和计算性能。
性能实测：数据说话
DeepSeek AI 团队在 H800 SXM5 GPU 上对 FlashMLA 进行了测试。结果令人印象深刻：
-
• 在内存受限配置下，FlashMLA 实现了高达 3000 GB/s 的内存带宽。 -
• 在计算受限配置下，FlashMLA 实现了高达 580 TFLOPS 的计算性能。
这些数据表明，FlashMLA 在充分利用 Hopper GPU 硬件性能方面，达到了非常高的水平。
深入剖析：FlashMLA 的技术细节
安装与快速上手
想要体验 FlashMLA 的强大性能？DeepSeek AI 团队提供了详细的安装和使用指南。
首先，你需要通过以下命令安装 FlashMLA：
python setup.py install安装完成后，你可以运行基准测试脚本来验证 FlashMLA 的性能：
python tests/test_flash_mla.py代码解读：flash_api.cpp 核心逻辑
FlashMLA 的核心代码位于 flash_api.cpp 文件中。这个文件包含了两个主要的函数：get_mla_metadata 和 mha_fwd_kvcache_mla。
get_mla_metadata 函数用于获取 MLA 内核的元数据，包括 tile 调度器的元数据和分割数量。这些元数据对于后续的计算至关重要。
mha_fwd_kvcache_mla 函数是 FlashMLA 的核心计算函数。它实现了多头注意力机制的前向传播，并支持 KV 缓存。该函数接受查询张量（q）、键缓存张量（kcache）、值缓存张量（vcache_）、序列长度信息、块表等输入，并返回输出张量和 softmax 的对数和。
以下是部分核心代码逻辑：
std::vector<at::Tensor>
mha_fwd_kvcache_mla(
at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const at::Tensor &kcache, // num_blocks x page_block_size x num_heads_k x head_size
c10::optional<const at::Tensor> &vcache_, // num_blocks x page_block_size x num_heads_k x head_size_v
const int head_size_v,
const at::Tensor &seqlens_k, // batch_size
const at::Tensor &block_table, // batch_size x max_num_blocks_per_seq
const float softmax_scale,
bool is_causal,
const at::Tensor &tile_scheduler_metadata, // num_sm_parts x TileSchedulerMetaDataSize
const at::Tensor &num_splits // batch_size + 1
) {
// ... 代码省略 ...
// 设置参数
Flash_fwd_mla_params params = {};
params.b = batch_size;
params.seqlen_q = seqlen_q;
params.cu_seqlens_k = seqlens_k.data_ptr<int>();
params.h = num_heads;
// ... 代码省略 ...
// 执行计算
run_mha_fwd_splitkv_mla<cutlass::bfloat16_t, 576>(params, stream);
// ... 代码省略 ...
}run_mha_fwd_splitkv_mla 函数是 FlashMLA 的核心计算函数，它基于 cutlass 库实现了高效的矩阵运算。具体来说，它使用了 cutlass 中的 bfloat16_t 类型和 576 作为模板参数，这表明 FlashMLA 针对 BF16 数据格式和特定的 head size 进行了优化。
技术要求与限制
FlashMLA 对硬件和软件有一定的要求：
-
• GPU： Hopper 架构的 GPU -
• CUDA： 12.3 及以上版本 -
• PyTorch： 2.0 及以上版本
站在巨人的肩膀上：FlashMLA 的灵感来源
DeepSeek AI 团队坦诚地表示，FlashMLA 的设计灵感来源于两个开源项目：
-
1. FlashAttention 2&3: 这是由 Dao-AILab 开发的快速、优化的注意力机制实现。FlashMLA 在其基础上进行了进一步的优化，特别是在 Hopper GPU 上的性能。 -
2. cutlass: 这是 NVIDIA 开源的一个 CUDA 库，提供了各种高性能的 GEMM（通用矩阵乘法）和其他计算原语。FlashMLA 利用 cutlass 库来实现高效的矩阵运算。
小团队如何挑战 AGI 领域？
DeepSeek AI 团队以小团队的身份，在 AGI 领域做出了一系列令人瞩目的成果。他们的开源策略，不仅体现了技术自信，也为整个行业带来了积极的影响。他们的行动证明，即使是资源有限的小团队，也能在尖端技术领域有所作为，甚至推动整个行业的发展。
FlashMLA 的开源，不仅仅是提供了一个工具，更是为整个 AI 社区注入了一剂强心针。它鼓励更多的开发者和研究者参与到开源项目中来，共同推动技术的进步。这种开放、协作的精神，正是 AI 技术快速发展的关键。
开源精神，共创未来
DeepSeek AI 团队的 #OpenSourceWeek 活动，不仅仅是一次技术展示，更是一次开源精神的宣扬。他们用实际行动证明，开源是推动技术进步的强大动力。
我们期待 DeepSeek AI 团队在未来能够继续为社区带来更多的惊喜。同时，我们也呼吁更多的研究者和开发者，积极参与到开源社区中来，共同为 AGI 的发展贡献力量。
推荐阅读
-
• FlashMLA GitHub 仓库: https://github.com/deepseek-ai/FlashMLA
（文：子非AI）