第一个项目与推理加速有关。北京时间周一上午 9 点,刚一上班(同时是硅谷即将下班的时候),DeepSeek 兑现了自己的诺言,开源了一款用于 Hopper GPU 的高效型 MLA 解码核:FlashMLA。
项目地址:https://github.com/deepseek-ai/FlashMLA
本地部署
在部署这个项目之前,需要准备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,
)
...