分享
简明科普:DeepSeek Infra开源周 全解析
输入“/”快速插入内容
🍵
简明科普:DeepSeek Infra开源周 全解析
用户5849
用户5849
2025年3月13日修改
本篇内容是对DeepSeek Infra开源周内容的一次费曼学习结果
阅读前最好能看过前置内容
万字长文:DeepSeek 647天铸就的登神长阶
。否则虽然也能看懂,但会无法建立全局视野,摄入的知识产生折扣。
阅读的时候可以先跳过每个章节的“术语说明”,看不懂了,再回去翻术语说明。不然没有上下文硬啃术语是挺难受的。
DAY1 FlashMLA
地址:
https://github.com/deepseek-ai/FlashMLA
,11.2K stars
术语说明
MLA
(Multi-Head Latent Attention),DeepSeek-V2中首次提出,用以取代他们在DeepSeek-67B中使用的GQA。是一种成本低,效果好的注意力算法。
Flash
,来自FlashAttention这个项目,Flash是闪光、快速,Attention是注意力。这是一个老牌的,专门优化注意力计算/内存效率的项目。
FlashMLA,
所以FlashMLA,就是DeepSeek开源的,借鉴了FlashAttention项目中的一些理念,然后专门针对他们
自研MLA
进行优化的CUDA内核。
所以,通俗来讲,MLA是MLA,FlashMLA是
“如何更好在机器上跑MLA”
的优化方法。
SM(Streaming Multiprocessor,流式多处理器),
这是GPU上的计算单元,H800上有132个SM。
Wrap,
是SM上线程调度的基本单位,你可以想象为流水线。所以GPU→SM→Wrap,DeepSeek很多精细化的工作是在Wrap层级上进行的。
项目特点
①
Hopper专用
Hopper是英伟达的显卡系列名称,换成耳熟能详的就是H800、H100。
也就是说这个项目如果在其他卡上直接运行,例如华为昇腾、英伟达的A100,都是行不通的。至于为什么,下面会说。
②
对KV缓存进行分块
每块大小64。分块后,整个计算和内存的效率会得到相当大的提升。
这就像以前以订单为单位去做仓库管理,有的订单10000个货品,有的订单100个货品,这时候可能出现某个卡车1订单装不下或者1订单装不满的情况,利用率很低。
这个时候平台推动了改革,不再以订单为单位做管理,而是跟踪到每个SKU。那么这个时候你的管理颗粒度上升,虽然会带来更多管理成本,但整体的物流仓管效率也会因为细粒度而得到巨额提升。
不过这个不是DeepSeek的独特发明,而是来自FlashAttention的理念。
③
极致利用共享内存
共享内存(Shared Memory)的访问速度很快,但容量很小,英伟达的Hopper系列,每个SM(GPU计算单元)上最大的共享内存为228K(
数据来源英伟达官网
)。
而DeepSeek在项目中将KV计算的中间结果都放在共享内存上了,每个SM单元下利用其中的224K(
此数据来源知乎ling ye
),从而实现了224/228=98.2%的利用率。
一方面,这极大利用了共享内存的高速特性,
但也将这个项目牢牢限定在Hopper系列上
,因为别的系列很难支持228K/SM的共享内存(例如A100仅有164KB)。
④ Wrap级别的精雕细琢
前面讲到DeepSeek将每个SM的共享内存利用得淋漓尽致,那么这里讲的就是他对每个SM计算、内存通信上的性能压榨。
DeepSeek将整个KV的计算分为两个Wrap Group,其中一个主要负责计算,一个主要负责加载内存,在64分块大小下,刚好每次完成一个分块的计算。
如下图所示,warp0负责Gemm1 这个矩阵的计算+Gemm2一半的计算。Wrap1则负责整个计算过程Q、K、V的内存加载搬运+Gemm2另一半的计算。
p.s,请记住Gemm这个概念,在两天后的第三个项目就会提到他。
⑤ 动态输入序列的适配
在实际的场景中,输入的序列是长度不一的,特别是R1类推理模型或阅读理解类任务,输入序列会很长。
在这个项目里DS采用了双线程缓冲的模式,会进行一个动态负载进行计算。如果短序列就采用计算优先模式,长序列就采用内存优先模式。(
此部分细节需要要看代码,我的参考来自ZOMI酱的解说
)
总结
最终,在H800 SXM5这个卡上,FlashMLA实现了内存带宽3000GB/S(上限是3.2TB,利用率已经接近90%了),计算浮点数580 TFLOPS的表现