速看！DeepSeek开源周第一弹：FlashMLA加速GPU计算！
速看！DeepSeek开源周第一弹：FlashMLA加速GPU计算！
AgenticAI
发布于 2025-03-18 16:45:24
发布于 2025-03-18 16:45:24
代码可运行
运行总次数：0
代码可运行
🚀DeepSeek开源周第一天，发布加速推理的基础设置FlashMLA，DeepSeek为 Hopper GPU 优化高效 MLA 解码内核，支持变长序列，并已投入生产使用。
✅ 支持 BF16 ✅ 分页 KV 缓存（块大小 64） ⚡ 在 H800 上实现 3000 GB/s 内存带宽 和 580 TFLOPS 计算性能！
快速开始
- 克隆项目
git clone https://github.com/deepseek-ai/FlashMLA.git- 安装
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,
)
...
- 最低要求
Hopper GPUs
CUDA 12.3 and above
PyTorch 2.0 and above本文参与 腾讯云自媒体同步曝光计划，分享自微信公众号。
原始发表：2025-02-24，如有侵权请联系 cloudcommunity@tencent.com 删除
评论
登录后参与评论
推荐阅读
目录
腾讯云开发者
