DeepSeek放出了开源周首日的重磅炸弹——FlashMLA。
这是DeepSeek专为英伟达Hopper GPU打造的高效MLA解码内核,特别针对变长序列进行了优化,目前已正式投产使用。
经实测,FlashMLA在H800 SXM5平台上(CUDA 12.6),在内存受限配置下可达最高3000GB/s,在计算受限配置下可达峰值580 TFLOPS。
开源地址:https://github.com/deepseek-ai/FlashMLA
团队在致谢部分表示,FlashMLA的设计参考了FlashAttention-2、FlashAttention-3以及CUTLASS的技术实现。
有网友对此表示,「DeepSeek王炸开局,FlashMLA是真正能加速AGI进程的」。
首先,需要打开终端,输入下面代码
安装setup.py文件:
这是一个基于Python的安装命令,用于编译和安装FlashMLA模块,确保其高效运行于特定硬件。
基准测试:
这段代码是一个测试脚本,用于验证FlashMLA的功能和性能,并与PyTorch的基准实现进行对比。
python tests/test_flash_mla.py
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
for i in range(num_layers):
...
o_i, lse_i = flash_mla_with_kvcache(
q_i, kvcache_i, block_table, cache_seqlens, dv,
tile_scheduler_metadata, num_splits, causal=True,
)
...
DeepSeek的成本涉及两项关键的技术:一个是MoE,一个就是MLA(多头潜注意力)。
其中,MLA的开发
耗时数月,可将每个查询KV缓存量减少93.3%,显著减少了推理过程中的内存占用(在训练过程也是如此)。
MLA架构需要一些巧妙的设计,因此实现的复杂性大大增加。而DeepSeek成功地将这些技术整合在一起,表明他们在高效语言模型训练方面走在了前沿
多头潜注意力(MLA)
KV缓存是Transforme模型中的一种内存机制,用于存储表示对话上下文的数据,从而减少不必要的计算开销。
随着对话上下文的增长,
KV缓存会
不断扩大,从而造成显著的内存限制。
通过大幅减少每次查询所需的KV缓存量,可以相应减少每次查询所需的硬件资源,从而降低运营成本。
与标准注意力机制相比,MLA将每次查询所需的KV缓存减少了约93.3%。
MLA这种全新多头潜注意力,可以将注意力机制的内存占用减少大约80%到90%,尤其有助于处理长上下文
此外,由于H20芯片比H100具有更高的内存带宽和容量,DeepSeek在推理工作负载方面获得了更多效率提升。
除了MLA,DeepSeek其他突破性进展还有哪些?
训练(前期和后期)
不是「下一个token预测」,而是「多token预测」
DeepSeek V3以前所未见的规模实现了多Token预测(MTP)技术,这些新增的注意力模块可以预测接下来的多个Token,而不是传统的单个Token。
这显著提高了训练阶段的模型性能,且这些模块可以在推理阶段移除。
这是一个典型的算法创新案例,实现了在更低计算资源消耗下的性能提升。
其他方面,虽然DeepSeek在训练中采用了FP8精度,但像全球一些顶尖的实验室已经采用这项技术相当长时间了。
DeepSeek V3采用了我们常见的「混合专家模型」(MoE)架构,个由多个专门处理不同任务的小型专家模型组成的大模型,展现出强大的涌现能力。