inclusionAI/cuLA

CUDA kernels for linear attention variants, written in CuTe DSL and CUTLASS C++.

解決する課題

cuLAは、線形アテンションバリアント向けに高性能なCUDAカーネルを提供します。これは、LLMにおける標準的な二次アテンションを置き換え、長文脈ワークロードをより効率的に処理するために設計されています。最新のNVIDIA GPU(BlackwellおよびHopper)をターゲットにしており、ハードウェアアクセラレーションを最大限に引き出します。

仕組み

このライブラリはCuTe DSLとCUTLASS C++を使用して記述されており、GLA、KDA (Kimi Delta Attention)、GDN、Lightning Attentionといった線形アテンションアルゴリズムの手動チューニング実装を提供します。flash-linear-attention (FLA) ライブラリのドロップインリプレースメントとして設計されており、同じインターフェースを共有しているため、ユーザーはcuLAカーネルに切り替えることでパフォーマンスを向上させることができます。

対象ユーザー

NVIDIA Hopper (SM90) および Blackwell (SM10X) GPU上で最大限のスループットと効率を必要とする、長文脈LLMを扱うAI研究者およびエンジニア。

特徴

  • ハードウェア最適化: NVIDIA Blackwell (SM10X) および Hopper (SM90) アーキテクチャ向けに特別にチューニングされています。
  • 大幅な高速化: FLAベースラインと比較して大幅なパフォーマンス向上を実現しており、Hopper上でのFlashKDAプリフィルでは最大7.56倍の高速化を達成しています。
  • 幅広いアルゴリズムサポート: Lightning AttentionやKDAなどのバリアントを実装し、プリフィルおよびデコードの両フェーズをサポートしています。
  • FLA互換: flash-linear-attention エコシステムとシームレスに統合できるように設計されています。

関連

  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト