DeepSeek开源周Day1，硬核发布…

作者
DeepSeek的开源周（OpenSourceWeek）开始了，Day1开源了FlashMLA，不到一天直接冲上6.8k star~
DeepSeek分享了他们的FlashMLA——为Hopper GPU开发的高效MLA解码内核，针对变长序列进行了优化，目前已投入生产。
✅ 支持BF16
✅ 分页KV缓存（块大小64）
⚡ 在H800上，内存受限时速度可达3000 GB/s，计算受限时速度可达580 TFLOPS

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,    )    ...
    https://github.com/deepseek-ai/FlashMLA

    （文：PaperAgent）

    欢迎分享

    发表评论