新聞｜ReplaySSM：混合模型解碼改存近期輸入，SSM 狀態每 16 步才寫回一次

重點

Princeton 的 Dao AI Lab 與 NVIDIA 在 10 月 12 日（美國時間）發表 ReplaySSM，改寫混合模型（SSM 層與注意力層交錯，例如 Nemotron 3、Qwen3.5）裡 SSM 層的解碼方式。
做法是把狀態存成一個檢查點，加上最近幾個 token 的輸入；緩衝累積到 L 筆才把狀態寫回 GPU 記憶體一次。論文用 L = 16。
batch 256 時，SSM kernel 快 1.46 到 1.80 倍，整個解碼步快 1.20 到 1.47 倍；batch 1 時沒有加速。
推測解碼不必再替每個草稿 token 存一份狀態快照。batch 512 時，兩個 120B 級模型的吞吐量是 vLLM 推測解碼的 1.85 倍與 2.10 倍。
輸出讀取與遞迴更新改用不同精度後，4 位元（NVFP4）狀態在 GPQA-Diamond 的平均損失從 13.05 分降到 1.77 到 2.36 分。

為什麼 SSM 層會拖慢解碼

狀態空間模型（state-space model，SSM）用固定大小的狀態摘要過去所有 token，取代注意力層那個隨上下文變長的 KV cache。照理說，這讓解碼變便宜。實際上每一步都要把整個狀態從 HBM（GPU 的高頻寬記憶體）讀進來，做一次衰減加秩一更新，再整個寫回去。計算量很少，搬運量很大，所以 kernel 卡在記憶體頻寬。

以 Mamba-2 為例，每個 token 的遞迴是：

S_t = \alpha_t S_{t-1} + v_t k_t^\top, \qquad o_t = S_t q_t

其中 $S_t$ 是狀態， $\alpha_t$ 是依資料而定的衰減， $v_t$ 、 $k_t$ 、 $q_t$ 是該 token 的輸入。

論文在 Nemotron-3-Super-120B-A12B-NVFP4、batch 256、單張 B300 上量測：上下文在 100K token 以內時，SSM 層是解碼步裡最大的延遲來源，比注意力層還慢。原因有兩個：

這個模型的 SSM 層數是注意力層的 5 倍。
一層 SSM 的 FP32 狀態每條序列 4 MiB，大小等於一層注意力在 8K token 時的 FP8 KV cache。解碼時 SSM 狀態要讀也要寫，KV cache 只要讀，所以按搬運位元組算，SSM 層要到大約 80K token 才被注意力追上。實測交叉點在 100K 左右，作者認為是注意力 kernel 被優化得比較徹底。

每步都寫回狀態還帶來另外兩個問題。推測解碼要在草稿被拒絕時回到最後接受的位置，SSM 狀態更新後很難倒回去。狀態量化之後，捨入誤差每一步都會重新進入遞迴，生成越長累積越多。

方法：存輸入，少存狀態

把遞迴在一個 $h$ 個 token 的窗口內展開：

S_t = \gamma_h S_{\mathrm{ckpt}} + \sum_{j=1}^{h} \gamma_{j+1:h}\, v_j k_j^\top

$S_{\mathrm{ckpt}}$ 是窗口開始前的狀態， $\gamma$ 是累積衰減。也就是說，目前的狀態可以從一個舊檢查點加上近期輸入 $(\alpha_j, v_j, k_j)$ 重建出來。ReplaySSM 就用這個形式存狀態：

一般步（non-flush）：把這一步的輸入記錄附加到緩衝區，從檢查點與緩衝算出輸出，不寫回狀態。
清空步（flush）：緩衝累積到 L 筆時，把這些輸入併進檢查點，寫回 HBM 一次，清空緩衝。

輸出與傳統解碼相同，差別只在浮點誤差。

在 L 步的窗口裡，兩種做法都要讀 L 次狀態，因為每步的輸出都需要它。傳統做法還要寫 L 次，ReplaySSM 只寫一次，另外讀寫的是比狀態小得多的輸入記錄，所以主要的狀態流量大約減半。L 越大，寫回攤得越薄，但每步要重播的輸入越多。四個模型在 batch 256 時都是 L = 16 最快。

