技術要約: DeltaLog

問題点: 線形アテンションのデコーディングにおける状態更新税 (State-Update Tax)

線形アテンションモデル（GDN、KDA、RWKV6など）は、二次的なプレフィックス計算やコンテキストが増大するKVキャッシュを、固定サイズの再帰的な状態更新に置き換えることで、ソフトマックス・アテンションに代わる魅力的な選択肢を提供します。これらのモデルにおいて、状態 $S_t$ はサイズ $d_k \times d_v$ の密な行列であり、シーケンス長に依存しません。

しかし、既存のデコーディング実装は、「状態更新税」と呼ばれる重大な非効率性に悩まされています。単一のトークンによって寄与される論理的な更新はコンパクト（サイズ $O(d_k + d_v)$ のベクトルで表される）であるにもかかわらず、標準的な「逐次的（eager）」な実装では、この更新を即座にフルサイズの密な状態へと実体化し、毎トークンの生成後に $[B, H, d_k, d_v]$ テンソル全体をHBM（高帯域幅メモリ）へ書き戻してしまいます。

• メモリトラフィックの不一致: バッチサイズ $B$ とヘッド数 $H$ に対して、実際に加えられる情報は $O(BH(d_k + d_v))$ ですが、物理的なトラフィックは $O(BH d_k d_v)$ 個の要素の読み書きを伴います。
• メモリ帯域のボトルネック: この逐次的な実体化は、算術強度を約0.5–1.0 FLOPs/byteに低下させます。これは、現代のGPU（例：H200の~208 FLOPs/byte）のマシンバランスを大幅に下回っています。その結果、再帰的な状態の維持がメモリトラフィックの主要な原因となり、特にバッチサイズや状態サイズが増大するにつれて、デコード・レイテンシを支配することになります。測定によれば、大規模バッチのシナリオでは、これがデコード・レイテンシの最大44.6%を消費することがあります。

手法: DeltaLog による遅延実体化 (Deferred Materialization)

本論文は、論理的な再帰と状態の物理的な実体化を切り離すデコーディング手法である DeltaLog を提案しています。DeltaLogは、モデルの重みや数学的な再帰を変更するものではなく、状態更新の物理的なスケジュールのみを変更します。

コアメカニズム: Base + Log

DeltaLogは、ライブな再帰状態を以下の2つのコンポーネントとして表現します：

1. 密なベース状態 ($\bar{S}$): 最後の「マージ」ポイントまでの累積履歴を格納する密な行列。
2. 限定的な更新ログ (Bounded Update Log): 最近のコンパクトな更新因子 ($a_t, b_t$) と、それに関連する対数減衰メタデータ ($\Lambda$) を格納するバッファ。

デコーディング・サイクルは2つのフェーズで動作します：

• アペンド・フェーズ (Append Phase - 通常時): $M-1$ ステップの間、システムはフルサイズの密な状態を書き戻しません。代わりに、コンパクトな更新因子と現在の対数減衰スナップショットをログ・バッファに追加（append）します。密なベース $\bar{S}$ は変更されません。
• マージ・フェーズ (Merge Phase - 定期的): $M$ ステップごとに、システムは累積された論理状態（ベース ＋ 再再生されたログのエントリ）を密なベース $\bar{S}$ に実体化し、ログをリセットして、遅延された更新をクリアします。

トラフィックモデル

密な書き戻しを遅延させることで、DeltaLogは頻繁なフル状態の書き込みを以下に置き換えます：

• コンパクトな因子 ($O(d_k + d_v)$) の軽量なアペンド。
• $M$ ステップに一度行われる、償却された密な書き込み。
• マージ・ステップにおけるログ・エントリの再生（replay）。

最適なマージ間隔 $M$ は、増加する再生トラフィックのコストと、密な書き込みの削減による節約とのバランスを取ります。

カーネル実装

GPU上でこれを効率的に実装するために、著者らはタイル分割カーネル分解 (tiled kernel decomposition) を採用しています。

