DeepSeek 开源周day1: FlashMLA
https://github.com/deepseek-ai/FlashMLA FlashMLA FlashMLA 是针对Hopper架构GPU优化的高效MLA解码内核,专为变长序列服务场景设计。 当前已发布特性: BF16支持 分页式kvcache(块大小为64) 快速入门 安装 python setup.py install 性能测试 python tests/test_flash_mla.py 在H800 SXM5(CUDA 12.6环境)上实现内存受限配置下3000GB/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架构GPU CUDA 12.3及以上版本 PyTorch 2.0及以上版本 致谢 FlashMLA的设计灵感来源于FlashAttention 2&3以及cutlass项目。