PaperlayerView exact revision on arXiv

Source-linked AI summary

ReplaySSM: Cache SSM Inputs, Not State

Ze-Wei Liou, Harris Nover, Amey Naik, Andrii Skliar, Igor Shovkun, Ido Shahaf, Tri Dao

arXiv:2610.13778v1cs.LGcs.AI

TL;DR

Eager SSM state management can bottleneck hybrid-model decoding and complicate rollback and quantized updates. ReplaySSM represents state with a checkpoint and recent-input buffer, separating output computation from checkpoint updates. At batch size 256, it improves standard decoding over tuned vLLM and raises speculative-decoding throughput on two 120B-class models.

  • Problem

    Eagerly reading and writing recurrent state at every step can bottleneck hybrid-model decoding, complicate speculative rollback, and feed quantization error into subsequent updates.

    Source: §3.1 · p. 4, §3.2 · p. 4, §9 · p. 11

  • Method

    ReplaySSM unrolls the recurrence into a checkpoint and bounded buffer of recent SSM inputs, separating output computation from checkpoint updates.

    Source: §1 · pp. 2–3

  • Results

    At batch size 256, standard decoding is 1.20–1.47× faster than tuned vLLM across four hybrid models, and speculative decoding reaches about 1.8× vLLM throughput on two 120B-class models.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Separate state precisions reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.

    Source: §1 · p. 3

    Every step still reads the checkpoint; skipping reads for heads with negligible cumulative decay remains a future direction.

    Source: §9 · p. 11

Abstract

from arXiv · show

State-space models (SSMs) maintain a fixed-size recurrent state, but reading and writing it at every decoding step can bottleneck hybrid models. This eager state management causes large state I/O with low arithmetic intensity, leaves speculative decoding no cheap rollback, and feeds state quantization error into subsequent recurrent updates. We introduce ReplaySSM, which unrolls the recurrence to represent the state as a checkpoint and a buffer of recent SSM inputs. This separates output computation from checkpoint updates. By updating the checkpoint periodically, ReplaySSM avoids state writes on most steps and roughly halves the dominant state memory traffic. Keeping recent inputs also makes speculative rollback a buffer-pointer operation and enables output-only verification without materializing intermediate states. The same separation allows a low-precision checkpoint to serve output computation on most steps while a higher-precision checkpoint serves recurrent updates, so rounding error in the low-precision copy never enters later updates. We implement ReplaySSM in vLLM for Mamba-2 and Gated DeltaNet. At serving batch size 256, it accelerates standard decoding by $1.20$--$1.47\times$ over a tuned vLLM baseline across four hybrid models, and delivers about $1.8\times$ the throughput of vLLM's speculative decoding on the two 120B-class models. Separate state precisions also reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.

1 INTRODUCTION

ReplaySSM represents SSM state as a checkpoint plus recent inputs, separating output computation from checkpoint updates to address state traffic, rollback, and quantization challenges.

1 INTRODUCTION

  • SSMs summarize history in a fixed-size recurrent state, but the per-step state reads and writes can bottleneck hybrid-model decoding.The introduction notes SSM layers can remain a bottleneck up to 100K tokens.Source: §1 · p. 1, Figure 1 · p. 2
  • ReplaySSM unrolls the recurrence into a checkpoint and bounded buffer of recent SSM inputs, computing outputs between periodic checkpoint flushes.Outputs match conventional decoding up to floating-point differences.Source: §1 · pp. 2–3
  • This representation amortizes checkpoint writes, roughly halves dominant state traffic, enables pointer-based speculative rollback and output-only verification, and keeps low-precision rounding out of recurrent updates.The low-precision checkpoint serves outputs while a higher-precision checkpoint carries recurrence across flushes.Source: §1 · pp. 2–3
  • At batch size 256, ReplaySSM speeds standard decoding by 1.20–1.47× over tuned vLLM across four hybrid models and delivers about 1.8× vLLM speculative-decoding throughput on two 120B-class models.Separate state precisions also reduce output KL divergence and downstream accuracy losses compared with a shared low-precision checkpoint.Source: §1 · p. 3

2 PRELIMINARIES

The preliminaries contrast fixed-size SSM recurrence with growing attention caches and introduce speculative decoding and state quantization as relevant settings.

