DeepSeek OpenSourceWeek第一弹:FlashMLA推理加速解码核分析

edwin99
edwin99
2025-02-24 07:19
29 阅读
0 评论
文章封面
目录
正在加载目录...

第一个项目与推理加速有关。北京时间周一上午 9 点,刚一上班(同时是硅谷即将下班的时候),DeepSeek 兑现了自己的诺言,开源了一款用于 Hopper GPU 的高效型 MLA 解码核:FlashMLA。

项目地址:https://github.com/deepseek-ai/FlashMLA

FlashMLA 是适用于 Hopper GPU 的高效 MLA 解码内核,针对可变长度序列服务进行了优化。目前已发布的内容有:BF16、块大小为64的分页kvcache。其速度非常快,在 H800 SXM5 GPU 上具有 3000 GB/s 的内存速度上限以及 580 TFLOPS 的计算上限。
 

本地部署

在部署这个项目之前,需要准备Hopper GPU,CUDA 12.3及以上版本、PyTorch 2.0及以上版本。开销大概是10美元以内/小时,一个小时拿来实验绰绰有余。

以下是安装、基准、启动用到的代码:

#安装
python setup.py install

#基准
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 //h_kv, h_kv)


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,
)
...

使用 CUDA 12.6,在 H800 SXM5 上,在内存绑定配置下实现高达 3000 GB/s,在计算绑定配置下实现 580 TFLOPS。
 
 
 
 

评论区 (0)

登录后参与评论

暂无评论,抢沙发吧!