MetaEra、FlashKDAをオープンソース化し、Kimiの線形推論速度を1.7〜2.2倍に向上

MEニュース：4月22日（UTC+8）、動察Beatingの監視によると、月之暗面はGitHub上でFlashKDAをオープンソース公開しました。これはNVIDIA HopperシリーズGPU（H100、H20など）向けにモデル推論を加速するためのツールで、MITライセンスです。このツールは、月之暗面が昨年Kimi Linear論文で提案した新しいアテンションメカニズム「KDA」を対象としています。従来のアテンションメカニズムでは、大規模モデルが長文を処理する際に計算量が長さの二乗に比例して増加しますが、線形アテンションはこのコストを線形増加に抑えます。KDAはこの方向性における改良版の一つです。Kimi Linearモデルの構造は、3層のKDAと1層の従来アテンションを交互に配置しています。以前から、Triton言語で書かれたKDAの実装がオープンソースライブラリ「flash-linear-attention」（略称：fla）に存在していました。FlashKDAは、これをNVIDIAの低レベルGPUライブラリCUTLASSを用いて再実装し、Hopper GPUの性能を最大限に引き出すように最適化されています。公式テストではH20で同様のフォワード計算を実行したところ、FlashKDAはTriton版より1.7～2.2倍高速でした。特に入力長が不均一でバッチ処理を行うシナリオでの加速効果が顕著です。ただし、公式は自社のTriton版とのみ比較しており、他の線形アテンション手法とは比較していません。今回のオープンソース化ではフォワード計算のみが公開されており、つまり「モデルの推論」のみが可能で、「モデルの学習」はまだ従来のTriton版を使用する必要があります。使用要件：Hopper以降のGPU（SM90アーキテクチャ以上）、CUDA 12.9以上、PyTorch 2.4以上。FlashKDAは新バックエンドとしてflaのアップストリームにマージされました（PR #852）。既存ユーザーは設定を1行変更するだけで切り替え可能です。（出典：BlockBeats）