2.1 STATE-SPACE MODELS

  • SSMs compress token history into a fixed-size state; Mamba-2 updates it with a data-dependent decay and rank-1 input update, then computes output from the state.Each step loads the prior state and stores the updated state.Source: §2.1 · p. 3
  • GDN and KDA modify stored information before adding new information; GDN retains the outer-product update form using a correction term computed from the previous state.GDN uses u_t in place of v_t, at the cost of computing u_t from the previous state.Source: §2.1 · p. 3

2.2 SPECULATIVE DECODING AND QUANTIZATION

  • Speculative decoding verifies drafted tokens together and retains only the accepted prefix, while quantization reduces storage through lower-precision values that introduce rounding error.Round-to-nearest can be biased; stochastic rounding is unbiased in expectation but has nonzero variance.Source: §2.2 · p. 3

3 MOTIVATION AND KEY INSIGHT

The section traces three costs of eager recurrent-state updates and introduces ReplaySSM’s checkpoint-and-input-buffer representation as the shared design response.

3 MOTIVATION AND KEY INSIGHT

  • SSMs compress an O(N)-token history into an O(1)-size state, but eagerly materializing each updated state creates low arithmetic intensity, expensive rollback, and propagating quantization error.These challenges arise from writing each recurrent update immediately.Source: §3 · p. 3

3.1 CHALLENGE 1: LARGE STATE I/O WITH LITTLE COMPUTE

  • For Nemotron-3-Super-120B-A12B-NVFP4 at batch size 256 on one B300, SSM updates remain the largest latency component up to 100K tokens.The model has five times as many SSM layers as attention layers, and their FP32 states move more bytes than attention KV caches until context reaches about 80K tokens.Source: §3.1 · p. 4

3.2 CHALLENGE 2: RECURRENT STATE HAS NO CHEAP ROLLBACK

  • Conventional speculative decoding stores a full-state snapshot per draft token, adding T state writes per step and reducing serving concurrency.Prior approaches that avoid snapshots still pay extra state traffic to recover the accepted state and are no more than a few percent faster at serving batch size.Source: §3.2 · p. 4

3.3 CHALLENGE 3: QUANTIZATION ERROR PROPAGATES THROUGH RECURRENCE

  • In low-precision decoding, the quantized state produces the current output and feeds the next update, so rounding error re-enters the recurrence and can accumulate with sequence length.Round-to-nearest adds systematic bias; stochastic rounding still propagates variance.Source: §3.3 · p. 4

3.4 REPLAYSSM: CACHE SSM INPUTS, NOT STATE

  • ReplaySSM stores a checkpoint plus a bounded buffer of recent inputs, computes outputs without state writes on non-flush steps, and folds buffered inputs into the checkpoint at flushes.Its outputs match conventional decoding up to floating-point differences; the design extends to standard decoding, speculative decoding, and low-precision state.Source: §3.4 · p. 4

4 REPLAYSSM FOR STANDARD DECODING

For standard decoding, ReplaySSM reconstructs state from a checkpoint and recent inputs, writing a full state only when the replay window flushes.

4 REPLAYSSM FOR STANDARD DECODING

  • Algorithm 1 appends each token’s inputs, reconstructs state, and computes output; only the flush branch stores a full state.For GDN, each buffer record stores the correction u_j in place of v_j, and the flush replays the same sum.Source: §4 · p. 5
  • Over an L-step window, ReplaySSM performs one full-state write instead of L, while both methods perform L full-state reads; the smaller buffer traffic makes dominant state traffic roughly halve.The standard-decoding speedup figure compares against tuned conventional vLLM at B = 256, L = 16, with FP32 state.Source: §4 · p. 5, Figure 2 · p. 6
  • A larger replay window amortizes checkpoint writes over more steps but increases buffered-input loads and applications; the implementation uses L = 16.The choice balances checkpoint writes against replay overhead.Source: §4 · p. 5

5 REPLAYSSM FOR SPECULATIVE DECODING

ReplaySSM carries recent inputs in a buffer so speculative decoding can roll back by discarding rejected records. It also verifies draft outputs without materializing intermediate recurrent states.

