ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding

來源原標題

ReplaySSM: Rethinking State Storage for Faster Hybrid-Model Decoding

作者

Michelle Horton

原始標籤

Agentic AI / Generative AI · AI Inference · Inference Performance · Nemotron · research · TensorRT-LLM · vLLM

原始連結

https://developer.nvidia.com/blog/replayssm-rethinking-state-storage-for-faster-hybrid-model-decoding/

摘要原文

Hybrid models interleave state-space model (SSM) layers with a few attention layers. NVIDIA Nemotron 3 , Qwen3.5 through Qwen3.8 , Kimi-K3 , and GLM-5.3-Flash are all built this way, with Mamba-2 , Gated DeltaNet (GDN) , or Kimi Delta Attention (KDA) as the SSM layer. SSM layers are designed to make decoding fast: they use a fixed-size state in place of a key-value (KV) cache that grows with context length. At serving batch sizes, however, the SSM layers, not attention, can be the latency bottleneck up to 100K tokens of context. SSM layers also complicate two techniques that reliably speed up attention decoding: speculative decoding and quantization. ReplaySSM makes SSM layers the fast part again. It represents the SSM state as a checkpoint plus a buffer of recent inputs. This separates output computation from state updates and addresses three challenges in current SSM decoding: state traffic, rollback for speculative decoding, and error accumulation in a quantized state. ReplaySSM is joint work between NVIDIA and the Dao AI Lab at Princeton. The paper, ReplaySSM: Cache SSM Inputs, Not State , details all three benefits and presents the full evaluation. It is available on arXiv, with code on GitHub . It is integrated into NVIDIA TensorRT-LLM, vLLM, and FlashInfer, and also runs in production in GLM-5.3-Flash and Kimi-K3. Figure 1. Conventional decoding loads and stores the full state at every step. ReplaySSM caches recent inputs, computes outputs from the checkpoint plus buffer, and writes the checkpoint only when the buffer fills Why SSM layers can become a decoding bottleneck Take the Mamba-2 layer used in NVIDIA Nemotron 3 Ultra as an example. It decodes one token with the following recurrence: S t = α t S t − 1 + v t k t ⊤ , o t = S t q t , S_t = \alpha_t S_{t-1} + v_t k_t^\top, \qquad o_t = S_t q_t, where S t S_t is the state, α t \alpha_t is a data-dependent decay, and v t v_t , k t k_t , and q t q_t are the token’s inputs. Each step loads S t − 1 S_{t-1} from high-bandwidth memory (HBM), applies the decay and a rank-1 update, computes the output with q t q_t , and stores S t S_t back. GDN and KDA add a delta-rule correction, but their updates can be written in the same outer-product form, so what follows applies to them too. During decoding, the state does two jobs. It is read to produce this token’s output, and it is updated to become the next token’s state. Conventional decoding does both at once: each step updates the state, reads the output from it, and stores it for the next step. All three challenges come from this fusion. Figure 2. Three challenges of conventional SSM decoding and the ReplaySSM answer to each. Left: decode-step latency breakdown for Nemotron-3-Super-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU State memory traffic An SSM decoding step loads and stores the full state while performing little arithmetic on it, so the kernel is bound by HBM bandwidth, not compute. At serving batch sizes, that makes the SSM layers, not attention, the largest component of the decode step up to 100K tokens on Nemotron-3-Super-NVFP4 at batch size 256 on one NVIDIA Blackwell Ultra B300 GPU (Figure 2, left). Two factors drive this: Layer count: The model has five times as many SSM layers as attention layers. State size: One layer’s FP32 state is 4 MiB per sequence, while an FP8 KV entry is 0.5 KiB per token, so an SSM layer’s state is the size of an attention layer’s KV cache for 8K tokens. Rollback for speculative decoding Speculative decoding has a drafter propose T tokens, verifies them in one pass through the target model, and rolls back to the last accepted token. Attention naturally supports rollback by moving the KV cache pointer. But once an SSM has summarized the drafts into its state, the update has no easy inverse, so rollback becomes expensive. Serving engines such as vLLM work around this by storing a state snapshot after every draft token and picking the one that matches the accepted prefix. That is one read and T writes per step on a memory-bound path, and the snapshots cap how many requests fit in GPU memory. As a result, speculative decoding, a reliable speedup for Transformers, barely helps SSMs at serving batch sizes. State quantization The quantized state produces the output and is the input to the next update, so its rounding error re-enters the recurrence at every step. A KV cache has no such loop: a quantized entry is read many times but never written back into itself, which is why KV caches tolerate lower precision than SSM states. Qwen3.5-122B-A10B-NVFP4, for example, uses the NVIDIA NVFP4 4-bit floating-point format for its weights, with an FP8 KV cache and an FP32 SSM state. How ReplaySSM represents the state The SSM recurrence can be unrolled over a window of h tokens: S t = α t S t − 1 + v t k t ⊤ = γ h S c k p t + ∑ j = 1 h γ j + 1 : h v j k j ⊤ , S_t = \alpha_t S_{t-1} + v_t k_t^\top = \gamma_h\, S_{\mathrm{ckpt}} + \sum_{j=1}^{h} \gamma_{j+1:h}\, v_j k_j^\top, where S c k p t S_{\mathrm{ckpt}} is the state before the window and γ \gamma is the cumulative decay. Conventional decoding follows the left side step by step. ReplaySSM uses the right side: it represents the state as a checkpoint plus a small buffer of recent input records ( α j \alpha_j , v j v_j , k j k_j ) for Mamba-2. On most steps, which we call non-flush steps, it appends the current inputs to the buffer and computes the output from the checkpoint and the buffer without storing a state. When the buffer of L records fills, a flush step folds the buffered inputs into the checkpoint, writes the new checkpoint to HBM once, and clears the buffer. The outputs match conventional decoding up to floating-point differences. ReplaySSM’s approach can be summarized in three questions: 1 What does ReplaySSM do? It caches recent SSM inputs and represents the state as a checkpoint plus those inputs. 2 What does this change? It separates output computation from state updates: outputs are computed every step, while the checkpoint is updated once per window. 3 What does this achieve? It reduces state memory traffic, simplifies speculative rollback, and supports low-precision state reads without propagating their rounding error through recurrent updates. How we evaluated ReplaySSM We evaluate four hybrid models from two families, Nemotron 3 with Mamba-2 layers and Qwen3.5 with GDN layers, from 4B dense models on one NVIDIA H100 Tensor Core GPU to 120B-class mixture-of-experts (MoE) models on one NVIDIA Blackwell Ultra B300 GPU (Table 1). We implemented ReplaySSM in vLLM and compared it with vLLM’s own kernels at their best launch configuration. Performance runs keep the SSM state in FP32, the default precision for these models, and use a window of L = 16, the best-performing window size for all four models. Speculative decoding tests use each model’s multi-token prediction (MTP) head to generate draft tokens. Model SSM layer Weights GPU SSM / attention layers 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 Table 1. Evaluated models and hardware. Each model runs on one GPU without tensor parallelism Benefit 1: Reduce state memory traffic Over a window of L tokens, conventional decoding reads the state L times and writes it L times. ReplaySSM still reads the checkpoint on every step, because the output needs it, but writes it only once, during the flush. A non-flush step has one job—computing the output—and writes only the new input record, a few vectors instead of a full matrix. This roughly halves the dominant state traffic. At serving batch sizes, memory bandwidth limits the conventional kernel’s performance, so reducing state memory traffic speeds up decoding. At batch size 1, other sources of latency dominate, and ReplaySSM provides no speedup in these tests. At batch size 256, the kernel runs 1.46× to 1.