deepseek-ai/FlashMLA

FlashMLA – DeepSeekモデル向け高性能アテンションカーネル

概要 – FlashMLAは、DeepSeekの大規模言語モデル（最新の DeepSeek-V4.1-Flash モデルなど）で使用される高度に最適化されたアテンション演算子を実装したC++/CUDA（およびHuawei Ascend）ライブラリです。主な機能は以下の通りです。

• スパースアテンションカーネル（トークンレベルの「DeepSeek Sparse Attention」）：prefillフェーズ（モデルが初期プロンプトを処理する際）とdecodingフェーズ（一度に1トークンずつ生成する際）の両方に対応。これらのカーネルは FP8またはFP4量子化キー・バリューキャッシュ を使用し、精度を維持しながらメモリ帯域幅を劇的に削減します。
• デンスマルチヘッドアテンション (MHA) カーネル：標準的な（スパースではない）prefillおよびバックワードパス用。
• フュージョンカーネル：Q-norm、RoPE（回転位置エンコーディング）、コアアテンション計算、出力側の共役RoPE、そして最終的なFP8へのキャストを単一の起動で実行します。これにより、多数の小さなカーネルによるオーバーヘッドを排除します。
• NVIDIA GPU (SM100/SM103, CUDA 13.1+) および Huawei Ascend 950 NPU (CANN 9.2+ および torch_npu 経由) をサポート。

重要性 – 大規模言語モデルの推論時間の大部分はアテンション演算に費やされます。FlashMLAは低精度 (FP8/FP4) でアテンションを書き直し、スパース性を活用することで、最新ハードウェア上で >1,000 TFlops を達成し、デバイスの理論上のピーク性能に迫ります。これにより、DeepSeek-V4.1のようなモデルの応答時間が短縮され、サービングコストが削減されます。

使用方法

1. インストール – リポジトリをクローンし、サブモジュールをプルして pip install -v . --no-build-isolation を実行します。ビルドスクリプトはCUDA環境かAscend環境かを自動検出し、適切な拡張機能をコンパイルします。
2. Prefill (スパース) – クエリテンソル、KVキャッシュ、および各クエリトークンがどのKVトークンを参照すべきかをカーネルに伝える3-D indices テンソルを使用して flash_mla_sparse_fwd(q, kv, indices, …) を呼び出します。
3. Decoding (スパース) – まず get_mla_metadata() でスケジューラのメタデータを取得し、生成ステップごとに flash_mla_with_kvcache(...) を繰り返し呼び出します。この関数はREADMEに記載されているFP8/FP4 KVキャッシュレイアウトを想定しています。
4. デンスMHA – 一般的な flash_attn パッケージと互換性のある flash_attn_varlen_func / flash_attn_varlen_qkvpacked_func APIを使用しますが、追加の head_dim_qk 引数が必要であること、および戻り値に出力とlog-sum-exp (lse) が含まれることに注意してください。
5. フュージョンカーネル – QプロジェクションとWvプロジェクションの重みをフュージョンカーネルが必要とするレイアウトに並べ替えた後（ヘルパーは fused_norm_rope_attn_rope_cast.permute_* にあります）、prefill(...) または decode(...) を呼び出して、後続のFP8-einsum用のFP8出力を取得します。

パフォーマンスのハイライト（リポジトリの報告に基づく）

• スパースprefill：NVIDIA B200 GPU (CUDA 13.3) で最大 1350 TFlops。Ascend 950 NPU で 410 TFlops (ハードウェアピークの約95%)。
• スパースdecoding：B200で最大 1024 TFlops、Ascend 950で 360 TFlops (ピークの約83%)。
• フュージョンnorm-RoPE-attn-RoPE-cast：B200で最大 1460 TFlops (prefill) および 950 TFlops (decoding) で、非フュージョン版より10〜15%向上。
• デンスMHA：B200でフォワード最大 1460 TFlops、バックワード最大 1000 TFlops。

技術的な詳細

• KVキャッシュフォーマット – FP8 (float8_e4m3) または FP4 (e2m1) 量化キャッシュのみサポート。レイアウトは生の量子化値とトークンごとのスケールファクターを格納し、オンザフライでの高速な逆量子化を可能にします。
• Indicesテンソル – 各クエリトークンがどのKV「ページ」を参照すべきかをエンコードする3-D整数テンソル。無効なエントリは -1 とマークされ、カーネルによって無視されます。
• レジスタスピル安全性 – CUDA拡張のビルド後、ビルドシステムはレジスタスピルをチェックします。FLASH_MLA_SKIP_REG_SPILL_CHECK=1 を設定しない限り、ビルドは失敗します。
• 互換性に関する注意 – 2026年9月のリリースでは、NVIDIA Hopper GPUおよび古いDeepSeekモデル (V3/V3.2/V4.0) のサポートが終了し、KVキャッシュフォーマットが変更されました。古い構成が必要なユーザーは、READMEにリンクされているレガシーコミットをチェックアウトする必要があります。

詳細情報 – リポジトリには、アルゴリズムと低レベルの最適化を説明する詳細なmarkdownレポート (docs/…) や、ハードウェアのパフォーマンスを検証するためのテストスクリプト (tests/*.py) が含まれています。

---

FlashMLAは、最先端のNVIDIA GPUとHuawei Ascend NPUの両方をターゲットとした、DeepSeekのLLM推論パイプラインにおけるパフォーマンスクリティカルなビルディングブロックです。