Dao-AILab/flash-attention

Fast and memory-efficient exact attention

解決する課題

FlashAttentionは、TransformerモデルにおけるAttentionメカニズムを最適化し、より高速でメモリ効率の高いものにします。標準的なAttention演算に伴う高いメモリおよび計算オーバーヘッドに対処し、特にGPUメモリ間でのデータ転送回数を削減（IO-awareness）することでこれを実現しています。

仕組み

メモリへのアクセスを最小限に抑える、最適化されたスケーリングドット積Attention（scaled dot product attention）を実装しています。このプロジェクトでは、NVIDIA（Ampere, Ada, Hopper, Blackwell）やAMD（ROCm）など、さまざまなGPUアーキテクチャ向けに最適化された複数のバージョン（FlashAttention, FlashAttention-2, 3, 4）を提供しています。因果マスク（causal masking）、スライディングウィンドウ局所Attention、マルチクエリ/グループクエリAttention（MQA/GQA）などの高度な機能をサポートしています。

対象ユーザー

大規模言語モデル（LLMs）のトレーニングやデプロイを行う機械学習エンジニアや研究者向けに設計されており、AttentionフェーズにおけるGPUスループットの最大化とメモリ消費量の削減を必要とする方に最適です。

ハイライト

  • マルチGPUアーキテクチャのサポート: NVIDIA Hopper (H100) および Blackwell (B200) GPU、ならびにAMD ROCm互換GPU向けに最適化されたカーネル。
  • メモリ効率: IOオーバーヘッドを削減し、順伝播および逆伝播の両方を高速化。
  • 推論の最適化: KVキャッシュをサポートしたインクリメンタルデコード用の専用関数を搭載。
  • 柔軟なAttentionパターン: 因果マスク、ALiBiスロープ、スライディングウィンドウAttentionをサポート。

関連

  • プロジェクト
    flagos-ai/FlagAttention

    FlagAttentionは、Tritonベースのライブラリで、高速かつメモリ効率の高いAttentionカーネル（FlashAttention、区分的、ページング、スパース、量子化、再帰的など）を提供し、多くのアクセラレータでの大規模モデルの学習と推論をサポートします。安定した密行列演算子は本番環境で使用可能で、Alpha版ではKVキャッシュのページング、ブロック・スパース・デコーディング、INT8量子化などの実験的機能を提供します。

  • プロジェクト
    HKUSTDial/flash-sparse-attention

    Flash Attentionのメモリ効率とスパース計算を組み合わせた、極めて長いシーケンスを扱うTransformerモデル向けの高性能トレーナブルスパースアテンション実装。

  • プロジェクト
    lucidrains/x-transformers

    x-transformersは、設定可能なTransformerモデル（エンコーダー・デコーダー、デコーダーのみ、エンコーダーのみ、ビジョン）を構築するためのPyTorchライブラリであり、多くの研究グレードの機能（フラッシュアテンション、メモリトークン、代替正規化、スパース/top-kアテンション、トーキングヘッド、GLUバリアントなど）を備えています。研究者や実務者が新しいアイデアを迅速にプロトタイプ化できるよう、簡潔で拡張性の高い実装を提供することを目的としています。

  • Dispatch
    本番環境でのLLM最適化: 精度、注意、アーキテクチャ

    Hugging Faceは、効率的なLLMデプロイのための主要な技術を概説し、低精度量子化、メモリ効率のためのFlash Attention、およびRoPE、ALiBi、MQA、GQAなどのアーキテクチャ最適化に焦点を当てています。

  • プロジェクト
    vllm-project/vllm

    vLLM は、大規模言語モデルの高スループット・低レイテンシーなサービングのためのオープンソースライブラリです。効率的な KV キャッシュ処理のための PagedAttention を導入し、多くの量子化形式をサポートし、高速アテンションカーネルを統合し、GPU、CPU、その他のアクセラレータで動作する OpenAI 互換 API を提供します。