vipshop/ffpa-attn

解決する問題

FFPA（Fast and Memory-Efficient Exact Attention）は、大きなヘッド次元（D > 256、最大 1024）を扱う際の標準的なアテンションメカニズムのパフォーマンスボトルネックに対処します。標準的な実装は、これらのスケールでメモリとレジスタのプレッシャーに悩まされ、大幅な速度低下を引き起こすことがよくあります。FFPA は、PyTorch のネイティブな Scaled Dot Product Attention（SDPA）と比較して大幅な高速化を実現する高度に最適化されたカーネルを提供します。

仕組み

FFPA は、アテンションを最適化するために 2 つの主要な技術を使用します。

• Split-D：このタイリング戦略は、ヘッド次元（D）軸を分割して SRAM 使用量を一定に保ち、メモリボトルネックを効果的に解決します。
• TiledMMA：このアプローチは、行列乗算の N 次元を分割してレジスタ使用量を最適化し、レジスタスピルを防ぎ、従来の実装よりもはるかに大きなヘッド次元をカーネルが処理できるようにします。

このライブラリは、CUDA、Triton、CuTe-DSL を含む複数のバックエンドをサポートし、1 行のコードで既存の PyTorch ワークフローに統合できます。

対象読者

このプロジェクトは、大規模モデルのトレーニングと推論に取り組む AI 研究者とエンジニアを対象としており、特に大きなヘッド次元を持つモデルや、高スループットの FP8/FP4 精度を必要とするモデルのアテンション層を最適化する必要がある人々に適しています。

ハイライト

• 高性能：標準の PyTorch SDPA と比較して 1.5 倍から 15 倍の高速化を実現。
• 大きなヘッド次元のサポート：精度を損なうことなく、ヘッド次元を最大 1024 までサポート。
• 量子化サポート：より高速な推論のための FP8 および FP4 アテンションカーネルの実験的サポートを含む。
• 柔軟な統合：torch.nn.functional.scaled_dot_product_attention をモンキーパッチするためのシンプルな API を提供し、即時のパフォーマンス向上を実現。
• 幅広い互換性：自己アテンション、GQA/MQA、クロスアテンション、因果/マスクアテンションなど、さまざまなアテンションタイプをサポート。