17370845950

刚刚,DeepSeek开源FlashMLA,推理加速核心技术,Star量飞涨中

deepseek开源高效型mla解码核flashmla,助力hopper gpu推理加速!上周五deepseek预告开源周计划,并于北京时间周一上午9点开源了首个项目——flashmla,一款针对hopper gpu优化的高效mla解码内核,仅上线45分钟便收获400+star!

☞☞☞AI 智能聊天, 问答助手, AI 智能搜索, 免费无限量使用 DeepSeek R1 模型☜☜☜

该项目Star数量持续飙升!

项目地址:https://www./link/a588b762d68fe60225d3de3c647a52b9

FlashMLA的核心在于优化可变长度序列处理,减少推理过程中的KV Cache,从而在有限硬件资源下实现更长上下文推理,显著降低推理成本。 目前已发布的版本支持BF16精度和64大小的分页kvcache,在H800 SXM5 GPU上内存速度上限达3000 GB/s,计算上限达580 TFLOPS。

使用条件:

  • Hopper GPU
  • CUDA 12.3及以上版本
  • PyTorch 2.0及以上版本

快速上手:

安装:

python setup.py install

基准测试:

python tests/test_flash_mla.py

在H800 SXM5上使用CUDA 12.6,内存绑定配置下可达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,
    )
    ...

FlashMLA开源后获得广泛好评,甚至有网友调侃“第五天会是AGI”。

这无疑是真正的开源精神的体现!