deepseek-ai/FlashMLA

解決的問題

FlashMLA 提供高度優化的注意力算子，旨在加速大型語言模型的推理與預填充（prefilling）階段，特別是為 DeepSeek-V3 與 V3.2 模型提供支持。它透過為密集（dense）與稀疏（sparse）注意力模式提供專門的實現，解決了 Multi-head Latent Attention (MLA) 的計算瓶頸。

工作原理

該函式庫實現了針對不同 GPU 架構（SM90 與 SM100）及運行模式量身定制的多种注意力算子：

• 稀疏注意力算子 (Sparse Attention Kernels)：這些算子實現了 DeepSeek Sparse Attention (DSA)，允許模型在預填充與解碼階段僅針對特定 Token（透過 indices 張量）計算注意力。解碼算子在進行 bfloat16 計算的同時，利用 FP8 KV cache 來減少記憶體開銷。
• 密集注意力算子 (Dense Attention Kernels)：為預填充與解碼提供標準的密集注意力實現。
• 硬體優化：算子針對 NVIDIA H800 與 B200 GPU 進行了優化，利用特定的架構特性實現高 TFlops 性能。

適用對象

本專案面向使用 DeepSeek 模型的 AI 研究人員與工程師，以及需要在高端 NVIDIA GPU 上實現最大硬體利用率與推理速度的 Multi-head Latent Attention (MLA) 實作者。

亮點

• 高吞吐量：在 B200 GPU 上，密集 MLA 解碼可達到高達 660 TFlops，稀疏 MLA 預填充可達到 1450 TFlops。
• FP8 KV Cache 支援：透過使用帶有縮放因子的量化 FP8 格式，減少稀疏解碼期間的記憶體占用。
• Token 級稀疏性：支援選擇性 Token 注意力，以優化預填充與解碼階段的計算。
• 廣泛的硬體支援：針對 SM90 與 SM100 架構進行了優化，並為其他各種 GPU 加速器提供了社群移植版本。