速看!DeepSeek开源周第一弹:FlashMLA加速GPU计算!

?DeepSeek在开源周的第一天推出了加速推理的基础设置FlashMLA。DeepSeek针对Hopper GPU优化了高效的MLA解码内核，支持变长序列，并且已经投入生产使用。

✅ 支持BF16 ✅ 分页KV缓存（块大小64） ⚡ 在H800上实现3000 GB/s的内存带宽和580 TFLOPS的计算性能！

快速开始克隆项目代码语言：javascript代码运行次数：0运行复制```javascriptgit clone https://www.php.cn/link/3dc20e2ac60a2551a4d3f08b03377def

安装代码语言：javascript代码运行次数：0运行复制```javascriptpython setup.py install

测试代码语言：javascript代码运行次数：0运行复制```javascriptpython tests/test_flash_mla.py

在H800 SXM5上，使用CUDA 12.6，在受内存限制的配置下可实现高达3000 GB/s的带宽，在受计算限制的配置下可达到580 TFLOPS的计算性能。
使用代码语言：javascript代码运行次数：0运行复制javascript 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,    )    ...

最低要求代码语言：javascript代码运行次数：0运行复制javascriptHopper GPUsCUDA 12.3 and abovePyTorch 2.0 and above