FlashQLA: CP-/Bwd フレンドリーな融合線形注意カーネル（GDN 用）
Qwen は FlashQLA をリリースしました。これは Gated Delta Network (GDN) レイヤーの最適化を目的とした高性能線形注意カーネルライブラリです。TileLang 上に構築されており、NVIDIA Hopper GPU 上の FLA Triton カーネルに比べて前方で 2〜3 倍、後方で 2 倍の速度向上を実現し、特に事前学習やエッジ側エージェント推論で効果を発揮します。
GDN チャンク化プリフィルの最適化
FlashQLA は、元の FLA 実装で見られる Gated Delta Network (GDN) のチャンク化プリフィルプロセスにおける 2 つの主要な効率ボトルネックに対処します。
- メモリバウンドカーネル: 標準フローは中間変数 ($W, U, S$) を高帯域幅メモリ (HBM) に何度も読み書きし、かなりのオーバーヘッドを生み出します。
- GPU 利用率の低さ: State Space Model (SSM) の状態が再帰的であるため、同時に実行できるスレッドブロック数が
batch_size * num_headsに制限されます。モデルが小さい、バッチが小さい、または Tensor Parallelism (TP) を使用しているシナリオでは、GPU の Streaming Multiprocessor (SM) がアイドル状態になります。
これらの相反する課題を解決するため、FlashQLA は小バッチシナリオで失敗する完全融合カーネルを回避し、代わりに前方計算を 2 つの融合カーネルに分割し、その間に Context Parallelism (CP) 前処理ステップを挿入します。
主な技術的イノベーション
ゲート駆動自動カード内コンテキスト並列化 (AutoCP)
FlashQLA は、TP、長シーケンス、ヘッド数が少ない設定下で SM の利用率を向上させる自動カード内 CP メカニズムを実装しています。最適な並列度 ($L = \lambda \sqrt{N}$) を決定するために数式モデルを使用し、ここで $N$ はチャンク数、$L$ は CP ランクあたりのチャンク数です。
さらにオーバーヘッドを削減するため、FlashQLA は GDN ゲートの指数減衰特性を利用します。ゲート $\alpha_i \in (0,1)$ を持つヘッドでは、過去状態の影響が指数的に減衰します。FlashQLA は「ウォームアップ」プロセス（通常 6〜8 チャンク）を用いて状態誤差をノイズフロア以下に下げ、コストの高い補正項 $M$ 行列計算を省略し、正確なサブシーケンス $S_0$ を直接取得します。
TileLang ワープ特化カーネル
FlashQLA は TileLang を利用してワープグループ特化カーネルを実装します。このアーキテクチャは同一 SM 内に 1 つのプロデューサー・ワープグループと 3 つのコンシューマー・ワープグループを配置し、共有メモリでデータをやり取りし、mbarriers で同期します。
- 前方パス: 3 つのコンシューマー・ワープグループがそれぞれ $V'$, $S$, $O$ を計算し、ピンポン構造で計算とメモリトラフィックをオーバーラップさせます。
- CP 前処理: 単一の融合カーネルが元の $M$ と $S$ の計算、および軽量なスライディングウィンドウ・ウォームアップ手法の両方を処理します。
- 後方パス: FlashQLA は
bwd_dv,bwd_dhu,bwd_dqkwg,bwd_wyを単一カーネルに融合します。オンチップリソースの制約により、メモリトラフィックを隠すために長い計算チェーンに依存し、マルチステージパイプラインは使用しません。
パフォーマンスベンチマーク
NVIDIA H200 GPU 上で FLA Triton および FlashInfer ベースラインと比較したベンチマーク結果は、特に Tensor Parallelism (TP) の度合いが高まるほど大幅な性能向上を示しています。
| モデル / TP | シーケンス長 | $h_{qk}$ | $h_v$ | FlashQLA | FlashInfer | FLA vs FLA vs FI |
|---|---|---|---|---|---|---|
| 397B/122B TP8 | 1x32768 | 2 | 8 | 0.310ms | 1.653ms | 2.95× |
| 397B/122B TP4 | 1x32768 | 4 | 16 | 0.486ms | 1.654ms | 2.57× |
| 27B TP2 | 1x32768 | 8 | 24 | 0.659ms | 1.616ms | 2.37× |
| 2B/0.8B TP1 | 1x32768 | 16 | 16 | 0.493ms | 1.640ms | 2.60× |
実装と要件
FlashQLA は、FLA のシグネチャと互換性のある高レベル API と、前方・後方パス用の低レベルエントリポイントを提供します。
システム要件:
- ハードウェア: NVIDIA SM90 (Hopper)
- ソフトウェア: CUDA 12.8+, PyTorch 2.8+
コードとベンチマークは github.com/QwenLM/FlashQLA で入手可能です。