5.1 CHEAP ROLLBACK WITH CACHED INPUTS

  • Cached inputs make rollback a buffer-pointer update: after verification, ReplaySSM commits only the accepted prefix and overwrites rejected records on the next append.A circular buffer holds committed records and drafts; flushes fold only committed inputs into the checkpoint.Source: §5.1 · p. 6
  • During verification, ReplaySSM appends all draft inputs before acceptance is known, then advances the committed buffer position by the accepted length.Only flush steps write the checkpoint; non-flush steps write small input records.Source: §5.1 · p. 6

5.2 OUTPUT-ONLY VERIFICATION

  • Output-only evaluation computes outputs from the checkpoint and buffered inputs, avoiding materialization of states that no later step reads.For Mamba-2, the method brackets the value-key-query product to compute the output without forming the full state; GDN admits the same form.Source: §5.2 · pp. 6–7

6 REPLAYSSM FOR LOW-PRECISION STATE

ReplaySSM separates the checkpoint used for recurrent updates from the copy used for output computation. This keeps rounding error in the low-precision copy out of later recurrent updates.

6 REPLAYSSM FOR LOW-PRECISION STATE

  • The separation also quantizes the state once per window rather than once per step, although a single low-precision checkpoint still feeds error into later updates.ReplaySSM's separate-precision design addresses that remaining feedback path.Source: §6 · p. 7
  • A high-precision checkpoint carries recurrent updates while a low-precision copy serves non-flush outputs, assigning different precisions to the two roles.Flushes reconstruct from the high-precision checkpoint and independently cast the result into both formats.Source: §6 · p. 7
  • The low-precision copy is never read back into the recurrence, so its rounding error does not accumulate across updates.Rounding error remains in the high-precision path, and non-flush outputs still depend on the approximate copy.Source: §6 · p. 7

7 EVALUATION

ReplaySSM improves standard decoding across four hybrid models and reduces verification latency in speculative decoding. Separate state precisions also reduce output divergence and downstream accuracy losses relative to a shared low-precision checkpoint.

7.1 SETUP

  • The evaluation covers four hybrid models from two families, compares against tuned conventional vLLM decoding, and uses FP32 state for performance experiments.Low-precision experiments measure teacher-forced output KL on PG-19 and generated-answer accuracy on MMLU-Pro, GPQA-Diamond, and Math-L5.Source: §7.1 · pp. 7–8

7.2 STANDARD DECODING

  • 1.20–1.47× whole-decode-step speedup over tuned vLLM at B = 256 and L = 16 across four models.The SSM-kernel speedup is 1.46–1.80×; the reported end-to-end gain is bounded by the SSM share of decode time.Source: §7.2 · p. 8

7.3 SPECULATIVE DECODING

  • At T = 8, conventional verification reaches 4.66× a standard-decoding step, while ReplaySSM verifies eight tokens for 0.89–1.21× of a standard step, flushes included.The ReplaySSM range spans all drafts rejected to all accepted.Source: §7.3 · p. 9

7.4 LOW-PRECISION STATE

  • At 128K tokens, NVFP4 reads with FP32 recurrence reduce output KL from 0.136 to 0.0038 versus a shared NVFP4 checkpoint.FP16+SR recurrence tracks FP32 recurrence closely; the remaining error reflects the low-precision read copy.Source: §7.4 · p. 9
  • NVFP4 reads with FP16+SR recurrence cut mean accuracy losses to 2.36 points on GPQA-Diamond, 0.15 on MMLU-Pro, and 0.10 on Math-L5, versus 13.05, 5.88, and 10.92 points with shared NVFP4+SR.On Nemotron-3-Super-120B-A12B, the two longest GPQA-Diamond quartiles lose 20.9 and 28.2 points with shared NVFP4+SR, compared with 2.0 and 3.4 with NVFP4 reads and FP8+SR recurrence.Source: §7.4 · pp. 9–10

8 RELATED WORK

Prior work optimizes chunkwise SSM computation for training and prefill, while ReplaySSM targets memory-bound autoregressive decoding. It also distinguishes its input-buffer state representation from prior rollback and state-quantization approaches.

