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

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

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

开源地址: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_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,     
 )

引用:

@misc{flashmla2025,       
 title={FlashMLA: Efficient MLA decoding kernel},        
 author={Jiashi Li},       
 year={2025},       
 publisher = {GitHub},       
 howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}}, 
}
声明:魔果智讯倡导尊重与保护知识产权。如发现本站文章存在版权等问题,烦请30天内提供版权疑问、身份证明、版权证明、联系方式等发邮件至moguoai@yeah.net!我们将及时沟通与处理。
0 条回复 A文章作者 M管理员
    暂无讨论,说说你的看法吧
个人中心
购物车
优惠劵
今日签到
有新私信 私信列表
搜索