MoonshotがFlashKDAをオープンソース化—Kimi Delta Attention向けCUTLASSカーネル、H20でTriton基準比最大2.22倍

github.com/MoonshotAI/FlashKDA

今週は K2.6 を扱うさまざまなルーティング層がどう動くか比較していて、OpenRouter、Together、Orq を見ていました。その調査中に、K2.6 の活動と一緒に Moonshot が出した FlashKDA に出会いました。注目されずに見過ごされているようですが、モデルのリリースとは別に、カーネル作業そのものが本当に面白いのでここで共有します。

それが何か。Kimi Delta Attention の順伝播カーネルに対する、CUTLASS の C++ 実装です。これは Kimi Linear 論文の線形アテンションのバリアントになります。FLA の pull request #852 経由で、バックエンドとして flash-linear-attention に組み込まれるので、すでに KDA ベースのモデルで FLA を使っている人は、バックエンド層で FlashKDA にルーティングできます。

彼らの H20 ベンチマークの数値（FLA 既存の Triton 経路に対して計測）：

T=8192、H=96、D=128、固定長シーケンスで 1.72x。混在 seq_lens の可変長で 1.95x。1024x8 の一様な可変長で 2.22x。

なぜ重要か。KDA のような線形アテンションのアーキテクチャは、シーケンス長に対して線形にスケールすることを約束しますが、その約束が成り立つのは、カーネル実装が実際にハードウェア効率的である場合だけです。FLA の Triton 経路はリファレンスで機能していますが、ホッパーのメモリアクセスパターン向けにチューニングされた CUTLASS によって、理論的なコストモデルと実際に GPU で見える性能のギャップを埋められます。

要件は SM90 以上、CUDA 12.9 以上、PyTorch 2.4 以上です。MIT ライセンス。

率直に挙げるべき制約として、ベンチマークは順伝播のみで、すべて数値は H20 上です。H20 は中国向けのホッパー派生なので、H100 や Blackwell での絶対値は異なるはずです。相対的なスピードアップは方向性としては似ると思われますが、まだその数値を投稿した人はいません。

ここで誰かが H100 で試したか、あるいは逆伝播パスのカーネルがいつ出そうか、考えがあるか気になります。順伝播のみという話だと、現状の学習用途には制限があります。