Chunkwise SSM algorithms and prefix caching

  • Chunkwise SSM algorithms optimize training and prefill throughput, whereas ReplaySSM targets decoding traffic; prefix caching addresses repeated prefill and can complement it.Marconi adds prefix caching for hybrid models.Source: §8 · p. 10

SSM speculative decoding

  • Prior speculative SSM methods use per-token snapshots or replay accepted tokens, incurring extra state traffic; ReplaySSM instead rolls back by moving a buffer pointer.The cited comparison describes snapshot writes as a concurrency cap at serving batch sizes.Source: §8 · p. 10
  • ReplaySSM's checkpoint plus recent-input buffer represents the state itself; any buffer prefix gives the state after that many accepted tokens, so rollback needs no recurrence replay.Most steps read the state once and write nothing; the checkpoint is written only at a flush.Source: §8 · p. 10

SSM state quantization

  • Earlier state-quantization approaches still feed a quantized state into the next update; ReplaySSM keeps recurrence on the high-precision checkpoint and reads the low-precision copy only for outputs.The passage notes that decoupled scales from Q-Mamba could further reduce read error.Source: §8 · p. 11

9 CONCLUSION AND FUTURE WORK

ReplaySSM reduces recurrent-state traffic and provides pointer-based rollback with output-only verification. The appendix compares speculative state-recovery methods at serving batch size and describes their per-step state traffic.

9 CONCLUSION AND FUTURE WORK

  • ReplaySSM roughly halves dominant state traffic, makes speculative rollback a buffer-pointer move, and keeps low-precision read-checkpoint error out of later updates.In vLLM, standard decoding is 1.20–1.47× faster and speculative decoding reaches about 1.8× vLLM throughput on two 120B-class models.Source: §9 · p. 11

A SPECULATIVE DECODING COMPARED WITH PRIOR METHODS

  • At B = 256, the comparison uses one B300, FP32 state, one draft chain per request, and kernels at their best swept launch configurations.The batch holds 1.07 GB of state per layer; STree is measured as its algorithm without the released path’s two state copies.Source: A SPECULATIVE DECODING COMPARED WITH… · p. 15

A.1 PER-STEP OPERATIONS AND STATE TRAFFIC

  • ReplaySSM reads the checkpoint at each step but writes it only on a flush; vLLM instead stores a full state snapshot for each of the T tokens.ReplaySSM uses 1R on a non-flush step and 1R + 1W on a flush; vLLM uses 1R + TW, independent of acceptance.Source: §A.1 · pp. 15–16

A.2 RESULTS

At serving batch size, state recovery determines whether speculative decoding pays off against vLLM snapshots. ReplaySSM is faster than the compared prior methods, although moving fewer state bytes alone does not ensure lower latency.

A.2 RESULTS

  • 3.3–3.9× faster than STree and 2.5–3.6× faster than The Mamba in the Llama when every token is committed; with one token committed, the gains are 3.5–5.3× and 2.3–5.1×.Neither prior method beats vLLM’s snapshots at this batch size.Source: §A.2 · p. 16
  • At serving batch size, state traffic is the bottleneck and recovery cost decides whether speculation pays off; none of the prior methods beats vLLM snapshots.Source: §A.2 · p. 16
  • At T ≥4, STree and The Mamba in the Llama move 1.6–2.9× fewer state bytes than vLLM’s kernel, yet are not faster.Their serial recovery recurrences limit occupancy and leave much of DRAM bandwidth idle, increasingly so as T grows.Source: §A.2 · p. 16

B GDN DERIVATIONS AND ALGORITHMS

The appendix derives and implements ReplaySSM for Gated DeltaNet (GDN), including speculative verification and separate checkpoint precisions. It also analyzes state traffic, output-only computation, and replay-window choices.

B.1 OUTER-PRODUCT AND UNROLLED FORMS

  • GDN’s correction u_t is the part of v_t not already held by the decayed state, scaled by β_t; caching corrections lets ReplaySSM reuse that state-dependent computation.The unrolled form reconstructs the state from the checkpoint and cached records.Source: §B.1 · p. 17

B.2 STANDARD DECODING

  • GDN standard decoding reconstructs the state before the token because its correction depends on S_{t−1}k_t, then appends the correction record and advances the state.Source: §B.2 · p. 17

