flagos-ai/FlagAttention

A collection of memory efficient attention operators implemented in the Triton language.

FlagAttention – 高速でメモリ効率の高いAttentionカーネル

概要 – FlagAttentionは、Triton（高性能GPUカーネルを記述するための言語）で記述されたカスタムAttention演算子のコレクションを提供するPythonライブラリです。これらのカーネルは、一般的なFlashAttentionアルゴリズムと同じ機能を実装し、さらにスパース、ページングキャッシュ、区分的、量子化、再帰的Attentionなどの多くの特殊なバリエーションを追加しています。これらは、さまざまなアクセラレータ（NVIDIA GPU、AMD、Ascendなど）上で大規模言語モデルやその他のTransformerスタイルのネットワークを学習および推論するために設計されています。

主要な演算子ファミリー

ファミリー	一般的なユースケース	APIエントリーポイント	勾配サポート	成熟度
FlashAttention	標準的な密行列Attention (scaled-dot-product)	`flag_attn.flash_attention`	順伝播 + 逆伝播	安定
Piecewise Attention	トークン距離に基づく2つのQ/Kペアの切り替え (例: NLPE)	`flag_attn.piecewise_attention`	順伝播 + 逆伝播	安定
Split-KV FlashAttention	長いKVデコーディング、低いクエリ並列度	`flag_attn.flash_attention_split_kv`	順伝播のみ	Alpha
Paged Attention	ページングされたKVキャッシュ上のシングルトークンクエリ (デコーダー専用モデルで一般的)	`flag_attn.paged_attention`	順伝播のみ	Alpha
MiniMax M3 Sparse Attention	128トークンページによるブロック・スパース・プリフィル/デコード	`flag_attn.minimax_m3_*`	推論のみ	Alpha
Chunk GLA / GDN2 / KDA	長いシーケンスのための再帰的ゲート線形またはデルタルールAttention	`flag_attn.chunk_*`	順伝播 (+ GLAは逆伝播も)	Alpha
SageAttention	ブロックごとのINT8量子化Q/K (オプションのマスク付き)	`flag_attn.sage_attention.forward`	順伝播のみ	Alpha
Parallel NSA	ネイティブ・スパースAttention + 圧縮	`flag_attn.parallel_nsa.parallel_nsa`	順伝播 + 逆伝播	Alpha

インストール (クイックガイド)

# Clone the repo
git clone https://github.com/flagos-ai/FlagAttention.git
cd FlagAttention

# Install the editable package with Triton and test extras
pip install -e "[triton,test]"

アクセラレータに対応したPyTorchビルド (例: torch==2.x+cu118) が既にインストールされている必要があります。 正しいビルドはGPUベンダーによって異なるため、このパッケージは特定のPyTorchホイールに依存しません。

最小限の使用例 – FlashAttention

import torch
from flag_attn import flash_attention

B, Hq, Hkv, M, N, D = 2, 16, 4, 2048, 4096, 128
q = torch.randn(B, Hq, M, D, device='cuda', dtype=torch.float16, requires_grad=True)
k = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)
v = torch.randn(B, Hkv, N, D, device='cuda', dtype=torch.float16, requires_grad=True)

out = flash_attention(q, k, v, causal=True)
out.sum().backward()   # gradient flow works

関数シグネチャはオリジナルのFlashAttention APIを反映しており、ドロップアウト、ログ正規化項の返却、合計Attention質量、Philox RNG状態のためのオプションフラグが追加されています。

演算子の使い分け

シナリオ	推奨演算子
単一GPUでの標準的な密行列学習または推論	FlashAttention (安定)
距離閾値後にAttention行列を切り替える必要がある場合 (例: 位置埋め込みのトリック)	Piecewise Attention (安定)
クエリが少ない非常に長いKVキャッシュ (デコーディング)	Split-KV FlashAttention (alpha)
ページングされたKVキャッシュを保持するデコーダー専用モデル (例: LLMサービング)	Paged Attention (alpha)
vLLMスタイルのトークン・スパース・デコーディングなどのブロック・スパースモデル	MiniMax M3 (alpha)
ストリーミングまたは非常に長いシーケンスのための再帰的線形Attention	Chunk GLA / GDN2 / KDA (alpha)
メモリ帯域幅を節約するためにブロックごとのINT8量子化を行いたい場合	SageAttention (alpha)
オプションの圧縮を伴うネイティブ・スパースAttentionが必要な場合	Parallel NSA (alpha)

成熟度とテスト

リポジトリは conf/operators.yaml で演算子の成熟度を追跡しています。FlashAttention と Piecewise Attention のみが安定とマークされており、その他はすべて alpha であり、APIシグネチャが変更される可能性があります。
完全なpytestスイート (pytest tests) がFP32 PyTorchリファレンスに対して数値精度を検証します。精度許容誤差はFP32誤差の約2倍で、小さな絶対マージンが設定されています。
ベンチマークは benchmark/ に提供されており、ヘルパースクリプト tools/run_tests.py または直接実行 (例: python benchmark/flash_benchmark.py) で実行できます。結果はFlagGemsプラットフォームと互換性のあるJSONとして出力されます。

サポートされているハードウェアとバックエンド

プライマリ: NVIDIA CUDA GPU (Triton経由)。このライブラリには、AMD、Ascend、Enflame、MetaX、MThreads、Intelなどのベンダー固有のバックエンドも含まれています。バックエンドの選択は検出されたPyTorchデバイスに基づいて自動的に行われますが、すべての演算子がすべてのベンダーで実装されているわけではありません。詳細は conf/operators.yaml および src/flag_attn/runtime/backend ディレクトリを参照してください。
CPU: パッケージをインポートしてメタデータを検査することはできますが、GPU互換のランタイムなしではカーネルを実行できません。

詳細情報

ドキュメント – README自体に完全なAPIリファレンス、インストールノート、および examples/ 下のサンプルスクリプトが含まれています。
テストとベンチマーク – 具体的な使用パターンについては tests/ を、パフォーマンス数値については benchmark/ を参照してください。
ベンダー固有の拡張 – 低レベルの実装については src/flag_attn/runtime/backend/_<vendor> を参照してください。

まとめ

FlagAttentionは、研究者やエンジニアに幅広いTransformerワークロードのための高速でメモリを節約するAttentionカーネルを提供する、実用的な本番グレードのライブラリです。その安定したFlashAttentionおよびPiecewiseAttention APIは、学習ループ内のデフォルトのPyTorch Attentionを置き換えることができ、Alpha演算子は高度なLLMサービングや研究のために、ページングキャッシュ、ブロック・スパース・デコーディング、量子化カーネル、再帰的線形Attentionなどの実験的機能を提供します。