FlashMLA
GitHub - deepseek-ai/FlashMLA
FlashMLA 是適用于 Hopper GPU 的高效 MLA 解碼內核，針對可變長度序列服務進行了優化。
當前發布：
- BF16、FP16
- 塊大小為 64 的分頁 kvcache
快速開始
安裝
python setup.py install
基準
python tests/test_flash_mla.py
使用 CUDA 12.8，在 H800 SXM5 上實現高達 3000 GB/s 的內存綁定配置和 580 TFLOPS 的計算綁定配置。
用法
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,)...
要求
- Hopper GPU
- CUDA 12.3 及更高版本
- 但我們強烈建議使用 12.8 或更高版本以獲得最佳性能
- PyTorch 2.0 及更高版本
確認
FlashMLA 的靈感來自?FlashAttention 2&3?和彎刀項目。
引文
@misc{flashmla2025,title={FlashMLA: Efficient MLA decoding kernels},author={Jiashi Li},year={2025},publisher = {GitHub},howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}
