DeepSeek 代码库开源进度 1/5:为 Hopper GPU 优化的高效 MLA 解码内核 FlashMLA

DeepSeek 今日启动“开源周”,首个开源的代码库为 FlashMLA—— 针对 Hopper GPU 优化的高效 MLA 解码内核,专为处理可变长度序列而设计。据介绍,FlashMLA 的灵感来自 FlashAttention 2&3 和 cutlass 项目。

DeepSeek 代码库开源进度 1/5:为 Hopper GPU 优化的高效 MLA 解码内核 FlashMLA

IT之家附开源地址:https://github.com/deepseek-ai/FlashMLA

官方对其的介绍大意如下:

需求:

  • Hopper GPU

  • CUDA 12.3 及以上版本

  • PyTorch 2.0 及以上版本

安装:

python setup.py install

benchmark:

python tests/test_flash_mla.py

使用 CUDA 12.6,H800 SXM5 在内存受限配置下可达 3000 GB/s 带宽,在计算受限配置下可达 580 TFLOPS 算力。

用法:

from flash_mla import get_mla_metadata flash_mla_with_kvcachetile_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,    )

引用:

@misc{flashmla2025,      title={FlashMLA: Efficient MLA decoding kernel},       author={Jiashi Li},      year={2025},      publisher = {GitHub},      howpublished = {url{https://github.com/deepseek-ai/FlashMLA}},}

免责声明:文章内容来自IT之家

本站不对其内容的真实性、完整性、准确性给予任何担保、明示、暗示和承诺,本文仅供读者参考!

数码迷尊重原作者的辛勤劳动并致力于保护原著版权以及相关的知识产权,所转载的文章,其版权归原作者所有。

如本文内容影响到您的合法权益(内容、图片等),请通过邮箱5937331#qq.com联系我们,我们将第一时间回复处理。

(0)
上一篇 2025年2月25日 上午8:38
下一篇 2025年2月25日 上午8:38

相关推荐