Sol-Attn、online softmax内で注意を動的にスパース化し、動画生成をエンドツーエンドで2倍超高速化

NVIDIAの研究チームは、高解像度動画の生成時に拡散Transformerで生じる注意コストの削減を目指すSol-Attnを公開した。従来の動的スパース注意では、通常、先にproxy score mapを作成し、固定比率または累積確率に基づいてkey-valueブロックを選択する。前者はqueryの難易度に応じて計算量を調整できず、後者はワークロードの不均衡を招く可能性がある。さらに、どちらにもルーティングと中間テンソルのコストが伴う。

Sol-Attnは、ルーティング、スパース計算、誤差補正を1回のonline softmaxに統合する。中核となる手法では、完全なスコアマップを事前に実体化せず、チップ上でブロック単位にproxy scoreとしきい値を比較する。しきい値を下回るブロックを単純にゼロにするのではなく、そのスコア列を再利用して、ロングテールの注意寄与を近似する。これにより、queryごとに異なるブロック予算を使用できると同時に、従来のkeep-or-drop方式でスパース率を高めた際に映像品質が急速に損なわれる問題を回避する。

著者らは、同一のpromptとseedを用いて密な注意パイプラインと比較した。720pのtext-to-videoでは、Wan 2.1が2.02倍、Hunyuan 1.5が2.12倍のエンドツーエンド高速化を達成した。LTX 2.3による1080pの2段階生成では1.90倍、Berniniの動画編集では2.34倍となった。さらに、カーネル融合とキャッシュを組み合わせてSol-Engineを構成した場合、公式報告では最大約5倍に達する。ただし、この数値をSol-Attn単独の技術によるものとみなすことはできない。

コードはSanaリポジトリの`sol-engine`ブランチで公開されており、エンジニアリングチームはカーネルとパイプラインの統合を確認できる。ただし、現時点のデモは主にNVIDIAによる自己評価と視覚的な事例に基づいており、複数のGPU、異なる動画長、高ダイナミックなシーンを対象とするブラインドテストはまだ不足している。B200向けカーネルも引き続き最適化中だ。今後は、ピークメモリ使用量、バッチスループット、しきい値が低い場合の実行時間の一貫性、ならびに主要な推論フレームワークへ統合できるかどうかを注視する必要がある。