NVIDIA ReplaySSM 改以檢查點加近期輸入儲存 SSM 狀態，改善混合模型解碼並整合三套推理框架

NVIDIA 發表 ReplaySSM，與 Princeton Dao AI Lab 共同提出以「狀態檢查點＋近期輸入緩衝區」取代逐步儲存完整 SSM 狀態的方法，目標是改善混合模型的解碼效能。ReplaySSM 已整合進 NVIDIA TensorRT-LLM、vLLM 與 FlashInfer，並在 GLM-5.3-Flash 和 Kimi-K3 的生產環境運作。 混合模型會交錯使用 state-space model（SSM）層與少量 attention 層，NVIDIA Nemotron 3、Qwen3.5 至 Qwen3.8、Kimi-K3 和 GLM-5.3-Flash 都採用這類架構，SSM 層則包括 Mamba-2、Gated DeltaNet（GDN）或 Kimi Delta Attention（KDA）。SSM 原本以固定大小的狀態取代會隨上下文長度增加的 KV cache，但在服務批次大小下，SSM 層可能在最多 100K tokens 的上下文中成為解碼瓶頸。 以 NVIDIA Nemotron 3 Ultra 使用的 Mamba-2 為例，傳統解碼每產生一個 token，都要從高頻寬記憶體（HBM）載入完整狀態、套用衰減與更新、計算輸出，再把更新後的狀態寫回。以 batch size 256、單張 NVIDIA Blackwell Ultra B300 GPU 執行 Nemotron-3-Super-NVFP4 時，SSM 層在最多 100K tokens 的上下文中是解碼步驟的主要部分；原因包括模型的 SSM 層數是 attention 層的五倍，以及單一序列每層 FP32 狀態為 4 MiB，相當於 attention 層 8K tokens 的 KV cache 大小。 ReplaySSM 將狀態表示為一個檢查點和近期輸入記錄。在非清空步驟中，系統把目前輸入加入緩衝區，根據檢查點與緩衝區計算輸出，但不儲存完整狀態；當緩衝區累積至窗口大小 L 時，清空步驟才會把輸入合併進檢查點、向 HBM 寫入一次新檢查點並清除緩衝區。NVIDIA 表示，這種做法在浮點數差異之外，可得到與傳統解碼相同的輸出。 這種狀態表示法針對三項問題：減少每一步載入與儲存完整狀態造成的記憶體流量；讓推測解碼能透過捨棄近期輸入記錄來回滾，而不必為每個草稿 token 儲存狀態快照；以及讓低精度狀態讀取的捨入誤差不會在遞迴更新中持續累積。NVIDIA 指出，這些問題使推測解碼與量化在 SSM 上較難像在 Transformer attention 上有效運作。 ReplaySSM 的論文《ReplaySSM: Cache SSM Inputs, Not State》已公開於 arXiv，程式碼則放在 GitHub。NVIDIA 以四個混合模型進行評估，涵蓋採用 Mamba-2 的 Nemotron 3 與採用 GDN 的 Qwen3.5，模型規模從單張 NVIDIA H100 Tensor Core GPU 上的 4B dense 模型，到單張 NVIDIA Blackwell Ultra B300 GPU 上的 120B 級 MoE 模型；測試使用 FP32 SSM 狀態與窗口大小 L=16，並與 vLLM 在最佳啟動設定下的原生 kernels 比較。