deepseek-ai/FlashMLA

解決する課題

FlashMLA は、大規模言語モデル（LLM）の推論およびプリフィル段階を加速するために設計された、高度に最適化されたアテンションカーネルを提供し、特に DeepSeek-V3 および V3.2 モデルを強力にサポートします。密な（dense）およびスパースな（sparse）アテンションパターンの両方に対して特化した実装を提供することで、Multi-head Latent Attention (MLA) の計算上のボトルネックを解消します。

仕組み

このライブラリは、異なる GPU アーキテクチャ（SM90 および SM100）と動作モードに合わせて調整された、さまざまなアテンションカーネルを実装しています：

• Sparse Attention Kernels: これらは DeepSeek Sparse Attention (DSA) を実装しており、プリフィルおよびデコーディング段階の両方において、特定のトークン（indices テンソル経由）に対してのみアテンションを計算することを可能にします。デコーディングカーネルは、bfloat16 で計算を行いながら、メモリオーバーヘッドを削減するために FP8 KV キャッシュを利用します。
• 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 キャッシュのサポート: スケール係数を用いた量子化 FP8 フォーマットを使用することで、スパースデコーディング中のメモリフットプリントを削減。
• トークンレベルのスパース性: プリフィルおよびデコーディング段階の計算を最適化するために、選択的なトークンアテンションをサポート。
• 幅広いハードウェアサポート: SM90 および SM100 アーキテクチャ向けに最適化されており、その他の様々な GPU アクセラレータ向けのコミュニティポートも提供されています。