前言:

DeepSeek 兑现了自己的诺言,开源了一款用于 Hopper GPU 的高效型 MLA 解码核:FlashMLA。

项目地址:https://github.com/deepseek-ai/FlashMLA

1:FlashMLA 是什么呀?

MLA是DeepSeek大模型的重要技术创新点,主要就是减少推理过程的KV Cache,从而实现在更少的设备上推理更长的Context,极大地降低推理成本。DeepSeek 直接开源了该核心技术的改进版本,可以说是诚意满满

FlashMLA适用于 Hopper GPU 的高效 MLA 解码内核,针对可变长度序列服务进行了优化。

目前已发布的内容有:

 BF16

块大小为 64 的分页 kvcache

其速度非常快,在 H800 SXM5 GPU 上具有 3000 GB/s 的内存速度上限以及 580 TFLOPS 的计算上限。

2:部署

在部署这个项目之前,你需要的有:

  •  Hopper GPU

  •  CUDA 12.3 及以上版本

  •  PyTorch 2.0 及以上版本

2.1安装

使用 CUDA 12.6,在 H800 SXM5 上,在内存绑定配置下实现高达 3000 GB/s,在计算绑定配置下实现 580 TFLOPS。

2.2 使用

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,)...

你可以试试!

最后,还是那句话:这才是真正的 open source 的精神所在!

Logo

欢迎加入DeepSeek 技术社区。在这里,你可以找到志同道合的朋友,共同探索AI技术的奥秘。

更多推荐