Gated DeltaNet（GDN）與 Kimi Delta Attention（KDA）用的是 delta rule，更新前會先修改既有內容。論文把它改寫成同樣的外積形式，緩衝改存修正項 $u_j$ 。實作涵蓋 Mamba-2 與 GDN，以 Triton kernel 寫在 vLLM 裡。

標準解碼的結果

模型	SSM 類型	權重	GPU	SSM／注意力層數
Nemotron-3-Nano-4B	Mamba-2	BF16	1×H100	21／4
Nemotron-3-Super-120B-A12B	Mamba-2	NVFP4	1×B300	40／8
Qwen3.5-4B	GDN	BF16	1×H100	24／8
Qwen3.5-122B-A10B	GDN	NVFP4	1×B300	36／12

對照組是 vLLM 調到最佳啟動設定的原生 kernel，比 Mamba-2 出廠預設值還快。在 batch 256、L = 16、FP32 狀態下：

SSM kernel 快 1.46 到 1.80 倍（含清空步成本）。
整個解碼步快 1.20 到 1.47 倍。
batch 1 時 kernel 受延遲限制，端到端只差幾個百分點，等於沒有加速。

加速隨 batch 變大而增加，直到傳統 kernel 吃滿頻寬為止。端到端的上限取決於 SSM 在解碼步裡的占比，注意力與模型其他部分沒有改動。

推測解碼：回滾只要移動指標

推測解碼先由小模型或 MTP head 猜出 T 個 token，再讓目標模型一次驗證，保留猜對的前綴。注意力層回滾只要把 KV cache 的指標往回移。SSM 狀態一旦吸收了草稿就倒不回去，所以 vLLM 的做法是每個草稿 token 存一份狀態快照，每步變成 1 次讀加 T 次寫，快照還會占滿記憶體、壓低能同時服務的請求數。先前的 Mamba in the Llama、Snakes and Ladders、STree 避開了逐 token 快照，但仍要額外搬運狀態，在服務級 batch 下都只比快照快幾個百分點。

ReplaySSM 的做法：

驗證前把 T 個草稿的輸入全部附加到緩衝區，驗證後依接受長度移動指標。被拒絕的記錄留在指標之後，下一次附加會直接覆蓋。
緩衝區是容量 L + T 的環形結構，清空步只併入已確認的記錄，不必等驗證結果。
一般步只需要輸出，不需要狀態，所以可以換一種算法。單一更新時 $o = (v k^\top) q = v (k^\top q)$ ，右式只算一個純量再乘向量，不必組出 $d_v \times d_k$ 的狀態矩陣。T 個草稿時寫成：

O = S_{\mathrm{ckpt}}\, Q D + V \bigl( W \odot (K^\top Q) \bigr)

$D$ 是草稿位置的累積衰減， $W$ 是帶因果遮罩的衰減矩陣。檢查點只讀一次，其餘是兩個小矩陣乘法，可以用 Tensor Core。這跟 prefill 用的分塊（chunkwise）算法是同一種形式，緩衝加草稿就是那個 chunk。

實驗用兩個 120B 級模型，各自的 MTP head 產生草稿，T = 4，GSM8K 題目，每個請求生成 256 個 token：

指標	Nemotron-3-Super	Qwen3.5-122B
batch 512 吞吐量，相對 vLLM 推測解碼	1.85×	2.10×
batch 512 吞吐量，相對標準解碼	2.16×	1.90×

兩者的接受長度相同。傳統推測解碼過了 batch 128 就不再變快，Qwen3.5-122B 在 batch 512 時甚至比不用推測解碼還慢；ReplaySSM 則隨 batch 持續成長。

驗證 kernel 本身的差距更明顯。在 Nemotron-3-Super、batch 256 下，傳統驗證 8 個 token 的延遲是一次標準解碼步的 4.66 倍；ReplaySSM 是 0.89 到 1.21 倍，範圍取決於草稿全被拒絕或全被接受。

低精度狀態：讀取用低精度，遞迴用高精度

