ある開発者が、二次関数的なMHAのボトルネックを回避するために設計された、融合Triton/CUDA状態蓄積カーネルを特徴とする減衰ゲート型O(N)因果線形アテンションアーキテクチャをオープンソースとして公開しました。

  • コンテキストスケーリング: NVIDIA T4上で<185MBのVRAMを使用して128Kのコンテキストシーケンスを108.78msで実行。
  • カーネル加速: 融合Tritonカーネルにより、標準的なPyTorch実装と比較して3.60倍の高速化を実現。
  • 精度保護: 長期コンテキスト学習中のFP16/BF16勾配オーバーフローを防ぐために、FP32マスターアキュムレータとRMS正規化を使用。
  • 検索精度: 1.89xのSNRで16K Needle-In-A-Haystack (NIAH)ベンチマークをパス。

著者は、このモデルをHugging Face Transformersに統合するか、既存の線形モデルとのベンチマーク比較を行うことについてフィードバックを求めています。