• 値タイル・パラレリズム (Value-Tile Parallelism): 状態は値の次元 ($d_v$) に沿って分割されます。各ワークユニットは、密なベースのスライスと、対応する値側のログのシャードを処理します。
• 共有メタデータ: キー側の因子と対数減衰メタデータは、冗長な計算を避けるために値タイル間で共有されます。
• 段階的実行: カーネルは、共有の「ReplayPrep」ステージ（スロット・ヘッドごとに係数を一度だけ計算）と、値タイル用の並列な「AppendDecode」ステージに分割されています。

この手法は、異なる減衰粒度と更新セマンティクスを持つ3つのオペレータに対してインスタンス化されています：

• GDN (Gated DeltaNet): ヘッド単位のスカラー減衰。Delta-ruleによる補正を使用。
• KDA (Kimi Delta Attention): キー単位の減衰。Delta-ruleによる補正を使用。
• RWKV6: キー単位の減衰。生の値の更新（Delta-ruleによる補正なし）を使用。

主な貢献

1. 状態更新税の特定: 本論文は、線形アテンションのサービングにおける、逐次的な密な実体化によって引き起こされる1トークンあたりのオーバーヘッドを特性化し、それが論理的な再帰コストとは異なるメモリ帯域幅のボトルネックであることを特定しました。
2. DeltaLog 実行戦略: 状態を密なベースとコンパクトな更新の限定的なログとして表現する、遅延実体化のアプローチを提供し、ほとんどのステップでフル状態の書き戻しではなく、軽量なアペンドを可能にします。
3. システム実装と評価: vLLMベースのサービング・スタックに統合されたGDN、KDA、RWKV6の完全な実装を提供し、この最適化が（浮動小数点精度の範囲内で）厳密な算術的等価性を維持しながら、パフォーマンスを向上させることを実証しました。

結果

実験は、NVIDIA H200およびRTX 4090 GPUを用い、様々なバッチサイズとコンテキスト長にわたって実施されました。

• カーネルレベルのパフォーマンス:
  • DeltaLogは、高度に最適化されたFast Linear Attention (FLA) ベースラインと比較して、再帰的な状態更新カーネルを 1.19倍から1.86倍 加速させました。
  • 再帰的な状態の書き込みトラフィックを最大 7.83倍 削減しました。
  • GDNは、コンパクトなスカラー減衰メタデータにより最も大きな恩恵を受け、RWKV6とKDAも、より重いキー単位のメタデータがあるにもかかわらず、顕著な利得を示しました。
• エンドツーエンドのサービング・パフォーマンス:
  • プロトタイプ・サービング・スタックに統合されたDeltaLogは、Qwen3.6-35BやKimi-Linear-48Bといったモデルにおいて、密な再帰ベースラインに対して 1.05倍から1.20倍 のエンドツーエンドの速度向上を実現しました。
  • 再帰的な更新はフルデコードステップ（MLP、正規化、スケジューリングを含む）の構成要素の一つに過ぎないため、カーネルレベルの改善に比べると、アムダールの法則により向上幅は小さくなっています。
• 数値的な正確性:
  • 本手法は、逐次的なデコーディングとの数値的な等価性を維持しており、テストされたすべてのオペレータにおいて、最大出力相対誤差は標準的なリポジトリの許容範囲（例：< 0.005）内に収まっています。

重要性

本論文は、再帰的な状態の実体化が、特に大規模バッチのサービングにおいて、線形アテンションのデコーディングにおける、見落とされがちな重要なボトルネックであることを主張しています。アルゴリズムの再帰自体ではなく、状態更新の物理的なスケジュールを最適化することで、モデルのアーキテクチャを変更することなく、大幅な効率向上を実現できることを示しています。

DeltaLogは、線形アテンションがKVキャッシュのコンテキスト長依存性を排除する一方で、新たな固定状態の維持コストを導入していることを浮き彫りにしました。DeltaLogは、この「状態更新税」に対する実用的な解決策を提供し、密な書き戻しを遅延させることが、再帰的な線形アテンションモデルの推論スループットを向上させるための有効な戦略であることを示しています。著者らは、この恩恵は、再帰的な状態の書き戻しがメモリトラフィックの大部分を占める領域（大規模バッチ、大規模な状態サイズ）で最も顕著になることを指摘しています。