[RFC] Mamba2: exact-replay decode so that prefill, chunked prefill and decode produce identical bits

[NolenLiang]

1. Problem

Mamba2 layers run two different kernels: the chunked scan (mamba_chunk_scan_combined_varlen, SSD) for prefill and the recurrent update (selective_state_update, SSU) for decode. Their outputs are not bit-identical, so a request's Mamba-layer outputs depend on how it was scheduled (prefill vs decode, chunked-prefill split points, preemption and recompute). This blocks batch-invariant / deterministic serving for hybrid models (#54993 discussion; #38561 closed on the same wall) and blocks exact train-inference matching for RL, where the trainer computes the whole sequence with the chunked form while sampled tokens came from the recurrent form.

2. Evidence (direct kernel calls, GB200, vLLM 0.28.0, fp64 reference; realistic Mamba2 init A∈[-16,-1], softplus(dt_bias)∈[1e-3,1e-1]; one seed, two shapes: 32 heads/1 group and 64 heads/8 groups, headdim 64, dstate 128, chunk 256, 2048 tokens)

1. SSD vs SSU differ from token 0. Mean |Δ| is 0.07% (fp32 SSM cache) to 0.22% (bf16 cache) of mean |output|; 22–45% of output elements change bits; max |Δ| 0.125 on outputs of magnitude up to ~30. Against the fp64 reference both kernels sit near the bf16 output-rounding floor (floor 0.141%; SSD 0.166–0.169%; SSU 0.141% with fp32 state, 0.245–0.269% with bf16 state, growing with position). Neither is "wrong"; they are two roundings of the same math.
2. SSU is batch-invariant at these shapes: batch 1/4/16/64/256 give bit-identical outputs and states although try_get_optimal_ssm_config selects different (BLOCK_SIZE_M, num_warps).
3. SSD is prefix-invariant (token t's output does not depend on later tokens; prefixes 1..511 all bitwise equal to the 512-token run) and packing-invariant (varlen packing with another sequence changes no bits).
4. Chunked-prefill splits: aligned splits (multiples of the chunk size) are bitwise-invariant only with mamba_ssm_cache_dtype=float32. With the default bf16 cache, [256|257] already differs at token 512 and every token after the next boundary differs (88/88), because _state_passing_fwd carries the inter-chunk state in fp32 inside one call and rounds only on store, while a resumed call starts from the rounded value; the slowest heads keep 42–73% of their state across a 256-token chunk. Unaligned splits change every token after the split (97–98% of elements bit-identical, max |Δ| 0.031 on outputs of mean magnitude ≈1).
5. Replaying the current partial chunk through the SSD kernels from the fp32 boundary state reproduces the single-shot prefill bitwise: 64/64 steps inside a chunk and 344/344 steps across a chunk boundary (bf16 cache: 0/88 after the boundary).
6. causal_conv1d_fn (prefill) and causal_conv1d_update (decode) are already bit-identical (1792/1792 tokens, two widths); [Bugfix][Kernel] Promote BF16 causal-conv operands before accumulation #52905 changes both symmetrically.
7. Engine-level (prototype, vLLM v0.28.0 + patch, AntonV/mamba2-130m-hf, 24 Mamba2 layers, bf16, fp32 SSM cache, eager, TP1, one request; random token stream, P=600 prompt + K=300 forced continuation; full raw-logits rows compared with torch.equal). Baseline: step-by-step decode differs from single-shot prefill in 299/299 rows (5.6% of logit elements bit-identical, mean |Δ| 1.3 logits); unaligned chunked prefill (max_num_batched_tokens=100) differs in 800/900 rows from the first split point. With exact replay (with and without VLLM_BATCH_INVARIANT=1): 0/299 decode rows differ (168 inside the first partial chunk, 131 after crossing the boundary at 768), 0/900 chunked-prefill rows differ; max |Δ| = 0. Single-shot prefill logits are unchanged by the mode. VLLM_BATCH_INVARIANT=1 without exact replay fails closed at engine init (batch_invariant mode is not supported for MAMBA2_ATTN). Cost in this launch-bound setting: 300 decode steps took 6.5 s (SSU) vs 11.9 s (replay), ≈ +0.75 ms per layer per step, consistent with the direct-kernel timing.

3. Proposal: exact-replay decode

State representation per sequence and layer = fp32 boundary state (the SSM state at the last multiple of chunk_size, i.e. the existing SSM cache with mamba_ssm_cache_dtype=float32) + partial-chunk input buffer holding x, dt, B, C for the tokens since that boundary (at most chunk_size tokens).

• Decode step: one mamba_chunk_scan_combined_varlen call over all decoding sequences, each sequence = its buffered partial chunk plus the new token, initial_states = boundary state, chunk metadata built with num_computed = boundary position (the existing varlen API supports this directly). The output of the last token of each sequence is the layer output; the new token's inputs are appended to the buffer; when a chunk fills, the returned final state becomes the new boundary state and the buffer is emptied.
• Prefill / chunked prefill / recompute after preemption: any SSD call for a sequence starts at its boundary: the buffered tokens are prepended, their outputs discarded, the trailing partial chunk's inputs written to the buffer, and the boundary state taken from the last full chunk (return_intermediate_states, as prefix caching already does).
• Because SSD is prefix- and packing-invariant at the kernel level (Evidence 3), the design goal is that prefill, chunked prefill (aligned or not), decode and post-preemption recompute produce the same bits regardless of decode batch composition. The engine-level prototype has verified this for a single request (Evidence 7) and the kernel-level test covers mixed-offset multi-request batches; engine-level mixed-offset batches are not yet tested.
• Train-inference exactness additionally needs the trainer to run the same chunk size and the same SSD kernels/configs (the Megatron-LM "mLite impl=vllm" pattern).

Relationship to existing code: ReplaySSM (--use-replayssm) already keeps a 16-slot ring of x/dt/B per sequence and flushes the state every 16 steps to save the per-step full-state store. The prototype does not extend that ring: ReplaySSM's dt_cache stores softplus(dt + dt_bias) whereas the SSD kernels take raw dt, and its cursor counts steps since the last flush rather than the position in the chunk grid. Exact replay therefore adds four separate token-major cache tensors (x, raw dt, B, C; chunk_size entries per slot) alongside conv_state/ssm_state, and the two modes are mutually exclusive.

3b. Relationship to the fixed-chunking scheduler work (#54993) and packed-history recovery

#54993 fixes the physical SSD chunking of prefill in the scheduler; its author's follow-up prototype (#27433, 2026-09-03) repairs the preemption gap by making recovery reproduce the SSU decode numerics: generated history is replayed with multi-token SSU and rounded to the cache dtype after every token (28/28 recoveries bit-identical to uninterrupted decode on Granite 4.0 H 350M; 64 ms per 255-token recovery). That guarantees decode == recovered decode at no per-step cost. Exact replay guarantees the mirror property — decode reproduces the SSD prefill numerics — so that prefill(prefix + token) == decode(token), which is what tests/v1/determinism/test_batch_invariance.py::test_decode_logprobs_match_prefill_logprobs checks and what training/inference alignment needs (the trainer scores generated tokens with the chunked scan). The two are complementary: exact replay is a separate opt-in that can sit on top of fixed chunking (the combination has not been tested); for serving-only determinism the packed-history route is the cheaper choice.

4. Requirements and costs

• mamba_ssm_cache_dtype=float32 is required (Evidence 4–5); the mode fails closed otherwise.
• The chunk size must be identical for prefill and decode (and for the trainer if train-inference exactness is the goal).
• Compute: at batch 64 with 128-token partial chunks, one SSD replay call took 0.80–1.05 ms versus 0.37–0.40 ms for one SSU step (Python-level timing, SSM scan only; kernel-level and end-to-end numbers are part of the prototype).
• Memory (32 heads/1 group, headdim 64, dstate 128, 27 Mamba layers, TP1, chunk 256): the partial-input buffer is ≈1.16 MiB per layer per sequence (≈31 MiB total); the fp32 boundary state adds 1 MiB per layer (≈27 MiB). Total resident ≈58 MiB per sequence; increment over the default bf16 state (13.5 MiB) ≈45 MiB per sequence, which scales linearly with concurrent sequences and must be budgeted in the KV-cache planner. A smaller shared chunk size (e.g. 64) cuts the buffer 4× at the cost of prefill efficiency and requires the trainer to use the same chunk size.

5. Out of scope for the first version

Speculative decoding (the num_accepted_tokens SSU path), prefix caching modes all/align beyond the aligned-boundary argument above, KV connectors, CUDA graphs for the decode replay call (data-dependent chunk metadata), TP > 1, RMSNormGated and projection GEMM invariance (covered by the generic batch-invariant work).

6. Phase-1 prototype (default-off, local)

Scope: TP1, eager mode, synchronous scheduling, one model, a fixed continuation token stream (forced decoding), no spec decode / prefix cache / KV connector / CUDA graph.

Validation: for the same token stream compare (a) single-shot prefill raw logits with (b) step-by-step exact-replay decode raw logits and (c) chunked prefill with unaligned split points and a forced preemption+recompute. Gate: torch.equal on the compared raw logits, mismatch_count == 0, covering within-chunk, across-boundary and post-preemption tokens. Only after this passes: kernel-level profiling and the full cache/scheduler integration.

Status (2026-09-05): single-request phase 1a passed on Mamba2ForCausalLM (see Evidence 7; prototype branch: https://github.com/NolenLiang/vllm/tree/mamba2-exact-replay, rebased on main). Mixed-offset multi-request batches (unaligned prefill splits, resumes and interleaved decode in the same batches) are verified bitwise at the kernel level (tests/kernels/mamba/test_mamba2_exact_replay.py on that branch, contiguous and paged buffer layouts); still to verify before phase 1 is complete: the same at the engine level, and real scheduler preemption on a hybrid model. Scheduler-driven preemption cannot be triggered on a pure Mamba model (one state slot per request), so the post-preemption path is covered by the unaligned chunked-prefill resume, which runs the same num_computed-based code; real preemption is to be exercised on a hybrid model (Nemotron-H). Implementation notes: the partial-chunk buffers are separate cache tensors (x, raw dt, B, C; token-major; chunk_size long) rather than the ReplaySSM ring, because dt_cache stores softplus(dt + bias) and its cursor counts flush steps, not chunk positions; buffers live in the paged mamba cache, so slot and position are indexed separately; models opt in via a SupportsMambaExactReplay protocol so the model-level and layer-level page-size computations stay consistent; Mamba2AttentionBackend.supports_batch_invariance() is true only in this mode.

7. Open questions

1. Interaction of the partial-chunk buffers with mamba_block_size and the prefix-caching modes (all / align), which are out of scope for the first version.
2. Resolved in the prototype: the mode is a separate cache option (--mamba-exact-replay); VLLM_BATCH_INVARIANT=1 accepts the Mamba2 backend only when it is set.
3. Autotune: the SSD sub-kernels use @triton.autotune keyed on shapes; run-to-run determinism across processes needs pinned configs in aligned mode.

CC @SyaOtiLan @jzakrzew (from the #54993 discussion)