FlashMLA：HopperGPU向けMLAデコード・カーネルの最適化（DeepSeek Open Source Week 1日目）

はじめに

FlashMLAは、DeepSeek AIによって開発された効率的なMLA（Multi-head Latent Attention）デコーディングカーネルで、NVIDIA Hopper Architecture GPU向けに最適化されており、可変長シーケンス処理のパフォーマンスを向上させるように設計されています。このプロジェクトはGitHubでオープンソース化されており、開発者は無料で利用できる。BF16精度計算とページドKVキャッシング（ブロックサイズ64）をサポートし、H800 SXM5で優れた性能を発揮し、メモリ集約型構成では最大3000GB/秒の帯域幅、計算集約型構成では最大580TFLOPSを実現する。 FlashMLAはFlashAttention 2&3とCutlassプロジェクトに触発された。DeepSeek AIは、このオープンソースプロジェクトを通じてAI技術の革新性を示し、多くの注目を集めている。

機能一覧

• 効率的なMLAデコーディングHopper GPU用に最適化され、可変長シーケンスの処理を大幅に高速化。
• BF16の精度をサポート半精度浮動小数点演算を利用することで、精度を維持しながら計算効率を向上。
• ページング KV キャッシングブロックサイズ64のページングメカニズムにより、メモリを効率的に管理し、推論パフォーマンスを向上させる。
• 高性能H800 GPUで最大3000GB/秒のメモリ帯域幅と580TFLOPSの演算能力を提供します。
• オープンソース開発者によるカスタマイズや統合をサポートするため、完全なソースコードが提供されます。

ヘルプの使用

設置プロセス

FlashMLAはGitHubをベースとしたオープンソースプロジェクトで、使用する前に環境が要件を満たしていることを確認し、インストールを完了する必要があります。詳しい手順は以下の通りです：

1.環境への備え

• オペレーティングシステムLinux システムのサポート (Ubuntu 20.04 以上を推奨)。
• ハードウェア要件NVIDIA Hopper Architecture GPU（H800 SXM5など）が必要です。
• ソフトウェア依存::
  • CUDA 12.6以上（インストール方法はNVIDIAのウェブサイトを参照）。
  • PyTorch 2.0以上（経由推奨 pip install torch (インストール）。
  • Python 3.8以上。
• 検査ツールGitHubからコードをダウンロードするために、Gitがインストールされていることを確認する。

2.ソースコードのダウンロード

1. ターミナルを開き、以下のコマンドを入力してFlashMLAリポジトリをクローンします：

   git clone https://github.com/deepseek-ai/FlashMLA.git

2. プロジェクト・カタログにアクセスする：

   cd FlashMLA
   

3.依存関係のインストール

プロジェクトはPyTorchとCUDAに依存しており、以下のコマンドでインストールできる：

pip install -r requirements.txt

そうでなければ 要件.txt ファイルで、PyTorchがインストールされていることを確認するのは簡単です：

pip install torch torchvision

CUDAが利用可能であることを確認する：

python -c "import torch; print(torch.cuda.is_available())"

輸出 真 環境設定に成功したことを示す。

4.コンパイルとテスト

FlashMLAにはコンパイル済みのCUDAプラグインが用意されていますが、ローカルのCUDAバージョンに合わせてください：

1. ソース・ディレクトリに移動し、コンパイル・スクリプトを実行する（もしあれば）：

   python setup.py インストール
   

2. インストールが成功するかテストし、サンプルコードを実行してください：

   python example.py
   

エラーが報告されなければ、インストールは完了です。

使用方法

FlashMLAの核となる機能は、AIモデル推論タスクのための効率的なMLA解読サポートを提供することである。以下はその手順である：

機能1：FlashMLAのロードと実行

1. インポートモジュール::
   PythonスクリプトでのFlashMLAコア関数の紹介：

   from flash_mla import get_mla_metadata, flash_mla_with_kvcache
   

2. データ入力の準備::
   • キャッシュ・セキュレンスKVキャッシュのシーケンス長を定義する。
   • q_iクエリテンソル。
   • kvcache_iKVのキャッシュデータ。
   • ブロックテーブルページングキャッシュのブロックテーブル。
3. メタデータの取得::

   tile_scheduler_metadata, num_splits = get_mla_metadata(cache_seqlens, s_q * h_q // h_kv, h_kv)
   

4. ランニングデコード::

   o_i, lse_i = flash_mla_with_kvcache(q_i, kvcache_i, block_table, cache_seqlens, dv, tile_scheduler_metadata, num_splits, causal=True)
   

   輸出 o_i はデコード結果である。lse_i は対数和の値である。

機能2：可変長シーケンス処理の最適化

• 取る動的な長さの入力シーケンスを扱う場合、FlashMLAはKVキャッシュをページングすることでメモリフットプリントを削減する。
• リグ::
  1. ページング・パラメーターの設定：ブロック・サイズは64で固定されており、以下のように調整できる。 キャッシュ・セキュレンス シーケンスの長さを制御する。
  2. 実行時に指定する causal=真(c) 因果的注意メカニズムが有効であることを確認する。
• 効果H800のメモリ帯域幅は3000GB/sで、大規模な推論タスクに対応。

機能3：パフォーマンス・テスト

• 試験方法::
  1. サンプルスクリプトを編集する（例 example.py)、入力データのサイズが大きくなる。
  2. パフォーマンスを記録するには、以下のコードを使用する：

     インポート時間
     start = time.time()
     # デコードコードの実行
     o_i, lse_i = flash_mla_with_kvcache(...)
     print(f "時間: {time.time() - start} 秒")
     

• 期待される結果メモリ負荷の高いタスクでは約3000GB/s、演算負荷の高いタスクでは約580TFLOPS。

ほら

• ハードウェア互換性Hopper GPUのみがサポートされ、H800または同等のGPUを推奨します。
• デバッグのヒントCUDAのエラーが発生した場合は、バージョンの一致を確認するか、GitHub Issuesでコミュニティのサポートを求めてください。
• 生産環境既存のモデル推論プロセスに直接統合し、入力データ形式がFlashMLA要件と整合していることを保証します。

以上の手順で、ユーザーはすぐにFlashMLAを使い始めることができ、その効率的なデコードによってもたらされるパフォーマンスの向上を楽しむことができる。完全なコードとドキュメントはGitHubリポジトリで見ることができます。