B.3 SPECULATIVE DECODING

  • GDN speculative verification obtains draft corrections with one T × T triangular solve, then computes outputs without materializing intermediate states.Draft records follow committed records; accepted drafts are committed by moving the buffer pointer.Source: §B.3 · p. 18

B.4 LOW-PRECISION STATE

  • Folding cached GDN corrections computed from a low-precision checkpoint can carry rounding error into later windows; re-deriving corrections from cached raw inputs and the high-precision checkpoint keeps that error out of the recurrence.The correction-folding approach is stable only under a data-dependent decay condition.Source: §B.4 · p. 19

C.1 ARITHMETIC INTENSITY / C.2 SSM AND ATTENTION TRAFFIC

  • Conventional decoding stays under 1 FLOP/byte, while calculated SSM/attention traffic crosses at 81,920 tokens and the measured crossover is near 100K.Source: §C.1 · p. 21, §C.2 · p. 21

C.3 MEMORY TRAFFIC

  • For Nemotron-3-Super-120B-A12B with L = 16, ReplaySSM moves 0.58× conventional traffic, a 1.7× reduction.The checkpoint term alone yields a traffic ratio just above one half; buffer traffic grows quadratically with L.Source: §C.3 · p. 22

C.4 FLOPS FOR OUTPUT-ONLY COMPUTATION

  • For T = 4 and h = 12, output-only verification needs over 4× fewer FLOPs than materializing states.Its remaining dvdk-scale term is one checkpoint readout against the T queries.Source: §C.4 · p. 23

C.5 REPLAY-WINDOW TRADEOFF

  • For Nemotron, the traffic-only window prediction is L* ≈13 for standard decoding and L* ≈24 for T = 4 with mean acceptance 3.5; the sweep favors L = 16 across four models.The prediction is approximate because FLOPs and on-chip buffer storage also grow with L.Source: §C.5 · p. 23

C.6 QUANTIZATION ERROR PROPAGATION

Rounding errors propagate through recurrent updates, with their accumulation depending on decay and rounding behavior. ReplaySSM reduces this accumulation, while the analysis predicts scaling rather than error magnitude.

C.6 QUANTIZATION ERROR PROPAGATION

  • Stochastic rounding is unbiased in expectation but retains variance, whereas round-to-nearest can introduce systematic bias.Stochastic-rounding variance is p(1 − p)δ², at most δ²/4; round-to-nearest error is at most δ/2.Source: §C.6 · pp. 23–24
  • For α < 1, accumulated deviation converges to b/(1 − α) in mean and σ²/(1 − α²) in variance, with slow-decay heads accumulating the most error.The analysis assumes independent step errors with fixed bias and variance; both limits grow without bound as α approaches 1.Source: §C.6 · p. 24
  • With separate precisions, low-precision rounding never enters the recurrence, making non-flush output error a single rounding independent of α and sequence length.A high-precision checkpoint drives recurrence updates while the low-precision copy is read only for outputs.Source: §C.6 · p. 24
  • The analysis predicts how error scales with decay and window size, not its magnitude, and assumes independent errors with fixed bias and variance.Real rounding errors depend on rounded values and can partly cancel; the paper therefore measures output error using KL divergence against FP32-state decoding.Source: §C.6 · p. 24

D SERVING IMPLEMENTATION

ReplaySSM’s serving implementation adjusts buffer flushing for speculative decoding and handles row-varying flush and commit decisions within captured CUDA graphs.

D SERVING IMPLEMENTATION

  • ReplaySSM is implemented in vLLM for Mamba-2 and GDN, though its design does not depend on vLLM.The implementation appendix describes buffer layout, paged-state placement, and cursor bookkeeping.Source: D SERVING IMPLEMENTATION · p. 24

D.1 FLUSH RULE UNDER SPECULATIVE DECODING

  • The speculative flush rule triggers when h + T > L, preventing new records from overwriting committed records in a ring of capacity L + T.This early flush can occur even when partial acceptance would have left room; the interval is roughly L/a steps, with higher acceptance flushing more often.Source: §D.1 · pp. 24–25

D.2 CUDA GRAPHS

  • Captured CUDA graphs handle row-varying flush and commit decisions at runtime, with a device-side commit kernel updating each row’s buffer pointer and flush flags.This avoids reading accepted counts back to the host, which would require synchronization at every step.Source: §D.2 · p. 25

