北京时间 2024 年 2 月 26 日讯，国内人工智能公司 DeepSeek 今日正式开源其用于 Hopper GPU 的高效 MLA 解码内核 FlashMLA。该项目一经发布便迅速引爆 AI 社区，短短 45 分钟内 GitHub Star 数量突破 400，并持续呈现爆发式增长。这一举动被业内视为 DeepSeek 在大模型推理加速领域迈出的重要一步，彰显了其拥抱开源、推动 AI 技术发展的决心。
背景：大模型推理成本居高不下，KV Cache 成瓶颈
近年来，大型语言模型（LLM）在自然语言处理、机器翻译、文本生成等领域取得了显著进展。然而，随着模型规模的不断扩大，其推理成本也呈指数级增长，成为制约 LLM 应用普及的关键因素。
在 LLM 推理过程中，KV Cache（Key-Value Cache）用于存储先前计算得到的 Key 和 Value 向量，以便在后续的解码步骤中重复使用，从而避免重复计算，提高推理效率。然而，随着上下文长度的增加，KV Cache 的大小也会线性增长，占用大量的 GPU 内存，甚至导致 OOM（Out Of Memory）错误，严重影响推理性能和成本。
为了解决这一问题，研究人员提出了多种 KV Cache 优化技术，包括量化、剪枝、蒸馏等。其中，MLA（Multi-Query Attention）是一种有效的减少 KV Cache 大小的技术，通过在多个 Query 之间共享 Key 和 Value 向量，从而降低 KV Cache 的存储需求。
DeepSeek MLA 技术：突破推理瓶颈的关键
DeepSeek 在大模型领域深耕已久，MLA 技术是其重要的技术创新点。通过减少推理过程中的 KV Cache，MLA 使得 DeepSeek 的大模型能够在更少的设备上推理更长的 Context，极大地降低了推理成本。
此次 DeepSeek 开源的 FlashMLA 是 MLA 技术的改进版本，专门针对 Hopper GPU 进行了优化，进一步提升了推理效率。
FlashMLA：Hopper GPU 上的高效 MLA 解码内核
FlashMLA 是一款适用于 Hopper GPU 的高效 MLA 解码内核，针对可变长度序列服务进行了优化。其主要特点包括：
- 高性能： FlashMLA 充分利用 Hopper GPU 的硬件特性，实现了极高的内存带宽和计算吞吐量。在 H800 SXM5 GPU 上，FlashMLA 能够实现高达 3000 GB/s 的内存速度上限以及 580 TFLOPS 的计算上限。
- 低延迟： FlashMLA 针对可变长度序列服务进行了优化，能够有效地处理不同长度的输入序列，降低推理延迟。
- 易于使用： FlashMLA 提供了简洁易用的 API，方便开发者将其集成到现有的 LLM 推理框架中。
目前，FlashMLA 已经发布了 BF16 块大小为 64 的分页 kvcache 实现。
技术细节：深入剖析 FlashMLA 的核心机制
FlashMLA 的高性能和低延迟得益于其精巧的设计和优化。以下是一些关键的技术细节：
- Hopper GPU 优化： FlashMLA 充分利用了 Hopper GPU 的 Tensor Core 和 Transformer Engine 等硬件加速单元，实现了高效的矩阵乘法和注意力计算。
- 内存访问优化： FlashMLA 采用了多种内存访问优化技术，包括数据预取、缓存优化、内存对齐等，最大限度地减少了内存访问延迟。
- 并行计算优化： FlashMLA 采用了细粒度的并行计算策略，将计算任务分解成多个小的子任务，并在多个 GPU 核心上并行执行，从而提高了计算效率。
- 可变长度序列处理： FlashMLA 针对可变长度序列服务进行了优化，采用了动态规划和填充等技术，有效地处理不同长度的输入序列。
部署与使用：快速上手 FlashMLA
要部署和使用 FlashMLA，需要满足以下条件：
- Hopper GPU
- CUDA 12.3 及以上版本
- PyTorch 2.0 及以上版本
安装步骤如下：
python setup.py install
基准测试：
python tests/test_flash_mla.py
使用示例：
“`python
from flashmla import getmlametadata, flashmlawithkvcache
tileschedulermetadata, numsplits = getmlametadata(cacheseqlens, sq * hq // hkv, hkv)
for i in range(numlayers):
…
oi, lsei = flashmlawithkvcache(qi, kvcachei, blocktable, cacheseqlens, dv,
tileschedulermetadata, num_splits, causal=True)
…
“`
社区反响：好评如潮，期待 AGI
FlashMLA 的开源引起了 AI 社区的广泛关注和热烈讨论。许多开发者表示，FlashMLA 能够有效地降低 LLM 的推理成本，提高推理效率，为 LLM 的应用普及带来了新的希望。
一些网友在社交媒体上表达了对 FlashMLA 的赞赏和期待：
- “DeepSeek 这次开源真是诚意满满，FlashMLA 解决了 LLM 推理成本高的痛点，期待更多的开源项目。”
- “FlashMLA 的性能太强了，在 H800 上能够达到 3000 GB/s 的内存速度上限，简直是推理加速神器。”
- “DeepSeek 这是要搞事情啊，听说第五天会开源 AGI，拭目以待！”
DeepSeek 的开源战略：拥抱开放，共建 AI 生态
FlashMLA 的开源是 DeepSeek 开源战略的重要组成部分。DeepSeek 表示，将持续拥抱开源，积极参与 AI 社区的建设，与开发者共同推动 AI 技术的发展。
DeepSeek 的开源战略具有以下几个方面的意义：
- 加速技术创新： 开源能够吸引更多的开发者参与到项目中，共同贡献代码、发现 Bug、提出改进建议，从而加速技术创新。
- 提高代码质量： 开源的代码经过了广泛的测试和审查，能够有效地提高代码质量和可靠性。
- 降低开发成本： 开源能够减少重复开发，降低开发成本，提高开发效率。
- 促进技术普及： 开源能够促进技术的普及和应用，让更多的人受益于 AI 技术的发展。
未来展望：FlashMLA 的发展前景
FlashMLA 作为一款高性能、低延迟的 MLA 解码内核，具有广阔的应用前景。未来，FlashMLA 将在以下几个方面得到进一步发展：
- 支持更多的 GPU： FlashMLA 目前只支持 Hopper GPU，未来将扩展到支持更多的 GPU，包括 NVIDIA 的其他型号 GPU 以及 AMD 和 Intel 的 GPU。
- 支持更多的 LLM： FlashMLA 目前只支持 DeepSeek 的 LLM，未来将扩展到支持更多的 LLM，包括 GPT、BERT、Transformer 等。
- 提供更多的优化选项： FlashMLA 将提供更多的优化选项，方便开发者根据自己的需求进行定制。
- 集成到更多的推理框架中： FlashMLA 将集成到更多的推理框架中，方便开发者将其应用到实际场景中。
结论：开源 FlashMLA，DeepSeek 引领 AI 推理加速新篇章
DeepSeek 开源 FlashMLA 是 AI 领域的一项重要事件。FlashMLA 作为一款高性能、低延迟的 MLA 解码内核，能够有效地降低 LLM 的推理成本，提高推理效率，为 LLM 的应用普及带来了新的希望。
DeepSeek 的开源战略彰显了其拥抱开放、推动 AI 技术发展的决心。相信在 DeepSeek 和广大开发者的共同努力下，FlashMLA 将在未来得到更广泛的应用和发展，为 AI 领域带来更多的创新和突破。
参考文献：
- DeepSeek FlashMLA GitHub 仓库：https://github.com/deepseek-ai/FlashMLA
- 机器之心相关报道：https://www.jiqizhixin.com/articles/2024-02-24-deepseek-flashmla
致谢：
感谢 DeepSeek 团队为 AI 社区做出的贡献，期待 DeepSeek 在未来能够带来更多的开源项目，共同推动 AI 技术的发展。
Views: 0