当前位置: 首页 » 资讯 » 新科技 » 正文

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

IP属地 北京 编辑:唐云泽 IT之家 时间:2025-02-24 10:33:29

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

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

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

需求:

Hopper GPU

CUDA 12.3 及以上版本

PyTorch 2.0 及以上版本

安装: python setup.py installbenchmark: 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}},}

免责声明:本网信息来自于互联网,目的在于传递更多信息,并不代表本网赞同其观点。其内容真实性、完整性不作任何保证或承诺。如若本网有任何内容侵犯您的权益,请及时联系我们,本站将会在24小时内处理完毕。