E.1 EVALUATION SETUP

The evaluation covers four hybrid models on single GPUs, comparing ReplaySSM with tuned vLLM kernels under standard and speculative decoding and testing low-precision checkpoints. Results examine throughput, output KL, and downstream accuracy across batch sizes, replay windows, and generation lengths.

E.2.1 BATCH-SIZE SWEEP

  • ReplaySSM standard decoding gains are concentrated at serving batches, where conventional kernels become memory-bound; at B = 1, speedups are below parity.The tuned baseline reaches 84–86% of peak DRAM bandwidth from B ≈128, while long-sequence end-to-end gains shrink as attention dominates.Source: §E.2.1 · p. 26, Figure 8 · p. 27

E.2.2 REPLAY-WINDOW SWEEP

  • At B = 256, L = 16 gives the lowest amortized kernel latency across all four models, balancing flush frequency against per-step buffer work.The traffic analysis predicts an optimum near L = 13, making 16 the nearest power of two.Source: §E.2.2 · p. 27, Figure 9 · p. 27

E.3.1 SERVING THROUGHPUT ACROSS BATCH SIZE AND VERIFICATION LENGTH

  • At larger batches, ReplaySSM speculative throughput grows with both batch size and verification length, while conventional speculation saturates near B ≈128 and can fall below standard decoding at large T.At B = 1, the methods are close: 0.88–1.17× on Nemotron-3-Super-120B-A12B-NVFP4 and 0.87–1.12× on Qwen3.5-122B-A10B-NVFP4; GDN kernel latency is 2.09× a standard step at T = 2 and 6.40× at T = 8 for conventional decoding.Source: §E.3.1 · p. 28, Table 8 · p. 28, Table 9 · p. 28

E.4.2 LATENCY OF LOW-PRECISION READS

  • With the FP16+SR recurrence checkpoint fixed, NVFP4 read checkpoints make non-flush decoding kernels 3.83× faster on Mamba-2 and 2.76× faster on GDN than FP32 reads at B = 256.Across a full window, flushes still read and write the higher-precision checkpoint, so gains are smaller.Source: §E.4.2 · p. 29, Figure 11 · p. 29

E.4.3 OUTPUT KL ACROSS MODELS

  • At the last plotted generation length, NVFP4 reads with FP32 recurrence show 13–29× lower output KL than a shared NVFP4 checkpoint.Separate recurrence reduces KL growth rather than eliminating it; the two 120B-class models also have higher numerical floors than the 4B models.Source: §E.4.3 · p. 30, Figure 12 · p. 30

E.4.4 WINDOW-SIZE SENSITIVITY OF KL

  • From L = 1 to 64, shared FP8 and NVFP4 checkpoint KL falls 12× and 23× at N = 1K, while separate configurations barely change across L = 4 to 64.A longer window reduces shared-checkpoint rounding events; separate read checkpoints are regenerated from the unquantized recurrence and are not fed back.Source: §E.4.4 · p. 31, Figure 13 · p. 31

E.4.5 STOCHASTIC ROUNDING AGAINST ROUND-TO-NEAREST

  • At N = 1K, round-to-nearest yields 0.40–0.51× the stochastic-rounding error for NVFP4, but its error can overtake stochastic rounding as generation length grows.On the two GDN models, round-to-nearest makes FP8 and NVFP4 states non-finite by 128K and 64K tokens, respectively, while stochastic rounding stays finite.Source: §E.4.5 · p. 31, Table 11 · p. 32

E.4.6 DOWNSTREAM ACCURACY

  • Separate checkpoints keep the three larger models within 2 points of FP32 across benchmarks, while shared NVFP4 losses grow with generation length and can include truncation.On RULER-128K, every configuration stays within 0.3 points of FP32, including shared NVFP4; within-benchmark length bins show growing shared-NVFP4 gaps, while NVFP4 reads with FP8+SR recurrence stay within the FP32 interval in 41 of 48 bins.Source: §E.4.6 · pp. 32–33, Table 12 · p. 33, Table 13 · p. 34

Related papers

Loading 2610.13778v1…