ReplaySSM：ハイブリッドモデルの高速デコーディングを実現する状態ストレージ再考
ハイブリッドモデル推論のボトルネックとReplaySSMの核心
大規模言語モデル（LLM）の推論において、ハイブリッドモデル（TransformerとState Space Model (SSM) レイヤーを組み合わせたモデル、例えばMamba-2やGated DeltaNetに基づくもの）は、その独特のアーキテクチャに起因するパフォーマンスの課題を抱えています。特に、SSMレイヤーを持つモデルのデコーディングプロセスは、メモリ帯域幅に制約されることが大きなボトルネックとなっています。これは、各ステップで大きな再帰状態（recurrent state）が高帯域幅メモリ（HBM）から読み出され、更新され、そして書き戻されるという操作が頻繁に発生するためです。特にバッチサイズが小さい場合、この状態の書き戻しが全体のパフォーマンスを支配します。
ReplaySSMは、この根本的な問題に対処するために、状態ストレージのアプローチを再考します。従来のSSMデコーディングが毎ステップで再帰状態をHBMに書き戻すのに対し、ReplaySSMは状態そのものを保存するのではなく、最近の入力（例：Mamba-2の場合、各ステップのペアと減衰因子）を小さな層ごとのリングバッファにキャッシュします。 この手法は、FlashAttentionが計算とI/Oのトレードオフを行ったのと同様に、少量の再計算を許容することでHBMへのトラフィックを大幅に削減します。 状態は必要に応じて、バッファされた入力から「リプレイ」によって再構築されます。これにより、毎ステップの状態更新と書き戻しを回避し、メモリトラフィックを実質的に半減させることができます。
ReplaySSMによる高速化メカニズムと性能ベンチマーク
ReplaySSMの核となる高速化メカニズムは、主に2つのパスで機能します。
-
標準デコーディングの最適化: ReplaySSMでは、フル状態の書き戻しはバッファが満杯になったときにのみ定期的に「フラッシュ」として実行されます。このフラッシュコストはバッファウィンドウ全体に償却されるため、毎ステップのHBMアクセスが大幅に削減されます。これにより、SSMカーネルはメモリバウンドな特性を克服し、レイテンシが直接改善されます。 NVIDIA H100およびB300 GPU上でvLLMとCUDA Graphを有効にした評価では、ReplaySSMは標準デコーディングにおいて、エンドツーエンドで最大1.48倍（大規模なMoEモデルでは1.43倍）の速度向上を達成しました。
-
推測デコーディング（Speculative Decoding）の強化: 推測デコーディングでは、ドラフトトークンが拒否された場合、SSMの状態をロールバックする必要があり、これは通常、コストのかかるフル状態のチェックポイント作成を伴います。ReplaySSMの入力リングバッファは、この課題に対してエレガントな解決策を提供します。リングバッファとカーソルを使用することで、ロールバックは単なるポインタの移動と必要な入力のリプレイによって安価に実行できます。 この革新により、vLLMの既存の実装が特定のバッチサイズで標準デコーディングを下回るのに対し、ReplaySSMは推測デコーディングにおいて1.87〜1.96倍の速度向上を実現しています。さらに、固定メモリ予算下でvLLMの推測パスよりも3.0〜3.3倍多くの同時リクエストを処理できることも示されています。
Nemotron-3 (Mamba-2ベース) およびQwen3.5 (GDNベース) ファミリーの4Bから550Bまでの幅広いモデルサイズで評価が行われ、ReplaySSMの有効性が実証されています。
推測デコーディングにおける革新と実装詳細
ReplaySSMは、SSMの再帰更新が線形であるという特性を利用しています。これにより、状態はチェックポイントと、直近のLステップ分の（d, k, g）レコードを含む小さなリングバッファからオンデマンドで再構築可能です。 フル状態を非フラッシュステップで具体化することなく、バッファされた入力から直接出力を読み出す「出力専用ルート」もサポートされます。 これは、メモリバウンドなカーネルにとって帯域幅を最適化するものであり、GDNデコーディングがバッチサイズ64以上でピークHBM帯域幅の70-79%で動作する場合、ステップあたりの状態トラフィックを半減させることは、ほぼ直接的な速度向上につながります。
ReplaySSMの実装はvLLM上で行われ、Mamba-2のselective_state_updateやGated DeltaNetのfused_recurrentといった既存のデコードパスと並行して、オプトインのTritonカーネルとして追加されています。 ロールバックがポインタ移動に限定されるリングバッファは、ツリーベースの推測デコーディングにおいても重要であり、この場合、受け入れられたトークンはバッファ内で連続しているとは限りません。 このアプローチは、SSM計算自体を変更するものではなく、同じ累積状態再帰を維持しつつ、フル再帰状態が具体化されメモリに書き込まれるタイミングを変更することで、異なるメモリパスを確立しています。
開発者・エンジニア視点での考察
-
ハイブリッドモデルの推論スケーラビリティ向上への道: ReplaySSMは、SSMベースのハイブリッドモデルの推論がHBM帯域幅に制約されるという根本的な課題に直接対処しています。これは、FlashAttentionがTransformerのKVキャッシュアクセスを最適化したのと同様に、SSMの主要なボトルネックを解消し、より大規模で複雑なハイブリッドモデルの展開を現実的にします。開発者は、この技術の統合により、メモリ制約に起因するパフォーマンス上の懸念を軽減し、より高いスループットと低いレイテンシを期待できます。
-
推測デコーディングのSSMモデルへの適用可能性拡大: SSMモデルにおける推測デコーディングの課題であった「状態ロールバックの複雑性とコスト」を、入力リプレイとポインタベースのリングバッファで解決したことは画期的です。これにより、SSMベースのモデルでもTransformerと同様に、またはそれ以上に効率的な推測デコーディング戦略を適用できるようになり、リアルタイム対話システムや低レイテンシが求められるアプリケーションでのSSMの採用が加速するでしょう。
-
既存の推論フレームワークへの統合とカスタムカーネル開発の重要性: ReplaySSMがvLLMの既存のデコードパスにTritonカーネルとして追加されていることは、最先端のモデル最適化がしばしばカスタムGPUカーネル開発を伴うことを示しています。これは、AI開発者にとって、高レベルなフレームワークの利用に加えて、パフォーマンス critical な部分でCUDAやTritonのような低レベルなプログラミングを通じてハードウェアに最適化された実装を行うスキルの重要性を再認識させるものです。
Source / 元記事
この記事について
この記事は、公開されているニュース、論文、公式発表、RSSフィードなどをもとに、AIが要約・補足調査・考察を行って作成しています。
元記事の完全な翻訳・逐語的な要約ではなく、AIによる背景説明や開発者向けの考察を含みます。
重要な技術仕様・価格・提供状況などは、必ず元記事または公式情報をご確認ください。
