ReplaySSM：状態ではなく入力をキャッシュすることでハイブリッドモデルのデコードを高速化

NVIDIAとプリンストン大学のDao AI Labは、ハイブリッド大規模言語モデルのデコードフェーズを最適化する新手法「ReplaySSM」を発表しました。2026年10月に公開されたこのアプローチは、推論中に状態空間モデル（SSM）が内部状態をどのように保存・更新するかを再考するものです。直近の入力トークンをキャッシュし、完全な状態マトリクスを継続的に上書きするのではなく、この手法によりメモリ帯域幅の使用量が大幅に削減され、より効率的な投機的デコードが可能になります。

何が起きたのか

SSM層と従来のアテンション層を組み合わせたハイブリッドモデルは、NVIDIA Nemotron 3、Qwen3.5、Kimi-K3などの現代的なアーキテクチャで一般的になりつつあります。これらのモデルでは、SSMを用いて長いコンテキストを効率的に処理し、アテンション機構の増大するキーバリュースタイルキャッシュを固定サイズの状態に置き換えています。しかし、サービング時のバッチサイズが高い場合、大きな状態マトリクスを高帯域幅メモリから頻繁に読み書きする必要があるため、SSM層が主要なレイテンシボトルネックになることがよくあります。

研究チームは、従来のデコードでは出力計算と状態更新が単一のステップに融合されており、過度なメモリトラフィックが発生することを特定しました。ReplaySSMは、状態をチェックポイントと直近入力のバッファとして表現することで、これらの操作を分離します。これにより、システムは完全な状態をメモリに書き戻すことなく出力を計算でき、入力バッファがいっぱいになったときのみチェックポイントを更新します。この手法は現在、TensorRT-LLM、vLLM、FlashInferなどの主要な推論フレームワークに統合されています。

仕組み

標準的なSSMデコードでは、各トークン生成ステップで以前の状態を読み込み、減衰と更新を適用し、出力を計算し、新しい状態を保存します。データ移動に対して演算強度が低いため、このプロセスはメモリバウンドとなります。ReplaySSMは、小さなトークンウィンドウ上で再帰を展開することでこの状況を変えます。毎ステップで状態を更新する代わりに、現在のトークンの入力ベクトルをバッファに追加します。出力は、保存されたチェックポイントとバッファ内の入力を使用して計算されます。

バッファが事前定義されたサイズに達すると、フラッシュ操作がすべてのバッファ内入力をチェックポイントに折りたたみ、一度だけメモリに書き込みます。これにより、状態書き込みの回数が約半分になります。複数のドラフトトークンを生成・検証する投機的デコードでは、ReplaySSMはロールバックを簡素化します。各ドラフトトークンごとに完全な状態のスナップショットを保存する代わりに、システムは単にドラフト入力をバッファに追加します。ドラフトが拒否された場合、バッファポインタがリセットされ、高コストな状態復元を行うことなく未コミットの入力を破棄します。

主要な詳細

• ReplaySSMは、直近入力のバッファが埋まった後にのみチェックポイントを書き込むことで、状態メモリトラフィックを約50%削減します。
• NVIDIA Blackwell Ultra B300 GPUにおいてバッチサイズ256の場合、SSMデコードカーネルは従来手法と比較して1.46倍から1.80倍高速に動作します。
• ロールバック時に状態スナップショットの保存と復元が不要になるため、120Bクラスモデルでの投機的デコードスループットは1.85倍から2.10倍向上します。
• この手法は、再帰用に高精度チェックポイントを保持しつつ、読み取りには低精度コピーを使用することで誤差蓄積を制限し、低精度状態ストレージをサポートします。
• 評価対象モデルにはNemotron-3およびQwen3.5のバリアントが含まれており、Mamba-2やGated DeltaNetなど異なるSSMアーキテクチャで一貫した性能向上を示しています。
• この手法は、GLM-5.3-FlashやKimi-K3などの本番システムですでに展開されています。

なぜ重要なのか

大規模ハイブリッドモデルの構築・提供を行うエンジニアにとって、メモリ帯域幅はしばしばスループットの制約要因となります。SSM層における支配的なメモリトラフィックを削減することで、ReplaySSMはトークンごとのレイテンシを増加させることなく、より高いバッチサイズを可能にします。これは、H100やB300といった高価なGPUでのハードウェア利用率向上を実現するため、コスト効率の良いサービングにとって極めて重要です。性能向上はバッチサイズに応じてスケールするため、高並行性の本番環境で特に価値があります。

さらに、ReplaySSMはSSMベースのモデルにおける投機的デコードを実用可能なものにします。従来は、ロールバックのための状態スナップショット保存のコストが、これらのアーキテクチャにおける投機による利点を相殺していました。入力バッファをキーバリューキャッシュポインタのように扱うことで、この手法はマルチトークン予測の効率性を回復させます。これにより、開発者は積極的な投機戦略を用いてレイテンシをさらに削減でき、SSMのパフォーマンスをハードウェアの理論的限界に近づけることができます。

やれること

• arXiv上のReplaySSM論文全文を読んで、数学的導出と詳細な評価指標を理解してください。
• GitHub上のオープンソース実装を検討し、独自プロジェクトでのキャッシュメカニズムを試してみてください。
• TensorRT-LLMまたはvLLMの最新バージョンにアップグレードして、互換性のあるハイブリッドモデルでReplaySSMサポートを有効にしてください。
• 現在のSSMワークロードをより高いバッチサイズでベンチマークし、メモリトラフィック削減による潜在的なスループット向上を定量化してください。
• ReplaySSMはロールバックに関連する以前のメモリペナルティを排除するため、より長いドラフト長を持つ投機的デコード設定を検討してください。
• 低精度状態を使用する場合、量子化エラーを監視してください。チェックポイントとバッファの分離により、FP32レベルに近い精度が維持されることに注意してください。