NVIDIAとプリンストン大学、ハイブリッドモデルの推論を高速化する「ReplaySSM」を発表──状態保存の再設計でメモリアクセスを半減

要点

• 米NVIDIAと米プリンストン大学のDao AI Labは、状態空間モデル（SSM）とAttention機構を組み合わせたハイブリッドモデルのデコーディング処理を高速化する新手法「ReplaySSM」を発表した。

• 毎ステップで状態全体を更新・保存する従来手法を改め、状態をチェックポイントと直近の入力バッファに分離して保持することで、出力計算と状態更新を切り離した。

• 状態メモリのトラフィックを約半分に削減し、NVIDIA Blackwell Ultra B300 GPU上でのSSMデコーディングカーネルの速度を最大1.80倍に向上させた。

• 投機的デコーディングにおけるロールバックをポインタ操作で簡素化し、120Bクラスのモデルにおいて推論スループットを最大2.10倍に引き上げた。

• 状態更新と読み出しの精度を分けることで量子化誤差の累積を防ぐ設計を取り入れており、TensorRT-LLMなどの推論エンジンや商用モデルの本番環境でもすでに導入されている。

• 米NVIDIAと米プリンストン大学のDao AI Labは2026年10月12日、大規模言語モデルなどの推論処理を大幅に高速化する技術「ReplaySSM」を発表した。研究グループは論文「ReplaySSM: Cache SSM Inputs, Not State」をarXivで公開するとともに、GitHub上で実装コードを提供している。状態空間モデル（SSM）とAttention機構を組み合わせた「ハイブリッドモデル」において、内部状態の保持方法を刷新することでメモリ転送の負荷を半減させるアプローチだという。

ハイブリッドモデルで顕在化していたSSM層の推論ボトルネック

近年、長大な文脈を効率的に扱うアーキテクチャとして、SSM層と一部のAttention層を交互に配置した「ハイブリッドモデル」の採用が広がっている。NVIDIAの「Nemotron 3」や「Nemotron 3 Ultra」をはじめ、「Qwen3.5」から「Qwen3.8」、「Kimi-K3」、「GLM-5.3-Flash」などがこの構造を採用しており、SSM層の実装には「Mamba-2」「Gated DeltaNet（GDN）」「Kimi Delta Attention（KDA）」などが用いられている。

SSM（状態空間モデル）は、入力系列を固定サイズの内部状態に圧縮・更新しながら処理する手法であり、文脈長に応じて増大するKVキャッシュ（Attention層が過去トークンの計算結果を保持するメモリ領域）を不要にすることで高速なデコーディングを実現する目的で設計された。しかし実際の大規模な推論サーバー環境では、最大10万トークン（100K）におよぶ文脈長において、Attention層ではなくSSM層こそがレイテンシのボトルネックになる現象が生じていたという。

従来のデコーディングでは、1トークンを生成するステップごとに、広帯域幅メモリ（HBM：GPUに搭載される超高速メモリ）から直前の状態を読み出して更新し、出力を計算した上で新しい状態をHBMへ書き戻す必要があった。この「出力の計算」と「次ステップへの状態更新」が一体化していたことが毎ステップの膨大なメモリアクセスを生み、性能低下や高速化手法の適用難航を招いていたと著者らは指摘している。

状態ではなく「入力」をキャッシュして計算と更新を分離

ReplaySSMは、この状態の保存方式を根本から再設計した技術だ。従来の全状態を毎ステップ更新・保存する手法とは異なり、状態を「特定時点のチェックポイント」と「直近の入力を保持するバッファ」の組み合わせとして表現する。

この仕組みにより、各ステップではチェックポイントとバッファ内の入力から出力を計算し、バッファが満杯になった段階でのみチェックポイントをHBMへ書き込む方式へと変更された。出力計算と状態更新を分離したことで、毎ステップ発生していたHBMへの状態書き戻しが不要となり、状態メモリのトラフィックを約半分に削減することに成功したとしている。

メモリ削減に加え、投機的デコーディングと量子化の課題も解消

このアプローチにより、メモリアクセスの削減だけでなく、従来のSSM層が抱えていた推論高速化技術との両立という課題も解消されたという。論文では以下の3つの改善が報告されている。

第1に、ハードウェア実行速度の向上だ。NVIDIA Blackwell Ultra B300 GPUにおいて、バッチサイズ256の条件下で測定したところ、SSMデコーディングカーネルの実行速度が1.46倍から1.80倍に向上することが確認された。

第2に、投機的デコーディング（小型モデルで候補トークンを先行生成し大型モデルで一括検証する高速化手法）におけるロールバック処理の簡素化だ。検証で棄却されたトークンを取り消す際、従来は状態全体のスナップショットを保存しておく必要があったが、ReplaySSMではバッファにドラフト入力を追加しポインタを戻すだけでロールバックが完了する。この結果、120B（1200億パラメータ）クラスのモデルにおいて、投機的スループットが1.85倍から2.10倍に向上した。

第3に、量子化（数値データのビット精度を下げてメモリを節約する手法）時の誤差蓄積の抑制だ。従来のSSMでは再帰処理のたびに量子化誤差が蓄積していく問題があったが、ReplaySSMでは再帰的な状態更新用には高精度のチェックポイントを維持し、出力の読み出し用には低精度のコピーを利用する構造をとる。難関ベンチマーク「GPQA-Diamond」による評価では、検証した4モデル中3モデルでFP32（単精度浮動小数点数）実行時との精度差を約2パーセントポイント以内に維持できたとしている。

主要推論エンジンへの統合と商用環境での先行導入

ReplaySSMはすでに実用フェーズに入っている。本手法はNVIDIAの推論エンジン「TensorRT-LLM」をはじめ、「vLLM」や「FlashInfer」といった主要なオープンソース推論フレームワークに統合された。

さらに研究用途にとどまらず、「GLM-5.3-Flash」や「Kimi-K3」の商用本番環境（プロダクション）でもすでに導入され、実際に稼働していることが明かされている。