inclusionAI/cuLA

解決する課題

cuLAは、長文コンテキストのLLMワークロードにおける線形アテンションメカニズムのパフォーマンスのボトルネックに対処します。線形アテンションは、標準的なアテンションの二次的な複雑さを線形時間の更新に削減しますが、これらを最新のNVIDIAアーキテクチャの計算能力を最大限に活用するために効率的に実装するには、深いハードウェアレベルの最適化が必要です。

仕組み

cuLAは、CuTe DSLとCUTLASS C++を使用して記述された、手動でチューニングされたCUDAカーネルを提供します。これには、GLA、KDA、GDN、Lightning Attentionを含む、線形アテンションの高性能なバリエーションが含まれます。これらのカーネルは、NVIDIA Blackwell (SM10X) および Hopper (SM90) GPU向けに特別に最適化されており、WGMMAデータフローやpersistent packed schedulingなどの機能を利用して、prefillおよびdecodeフェーズの両方でスループットを最大化します。

対象者

NVIDIA HopperまたはBlackwellハードウェア上で、長文コンテキストLLMを扱う際に最大限の推論およびトレーニングパフォーマンスを必要とするAI研究者やエンジニア向けに設計されています。また、flash-linear-attention (FLA) ライブラリの高性能バックエンドとしても意図されています。

ハイライト

• ハードウェア最適化: NVIDIA Blackwell (SM10X) および Hopper (SM90) アーキテクチャ向けに特別にチューニングされています。
• 大幅な高速化: FLAのTriton実装と比較して大幅なパフォーマンス向上を実現しており、HopperでのFlashKDA prefillにおいて最大7.56倍の高速化を実証しています。
• 幅広いアルゴリズムサポート: KDA (Kimi Delta Attention) や Lightning Attention など、複数の線形アテンションのバリエーションを実装しています。
• シームレスな統合: 単純なインポートの変更だけで、FLAカーネルのドロップイン・リプレイスメントとして使用できるように設計されています。