狀態量化後，同一份狀態既用來算輸出，也是下一步更新的起點，捨入誤差會一路累積。KV cache 沒有這個迴圈，量化後的項目會被讀很多次，但從不寫回自己，所以比 SSM 狀態更耐低精度。這也是為什麼 Qwen3.5-122B-A10B-NVFP4 的權重用 4 位元、KV cache 用 FP8，SSM 狀態卻維持 FP32。

ReplaySSM 的分工讓兩件事可以用不同精度：

高精度檢查點只負責遞迴，清空步從它出發。
低精度副本只給一般步讀取、計算輸出，清空時從高精度結果另外轉出，永遠不會再被讀回遞迴。

讀 NVFP4 狀態時，Mamba-2 的一般步 kernel 比讀 FP32 快 3.8 倍。GDN 要多一個調整：一般步算出的修正項來自低精度副本，所以另外快取 $v$ ，在清空時用高精度狀態重算修正項。

輸出偏差用 KL divergence 量測（PG-19 文字、teacher forcing）。在 Qwen3.5-4B 生成到 128K token 時，NVFP4 讀取搭配 FP32 遞迴，KL 從 0.136 降到 0.0038，約 36 倍。遞迴改用 FP16 加隨機捨入，結果跟 FP32 很接近。

下游準確率是四個模型的平均變化（百分點，FP32 為基準）：

狀態設定	MMLU-Pro	GPQA-Diamond	Math-L5
FP32 準確率（基準）	77.6%	73.7%	92.3%
共用 FP8	−0.24	−1.05	+0.01
共用 NVFP4	−5.88	−13.05	−10.92
NVFP4 讀取／FP16 遞迴	−0.15	−2.36	−0.10
NVFP4 讀取／FP8 遞迴	−0.40	−1.77	−0.38

損失集中在長生成。Nemotron-3-Super 在 GPQA-Diamond 最長四分之一題目的生成長度中位數是 61K token，MMLU-Pro 與 Math-L5 約 8K。在較長的兩個四分位，共用 NVFP4 比 FP32 差 20.9 與 28.2 分；NVFP4 讀取搭配 FP8 遞迴只差 2.0 與 3.4 分。

跟既有做法的關係

分塊算法：Mamba-2、GDN、KDA 的分塊算法都是為訓練與 prefill 的計算吞吐量設計。ReplaySSM 處理的是受記憶體限制的逐 token 解碼，一般步直接沿用這些分塊算法。
前綴快取：Marconi 這類做法省的是重複的 prefill，ReplaySSM 省的是解碼時的狀態流量，兩者可以並用。
週期性檢查點：Nemotron 3 Ultra 的技術報告已經試過每隔一段才量化一次狀態，可以減少捨入次數，但誤差仍會進入後續更新。ReplaySSM 多出的是讀取與遞迴分用兩種精度。

限制與待確認

加速只在大 batch 時出現，對單一使用者、batch 1 的本地推論沒有幫助。
論文實作只涵蓋 Mamba-2 與 GDN，KDA 只在數學形式上說明可以套用。推測解碼只測了兩個 120B 級模型、GSM8K 題目、T = 4。
NVFP4 讀取仍有固定的讀取誤差，KL 是 FP32 下限的好幾倍。最小的 Nemotron-3-Nano-4B 在 GPQA-Diamond 仍掉 5.6 到 6.7 分。
部署狀況說法不一。NVIDIA 部落格說已整合進 TensorRT-LLM、vLLM、SGLang、FlashInfer，GLM-5.3-Flash 也已用它服務；GitHub 的研究版 README 則寫 vLLM 的正式版本仍透過 RFC 與 pull request 上游合併中。Kimi-K3 在驗證時也快取草稿的 SSM 輸入，但它是把 ReplaySSM 列為同期做法，不是直接採用。
數字依版本略有出入。論文摘要寫推測解碼「約 1.8 倍」，對應 batch 256；1.85 與 2.10 倍是 batch 512 的結果。Tri Dao 個人部落格的介紹文寫的是標準解碼最高 1.48 倍、推測解碼 1.87 到 1.96 倍。本文以論文正文為準。