@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 866fea2b99

[Kernel] ReplaySSM: cache SSM inputs for faster Mamba2 standard decode (#48018)

Signed-off-by: Johnny-Liou <a897111@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: tomeras91 <57313761+tomeras91@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Johnny-Liou committed Jul 24, 2026 at 12:39 UTC 866fea2b9900bf49d552c205d2eaac4716fb63ac

23 files changed +2954 -31

benchmarks/replayssm/e2e_decode_speedup.py new

+267

@@ -0,0 +1,267 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""End-to-end autoregressive decode benchmark: ReplaySSM vs the standard SSM kernel.

4 +

5 +Loads a hybrid Mamba2 model, replicates one prompt across the batch, and times a

6 +long greedy decode (CUDA graphs on) once with the standard kernel and once with

7 +ReplaySSM, then reports the per-step / throughput speedup. The two modes run in

8 +separate subprocesses so each gets a clean CUDA context.

9 +

10 +The FlashInfer FP4-MoE autotuner is disabled by default (it is unstable under

11 +CUDA-graph capture on the pre-release Blackwell FP4 path); pass

12 +--no-disable-flashinfer-autotune for non-FP4 models.

13 +

14 +Examples:

15 + python e2e_decode_speedup.py --model-id nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16

16 + python e2e_decode_speedup.py --dtype auto --buffer-len 16 \

17 + --model-id nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 # B300 NVFP4

18 +"""

19 +

20 +import argparse

21 +import json

22 +import os

23 +import subprocess

24 +import sys

25 +import time

26 +

27 +DEFAULT_PROMPT = "My cat wrote all this CUDA code for a new language model and"

28 +

29 +MODE_LABEL = {"standard": "standard", "replayssm": "ReplaySSM"}

30 +

31 +

32 +def parse_args():

33 + p = argparse.ArgumentParser(

34 + description="E2E decode speedup: ReplaySSM vs the standard SSM kernel."

35 + )

36 + p.add_argument("--model-id", default="nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16")

37 + p.add_argument("--prompt", default=DEFAULT_PROMPT)

38 + p.add_argument("--batch-size", type=int, default=256)

39 + p.add_argument("--num-steps", type=int, default=1000)

40 + p.add_argument("--warmup-steps", type=int, default=128)

41 + p.add_argument("--repeats", type=int, default=1)

42 + p.add_argument(

43 + "--buffer-len", type=int, default=16, help="ReplaySSM input-buffer length."

44 + )

45 + p.add_argument(

46 + "--dtype",

47 + default="bfloat16",

48 + choices=["bfloat16", "float16", "float32", "auto"],

49 + )

50 + p.add_argument("--gpu-memory-utilization", type=float, default=0.9)

51 + p.add_argument("--max-model-len", type=int, default=None)

52 + p.add_argument(

53 + "--disable-flashinfer-autotune",

54 + action=argparse.BooleanOptionalAction,

55 + default=True,

56 + help="Disable the FlashInfer FP4-MoE autotuner (default: on). "

57 + "It is unstable under CUDA-graph capture on the "

58 + "pre-release Blackwell FP4 path; pass "

59 + "--no-disable-flashinfer-autotune for non-FP4 models.",

60 + )

61 + p.add_argument(

62 + "--mamba-ssm-cache-dtype",

63 + default="auto",

64 + choices=["auto", "float32", "float16", "bfloat16"],

65 + help="SSM state dtype (both modes). 'auto' = config-driven; "

66 + "'float32' = fp32 state, 'bfloat16' = s16 state.",

67 + )

68 + p.add_argument(

69 + "--baseline-ssm-config",

70 + default="",

71 + help="Pin the STANDARD baseline's SSM launch config as "

72 + "'bsm,nw' via override_ssm_config (forces the in-process "

73 + "engine so the override reaches the kernel). Empty = off.",

74 + )

75 + p.add_argument(

76 + "--worker",

77 + choices=["standard", "replayssm"],

78 + default=None,

79 + help=argparse.SUPPRESS,

80 + )

81 + return p.parse_args()

82 +

83 +

84 +def resolve_max_model_len(args) -> int:

85 + if args.max_model_len is not None:

86 + return args.max_model_len

87 + return args.num_steps + 256

88 +

89 +

90 +def run_worker(args):

91 + # override_ssm_config is a module global; it only reaches the model if the

92 + # engine runs in-process (default V1 spawns a separate EngineCore). Force it.

93 + if args.worker == "standard" and args.baseline_ssm_config:

94 + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"

95 +

96 + import torch

97 +

98 + from vllm import LLM, SamplingParams

99 +

100 + mode = args.worker

101 + max_model_len = resolve_max_model_len(args)

102 +

103 + llm_kwargs = dict(

104 + model=args.model_id,

105 + tensor_parallel_size=1,

106 + dtype=args.dtype,

107 + max_model_len=max_model_len,

108 + trust_remote_code=True,

109 + enable_prefix_caching=False,

110 + enable_chunked_prefill=False,

111 + max_num_seqs=args.batch_size,

112 + max_num_batched_tokens=max(max_model_len, args.batch_size * 64),

113 + enforce_eager=False,

114 + disable_log_stats=True,

115 + gpu_memory_utilization=args.gpu_memory_utilization,

116 + # SSM state dtype (applies to both standard and ReplaySSM).

117 + mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,

118 + )

119 + if args.disable_flashinfer_autotune:

120 + # FP4-MoE autotuner is unstable under CUDA-graph capture on Blackwell;

121 + # re-enable (--no-disable-flashinfer-autotune) only for non-FP4 models.

122 + llm_kwargs["kernel_config"] = {"enable_flashinfer_autotune": False}

123 + if mode == "replayssm":

124 + llm_kwargs.update(use_replayssm=True, replayssm_buffer_len=args.buffer_len)

125 +

126 + _ssm_cm = None

127 + if mode == "standard" and args.baseline_ssm_config:

128 + from vllm.model_executor.layers.mamba.ops.mamba_ssm import override_ssm_config

129 +

130 + _bsm, _nw = (int(x) for x in args.baseline_ssm_config.split(","))

131 + _ssm_cm = override_ssm_config((_bsm, _nw))

132 + _ssm_cm.__enter__() # active through LLM() graph capture + decode

133 + print(

134 + f"[{mode}] override_ssm_config -> (BLOCK_SIZE_M={_bsm}, num_warps={_nw})",

135 + flush=True,

136 + )

137 +

138 + llm = LLM(**llm_kwargs)

139 + prompts = [args.prompt] * args.batch_size

140 +

141 + def timed_generate(n_tokens):

142 + sp = SamplingParams(

143 + n=1,

144 + temperature=0.0,

145 + ignore_eos=True,

146 + min_tokens=n_tokens,

147 + max_tokens=n_tokens,

148 + )

149 + if torch.accelerator.is_available():

150 + torch.accelerator.synchronize()

151 + t0 = time.perf_counter()

152 + outs = llm.generate(prompts, sp, use_tqdm=False)

153 + if torch.accelerator.is_available():

154 + torch.accelerator.synchronize()

155 + elapsed = time.perf_counter() - t0

156 + produced = min(len(o.outputs[0].token_ids) for o in outs)

157 + assert produced == n_tokens, f"expected {n_tokens} tokens, got {produced}"

158 + return elapsed

159 +

160 + timed_generate(args.warmup_steps)

161 +

162 + best = None

163 + for _ in range(args.repeats):

164 + elapsed = timed_generate(args.num_steps)

165 + tok_s = args.batch_size * args.num_steps / elapsed

166 + per_step_ms = elapsed / args.num_steps * 1e3

167 + print(

168 + f"[{mode}] {elapsed:.3f}s {tok_s:,.0f} tok/s {per_step_ms:.3f} ms/step",

169 + flush=True,

170 + )

171 + if best is None or elapsed < best["elapsed_s"]:

172 + best = {

173 + "mode": mode,

174 + "elapsed_s": elapsed,

175 + "tok_s": tok_s,

176 + "per_step_ms": per_step_ms,

177 + }

178 +

179 + print("RESULT_JSON " + json.dumps(best), flush=True)

180 + if _ssm_cm is not None:

181 + _ssm_cm.__exit__(None, None, None)

182 +

183 +

184 +def run_one_mode(args, mode) -> dict:

185 + cmd = [

186 + sys.executable,

187 + __file__,

188 + "--worker",

189 + mode,

190 + "--model-id",

191 + args.model_id,

192 + "--prompt",

193 + args.prompt,

194 + "--batch-size",

195 + str(args.batch_size),

196 + "--num-steps",

197 + str(args.num_steps),

198 + "--warmup-steps",

199 + str(args.warmup_steps),

200 + "--repeats",

201 + str(args.repeats),

202 + "--buffer-len",

203 + str(args.buffer_len),

204 + "--dtype",

205 + args.dtype,

206 + "--gpu-memory-utilization",

207 + str(args.gpu_memory_utilization),

208 + "--mamba-ssm-cache-dtype",

209 + args.mamba_ssm_cache_dtype,

210 + "--baseline-ssm-config",

211 + args.baseline_ssm_config,

212 + ]

213 + cmd.append(

214 + "--disable-flashinfer-autotune"

215 + if args.disable_flashinfer_autotune

216 + else "--no-disable-flashinfer-autotune"

217 + )

218 + if args.max_model_len is not None:

219 + cmd += ["--max-model-len", str(args.max_model_len)]

220 +

221 + result = None

222 + proc = subprocess.Popen(

223 + cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1

224 + )

225 + for line in proc.stdout:

226 + sys.stdout.write(line)

227 + sys.stdout.flush()

228 + if line.startswith("RESULT_JSON "):

229 + result = json.loads(line[len("RESULT_JSON ") :])

230 + proc.wait()

231 + if proc.returncode != 0:

232 + raise RuntimeError(f"mode '{mode}' worker exited with {proc.returncode}")

233 + if result is None:

234 + raise RuntimeError(f"mode '{mode}' produced no RESULT_JSON line")

235 + return result

236 +

237 +

238 +def main():

239 + args = parse_args()

240 + if args.worker is not None:

241 + run_worker(args)

242 + return

243 +

244 + print(

245 + f"model={args.model_id} batch_size={args.batch_size} "

246 + f"steps={args.num_steps} buffer_len={args.buffer_len} dtype={args.dtype}"

247 + )

248 +

249 + std = run_one_mode(args, "standard")

250 + fla = run_one_mode(args, "replayssm")

251 + speedup = std["per_step_ms"] / fla["per_step_ms"]

252 +

253 + print()

254 + header = f"{'mode':<10}{'ms/step':>12}{'tok/s':>16}{'wall (s)':>12}"

255 + print(header)

256 + print("-" * len(header))

257 + for r in (std, fla):

258 + print(

259 + f"{MODE_LABEL[r['mode']]:<10}{r['per_step_ms']:>12.3f}"

260 + f"{r['tok_s']:>16,.0f}{r['elapsed_s']:>12.3f}"

261 + )

262 + print("-" * len(header))

263 + print(f"speedup (standard / ReplaySSM, per step): {speedup:.3f}x")

264 +

265 +

266 +if __name__ == "__main__":

267 + main()

tests/kernels/mamba/test_replayssm_prefill_decode_equivalence_mamba2.py new

+244

@@ -0,0 +1,244 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Prefill == decode equivalence for the Mamba2 ReplaySSM kernels.

4 +

5 +The SSM recurrence is path-independent: vLLM's chunked prefill (the SSD kernel

6 +``mamba_chunk_scan_combined_varlen``) and step-by-step decode must produce the

7 +same per-position outputs and final state over a sequence. This file feeds one

8 +set of inputs through the production dt flow (raw dt + softplus + a per-head

9 +dt_bias, applied inside each kernel) to:

10 +

11 + * the exact fp32 step recurrence (``selective_state_update_ref``) -- the

12 + ground truth,

13 + * the chunked prefill kernel,

14 + * the baseline decode kernel,

15 + * the ReplaySSM output_only decode kernel,

16 +

17 +and checks all of them agree. Prefill (a chunked scan) and decode (a step

18 +recurrence) are different code paths, so they differ numerically: the chunked

19 +scan carries ~2e-2 (fp32) / ~4e-2 (bf16) vs the exact recurrence, far above the

20 +near-exact decode. We therefore anchor every path on the exact recurrence at

21 +SSD-level tolerances (the same regime as ``test_mamba_ssm_ssd.py``), which the

22 +chunked scan sets, keyed off the activation dtype.

23 +

24 +State and activation/buffer precision are swept independently, including the

25 +fp32-state + bf16-activation production config.

26 +"""

27 +

28 +import pytest

29 +import torch

30 +

31 +from tests.kernels.mamba.utils import selective_state_update_ref

32 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update

33 +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501

34 + selective_state_update_replayssm_output_only,

35 +)

36 +from vllm.model_executor.layers.mamba.ops.ssd_combined import (

37 + mamba_chunk_scan_combined_varlen,

38 +)

39 +from vllm.utils.torch_utils import set_random_seed

40 +from vllm.v1.attention.backends.mamba2_attn import compute_varlen_chunk_metadata

41 +

42 +

43 +def _prefill_tolerances(act_dtype: torch.dtype) -> tuple[float, float]:

44 + # The chunked prefill scan, not the decode, sets these: it carries ~2e-2

45 + # (fp32) / ~6e-2 (bf16) vs the exact recurrence, while ReplaySSM decode is

46 + # near-exact (~2e-6 fp32). Keyed off the activation dtype (the outputs are

47 + # in act_dtype). Same regime as test_mamba_ssm_ssd.py.

48 + if act_dtype == torch.float32:

49 + return 1e-2, 3e-2

50 + return 6e-2, 1e-1

51 +

52 +

53 +def _run_prefill_decode_equivalence(

54 + *,

55 + state_dtype: torch.dtype,

56 + act_dtype: torch.dtype,

57 + nheads: int,

58 + headdim: int,

59 + ngroups: int,

60 + dstate: int,

61 + seqlen: int,

62 + chunk_size: int,

63 + max_cache_len: int,

64 + seed: int = 0,

65 +) -> None:

66 + """Prefill the whole sequence and decode it step by step; check both match

67 + the exact fp32 recurrence (and each other). All paths use the production dt

68 + flow (raw dt + softplus + a per-head dt_bias), so this also checks prefill

69 + and decode apply the softplus/bias preprocessing consistently. ``state_dtype``

70 + is the recurrent-state precision; ``act_dtype`` the activation/buffer one."""

71 + device = "cuda"

72 + rtol, atol = _prefill_tolerances(act_dtype)

73 + set_random_seed(seed)

74 +

75 + # Production dt flow: raw dt + a per-head dt_bias (~-4 keeps softplus(dt +

76 + # bias) small and well-conditioned). dt_bias is (nheads,) for prefill and

77 + # (nheads, headdim) for the decode kernels/reference.

78 + A = -torch.exp(torch.rand(nheads, device=device, dtype=act_dtype))

79 + dt = torch.randn(seqlen, nheads, device=device, dtype=act_dtype)

80 + dt_bias = torch.rand(nheads, device=device, dtype=act_dtype) - 4

81 + dt_bias_hd = dt_bias.view(nheads, 1).expand(nheads, headdim)

82 + X = torch.randn(seqlen, nheads, headdim, device=device, dtype=act_dtype)

83 + B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype)

84 + C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=act_dtype)

85 + A_bcast = A.view(nheads, 1, 1).expand(nheads, headdim, dstate)

86 +

87 + # Chunked prefill over the whole sequence (implicit batch=1, varlen). The

88 + # kernel always returns the final state in fp32, so no state_dtype plumbing.

89 + cu_seqlens = torch.tensor((0, seqlen), device=device).cumsum(0).to(torch.int32)

90 + cu_chunk_seqlens, last_chunk_indices, seq_idx = compute_varlen_chunk_metadata(

91 + cu_seqlens, chunk_size

92 + )

93 + y_prefill = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype)

94 + final_state_prefill = mamba_chunk_scan_combined_varlen(

95 + X,

96 + dt,

97 + A,

98 + B,

99 + C,

100 + chunk_size,

101 + cu_seqlens=cu_seqlens,

102 + cu_chunk_seqlens=cu_chunk_seqlens,

103 + last_chunk_indices=last_chunk_indices,

104 + seq_idx=seq_idx,

105 + out=y_prefill,

106 + D=None,

107 + dt_bias=dt_bias,

108 + dt_softplus=True,

109 + )

110 +

111 + # Step paths: exact fp32 recurrence (ground truth), baseline, ReplaySSM.

112 + # State follows state_dtype; caches follow act_dtype (dt_cache is fp32).

113 + state_ref = torch.zeros(

114 + 1, nheads, headdim, dstate, device=device, dtype=torch.float32

115 + )

116 + state_base = torch.zeros(

117 + 1, nheads, headdim, dstate, device=device, dtype=state_dtype

118 + )

119 + state_dec = torch.zeros(

120 + 1, nheads, headdim, dstate, device=device, dtype=state_dtype

121 + )

122 + x_cache = torch.zeros(

123 + 1, nheads, max_cache_len, headdim, device=device, dtype=act_dtype

124 + )

125 + dt_cache = torch.zeros(1, nheads, max_cache_len, device=device, dtype=torch.float32)

126 + B_cache = torch.zeros(

127 + 1, ngroups, max_cache_len, dstate, device=device, dtype=act_dtype

128 + )

129 + bc_pre = torch.empty(1, ngroups, max_cache_len, device=device, dtype=torch.float32)

130 + write_pos = torch.zeros(1, dtype=torch.int32, device=device)

131 + # No skip connection (D=0) on any path here; the D!=0 path is covered by the

132 + # standard-decode suite. The baseline kernel needs a D tensor, not None.

133 + D_zero = torch.zeros(nheads, headdim, device=device)

134 +

135 + y_ref = torch.empty(seqlen, nheads, headdim, device=device, dtype=torch.float32)

136 + y_base = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype)

137 + y_dec = torch.empty(seqlen, nheads, headdim, device=device, dtype=act_dtype)

138 + for t in range(seqlen):

139 + dt_t = dt[t].view(1, nheads, 1).expand(1, nheads, headdim)

140 + is_flush = write_pos == max_cache_len - 1

141 +

142 + y_ref[t] = selective_state_update_ref(

143 + state_ref,

144 + X[t : t + 1].float(),

145 + dt_t.float(),

146 + A_bcast.float(),

147 + B[t : t + 1].float(),

148 + C[t : t + 1].float(),

149 + dt_bias=dt_bias_hd.float(),

150 + dt_softplus=True,

151 + )[0]

152 +

153 + out_b = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype)

154 + selective_state_update(

155 + state_base,

156 + X[t : t + 1],

157 + dt_t,

158 + A_bcast,

159 + B[t : t + 1],

160 + C[t : t + 1],

161 + D=D_zero,

162 + dt_bias=dt_bias_hd,

163 + dt_softplus=True,

164 + out=out_b,

165 + )

166 + y_base[t] = out_b[0]

167 +

168 + out_d = torch.empty(1, nheads, headdim, device=device, dtype=act_dtype)

169 + common = dict(

170 + dt_bias=dt_bias_hd,

171 + dt_softplus=True,

172 + x_cache=x_cache,

173 + dt_cache=dt_cache,

174 + B_cache=B_cache,

175 + write_pos=write_pos,

176 + is_flush=is_flush,

177 + max_cache_len=max_cache_len,

178 + out=out_d,

179 + )

180 + selective_state_update_replayssm_output_only(

181 + state_dec,

182 + X[t : t + 1],

183 + dt_t,

184 + A_bcast,

185 + B[t : t + 1],

186 + C[t : t + 1],

187 + bc_pre=bc_pre,

188 + **common,

189 + )

190 + y_dec[t] = out_d[0]

191 +

192 + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1)

193 +

194 + # Every path computes the same recurrence; anchor each on the fp32 truth.

195 + torch.testing.assert_close(y_prefill.float(), y_ref, rtol=rtol, atol=atol)

196 + torch.testing.assert_close(y_base.float(), y_ref, rtol=rtol, atol=atol)

197 + torch.testing.assert_close(y_dec.float(), y_ref, rtol=rtol, atol=atol)

198 + # Headline: ReplaySSM decode matches the chunked prefill directly.

199 + torch.testing.assert_close(y_dec.float(), y_prefill.float(), rtol=rtol, atol=atol)

200 + # Final state too (the recurrence ends in state_ref after the loop).

201 + torch.testing.assert_close(

202 + final_state_prefill[0].float(), state_ref[0], rtol=rtol, atol=atol

203 + )

204 +

205 +

206 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

207 +@pytest.mark.parametrize(

208 + "precision",

209 + # fp32 state is the default; bf16/fp16 are reduced-footprint configs. fp16

210 + # appears as an activation dtype (s32_afp16, sfp16_afp16) and as a finer-

211 + # mantissa state under bf16 activations (sfp16_a16).

212 + [

213 + pytest.param((torch.float32, torch.float32), id="s32_a32"),

214 + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"),

215 + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"),

216 + pytest.param((torch.float32, torch.float16), id="s32_afp16"),

217 + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"),

218 + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"),

219 + ],

220 +)

221 +@pytest.mark.parametrize(

222 + "geometry", # (nheads, headdim, dstate, ngroups)

223 + [

224 + pytest.param((8, 64, 64, 2), id="small"),

225 + pytest.param((96, 80, 128, 8), id="nano4b"),

226 + ],

227 +)

228 +def test_replayssm_prefill_decode_equivalence(

229 + precision: tuple[torch.dtype, torch.dtype],

230 + geometry: tuple[int, int, int, int],

231 +):

232 + state_dtype, act_dtype = precision

233 + nheads, headdim, dstate, ngroups = geometry

234 + _run_prefill_decode_equivalence(

235 + state_dtype=state_dtype,

236 + act_dtype=act_dtype,

237 + nheads=nheads,

238 + headdim=headdim,

239 + ngroups=ngroups,

240 + dstate=dstate,

241 + seqlen=16,

242 + chunk_size=8,

243 + max_cache_len=4,

244 + )

tests/kernels/mamba/test_replayssm_standard_decode_mamba2.py new

+673

@@ -0,0 +1,673 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Standard (autoregressive) decode correctness for the Mamba2 ReplaySSM kernels.

4 +

5 +ReplaySSM caches the recent SSM inputs ``(x, dt, B)`` in a small ring buffer and

6 +reconstructs / reads out the recurrent state on the fly, writing the full state

7 +back to HBM only when the buffer flushes. This file checks that, over a

8 +multi-step decode, the ReplaySSM output_only kernel reproduces the exact SSM

9 +recurrence.

10 +

11 +For every step we assert, against trusted oracles driven one token at a time:

12 +

13 + * the ReplaySSM output matches its pure-PyTorch reference, which models the

14 + kernel's exact arithmetic (including its bf16 reconstruction), at every

15 + precision,

16 + * when the state is fp32, the output also matches the baseline decode kernel

17 + (``selective_state_update``); at bf16 state the baseline legitimately

18 + differs -- it downcasts the fp32 state to bf16 every step, while ReplaySSM

19 + accumulates a whole buffer in fp32 and is the more accurate path,

20 + * when state and activations are both fp32, the output also matches the exact

21 + elementwise ``selective_state_update_ref``,

22 + * the cached inputs match the reference cache management,

23 + * the checkpoint state matches the reference (and the baseline at fp32 state).

24 +

25 +State and activation/buffer precision are swept independently. Nemotron-3

26 +defaults to fp32 SSM state (``mamba_ssm_cache_dtype=float32``), but bf16 state is

27 +also supported; the buffer dtype follows the activation dtype and ``dt_cache``

28 +is always fp32.

29 +"""

30 +

31 +import pytest

32 +import torch

33 +

34 +from tests.kernels.mamba.utils import (

35 + allocate_update_caches,

36 + selective_state_update_ref,

37 + selective_state_update_replayssm_output_only_ref,

38 +)

39 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update

40 +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501

41 + selective_state_update_replayssm_output_only,

42 +)

43 +from vllm.utils.torch_utils import set_random_seed

44 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

45 +

46 +

47 +def _tolerances(dtype: torch.dtype) -> tuple[float, float]:

48 + # fp32 demands true fp32 parity. A correct fp32 reconstruction (tl.dot with

49 + # input_precision="tf32x3"/"ieee") matches the elementwise baseline to

50 + # ~1e-5, while a TF32 reconstruction drifts to ~1e-2. atol=1e-3 sits between,

51 + # so it flags TF32 degradation yet passes a correct fp32 kernel. bf16 stays

52 + # loose (bf16 rounding dominates and is unaffected by the matmul precision).

53 + if dtype == torch.float32:

54 + return 1e-4, 1e-3

55 + return 6e-2, 2e-1

56 +

57 +

58 +def _tied_A(nheads: int, headdim: int, dstate: int, device: str) -> torch.Tensor:

59 + A = -torch.rand(nheads, device=device) - 1.0

60 + return A.view(nheads, 1, 1).expand(nheads, headdim, dstate)

61 +

62 +

63 +def _tied_dt(

64 + batch: int,

65 + nheads: int,

66 + headdim: int,

67 + device: str,

68 + dtype: torch.dtype,

69 +) -> torch.Tensor:

70 + dt = torch.randn(batch, nheads, device=device, dtype=dtype)

71 + return dt.unsqueeze(-1).expand(batch, nheads, headdim)

72 +

73 +

74 +def _tied_dt_bias(nheads: int, headdim: int, device: str) -> torch.Tensor:

75 + dt_bias = torch.rand(nheads, device=device) - 4.0

76 + return dt_bias.view(nheads, 1).expand(nheads, headdim)

77 +

78 +

79 +def _run_standard_decode(

80 + *,

81 + state_dtype: torch.dtype,

82 + act_dtype: torch.dtype,

83 + batch: int,

84 + nheads: int,

85 + headdim: int,

86 + ngroups: int,

87 + dstate: int,

88 + max_cache_len: int,

89 + num_steps: int,

90 + has_z: bool,

91 + dt_softplus: bool,

92 + use_dt_bias: bool,

93 + desync_write_pos: bool = False,

94 + seed: int = 0,

95 +) -> None:

96 + """Drive ``num_steps`` decode steps and check every path against the

97 + trusted recurrence. ``state_dtype`` is the recurrent-state precision;

98 + ``act_dtype`` is the activation/buffer precision."""

99 + device = "cuda"

100 + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32

101 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16)

102 + set_random_seed(seed)

103 +

104 + # One state copy per path; all start identical at the same precision.

105 + state0 = torch.randn(

106 + batch, nheads, headdim, dstate, dtype=state_dtype, device=device

107 + )

108 + state_anchor = state0.clone()

109 + state_baseline = state0.clone()

110 + state_cached = state0.clone()

111 + state_ref = state0.clone()

112 +

113 + A = _tied_A(nheads, headdim, dstate, device)

114 + dt_bias = _tied_dt_bias(nheads, headdim, device) if use_dt_bias else None

115 + D = torch.randn(nheads, headdim, device=device)

116 +

117 + # Caches follow the activation dtype (dt_cache is forced to fp32 inside).

118 + x_cache, dt_cache, B_cache, _ = allocate_update_caches(

119 + batch,

120 + nheads,

121 + ngroups,

122 + headdim,

123 + dstate,

124 + max_cache_len,

125 + device,

126 + act_dtype,

127 + act_dtype,

128 + )

129 + x_cache_ref, dt_cache_ref, B_cache_ref, _ = allocate_update_caches(

130 + batch,

131 + nheads,

132 + ngroups,

133 + headdim,

134 + dstate,

135 + max_cache_len,

136 + device,

137 + act_dtype,

138 + act_dtype,

139 + )

140 + bc_pre = torch.empty(

141 + batch, ngroups, max_cache_len, device=device, dtype=torch.float32

142 + )

143 +

144 + if desync_write_pos:

145 + # Rows start at different ring positions so they flush on different

146 + # steps, exercising per-row write-position handling.

147 + write_pos = (

148 + torch.arange(batch, device=device, dtype=torch.int32) % max_cache_len

149 + )

150 + else:

151 + write_pos = torch.zeros(batch, dtype=torch.int32, device=device)

152 +

153 + for _ in range(num_steps):

154 + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype)

155 + dt = _tied_dt(batch, nheads, headdim, device, act_dtype)

156 + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype)

157 + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype)

158 + z = torch.randn_like(x) if has_z else None

159 + is_flush = write_pos == max_cache_len - 1

160 +

161 + # Trusted recurrence (mutates state_anchor in place, returns output).

162 + out_anchor = selective_state_update_ref(

163 + state_anchor,

164 + x,

165 + dt,

166 + A,

167 + B,

168 + C,

169 + D=D,

170 + z=z,

171 + dt_bias=dt_bias,

172 + dt_softplus=dt_softplus,

173 + )

174 +

175 + # Upstream baseline decode kernel.

176 + out_baseline = torch.empty_like(x)

177 + selective_state_update(

178 + state_baseline,

179 + x,

180 + dt,

181 + A,

182 + B,

183 + C,

184 + D=D,

185 + z=z,

186 + dt_bias=dt_bias,

187 + dt_softplus=dt_softplus,

188 + out=out_baseline,

189 + )

190 +

191 + # ReplaySSM kernel under test + its pure-PyTorch reference.

192 + out_cached = torch.empty_like(x)

193 + common = dict(

194 + D=D,

195 + z=z,

196 + dt_bias=dt_bias,

197 + dt_softplus=dt_softplus,

198 + x_cache=x_cache,

199 + dt_cache=dt_cache,

200 + B_cache=B_cache,

201 + write_pos=write_pos,

202 + is_flush=is_flush,

203 + max_cache_len=max_cache_len,

204 + out=out_cached,

205 + )

206 + selective_state_update_replayssm_output_only(

207 + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common

208 + )

209 + out_ref = selective_state_update_replayssm_output_only_ref(

210 + state_ref,

211 + x,

212 + dt,

213 + A,

214 + B,

215 + C,

216 + D=D,

217 + z=z,

218 + dt_bias=dt_bias,

219 + dt_softplus=dt_softplus,

220 + x_cache=x_cache_ref,

221 + dt_cache=dt_cache_ref,

222 + B_cache=B_cache_ref,

223 + write_pos=write_pos,

224 + max_cache_len=max_cache_len,

225 + )

226 +

227 + # The reference models the kernel's exact arithmetic (including its bf16

228 + # reconstruction), so the kernel must match it tightly at every

229 + # precision. At fp32 this also flags any TF32 reconstruction drift.

230 + torch.testing.assert_close(out_cached, out_ref, rtol=rtol, atol=atol)

231 + # When the STATE is fp32 the baseline decode kernel is a valid oracle (it

232 + # does not downcast the state per step), so ReplaySSM must match it. At

233 + # bf16 state the baseline legitimately differs: it downcasts the fp32

234 + # state to bf16 every step, while ReplaySSM accumulates a whole buffer in

235 + # fp32 and is the MORE accurate path -- so it is not a tight oracle there.

236 + if state_dtype == torch.float32:

237 + torch.testing.assert_close(out_cached, out_baseline, rtol=rtol, atol=atol)

238 + # The exact elementwise reference is valid only when state AND

239 + # activations are fp32; otherwise it downcasts the state (at readout for

240 + # bf16 activations, or per step for bf16 state).

241 + if both_fp32:

242 + torch.testing.assert_close(out_cached, out_anchor, rtol=rtol, atol=atol)

243 +

244 + # Cached inputs match the reference cache management.

245 + torch.testing.assert_close(x_cache, x_cache_ref, rtol=rtol, atol=atol)

246 + torch.testing.assert_close(dt_cache, dt_cache_ref, rtol=rtol, atol=atol)

247 + torch.testing.assert_close(B_cache, B_cache_ref, rtol=rtol, atol=atol)

248 +

249 + # Checkpoint state at flush matches the reference (and, when the state is

250 + # fp32, the baseline kernel).

251 + if bool(is_flush.any()):

252 + torch.testing.assert_close(

253 + state_cached[is_flush], state_ref[is_flush], rtol=rtol, atol=atol

254 + )

255 + if state_dtype == torch.float32:

256 + torch.testing.assert_close(

257 + state_cached[is_flush],

258 + state_baseline[is_flush],

259 + rtol=rtol,

260 + atol=atol,

261 + )

262 +

263 + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1)

264 +

265 +

266 +# State/activation precisions. fp32 state is the default; bf16/fp16 are the

267 +# reduced-footprint configs. fp16 appears both as an activation dtype (fully-fp16

268 +# model sfp16_afp16, or fp16 act over fp32 state s32_afp16) and as a state dtype

269 +# under bf16 activations (sfp16_a16): fp16 has a finer mantissa than bf16 at the

270 +# same 2 bytes, so it is a more accurate state at no extra footprint. We still

271 +# skip fp16 state under fp32 activations (the unused low-state/high-act mix).

272 +_PRECISIONS = [

273 + pytest.param((torch.float32, torch.float32), id="s32_a32"),

274 + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"),

275 + pytest.param((torch.bfloat16, torch.bfloat16), id="s16_a16"),

276 + pytest.param((torch.float32, torch.float16), id="s32_afp16"),

277 + pytest.param((torch.float16, torch.float16), id="sfp16_afp16"),

278 + pytest.param((torch.float16, torch.bfloat16), id="sfp16_a16"),

279 +]

280 +# Small synthetic shapes for the full axis sweep (compile fast).

281 +_SMALL_GEOMETRIES = [

282 + pytest.param((8, 64, 64, 4), id="small"),

283 + pytest.param((4, 64, 16, 1), id="tiny"),

284 +]

285 +# Production Mamba2 shapes (nheads, headdim, dstate, ngroups), TP=1.

286 +_REAL_GEOMETRIES = [

287 + pytest.param((96, 80, 128, 8), id="nano4b"),

288 + pytest.param((128, 64, 128, 8), id="super120b"),

289 + pytest.param((256, 64, 128, 8), id="ultra550b"),

290 +]

291 +

292 +

293 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

294 +@pytest.mark.parametrize("precision", _PRECISIONS)

295 +@pytest.mark.parametrize("max_cache_len", [1, 4, 16])

296 +@pytest.mark.parametrize("geometry", _SMALL_GEOMETRIES)

297 +@pytest.mark.parametrize("has_z", [False, True])

298 +def test_replayssm_standard_decode_matches_reference(

299 + precision: tuple[torch.dtype, torch.dtype],

300 + max_cache_len: int,

301 + geometry: tuple[int, int, int, int],

302 + has_z: bool,

303 +):

304 + state_dtype, act_dtype = precision

305 + nheads, headdim, dstate, ngroups = geometry

306 + _run_standard_decode(

307 + state_dtype=state_dtype,

308 + act_dtype=act_dtype,

309 + batch=4,

310 + nheads=nheads,

311 + headdim=headdim,

312 + ngroups=ngroups,

313 + dstate=dstate,

314 + max_cache_len=max_cache_len,

315 + num_steps=2 * max_cache_len + 1,

316 + has_z=has_z,

317 + dt_softplus=True,

318 + use_dt_bias=True,

319 + )

320 +

321 +

322 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

323 +@pytest.mark.parametrize("precision", _PRECISIONS)

324 +@pytest.mark.parametrize("geometry", _REAL_GEOMETRIES)

325 +def test_replayssm_standard_decode_real_geometry(

326 + precision: tuple[torch.dtype, torch.dtype],

327 + geometry: tuple[int, int, int, int],

328 +):

329 + # Production Mamba2 shapes for the Nemotron-3 family (Nano-4B / Super-120B /

330 + # Ultra-550B), at the production buffer length (8). All three precisions,

331 + # including the bf16 state case.

332 + state_dtype, act_dtype = precision

333 + nheads, headdim, dstate, ngroups = geometry

334 + _run_standard_decode(

335 + state_dtype=state_dtype,

336 + act_dtype=act_dtype,

337 + batch=4,

338 + nheads=nheads,

339 + headdim=headdim,

340 + ngroups=ngroups,

341 + dstate=dstate,

342 + max_cache_len=8,

343 + num_steps=17,

344 + has_z=True,

345 + dt_softplus=True,

346 + use_dt_bias=True,

347 + )

348 +

349 +

350 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

351 +@pytest.mark.parametrize(

352 + "precision",

353 + [

354 + pytest.param((torch.float32, torch.float32), id="s32_a32"),

355 + pytest.param((torch.float32, torch.bfloat16), id="s32_a16"),

356 + pytest.param((torch.float32, torch.float16), id="s32_afp16"),

357 + ],

358 +)

359 +def test_replayssm_standard_decode_desync_write_pos(

360 + precision: tuple[torch.dtype, torch.dtype],

361 +):

362 + # Rows start at staggered ring positions, so they flush on different steps

363 + # and hold genuinely different cached histories.

364 + state_dtype, act_dtype = precision

365 + _run_standard_decode(

366 + state_dtype=state_dtype,

367 + act_dtype=act_dtype,

368 + batch=4,

369 + nheads=8,

370 + headdim=64,

371 + ngroups=4,

372 + dstate=64,

373 + max_cache_len=4,

374 + num_steps=12,

375 + has_z=True,

376 + dt_softplus=True,

377 + use_dt_bias=True,

378 + desync_write_pos=True,

379 + )

380 +

381 +

382 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

383 +@pytest.mark.parametrize("precision", _PRECISIONS)

384 +@pytest.mark.parametrize("with_padding", [False, True])

385 +def test_replayssm_standard_decode_with_batch_indices(

386 + precision: tuple[torch.dtype, torch.dtype],

387 + with_padding: bool,

388 +):

389 + # Sparse state allocation via state_batch_indices, with NULL_BLOCK_ID

390 + # padding rows. The pure-PyTorch references do not model the sparse

391 + # gather, so the anchor here is the upstream baseline decode kernel.

392 + state_dtype, act_dtype = precision

393 + device = "cuda"

394 + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32

395 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16)

396 + set_random_seed(0)

397 +

398 + batch = 3

399 + padding = 2 if with_padding else 0

400 + padded_batch = batch + padding

401 + total_state_slots = 16

402 + nheads = 4

403 + ngroups = 2

404 + headdim = 64

405 + dstate = 16

406 + max_cache_len = 4

407 + num_steps = 2 * max_cache_len

408 +

409 + state = torch.randn(

410 + total_state_slots, nheads, headdim, dstate, dtype=state_dtype, device=device

411 + )

412 + state_baseline = state.clone()

413 + state_cached = state.clone()

414 + state_before = state.clone()

415 +

416 + state_indices = (

417 + torch.randperm(total_state_slots - 1, device=device)[:batch] + 1

418 + ).to(torch.int32)

419 + state_batch_indices = torch.cat(

420 + [

421 + state_indices,

422 + torch.full((padding,), NULL_BLOCK_ID, dtype=torch.int32, device=device),

423 + ]

424 + )

425 + unused_states = torch.ones(total_state_slots, dtype=torch.bool, device=device)

426 + unused_states[state_indices] = False

427 +

428 + A = _tied_A(nheads, headdim, dstate, device)

429 + dt_bias = _tied_dt_bias(nheads, headdim, device)

430 + D = torch.randn(nheads, headdim, device=device)

431 + x_cache = torch.zeros(

432 + total_state_slots,

433 + nheads,

434 + max_cache_len,

435 + headdim,

436 + device=device,

437 + dtype=act_dtype,

438 + )

439 + dt_cache = torch.zeros(

440 + total_state_slots, nheads, max_cache_len, device=device, dtype=torch.float32

441 + )

442 + B_cache = torch.zeros(

443 + total_state_slots,

444 + ngroups,

445 + max_cache_len,

446 + dstate,

447 + device=device,

448 + dtype=act_dtype,

449 + )

450 + bc_pre = torch.empty(

451 + padded_batch, ngroups, max_cache_len, device=device, dtype=torch.float32

452 + )

453 + write_pos = torch.zeros(padded_batch, dtype=torch.int32, device=device)

454 +

455 + for _ in range(num_steps):

456 + x = torch.randn(padded_batch, nheads, headdim, device=device, dtype=act_dtype)

457 + dt = _tied_dt(padded_batch, nheads, headdim, device, act_dtype)

458 + B = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype)

459 + C = torch.randn(padded_batch, ngroups, dstate, device=device, dtype=act_dtype)

460 + z = torch.randn_like(x)

461 + is_flush = write_pos == max_cache_len - 1

462 +

463 + out_baseline = torch.empty_like(x)

464 + selective_state_update(

465 + state_baseline,

466 + x,

467 + dt,

468 + A,

469 + B,

470 + C,

471 + D=D,

472 + z=z,

473 + dt_bias=dt_bias,

474 + dt_softplus=True,

475 + state_batch_indices=state_batch_indices,

476 + out=out_baseline,

477 + )

478 +

479 + out_cached = torch.full_like(x, 42)

480 + common = dict(

481 + D=D,

482 + z=z,

483 + dt_bias=dt_bias,

484 + dt_softplus=True,

485 + x_cache=x_cache,

486 + dt_cache=dt_cache,

487 + B_cache=B_cache,

488 + write_pos=write_pos,

489 + is_flush=is_flush,

490 + max_cache_len=max_cache_len,

491 + state_batch_indices=state_batch_indices,

492 + out=out_cached,

493 + )

494 + selective_state_update_replayssm_output_only(

495 + state_cached, x, dt, A, B, C, bc_pre=bc_pre, **common

496 + )

497 +

498 + torch.testing.assert_close(

499 + out_cached[:batch], out_baseline[:batch], rtol=rtol, atol=atol

500 + )

501 + if with_padding:

502 + assert torch.equal(

503 + out_cached[batch:], torch.full_like(out_cached[batch:], 42)

504 + )

505 +

506 + if bool(is_flush[:batch].all()):

507 + torch.testing.assert_close(

508 + state_cached[state_indices],

509 + state_baseline[state_indices],

510 + rtol=rtol,

511 + atol=atol,

512 + )

513 +

514 + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1)

515 +

516 + assert torch.equal(state_cached[unused_states], state_before[unused_states])

517 + assert torch.equal(state_baseline[unused_states], state_before[unused_states])

518 +

519 +

520 +# Geometries with (nheads, ngroups) both divisible by the tp below.

521 +_TP_GEOMETRIES = [

522 + pytest.param((8, 64, 64, 4), id="small"),

523 + pytest.param((96, 80, 128, 8), id="nano4b"),

524 +]

525 +

526 +

527 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Need CUDA device")

528 +@pytest.mark.parametrize("precision", _PRECISIONS)

529 +@pytest.mark.parametrize("geometry", _TP_GEOMETRIES)

530 +@pytest.mark.parametrize("tp", [2])

531 +def test_replayssm_standard_decode_tp_head_shard_equivalence(

532 + precision: tuple[torch.dtype, torch.dtype],

533 + geometry: tuple[int, int, int, int],

534 + tp: int,

535 +):

536 + """Tensor-parallel correctness at the kernel boundary (single GPU).

537 +

538 + The kernel has no cross-rank communication, so head sharding must be exactly

539 + separable: one run over all heads must equal concatenating ``tp`` independent

540 + per-rank runs on ``nheads // tp`` heads and ``ngroups // tp`` groups. This

541 + guards the per-rank divisors and group->head mapping the state-shape wiring

542 + relies on. The real TP1==TP2 engine check lives in the v1/e2e suite."""

543 + state_dtype, act_dtype = precision

544 + nheads, headdim, dstate, ngroups = geometry

545 + assert nheads % tp == 0 and ngroups % tp == 0

546 + device = "cuda"

547 + both_fp32 = state_dtype == torch.float32 and act_dtype == torch.float32

548 + rtol, atol = _tolerances(torch.float32 if both_fp32 else torch.bfloat16)

549 + set_random_seed(0)

550 +

551 + batch = 4

552 + max_cache_len = 4

553 + num_steps = 2 * max_cache_len + 1

554 + nh_s = nheads // tp

555 + ng_s = ngroups // tp

556 +

557 + # Shards slice the tied params; never .contiguous() -- that would drop the

558 + # stride-0 broadcast the kernel's TIE_HDIM asserts require.

559 + A = _tied_A(nheads, headdim, dstate, device)

560 + dt_bias = _tied_dt_bias(nheads, headdim, device)

561 + D = torch.randn(nheads, headdim, device=device)

562 +

563 + state_full = torch.randn(

564 + batch, nheads, headdim, dstate, dtype=state_dtype, device=device

565 + )

566 + x_cache, dt_cache, B_cache, _ = allocate_update_caches(

567 + batch,

568 + nheads,

569 + ngroups,

570 + headdim,

571 + dstate,

572 + max_cache_len,

573 + device,

574 + act_dtype,

575 + act_dtype,

576 + )

577 + bc_pre = torch.empty(

578 + batch, ngroups, max_cache_len, device=device, dtype=torch.float32

579 + )

580 +

581 + # Per-rank shards, each seeded from the matching head slice so all start equal.

582 + shard_state = []

583 + shard_caches = []

584 + for r in range(tp):

585 + h0 = r * nh_s

586 + shard_state.append(state_full[:, h0 : h0 + nh_s].contiguous())

587 + xc, dtc, bc, _ = allocate_update_caches(

588 + batch,

589 + nh_s,

590 + ng_s,

591 + headdim,

592 + dstate,

593 + max_cache_len,

594 + device,

595 + act_dtype,

596 + act_dtype,

597 + )

598 + bcp = torch.empty(

599 + batch, ng_s, max_cache_len, device=device, dtype=torch.float32

600 + )

601 + shard_caches.append((xc, dtc, bc, bcp))

602 +

603 + write_pos = torch.zeros(batch, dtype=torch.int32, device=device)

604 + for _ in range(num_steps):

605 + x = torch.randn(batch, nheads, headdim, device=device, dtype=act_dtype)

606 + dt = _tied_dt(batch, nheads, headdim, device, act_dtype)

607 + B = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype)

608 + C = torch.randn(batch, ngroups, dstate, device=device, dtype=act_dtype)

609 + z = torch.randn_like(x)

610 + is_flush = write_pos == max_cache_len - 1

611 +

612 + out_full = torch.empty_like(x)

613 + selective_state_update_replayssm_output_only(

614 + state_full,

615 + x,

616 + dt,

617 + A,

618 + B,

619 + C,

620 + D=D,

621 + z=z,

622 + dt_bias=dt_bias,

623 + dt_softplus=True,

624 + x_cache=x_cache,

625 + dt_cache=dt_cache,

626 + B_cache=B_cache,

627 + bc_pre=bc_pre,

628 + write_pos=write_pos,

629 + is_flush=is_flush,

630 + max_cache_len=max_cache_len,

631 + out=out_full,

632 + )

633 +

634 + for r in range(tp):

635 + h0 = r * nh_s

636 + g0 = r * ng_s

637 + xc, dtc, bc, bcp = shard_caches[r]

638 + out_shard = torch.empty(

639 + batch, nh_s, headdim, device=device, dtype=act_dtype

640 + )

641 + selective_state_update_replayssm_output_only(

642 + shard_state[r],

643 + x[:, h0 : h0 + nh_s].contiguous(),

644 + dt[:, h0 : h0 + nh_s],

645 + A[h0 : h0 + nh_s],

646 + B[:, g0 : g0 + ng_s].contiguous(),

647 + C[:, g0 : g0 + ng_s].contiguous(),

648 + D=D[h0 : h0 + nh_s].contiguous(),

649 + z=z[:, h0 : h0 + nh_s].contiguous(),

650 + dt_bias=dt_bias[h0 : h0 + nh_s],

651 + dt_softplus=True,

652 + x_cache=xc,

653 + dt_cache=dtc,

654 + B_cache=bc,

655 + bc_pre=bcp,

656 + write_pos=write_pos,

657 + is_flush=is_flush,

658 + max_cache_len=max_cache_len,

659 + out=out_shard,

660 + )

661 +

662 + torch.testing.assert_close(

663 + out_full[:, h0 : h0 + nh_s], out_shard, rtol=rtol, atol=atol

664 + )

665 + if bool(is_flush.any()):

666 + torch.testing.assert_close(

667 + state_full[:, h0 : h0 + nh_s][is_flush],

668 + shard_state[r][is_flush],

669 + rtol=rtol,

670 + atol=atol,

671 + )

672 +

673 + write_pos = torch.where(is_flush, torch.zeros_like(write_pos), write_pos + 1)

tests/kernels/mamba/utils.py

+160

@@ -76,3 +76,163 @@ def selective_state_update_ref(

76 if not has_heads:

77 out = out.squeeze(1)

78 return out

79 +

80 +

81 +def selective_state_update_replayssm_output_only_ref(

82 + state: torch.Tensor,

83 + x: torch.Tensor,

84 + dt: torch.Tensor,

85 + A: torch.Tensor,

86 + B: torch.Tensor,

87 + C: torch.Tensor,

88 + D: torch.Tensor | None = None,

89 + z: torch.Tensor | None = None,

90 + dt_bias: torch.Tensor | None = None,

91 + dt_softplus: bool = False,

92 + x_cache: torch.Tensor | None = None,

93 + dt_cache: torch.Tensor | None = None,

94 + B_cache: torch.Tensor | None = None,

95 + write_pos: torch.Tensor | None = None,

96 + max_cache_len: int = 16,

97 +) -> torch.Tensor:

98 + """Pure-PyTorch cached-bc reference for validation."""

99 + has_heads = state.dim() > 3

100 + if state.dim() == 3:

101 + state = state.unsqueeze(1)

102 + if x.dim() == 2:

103 + x = x.unsqueeze(1)

104 + if dt.dim() == 2:

105 + dt = dt.unsqueeze(1)

106 + if A.dim() == 2:

107 + A = A.unsqueeze(0)

108 + if B.dim() == 2:

109 + B = B.unsqueeze(1)

110 + if C.dim() == 2:

111 + C = C.unsqueeze(1)

112 + if D is not None and D.dim() == 1:

113 + D = D.unsqueeze(0)

114 + if z is not None and z.dim() == 2:

115 + z = z.unsqueeze(1)

116 + if dt_bias is not None and dt_bias.dim() == 1:

117 + dt_bias = dt_bias.unsqueeze(0)

118 +

119 + batch, nheads, dim, dstate = state.shape

120 + assert x.shape == (batch, nheads, dim)

121 + assert dt.shape == x.shape

122 + assert A.shape == (nheads, dim, dstate)

123 + ngroups = B.shape[1]

124 + assert nheads % ngroups == 0, "nheads must be divisible by ngroups"

125 + assert B.shape == (batch, ngroups, dstate)

126 + assert C.shape == B.shape

127 +

128 + ratio = nheads // ngroups

129 +

130 + dt_val = dt[:, :, 0].float()

131 + if dt_bias is not None:

132 + dt_val = dt_val + dt_bias[:, 0].float()

133 + if dt_softplus:

134 + dt_val = F.softplus(dt_val)

135 + A_val = A[:, 0, 0].float()

136 + C_heads = C.repeat_interleave(ratio, dim=1)

137 + out = torch.empty(batch, nheads, dim, device=x.device, dtype=torch.float32)

138 +

139 + assert x_cache is not None

140 + assert dt_cache is not None

141 + assert B_cache is not None

142 + assert write_pos is not None

143 +

144 + for b in range(batch):

145 + cache_len = int(write_pos[b].item())

146 + is_flush = cache_len == max_cache_len - 1

147 + n_steps = cache_len + 1

148 +

149 + dt_all = torch.zeros(nheads, n_steps, device=x.device, dtype=torch.float32)

150 + if cache_len > 0:

151 + dt_all[:, :cache_len] = dt_cache[b, :, :cache_len]

152 + dt_all[:, cache_len] = dt_val[b]

153 +

154 + cumsum = torch.cumsum(dt_all, dim=-1)

155 + total = cumsum[:, -1]

156 + dA_cumsum = A_val[:, None] * cumsum

157 + dA_total = A_val * total

158 + total_decay = torch.exp(dA_total)

159 + scale = dt_all * torch.exp(dA_total[:, None] - dA_cumsum)

160 +

161 + x_all = torch.zeros(nheads, dim, n_steps, device=x.device, dtype=x.dtype)

162 + if cache_len > 0:

163 + x_all[..., :cache_len] = x_cache[b, :, :cache_len, :].permute(0, 2, 1)

164 + x_all[..., cache_len] = x[b]

165 +

166 + B_all = torch.zeros(ngroups, n_steps, dstate, device=B.device, dtype=B.dtype)

167 + if cache_len > 0:

168 + B_all[:, :cache_len, :] = B_cache[b, :, :cache_len, :]

169 + B_all[:, cache_len, :] = B[b]

170 +

171 + B_heads = B_all.repeat_interleave(ratio, dim=0)

172 + C_heads_b = C_heads[b]

173 +

174 + if is_flush:

175 + B_scaled = (B_heads.float() * scale[:, :, None]).to(B_heads.dtype)

176 + delta = torch.einsum("hdk,hkn->hdn", x_all.float(), B_scaled.float())

177 + state_new = state[b].float() * total_decay[:, None, None] + delta

178 + state[b].copy_(state_new.to(state.dtype))

179 + out[b] = torch.einsum("hdn,hn->hd", state_new, C_heads_b.float())

180 + else:

181 + checkpoint_out = torch.einsum(

182 + "hdn,hn->hd", state[b].float(), C_heads_b.float()

183 + )

184 + checkpoint_out = checkpoint_out * total_decay[:, None]

185 + BC = torch.einsum("hkn,hn->hk", B_heads.float(), C_heads_b.float())

186 + cache_out = torch.einsum("hdk,hk->hd", x_all.float(), scale * BC)

187 + out[b] = checkpoint_out + cache_out

188 + x_cache[b, :, cache_len, :] = x[b]

189 + dt_cache[b, :, cache_len] = dt_val[b]

190 + B_cache[b, :, cache_len, :] = B[b]

191 +

192 + if D is not None:

193 + out = out + (x.float() * D[None]).to(out.dtype)

194 + if z is not None:

195 + out = out * F.silu(z.float())

196 + out = out.to(x.dtype)

197 + if not has_heads:

198 + out = out.squeeze(1)

199 + return out

200 +

201 +

202 +def allocate_update_caches(

203 + batch: int,

204 + nheads: int,

205 + ngroups: int,

206 + dim: int,

207 + dstate: int,

208 + max_cache_len: int,

209 + device: torch.device,

210 + x_dtype: torch.dtype,

211 + B_dtype: torch.dtype,

212 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

213 + """Allocate dense reference caches for standalone validation."""

214 + x_cache = torch.zeros(

215 + batch,

216 + nheads,

217 + max_cache_len,

218 + dim,

219 + device=device,

220 + dtype=x_dtype,

221 + )

222 + dt_cache = torch.zeros(

223 + batch,

224 + nheads,

225 + max_cache_len,

226 + device=device,

227 + dtype=torch.float32,

228 + )

229 + B_cache = torch.zeros(

230 + batch,

231 + ngroups,

232 + max_cache_len,

233 + dstate,

234 + device=device,

235 + dtype=B_dtype,

236 + )

237 + write_pos = torch.zeros(batch, dtype=torch.int32, device=device)

238 + return x_cache, dt_cache, B_cache, write_pos

tests/models/test_registry.py

+16

@@ -134,6 +134,22 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda):

134 )

135

136

137 +@create_new_process_for_each_test()

138 +@pytest.mark.parametrize(

139 + "model_arch,supported",

140 + [

141 + # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today.

142 + ("NemotronHForCausalLM", True),

143 + ("Mamba2ForCausalLM", False),

144 + ("Zamba2ForCausalLM", False),

145 + ],

146 +)

147 +def test_registry_supports_replayssm(model_arch, supported):

148 + model_info = ModelRegistry._try_inspect_model_cls(model_arch)

149 + assert model_info is not None

150 + assert model_info.supports_replayssm is supported

151 +

152 +

153 def test_lazy_modelinfo_package_hash_includes_submodules(tmp_path):

154 package_dir = tmp_path / "model_package"

155 package_dir.mkdir()

tests/v1/attention/test_mamba_update_block_table.py

@@ -42,6 +42,8 @@ def _make_vllm_config(

42 cache_config=SimpleNamespace(

43 block_size=block_size,

44 mamba_cache_mode="all",

45 + use_replayssm=False,

46 + replayssm_buffer_len=16,

47 ),

48 compilation_config=SimpleNamespace(

49 cudagraph_mode=CUDAGraphMode.FULL,

tests/v1/attention/test_replayssm_metadata_builder.py new

+256

@@ -0,0 +1,256 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Mamba2 ReplaySSM decode write-position derivation in

4 +BaseMambaAttentionMetadataBuilder: write_pos and is_flush computed from the

5 +per-request ring origin (replayssm_decode_base) and num_computed.

6 +"""

7 +

8 +from dataclasses import dataclass

9 +

10 +import pytest

11 +import torch

12 +

13 +from tests.v1.attention.utils import (

14 + BatchSpec,

15 + MockMambaBuilder,

16 + create_common_attn_metadata,

17 + create_vllm_config,

18 +)

19 +from vllm.v1.kv_cache_interface import MambaSpec

20 +

21 +BLOCK_SIZE = 16

22 +DEVICE = torch.device("cpu")

23 +

24 +

25 +@dataclass

26 +class ReplaySSMBuildCase:

27 + """A decode batch and its expected per-row write_pos / is_flush.

28 +

29 + num_computed = seq_len - query_len; write_pos =

30 + (num_computed - decode_base) % buffer_len; is_flush = write_pos ==

31 + buffer_len - 1 (or a forced one-token flush when num_computed < decode_base).

32 + """

33 +

34 + seq_lens: list[int]

35 + query_lens: list[int]

36 + is_prefilling: list[bool]

37 + decode_base: list[int]

38 + buffer_len: int

39 + expected_write_pos: list[int]

40 + expected_is_flush: list[int]

41 + mamba_cache_mode: str = "none"

42 +

43 +

44 +REPLAYSSM_BUILD_CASES = {

45 + # decode_base == num_prompt (fresh request).

46 + "fresh_decode": ReplaySSMBuildCase(

47 + seq_lens=[106],

48 + query_lens=[1],

49 + is_prefilling=[False],

50 + decode_base=[100],

51 + buffer_len=16,

52 + expected_write_pos=[5],

53 + expected_is_flush=[0],

54 + ),

55 + # decode_base > num_prompt anchors write_pos at the resume point.

56 + "resumed_reanchors_to_zero": ReplaySSMBuildCase(

57 + seq_lens=[106],

58 + query_lens=[1],

59 + is_prefilling=[False],

60 + decode_base=[105],

61 + buffer_len=16,

62 + expected_write_pos=[0],

63 + expected_is_flush=[0],

64 + ),

65 + # write_pos == buffer_len - 1 flushes.

66 + "flush_boundary": ReplaySSMBuildCase(

67 + seq_lens=[116],

68 + query_lens=[1],

69 + is_prefilling=[False],

70 + decode_base=[100],

71 + buffer_len=16,

72 + expected_write_pos=[15],

73 + expected_is_flush=[1],

74 + ),

75 + # Resumed request landing on a flush boundary.

76 + "resumed_flush_boundary": ReplaySSMBuildCase(

77 + seq_lens=[121],

78 + query_lens=[1],

79 + is_prefilling=[False],

80 + decode_base=[105],

81 + buffer_len=16,

82 + expected_write_pos=[15],

83 + expected_is_flush=[1],

84 + ),

85 + # Per-row write_pos / is_flush are independent.

86 + "mixed_rows": ReplaySSMBuildCase(

87 + seq_lens=[104, 106, 216],

88 + query_lens=[1, 1, 1],

89 + is_prefilling=[False, False, False],

90 + decode_base=[100, 105, 200],

91 + buffer_len=16,

92 + expected_write_pos=[3, 0, 15],

93 + expected_is_flush=[0, 0, 1],

94 + ),

95 + # write_pos wraps within the buffer (6 % 4 == 2).

96 + "small_buffer_wrap": ReplaySSMBuildCase(

97 + seq_lens=[112],

98 + query_lens=[1],

99 + is_prefilling=[False],

100 + decode_base=[105],

101 + buffer_len=4,

102 + expected_write_pos=[2],

103 + expected_is_flush=[0],

104 + ),

105 + # Single-token prefill-as-decode still in the prompt (num_computed <

106 + # decode_base): forced one-token flush.

107 + "leftover_prompt_one_token_flush": ReplaySSMBuildCase(

108 + seq_lens=[100],

109 + query_lens=[1],

110 + is_prefilling=[True],

111 + decode_base=[100],

112 + buffer_len=16,

113 + expected_write_pos=[0],

114 + expected_is_flush=[1],

115 + ),

116 + # Align mode (block_size 16). Past the first boundary the ring re-anchors at

117 + # the block start: num_computed 117 -> block_start 112, write_pos 5 (vs 1 in

118 + # none mode).

119 + "align_reanchor_past_boundary": ReplaySSMBuildCase(

120 + seq_lens=[118],

121 + query_lens=[1],

122 + is_prefilling=[False],

123 + decode_base=[100],

124 + buffer_len=16,

125 + expected_write_pos=[5],

126 + expected_is_flush=[0],

127 + mamba_cache_mode="align",

128 + ),

129 + # First-block boundary (num_computed+1 == 112) forces a flush even though

130 + # write_pos (11) != buffer_len - 1.

131 + "align_first_block_boundary_flush": ReplaySSMBuildCase(

132 + seq_lens=[112],

133 + query_lens=[1],

134 + is_prefilling=[False],

135 + decode_base=[100],

136 + buffer_len=16,

137 + expected_write_pos=[11],

138 + expected_is_flush=[1],

139 + mamba_cache_mode="align",

140 + ),

141 + # First step of a new block re-anchors write_pos to 0.

142 + "align_new_block_start_zero": ReplaySSMBuildCase(

143 + seq_lens=[113],

144 + query_lens=[1],

145 + is_prefilling=[False],

146 + decode_base=[100],

147 + buffer_len=16,

148 + expected_write_pos=[0],

149 + expected_is_flush=[0],

150 + mamba_cache_mode="align",

151 + ),

152 + # block_size % buffer_len == 0: a later boundary lands on write_pos ==

153 + # buffer_len - 1, so the boundary flush coincides with the natural flush.

154 + "align_boundary_coincides_natural_flush": ReplaySSMBuildCase(

155 + seq_lens=[128],

156 + query_lens=[1],

157 + is_prefilling=[False],

158 + decode_base=[100],

159 + buffer_len=16,

160 + expected_write_pos=[15],

161 + expected_is_flush=[1],

162 + mamba_cache_mode="align",

163 + ),

164 + # block_size % buffer_len != 0 (buffer_len 6): the boundary step still flushes

165 + # although write_pos (3) != buffer_len - 1.

166 + "align_unaligned_buffer_forces_flush": ReplaySSMBuildCase(

167 + seq_lens=[128],

168 + query_lens=[1],

169 + is_prefilling=[False],

170 + decode_base=[100],

171 + buffer_len=6,

172 + expected_write_pos=[3],

173 + expected_is_flush=[1],

174 + mamba_cache_mode="align",

175 + ),

176 + # Per-row independence in align mode: partial-block / new-block / boundary.

177 + "align_mixed_rows": ReplaySSMBuildCase(

178 + seq_lens=[105, 113, 112],

179 + query_lens=[1, 1, 1],

180 + is_prefilling=[False, False, False],

181 + decode_base=[100, 100, 100],

182 + buffer_len=16,

183 + expected_write_pos=[4, 0, 11],

184 + expected_is_flush=[0, 0, 1],

185 + mamba_cache_mode="align",

186 + ),

187 +}

188 +

189 +

190 +def _make_mamba_spec(buffer_len: int) -> MambaSpec:

191 + # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups).

192 + return MambaSpec(

193 + block_size=BLOCK_SIZE,

194 + shapes=(

195 + (1, 1),

196 + (1, 1, 1),

197 + (1, buffer_len, 1),

198 + (1, buffer_len),

199 + (1, buffer_len, 1),

200 + ),

201 + dtypes=(torch.float32,),

202 + )

203 +

204 +

205 +def _create_replayssm_builder(

206 + buffer_len: int, mamba_cache_mode: str = "none"

207 +) -> MockMambaBuilder:

208 + vllm_config = create_vllm_config(

209 + model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE

210 + )

211 + # Set the flags after construction to skip validate_mamba_cached_kernel

212 + # (it requires a Triton backend) on the mock model.

213 + vllm_config.cache_config.use_replayssm = True

214 + vllm_config.cache_config.replayssm_buffer_len = buffer_len

215 + vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode

216 + return MockMambaBuilder(

217 + _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE

218 + )

219 +

220 +

221 +def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase):

222 + batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens)

223 + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace(

224 + is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool),

225 + replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32),

226 + )

227 + return builder.build(0, common)

228 +

229 +

230 +@pytest.mark.parametrize(

231 + "case", REPLAYSSM_BUILD_CASES.values(), ids=REPLAYSSM_BUILD_CASES.keys()

232 +)

233 +def test_replayssm_write_pos(case: ReplaySSMBuildCase):

234 + builder = _create_replayssm_builder(case.buffer_len, case.mamba_cache_mode)

235 + meta = _build(builder, case)

236 +

237 + assert meta.write_pos_d is not None

238 + assert meta.is_flush_d is not None

239 + n = len(case.expected_write_pos)

240 + assert meta.write_pos_d[:n].tolist() == case.expected_write_pos

241 + assert meta.is_flush_d[:n].tolist() == case.expected_is_flush

242 +

243 +

244 +def test_resumed_request_differs_from_fresh():

245 + """Same token count, different decode_base: fresh (base 100) -> write_pos 5,

246 + resumed (base 105) -> write_pos 0."""

247 + builder = _create_replayssm_builder(16)

248 + batch = BatchSpec(seq_lens=[106, 106], query_lens=[1, 1])

249 + common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace(

250 + is_prefilling=torch.tensor([False, False]),

251 + replayssm_decode_base_cpu=torch.tensor([100, 105], dtype=torch.int32),

252 + )

253 + meta = builder.build(0, common)

254 +

255 + assert meta.write_pos_d.tolist()[:2] == [5, 0]

256 + assert meta.is_flush_d.tolist()[:2] == [0, 0]

tests/v1/e2e/test_replayssm_decode.py new

+138

@@ -0,0 +1,138 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Engine-level parity: ReplaySSM standard decode vs the baseline SSM kernel."""

4 +

5 +import pytest

6 +

7 +from vllm.v1.metrics.reader import Counter

8 +

9 +from ...models.utils import check_logprobs_close

10 +from ...utils import large_gpu_mark, multi_gpu_test

11 +

12 +# Mamba2 (Nemotron-3) hybrid.

13 +MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"

14 +MODELS = [

15 + pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)),

16 +]

17 +

18 +PROMPTS = [

19 + "The capital of France is",

20 + "Once upon a time, in a small village,",

21 +]

22 +

23 +

24 +def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1):

25 + # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a

26 + # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are

27 + # common-mode and only ReplaySSM varies.

28 + common = dict(

29 + max_model_len=1024,

30 + trust_remote_code=True,

31 + enable_prefix_caching=False,

32 + mamba_cache_mode="none",

33 + tensor_parallel_size=tensor_parallel_size,

34 + )

35 + with vllm_runner(model_name, **common) as llm:

36 + baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)

37 + with vllm_runner(

38 + model_name, use_replayssm=True, replayssm_buffer_len=16, **common

39 + ) as llm:

40 + replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)

41 +

42 + check_logprobs_close(

43 + outputs_0_lst=baseline,

44 + outputs_1_lst=replay,

45 + name_0="baseline",

46 + name_1="replayssm",

47 + )

48 +

49 +

50 +@pytest.mark.parametrize("model_name", MODELS)

51 +def test_replayssm_decode_matches_baseline(vllm_runner, model_name):

52 + _check_replayssm_parity(vllm_runner, model_name)

53 +

54 +

55 +@multi_gpu_test(num_gpus=2)

56 +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL])

57 +def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name):

58 + # Tensor-parallel correctness: ReplaySSM's caches and checkpoint state are

59 + # sharded per rank, so TP2 decode must still match the baseline at TP2.

60 + _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2)

61 +

62 +

63 +# Prefix spans several mamba blocks; prefix caching only reuses full blocks.

64 +_PC_SENTENCE = (

65 + "In a detailed survey of state space models, the authors compared many "

66 + "architectures across a wide range of long-context language tasks and "

67 + "measured their throughput, memory use, and accuracy in careful detail. "

68 +)

69 +_PC_PREFIX = _PC_SENTENCE * 120

70 +PREFIX_CACHING_PROMPTS = [

71 + _PC_PREFIX + "The most important conclusion was that",

72 + _PC_PREFIX + "Surprisingly, the experiments showed that",

73 + _PC_PREFIX + "The most important conclusion was that",

74 +]

75 +

76 +

77 +def _prefix_cache_hits(llm) -> int:

78 + return sum(

79 + m.value

80 + for m in llm.llm.get_metrics()

81 + if isinstance(m, Counter) and m.name == "vllm:prefix_cache_hits"

82 + )

83 +

84 +

85 +def _check_replayssm_prefix_caching_parity(

86 + vllm_runner, model_name, *, tensor_parallel_size=1

87 +):

88 + # align mode materializes the exact SSM state at each block boundary, so

89 + # ReplaySSM's cached prefixes must match the always-materialized baseline.

90 + common = dict(

91 + max_model_len=8192,

92 + trust_remote_code=True,

93 + enable_prefix_caching=True,

94 + enable_chunked_prefill=True,

95 + mamba_cache_mode="align",

96 + disable_log_stats=False, # required for llm.get_metrics()

97 + tensor_parallel_size=tensor_parallel_size,

98 + )

99 + with vllm_runner(model_name, **common) as llm:

100 + baseline = llm.generate_greedy_logprobs(

101 + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5

102 + )

103 + with vllm_runner(

104 + model_name, use_replayssm=True, replayssm_buffer_len=16, **common

105 + ) as llm:

106 + # Prime the cache, then measure, so cache hits are deterministic.

107 + llm.generate_greedy_logprobs(

108 + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5

109 + )

110 + replay = llm.generate_greedy_logprobs(

111 + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5

112 + )

113 + replay_hits = _prefix_cache_hits(llm)

114 +

115 + # Without real cache hits the cached path is never exercised.

116 + assert replay_hits > 0, (

117 + "ReplaySSM align-mode run produced no prefix-cache hits; the shared "

118 + "prefix may be shorter than one mamba block, so prefix caching is inert"

119 + )

120 + check_logprobs_close(

121 + outputs_0_lst=baseline,

122 + outputs_1_lst=replay,

123 + name_0="baseline_align_pc",

124 + name_1="replayssm_align_pc",

125 + )

126 +

127 +

128 +@pytest.mark.parametrize("model_name", MODELS)

129 +def test_replayssm_prefix_caching_matches_baseline(vllm_runner, model_name):

130 + _check_replayssm_prefix_caching_parity(vllm_runner, model_name)

131 +

132 +

133 +@multi_gpu_test(num_gpus=2)

134 +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL])

135 +def test_replayssm_prefix_caching_matches_baseline_tp2(vllm_runner, model_name):

136 + _check_replayssm_prefix_caching_parity(

137 + vllm_runner, model_name, tensor_parallel_size=2

138 + )

vllm/config/cache.py

+11

@@ -145,6 +145,17 @@ class CacheConfig:

145 - "align": only cache the mamba state of the last token of each scheduler step and

146 when the token is at position i * block_size.

147 """

148 + replayssm_buffer_len: int = Field(default=16, gt=0)

149 + """ReplaySSM history buffer length B: with use_replayssm, standard decode

150 + caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint

151 + state to HBM every B steps. Default 16."""

152 + use_replayssm: bool = False

153 + """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip

154 + the per-step full-state store, writing the checkpoint back only on flush.

155 + Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton

156 + mamba backend; standard (non-speculative) decode only. In align mode flushes

157 + are most efficient when mamba_block_size is a multiple of replayssm_buffer_len,

158 + but this is not required."""

159

160 # Will be set after profiling.

161 num_gpu_blocks: int | None = field(default=None, init=False)

vllm/config/model.py

@@ -1696,6 +1696,10 @@ class ModelConfig:

1696 def supports_mamba_prefix_caching(self) -> bool:

1697 return self._model_info.supports_mamba_prefix_caching

1698

1699 + @property

1700 + def supports_replayssm(self) -> bool:

1701 + return self._model_info.supports_replayssm

1702 +

1703 @property

1704 def use_mla(self) -> bool:

1705 return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE

vllm/config/vllm.py

+32 -1

@@ -39,7 +39,7 @@ from .kv_events import KVEventsConfig

39 from .kv_transfer import KVTransferConfig

40 from .load import LoadConfig

41 from .lora import LoRAConfig

42 -from .mamba import MambaConfig

42 +from .mamba import MambaBackendEnum, MambaConfig

43 from .model import ModelConfig

44 from .observability import ObservabilityConfig

45 from .offload import OffloadConfig

@@ -2308,6 +2308,37 @@ class VllmConfig:

2308 )

2309 return self

2310

2311 + @model_validator(mode="after")

2312 + def validate_mamba_cached_kernel(self) -> "VllmConfig":

2313 + if not self.cache_config.use_replayssm:

2314 + return self

2315 + # ReplaySSM adds a 3-tensor ring to the mamba state; only models that

2316 + # opt in (supports_replayssm) build a consistent shape on both the layer

2317 + # and config paths. Reject others so the mamba page size cannot desync.

2318 + if self.model_config is not None and not self.model_config.supports_replayssm:

2319 + raise ValueError(

2320 + "--use-replayssm is only supported for Nemotron-H models "

2321 + f"(got architecture {self.model_config.architecture!r})"

2322 + )

2323 + if self.cache_config.mamba_cache_mode == "all":

2324 + raise ValueError(

2325 + "--use-replayssm supports prefix caching only in align mode; "

2326 + "pass --mamba-cache-mode align"

2327 + )

2328 + if self.num_speculative_tokens > 0:

2329 + raise ValueError("--use-replayssm does not support speculative decoding")

2330 + if self.mamba_config.backend != MambaBackendEnum.TRITON:

2331 + raise ValueError("--use-replayssm requires --mamba-backend triton")

2332 + if (

2333 + self.kv_transfer_config is not None

2334 + and self.kv_transfer_config.is_kv_transfer_instance

2335 + ):

2336 + raise ValueError(

2337 + "--use-replayssm is incompatible with KV connectors "

2338 + "(P/D disaggregation, KV cache offload)"

2339 + )

2340 + return self

2341 +

2342

2343 _current_vllm_config: VllmConfig | None = None

2344 _current_prefix: str | None = None

vllm/engine/arg_utils.py

@@ -695,6 +695,8 @@ class EngineArgs:

695 mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")

696 prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit")

697 mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode

698 + replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len

699 + use_replayssm: bool = CacheConfig.use_replayssm

700

701 mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON

702 enable_mamba_cache_stochastic_rounding: bool = (

@@ -1201,6 +1203,10 @@ class EngineArgs:

1203 cache_group.add_argument(

1204 "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]

1205 )

1206 + cache_group.add_argument(

1207 + "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"]

1208 + )

1209 + cache_group.add_argument("--use-replayssm", **cache_kwargs["use_replayssm"])

1210 cache_group.add_argument(

1211 "--kv-offloading-size", **cache_kwargs["kv_offloading_size"]

1212 )

@@ -1910,6 +1916,8 @@ class EngineArgs:

1916 mamba_block_size=self.mamba_block_size,

1917 prefix_match_unit=self.prefix_match_unit,

1918 mamba_cache_mode=self.mamba_cache_mode,

1919 + replayssm_buffer_len=self.replayssm_buffer_len,

1920 + use_replayssm=self.use_replayssm,

1921 kv_offloading_size=self.kv_offloading_size,

1922 kv_offloading_backend=self.kv_offloading_backend,

1923 )

vllm/model_executor/layers/mamba/mamba_mixer2.py

+96 -24

@@ -32,6 +32,9 @@ from vllm.model_executor.layers.mamba.ops.causal_conv1d import (

32 causal_conv1d_update,

33 )

34 from vllm.model_executor.layers.mamba.ops.layernorm_gated import rms_norm_gated

35 +from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_output_only import ( # noqa: E501

36 + selective_state_update_replayssm_output_only,

37 +)

38 from vllm.model_executor.layers.mamba.ops.ssd_combined import (

39 mamba_chunk_scan_combined_varlen,

40 )

@@ -494,12 +497,27 @@ class MambaMixer2(MambaBase, PluggableLayer):

497 if prefix in compilation_config.static_forward_context:

498 raise ValueError(f"Duplicate layer name: {prefix}")

499 compilation_config.static_forward_context[prefix] = self

497 - # The tuple is (conv_state, ssm_state)

498 - self.kv_cache = (torch.tensor([]), torch.tensor([]))

500

501 self.model_config = model_config

502 self.cache_config = cache_config

503 self.prefix = prefix

504 + self.use_replayssm = (

505 + cache_config.use_replayssm if cache_config is not None else False

506 + )

507 + self.replayssm_buffer_len = (

508 + cache_config.replayssm_buffer_len

509 + if cache_config is not None and cache_config.use_replayssm

510 + else None

511 + )

512 + self.mamba_config = vllm_config.mamba_config

513 + if self.use_replayssm and self.num_heads % self.tp_size != 0:

514 + raise ValueError(

515 + "--use-replayssm requires tensor-parallel heads to divide evenly"

516 + )

517 + # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode

518 + # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache).

519 + _n_state = 5 if self.use_replayssm else 2

520 + self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state))

521

522 self.num_spec = vllm_config.num_speculative_tokens

523 if self.num_spec > 0:

@@ -591,7 +609,7 @@ class MambaMixer2(MambaBase, PluggableLayer):

609

610 # Triton's autotuner includes tensor dtypes in its cache key,

611 # so state_dtype must match what real inference uses.

594 - _, ssm_state_dtype = self.get_state_dtype()

612 + ssm_state_dtype = self.get_state_dtype()[1]

613

614 # SSD kernel autotune keys depend on dtype and head dimensions,

615 # not on sequence length or batch size, so a single shape suffices.

@@ -702,6 +720,10 @@ class MambaMixer2(MambaBase, PluggableLayer):

720 else self.kv_cache[0].transpose(-1, -2)

721 )

722 ssm_state = self.kv_cache[1]

723 + if self.use_replayssm:

724 + x_cache, dt_cache, B_cache = self.kv_cache[2:]

725 + else:

726 + x_cache = dt_cache = B_cache = None

727 has_initial_states_p = attn_metadata.has_initial_states_p

728 prep_initial_states = attn_metadata.prep_initial_states

729 chunk_size = attn_metadata.chunk_size

@@ -1027,37 +1049,78 @@ class MambaMixer2(MambaBase, PluggableLayer):

1049 # - mamba_cache_params.ssm_state's slots will be selected

1050 # using state_indices_tensor_d

1051 # NOTE: final output is an in-place update of out tensor

1030 - selective_state_update(

1031 - ssm_state,

1032 - hidden_states_d,

1033 - dt_d,

1034 - A_d,

1035 - B_d,

1036 - C_d,

1037 - D_d,

1038 - dt_bias,

1039 - dt_softplus=True,

1040 - state_batch_indices=state_indices_tensor_d_input,

1041 - dst_state_batch_indices=state_indices_tensor_d_output,

1042 - out=preallocated_ssm_out_d.view(num_decode_tokens, -1, self.head_dim),

1043 - num_accepted_tokens=num_accepted_tokens,

1044 - cu_seqlens=query_start_loc_d,

1045 - is_blackwell=self.is_blackwell,

1052 + preallocated_ssm_out_d = preallocated_ssm_out_d.view(

1053 + num_decode_tokens, -1, self.head_dim

1054 )

1055 + if self.use_replayssm:

1056 + assert self.replayssm_buffer_len is not None

1057 + selective_state_update_replayssm_output_only(

1058 + ssm_state,

1059 + hidden_states_d,

1060 + dt_d,

1061 + A_d,

1062 + B_d,

1063 + C_d,

1064 + D_d,

1065 + dt_bias,

1066 + dt_softplus=True,

1067 + x_cache=x_cache,

1068 + dt_cache=dt_cache,

1069 + B_cache=B_cache,

1070 + bc_pre=attn_metadata.bc_pre_scratch,

1071 + write_pos=attn_metadata.write_pos_d,

1072 + is_flush=attn_metadata.is_flush_d,

1073 + max_cache_len=self.replayssm_buffer_len,

1074 + state_batch_indices=state_indices_tensor_d_input,

1075 + out=preallocated_ssm_out_d,

1076 + # Stochastic Rounding for the vanilla decode path is read

1077 + # from mamba_config inside ssu_dispatch; the replay kernel

1078 + # isn't a dispatch backend, so pass it here.

1079 + enable_stochastic_rounding=(

1080 + self.mamba_config.enable_stochastic_rounding

1081 + ),

1082 + cache_philox_rounds=(

1083 + self.mamba_config.stochastic_rounding_philox_rounds

1084 + ),

1085 + )

1086 + else:

1087 + selective_state_update(

1088 + ssm_state,

1089 + hidden_states_d,

1090 + dt_d,

1091 + A_d,

1092 + B_d,

1093 + C_d,

1094 + D_d,

1095 + dt_bias,

1096 + dt_softplus=True,

1097 + state_batch_indices=state_indices_tensor_d_input,

1098 + dst_state_batch_indices=state_indices_tensor_d_output,

1099 + out=preallocated_ssm_out_d,

1100 + num_accepted_tokens=num_accepted_tokens,

1101 + cu_seqlens=query_start_loc_d,

1102 + is_blackwell=self.is_blackwell,

1103 + )

1104

1048 - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]:

1105 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:

1106 assert self.model_config is not None

1107 assert self.cache_config is not None

1051 - return MambaStateDtypeCalculator.mamba2_state_dtype(

1108 + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype(

1109 self.model_config.dtype,

1110 self.cache_config.mamba_cache_dtype,

1111 self.cache_config.mamba_ssm_cache_dtype,

1112 )

1113 + if self.use_replayssm:

1114 + return MambaStateDtypeCalculator.append_replayssm_ring(

1115 + base_dtype, self.model_config.dtype

1116 + )

1117 + return base_dtype

1118

1057 - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]:

1058 - return MambaStateShapeCalculator.mamba2_state_shape(

1119 + def get_state_shape(self) -> tuple[tuple[int, ...], ...]:

1120 + tp_world_size = get_tensor_model_parallel_world_size()

1121 + base_shape = MambaStateShapeCalculator.mamba2_state_shape(

1122 intermediate_size=self.intermediate_size,

1060 - tp_world_size=get_tensor_model_parallel_world_size(),

1123 + tp_world_size=tp_world_size,

1124 n_groups=self.n_groups,

1125 num_heads=self.num_heads,

1126 head_dim=self.head_dim,

@@ -1065,6 +1128,15 @@ class MambaMixer2(MambaBase, PluggableLayer):

1128 conv_kernel=self.conv_kernel_size,

1129 num_spec=self.num_spec,

1130 )

1131 + if self.use_replayssm:

1132 + assert self.replayssm_buffer_len is not None

1133 + return MambaStateShapeCalculator.append_replayssm_ring(

1134 + base_shape,

1135 + self.n_groups,

1136 + tp_world_size,

1137 + self.replayssm_buffer_len,

1138 + )

1139 + return base_shape

1140

1141 @property

1142 def mamba_type(self) -> MambaAttentionBackendEnum:

vllm/model_executor/layers/mamba/mamba_utils.py

+34

@@ -80,6 +80,18 @@ class MambaStateDtypeCalculator:

80 model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype

81 )

83 + @classmethod

84 + def append_replayssm_ring(

85 + cls,

86 + base_dtypes: tuple[torch.dtype, ...],

87 + model_dtype: ModelDType | torch.dtype,

88 + ) -> tuple[torch.dtype, ...]:

89 + """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple:

90 + ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``.

91 + """

92 + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype)

93 + return (*base_dtypes, activation_dtype, torch.float32, activation_dtype)

94 +

95 @classmethod

96 def _mamba_state_dtype(

97 cls,

@@ -186,6 +198,28 @@ class MambaStateShapeCalculator:

198 temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size)

199 return conv_state_shape, temporal_state_shape

200

201 + @classmethod

202 + def append_replayssm_ring(

203 + cls,

204 + base_shapes: tuple[tuple[int, ...], ...],

205 + n_groups: int,

206 + tp_world_size: int,

207 + replayssm_buffer_len: int,

208 + ) -> tuple[tuple[int, ...], ...]:

209 + """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a

210 + base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape

211 + ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended

212 + ``n_groups``.

213 + """

214 + local_nheads, head_dim, state_size = base_shapes[1]

215 + local_ngroups = divide(n_groups, tp_world_size)

216 + return (

217 + *base_shapes,

218 + (local_nheads, replayssm_buffer_len, head_dim),

219 + (local_nheads, replayssm_buffer_len),

220 + (local_ngroups, replayssm_buffer_len, state_size),

221 + )

222 +

223 @classmethod

224 def short_conv_state_shape(

225 cls,

vllm/model_executor/layers/mamba/ops/replayssm_config.py new

+66

@@ -0,0 +1,66 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Launch-config selection for the ReplaySSM Mamba2 output_only decode kernel.

4 +

5 +Mirrors ``mamba_ssm.py``: a hard-coded heuristic per kernel, plus an

6 +``override`` context manager for benchmarks/tests/config sweeps. Hardware is

7 +auto-detected (Blackwell vs not) so call sites need not thread it through.

8 +"""

9 +

10 +import functools

11 +from contextlib import contextmanager

12 +

13 +from vllm.platforms import current_platform

14 +from vllm.triton_utils import triton

15 +

16 +

17 +@functools.cache

18 +def _is_blackwell() -> bool:

19 + try:

20 + return current_platform.is_device_capability_family(100)

21 + except Exception:

22 + return False

23 +

24 +

25 +# Per-kernel overrides keyed by the kernel name passed to get_replayssm_config.

26 +_overrides: dict[str, tuple] = {}

27 +

28 +

29 +@contextmanager

30 +def override_replayssm_config(kernel: str, config: tuple):

31 + """Pin ``kernel``'s launch config for the duration of the context."""

32 + prev = _overrides.get(kernel)

33 + _overrides[kernel] = config

34 + try:

35 + yield

36 + finally:

37 + if prev is None:

38 + _overrides.pop(kernel, None)

39 + else:

40 + _overrides[kernel] = prev

41 +

42 +

43 +def _dstate_tile(dstate: int, tile: int) -> int:

44 + return max(16, min(tile, triton.next_power_of_2(dstate)))

45 +

46 +

47 +def _mamba2_output_only(dstate, L, is_blackwell):

48 + # (block_size_m, num_warps, nf_dstate_tile, fl_dstate_tile, num_stages);

49 + # decoupled dstate tiling, serving-batch optimum, dtype-independent per device.

50 + if is_blackwell:

51 + return 64, 1, _dstate_tile(dstate, 32), _dstate_tile(dstate, 64), 2

52 + return 16, 1, _dstate_tile(dstate, 64), _dstate_tile(dstate, 128), 2

53 +

54 +

55 +def get_replayssm_config(kernel: str, **shape) -> tuple:

56 + """Return the launch config for ``kernel`` (override > tuned default).

57 +

58 + kernel: "mamba2_output_only". ``shape`` carries the keying dims (dstate;

59 + ``L`` for the buffer length, default 16); hardware is auto-detected.

60 + """

61 + if kernel in _overrides:

62 + return _overrides[kernel]

63 + bw = _is_blackwell()

64 + if kernel == "mamba2_output_only":

65 + return _mamba2_output_only(shape["dstate"], shape.get("L", 16), bw)

66 + raise ValueError(f"unknown ReplaySSM kernel config key: {kernel}")

vllm/model_executor/layers/mamba/ops/selective_state_update_replayssm_output_only.py new

+709

@@ -0,0 +1,709 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# ruff: noqa: E501

4 +

5 +import torch

6 +

7 +from vllm.model_executor.layers.mamba.ops.mamba_ssm import convert_rs_fp16x2, softplus

8 +from vllm.model_executor.layers.mamba.ops.replayssm_config import get_replayssm_config

9 +from vllm.triton_utils import tl, triton

10 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID

11 +

12 +

13 +@triton.heuristics(

14 + {

15 + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]

16 + is not None

17 + }

18 +)

19 +@triton.heuristics(

20 + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}

21 +)

22 +@triton.jit

23 +def _replayssm_output_only_precompute_kernel(

24 + B_ptr,

25 + C_ptr,

26 + B_cache_ptr,

27 + write_pos_ptr,

28 + is_flush_ptr,

29 + bc_pre_ptr,

30 + state_batch_indices_ptr,

31 + null_block_id,

32 + # Matrix dimensions

33 + batch,

34 + ngroups,

35 + dstate,

36 + # Input strides

37 + stride_B_batch,

38 + stride_B_group,

39 + stride_B_dstate,

40 + stride_C_batch,

41 + stride_C_group,

42 + stride_C_dstate,

43 + # Cache strides

44 + stride_B_cache_batch,

45 + stride_B_cache_group,

46 + stride_B_cache_pos,

47 + stride_B_cache_dstate,

48 + stride_bc_pre_batch,

49 + stride_bc_pre_group,

50 + stride_bc_pre_pos,

51 + stride_state_indices_batch,

52 + stride_state_indices_T,

53 + # Meta-parameters

54 + MAX_CACHE_LEN: tl.constexpr,

55 + BLOCK_SIZE_K: tl.constexpr,

56 + # heuristic-computed

57 + BLOCK_SIZE_DSTATE: tl.constexpr,

58 + HAS_STATE_BATCH_INDICES: tl.constexpr,

59 +):

60 + pid_b = tl.program_id(axis=0)

61 + pid_g = tl.program_id(axis=1)

62 +

63 + # On flush steps the main kernel does not read bc_pre, so skip the work.

64 + is_flush = tl.load(is_flush_ptr + pid_b) != 0

65 + if is_flush:

66 + return

67 +

68 + if HAS_STATE_BATCH_INDICES:

69 + state_batch_idx = tl.load(

70 + state_batch_indices_ptr

71 + + pid_b * stride_state_indices_batch

72 + + 0 * stride_state_indices_T

73 + ).to(tl.int64)

74 + if state_batch_idx == null_block_id:

75 + return

76 + else:

77 + state_batch_idx = pid_b

78 +

79 + offs_k = tl.arange(0, BLOCK_SIZE_K)

80 + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)

81 +

82 + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64)

83 +

84 + B_ptr += pid_b * stride_B_batch + pid_g * stride_B_group

85 + C_ptr += pid_b * stride_C_batch + pid_g * stride_C_group

86 + B_cache_ptr += state_batch_idx * stride_B_cache_batch + pid_g * stride_B_cache_group

87 + bc_pre_ptr += pid_b * stride_bc_pre_batch + pid_g * stride_bc_pre_group

88 +

89 + B_cur = tl.load(

90 + B_ptr + offs_n * stride_B_dstate,

91 + mask=offs_n < dstate,

92 + other=0.0,

93 + )

94 + C = tl.load(

95 + C_ptr + offs_n * stride_C_dstate,

96 + mask=offs_n < dstate,

97 + other=0.0,

98 + )

99 + B_cache_ptrs = (

100 + B_cache_ptr

101 + + offs_k[:, None] * stride_B_cache_pos

102 + + offs_n[None, :] * stride_B_cache_dstate

103 + )

104 + B_cache = tl.load(

105 + B_cache_ptrs,

106 + mask=(offs_k[:, None] < write_pos) & (offs_n[None, :] < dstate),

107 + other=0.0,

108 + )

109 + B_all = tl.where(offs_k[:, None] == write_pos, B_cur[None, :], B_cache)

110 + bc = tl.sum(B_all.to(tl.float32) * C[None, :].to(tl.float32), axis=1)

111 +

112 + tl.store(

113 + bc_pre_ptr + offs_k * stride_bc_pre_pos,

114 + bc,

115 + mask=(offs_k <= write_pos) & (offs_k < MAX_CACHE_LEN),

116 + )

117 +

118 +

119 +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None})

120 +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None})

121 +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})

122 +@triton.heuristics(

123 + {

124 + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"]

125 + is not None

126 + }

127 +)

128 +@triton.heuristics(

129 + {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}

130 +)

131 +@triton.jit

132 +def _replayssm_output_only_kernel(

133 + # Pointers to matrices

134 + state_ptr,

135 + rand_seed_ptr,

136 + x_ptr,

137 + dt_ptr,

138 + dt_bias_ptr,

139 + A_ptr,

140 + B_ptr,

141 + C_ptr,

142 + D_ptr,

143 + z_ptr,

144 + out_ptr,

145 + x_cache_ptr,

146 + dt_cache_ptr,

147 + B_cache_ptr,

148 + bc_pre_ptr,

149 + write_pos_ptr,

150 + is_flush_ptr,

151 + state_batch_indices_ptr,

152 + null_block_id,

153 + # Matrix dimensions

154 + batch,

155 + nheads,

156 + dim,

157 + dstate,

158 + nheads_ngroups_ratio,

159 + # State strides

160 + stride_state_batch,

161 + stride_state_head,

162 + stride_state_dim,

163 + stride_state_dstate,

164 + # Input strides

165 + stride_x_batch,

166 + stride_x_head,

167 + stride_x_dim,

168 + stride_dt_batch,

169 + stride_dt_head,

170 + stride_dt_bias_head,

171 + stride_A_head,

172 + stride_B_batch,

173 + stride_B_group,

174 + stride_B_dstate,

175 + stride_C_batch,

176 + stride_C_group,

177 + stride_C_dstate,

178 + stride_D_head,

179 + stride_D_dim,

180 + stride_z_batch,

181 + stride_z_head,

182 + stride_z_dim,

183 + stride_out_batch,

184 + stride_out_head,

185 + stride_out_dim,

186 + # Cache strides

187 + stride_x_cache_batch,

188 + stride_x_cache_head,

189 + stride_x_cache_dim,

190 + stride_x_cache_pos,

191 + stride_dt_cache_batch,

192 + stride_dt_cache_head,

193 + stride_dt_cache_pos,

194 + stride_B_cache_batch,

195 + stride_B_cache_group,

196 + stride_B_cache_pos,

197 + stride_B_cache_dstate,

198 + stride_bc_pre_batch,

199 + stride_bc_pre_group,

200 + stride_bc_pre_pos,

201 + stride_state_indices_batch,

202 + stride_state_indices_T,

203 + # Meta-parameters

204 + DT_SOFTPLUS: tl.constexpr,

205 + MAX_CACHE_LEN: tl.constexpr,

206 + BLOCK_SIZE_M: tl.constexpr,

207 + BLOCK_SIZE_K_CACHE: tl.constexpr,

208 + BLOCK_SIZE_K_DOT: tl.constexpr,

209 + NF_DSTATE_TILE: tl.constexpr,

210 + NF_NDS: tl.constexpr,

211 + FL_DSTATE_TILE: tl.constexpr,

212 + FL_NDS: tl.constexpr,

213 + USE_RS_ROUNDING: tl.constexpr,

214 + PHILOX_ROUNDS: tl.constexpr,

215 + # heuristic-computed

216 + BLOCK_SIZE_DSTATE: tl.constexpr,

217 + HAS_DT_BIAS: tl.constexpr,

218 + HAS_D: tl.constexpr,

219 + HAS_Z: tl.constexpr,

220 + HAS_STATE_BATCH_INDICES: tl.constexpr,

221 +):

222 + pid_m = tl.program_id(axis=0)

223 + pid_b = tl.program_id(axis=1)

224 + pid_h = tl.program_id(axis=2)

225 +

226 + # Resolve the physical state slot for this decode row; skip padded rows.

227 + if HAS_STATE_BATCH_INDICES:

228 + state_batch_idx = tl.load(

229 + state_batch_indices_ptr

230 + + pid_b * stride_state_indices_batch

231 + + 0 * stride_state_indices_T

232 + ).to(tl.int64)

233 + if state_batch_idx == null_block_id:

234 + return

235 + else:

236 + state_batch_idx = pid_b

237 +

238 + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)

239 + offs_n_full = tl.arange(0, BLOCK_SIZE_DSTATE)

240 +

241 + # Buffer cursor (number of cached tokens so far) and the flush flag.

242 + write_pos = tl.load(write_pos_ptr + pid_b).to(tl.int64)

243 + is_flush = tl.load(is_flush_ptr + pid_b) != 0

244 +

245 + # Advance every pointer to this (row, head, group).

246 + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head

247 + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head

248 + dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head

249 + B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group

250 + C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group

251 + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head

252 + x_cache_ptr += state_batch_idx * stride_x_cache_batch + pid_h * stride_x_cache_head

253 + dt_cache_ptr += (

254 + state_batch_idx * stride_dt_cache_batch + pid_h * stride_dt_cache_head

255 + )

256 + B_cache_ptr += (

257 + state_batch_idx * stride_B_cache_batch

258 + + (pid_h // nheads_ngroups_ratio) * stride_B_cache_group

259 + )

260 + bc_pre_ptr += (

261 + pid_b * stride_bc_pre_batch

262 + + (pid_h // nheads_ngroups_ratio) * stride_bc_pre_group

263 + )

264 +

265 + # Current-token dt (+ bias, softplus), scalar A, and current x. C, the

266 + # checkpoint state S_0, and current-token B are read per dstate tile below.

267 + dt_cur = tl.load(dt_ptr).to(tl.float32)

268 + if HAS_DT_BIAS:

269 + dt_cur += tl.load(dt_bias_ptr + pid_h * stride_dt_bias_head).to(tl.float32)

270 + if DT_SOFTPLUS:

271 + dt_cur = tl.where(dt_cur <= 20.0, softplus(dt_cur), dt_cur)

272 + A = tl.load(A_ptr + pid_h * stride_A_head).to(tl.float32)

273 + x_cur = tl.load(x_ptr + offs_m * stride_x_dim, mask=offs_m < dim, other=0.0)

274 +

275 + if not is_flush:

276 + # Output-only route: read y without materializing the state, using the

277 + # precomputed k^T q products (`bc`):

278 + # y = total_decay * (S_0 q) + sum_j s_j (k_j^T q) v_j.

279 + # Then append the current token to the buffer.

280 + offs_k_cache = tl.arange(0, BLOCK_SIZE_K_CACHE)

281 + # dt over the window (history + current token), then the decay weights.

282 + dt_all_cache = tl.load(

283 + dt_cache_ptr + offs_k_cache * stride_dt_cache_pos,

284 + mask=offs_k_cache < write_pos,

285 + other=0.0,

286 + ).to(tl.float32)

287 + dt_all_cache = tl.where(offs_k_cache == write_pos, dt_cur, dt_all_cache)

288 + dA_cumsum_cache = A * tl.cumsum(dt_all_cache, axis=0)

289 + dA_total_cache = A * tl.sum(dt_all_cache, axis=0)

290 + total_decay_cache = tl.exp(dA_total_cache)

291 + scale_cache = dt_all_cache * tl.exp(dA_total_cache - dA_cumsum_cache)

292 + scale_cache = tl.where(offs_k_cache <= write_pos, scale_cache, 0.0)

293 +

294 + # Gather buffered x over the window (history + current token).

295 + x_all_cache_ptrs = (

296 + x_cache_ptr

297 + + offs_m[:, None] * stride_x_cache_dim

298 + + offs_k_cache[None, :] * stride_x_cache_pos

299 + )

300 + x_all_cache = tl.load(

301 + x_all_cache_ptrs,

302 + mask=(offs_m[:, None] < dim) & (offs_k_cache[None, :] < write_pos),

303 + other=0.0,

304 + )

305 + x_all_cache = tl.where(

306 + offs_k_cache[None, :] == write_pos, x_cur[:, None], x_all_cache

307 + )

308 +

309 + # Decayed checkpoint readout sum_n S_0(m,n) q(n), streamed over NF dstate

310 + # tiles so the (M, N) state slice is never held whole.

311 + offs_nt = tl.arange(0, NF_DSTATE_TILE)

312 + ck_acc = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32)

313 + for i in tl.static_range(NF_NDS):

314 + offs_n = i * NF_DSTATE_TILE + offs_nt

315 + nmask = offs_n < dstate

316 + st = tl.load(

317 + state_ptr

318 + + offs_m[:, None] * stride_state_dim

319 + + offs_n[None, :] * stride_state_dstate,

320 + mask=(offs_m[:, None] < dim) & nmask[None, :],

321 + other=0.0,

322 + )

323 + c_chunk = tl.load(

324 + C_ptr + offs_n * stride_C_dstate, mask=nmask, other=0.0

325 + ).to(tl.float32)

326 + ck_acc += tl.sum(st.to(tl.float32) * c_chunk[None, :], axis=1)

327 + checkpoint_out = ck_acc * total_decay_cache

328 + bc_cache = tl.load(

329 + bc_pre_ptr + offs_k_cache * stride_bc_pre_pos,

330 + mask=offs_k_cache <= write_pos,

331 + other=0.0,

332 + )

333 + cache_out = tl.sum(

334 + x_all_cache.to(tl.float32) * (scale_cache * bc_cache)[None, :], axis=1

335 + )

336 + out = checkpoint_out + cache_out

337 +

338 + # Append the current token (x, dt, B) into the buffer at write_pos.

339 + tl.store(

340 + x_cache_ptr + offs_m * stride_x_cache_dim + write_pos * stride_x_cache_pos,

341 + x_cur,

342 + mask=offs_m < dim,

343 + )

344 + if pid_m == 0:

345 + B_cur = tl.load(

346 + B_ptr + offs_n_full * stride_B_dstate,

347 + mask=offs_n_full < dstate,

348 + other=0.0,

349 + )

350 + tl.store(dt_cache_ptr + write_pos * stride_dt_cache_pos, dt_cur)

351 + tl.store(

352 + B_cache_ptr

353 + + write_pos * stride_B_cache_pos

354 + + offs_n_full * stride_B_cache_dstate,

355 + B_cur,

356 + mask=offs_n_full < dstate,

357 + )

358 + else:

359 + # Flush step: state route. Reconstruct the state from cached inputs,

360 + # S_t = total_decay * S_0 + sum_j s_j (v_j k_j^T), persist it as the new

361 + # checkpoint, then read y = S_t q -- streamed over FL dstate tiles.

362 + offs_k_dot = tl.arange(0, BLOCK_SIZE_K_DOT)

363 + dt_all_dot = tl.load(

364 + dt_cache_ptr + offs_k_dot * stride_dt_cache_pos,

365 + mask=offs_k_dot < write_pos,

366 + other=0.0,

367 + ).to(tl.float32)

368 + dt_all_dot = tl.where(offs_k_dot == write_pos, dt_cur, dt_all_dot)

369 + dA_cumsum_dot = A * tl.cumsum(dt_all_dot, axis=0)

370 + dA_total_dot = A * tl.sum(dt_all_dot, axis=0)

371 + total_decay_dot = tl.exp(dA_total_dot)

372 + scale_dot = dt_all_dot * tl.exp(dA_total_dot - dA_cumsum_dot)

373 + scale_dot = tl.where(offs_k_dot <= write_pos, scale_dot, 0.0)

374 +

375 + # Gather buffered x over the window (history + current token).

376 + x_all_dot_ptrs = (

377 + x_cache_ptr

378 + + offs_m[:, None] * stride_x_cache_dim

379 + + offs_k_dot[None, :] * stride_x_cache_pos

380 + )

381 + x_all_dot = tl.load(

382 + x_all_dot_ptrs,

383 + mask=(offs_m[:, None] < dim) & (offs_k_dot[None, :] < write_pos),

384 + other=0.0,

385 + )

386 + x_all_dot = tl.where(

387 + offs_k_dot[None, :] == write_pos, x_cur[:, None], x_all_dot

388 + )

389 + x_all_ty = x_all_dot.to(x_ptr.dtype.element_ty)

390 +

391 + # Distinct tile locals (_f) from the nf branch: differing tile widths

392 + # would force a shape-mismatched merge at the if/else exit.

393 + offs_nt_f = tl.arange(0, FL_DSTATE_TILE)

394 + out = tl.zeros([BLOCK_SIZE_M], dtype=tl.float32)

395 + for i in tl.static_range(FL_NDS):

396 + offs_n_f = i * FL_DSTATE_TILE + offs_nt_f

397 + nmask_f = offs_n_f < dstate

398 + # Gather buffered B over the window (history + current token).

399 + B_all_dot = tl.load(

400 + B_cache_ptr

401 + + offs_k_dot[:, None] * stride_B_cache_pos

402 + + offs_n_f[None, :] * stride_B_cache_dstate,

403 + mask=(offs_k_dot[:, None] < write_pos) & nmask_f[None, :],

404 + other=0.0,

405 + )

406 + B_cur_tile = tl.load(

407 + B_ptr + offs_n_f * stride_B_dstate, mask=nmask_f, other=0.0

408 + )

409 + B_all_dot = tl.where(

410 + offs_k_dot[:, None] == write_pos, B_cur_tile[None, :], B_all_dot

411 + )

412 + # tf32x3 keeps fp32 parity with the elementwise baseline (plain tf32 on

413 + # fp32 inputs drifts ~1e-2); bf16/fp16 inputs are unaffected by this flag.

414 + B_scaled = (B_all_dot.to(tl.float32) * scale_dot[:, None]).to(

415 + x_ptr.dtype.element_ty

416 + )

417 + delta_state = tl.dot(x_all_ty, B_scaled, input_precision="tf32x3")

418 + state_ptrs = (

419 + state_ptr

420 + + offs_m[:, None] * stride_state_dim

421 + + offs_n_f[None, :] * stride_state_dstate

422 + )

423 + state_mask = (offs_m[:, None] < dim) & nmask_f[None, :]

424 + st_f = tl.load(state_ptrs, mask=state_mask, other=0.0)

425 + state_new = st_f.to(tl.float32) * total_decay_dot + delta_state.to(

426 + tl.float32

427 + )

428 + if USE_RS_ROUNDING:

429 + # Stochastic-round fp32->fp16 (Blackwell cvt.rs), mirroring the

430 + # baseline. Only the flush step stores state, so this runs at 1/L

431 + # the baseline's per-step rate. Absolute per-element offsets seed

432 + # the RNG so each state element draws independently.

433 + rand_seed = tl.load(rand_seed_ptr)

434 + rand_offsets = (

435 + state_batch_idx * stride_state_batch

436 + + pid_h * stride_state_head

437 + + offs_m[:, None] * stride_state_dim

438 + + offs_n_f[None, :] * stride_state_dstate

439 + )

440 + if PHILOX_ROUNDS > 0:

441 + rand = tl.randint(rand_seed, rand_offsets, PHILOX_ROUNDS)

442 + else:

443 + rand = tl.randint(rand_seed, rand_offsets)

444 + tl.static_assert(

445 + state_ptrs.dtype.element_ty == tl.float16,

446 + "stochastic rounding requires an fp16 SSM state cache",

447 + )

448 + state_store = convert_rs_fp16x2(state_new, rand)

449 + else:

450 + state_store = state_new.to(st_f.dtype)

451 + tl.store(state_ptrs, state_store, mask=state_mask)

452 + c_chunk_f = tl.load(

453 + C_ptr + offs_n_f * stride_C_dstate, mask=nmask_f, other=0.0

454 + ).to(tl.float32)

455 + out += tl.sum(state_new * c_chunk_f[None, :], axis=1)

456 +

457 + # Skip connection (D) and output gate (z).

458 + if HAS_D:

459 + D_ptr += pid_h * stride_D_head

460 + D = tl.load(D_ptr + offs_m * stride_D_dim, mask=offs_m < dim, other=0.0).to(

461 + tl.float32

462 + )

463 + out += x_cur.to(tl.float32) * D

464 + if HAS_Z:

465 + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head

466 + z = tl.load(z_ptr + offs_m * stride_z_dim, mask=offs_m < dim, other=0.0).to(

467 + tl.float32

468 + )

469 + out *= z * tl.sigmoid(z)

470 +

471 + tl.store(out_ptr + offs_m * stride_out_dim, out, mask=offs_m < dim)

472 +

473 +

474 +def selective_state_update_replayssm_output_only(

475 + state: torch.Tensor,

476 + x: torch.Tensor,

477 + dt: torch.Tensor,

478 + A: torch.Tensor,

479 + B: torch.Tensor,

480 + C: torch.Tensor,

481 + D: torch.Tensor | None = None,

482 + dt_bias: torch.Tensor | None = None,

483 + z: torch.Tensor | None = None,

484 + dt_softplus: bool = False,

485 + x_cache: torch.Tensor | None = None,

486 + dt_cache: torch.Tensor | None = None,

487 + B_cache: torch.Tensor | None = None,

488 + bc_pre: torch.Tensor | None = None,

489 + write_pos: torch.Tensor | None = None,

490 + is_flush: torch.Tensor | None = None,

491 + max_cache_len: int = 16,

492 + state_batch_indices: torch.Tensor | None = None,

493 + null_block_id: int = NULL_BLOCK_ID,

494 + out: torch.Tensor | None = None,

495 + enable_stochastic_rounding: bool = False,

496 + cache_philox_rounds: int = 0,

497 +) -> torch.Tensor:

498 + """Cached-bc SSM update for vLLM's autoregressive Mamba2 decode path."""

499 + has_heads = state.dim() > 3

500 + if state.dim() == 3:

501 + state = state.unsqueeze(1)

502 + if x.dim() == 2:

503 + x = x.unsqueeze(1)

504 + if dt.dim() == 2:

505 + dt = dt.unsqueeze(1)

506 + if A.dim() == 2:

507 + A = A.unsqueeze(0)

508 + if B.dim() == 2:

509 + B = B.unsqueeze(1)

510 + if C.dim() == 2:

511 + C = C.unsqueeze(1)

512 + if D is not None and D.dim() == 1:

513 + D = D.unsqueeze(0)

514 + if z is not None and z.dim() == 2:

515 + z = z.unsqueeze(1)

516 + if dt_bias is not None and dt_bias.dim() == 1:

517 + dt_bias = dt_bias.unsqueeze(0)

518 + if out is not None and out.dim() == 2:

519 + out = out.unsqueeze(1)

520 + if state_batch_indices is not None and state_batch_indices.dim() == 1:

521 + state_batch_indices = state_batch_indices.unsqueeze(1)

522 +

523 + _, nheads, dim, dstate = state.shape

524 + batch = x.shape[0]

525 + assert x.shape == (batch, nheads, dim)

526 + assert dt.shape == x.shape

527 + assert A.shape == (nheads, dim, dstate)

528 + ngroups = B.shape[1]

529 + assert nheads % ngroups == 0, "nheads must be divisible by ngroups"

530 + assert B.shape == (batch, ngroups, dstate)

531 + assert C.shape == B.shape

532 + if D is not None:

533 + assert D.shape == (nheads, dim)

534 + if z is not None:

535 + assert z.shape == x.shape

536 + if dt_bias is not None:

537 + assert dt_bias.shape == (nheads, dim)

538 + assert out is not None and out.shape == x.shape

539 +

540 + assert A.stride(-1) == 0 and A.stride(-2) == 0, (

541 + "Cached kernel requires TIE_HDIM (A scalar per head)"

542 + )

543 + assert dt.stride(-1) == 0, "Cached kernel requires TIE_HDIM (dt scalar per head)"

544 + if dt_bias is not None:

545 + assert dt_bias.stride(-1) == 0, (

546 + "Cached kernel requires TIE_HDIM (dt_bias scalar per head)"

547 + )

548 +

549 + assert x_cache is not None

550 + assert dt_cache is not None

551 + assert B_cache is not None

552 + assert x_cache.shape[1:] == (nheads, max_cache_len, dim)

553 + assert dt_cache.shape[1:] == (nheads, max_cache_len)

554 + assert B_cache.shape[1:] == (ngroups, max_cache_len, dstate)

555 + assert write_pos is not None and write_pos.shape[0] >= batch

556 + assert write_pos.dtype == torch.int32

557 + assert is_flush is not None and is_flush.shape[0] >= batch

558 + assert is_flush.dtype in (torch.bool, torch.int8)

559 + assert bc_pre is not None

560 + assert bc_pre.shape[0] >= batch and bc_pre.shape[1] >= ngroups

561 + assert bc_pre.shape[2] == max_cache_len

562 + assert bc_pre.dtype == torch.float32

563 + if state_batch_indices is not None:

564 + assert state_batch_indices.shape[0] >= batch

565 + assert state_batch_indices.shape[1] >= 1

566 +

567 + block_size_k_cache = max(1, triton.next_power_of_2(max_cache_len))

568 + block_size_k_dot = max(16, block_size_k_cache)

569 + block_size_m, num_warps, nf_tile, fl_tile, num_stages = get_replayssm_config(

570 + "mamba2_output_only", dstate=dstate, L=max_cache_len

571 + )

572 + bs_dstate = triton.next_power_of_2(dstate)

573 + nf_dstate_tile = max(16, min(nf_tile, bs_dstate))

574 + nf_nds = triton.cdiv(bs_dstate, nf_dstate_tile)

575 + fl_dstate_tile = max(16, min(fl_tile, bs_dstate))

576 + fl_nds = triton.cdiv(bs_dstate, fl_dstate_tile)

577 +

578 + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads)

579 + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0)

580 + state_indices_strides = (

581 + (state_batch_indices.stride(0), state_batch_indices.stride(1))

582 + if state_batch_indices is not None

583 + else (0, 0)

584 + )

585 + rand_seed = (

586 + torch.randint(0, 2**32, (1,), device=state.device)

587 + if enable_stochastic_rounding

588 + else None

589 + )

590 +

591 + with torch.accelerator.device_index(x.device.index):

592 + # Both kernels always launch: the precompute kernel self-skips flush

593 + # rows per row (the branch can't be hoisted out under CUDA graphs), so

594 + # it is a no-op when every row is flushing.

595 + _replayssm_output_only_precompute_kernel[(batch, ngroups)](

596 + B,

597 + C,

598 + B_cache,

599 + write_pos,

600 + is_flush,

601 + bc_pre,

602 + state_batch_indices,

603 + null_block_id,

604 + batch,

605 + ngroups,

606 + dstate,

607 + B.stride(0),

608 + B.stride(1),

609 + B.stride(2),

610 + C.stride(0),

611 + C.stride(1),

612 + C.stride(2),

613 + B_cache.stride(0),

614 + B_cache.stride(1),

615 + B_cache.stride(2),

616 + B_cache.stride(3),

617 + bc_pre.stride(0),

618 + bc_pre.stride(1),

619 + bc_pre.stride(2),

620 + state_indices_strides[0],

621 + state_indices_strides[1],

622 + max_cache_len,

623 + block_size_k_cache,

624 + num_warps=2,

625 + )

626 + _replayssm_output_only_kernel[grid](

627 + state,

628 + rand_seed,

629 + x,

630 + dt,

631 + dt_bias,

632 + A,

633 + B,

634 + C,

635 + D,

636 + z,

637 + out,

638 + x_cache,

639 + dt_cache,

640 + B_cache,

641 + bc_pre,

642 + write_pos,

643 + is_flush,

644 + state_batch_indices,

645 + null_block_id,

646 + batch,

647 + nheads,

648 + dim,

649 + dstate,

650 + nheads // ngroups,

651 + state.stride(0),

652 + state.stride(1),

653 + state.stride(2),

654 + state.stride(3),

655 + x.stride(0),

656 + x.stride(1),

657 + x.stride(2),

658 + dt.stride(0),

659 + dt.stride(1),

660 + dt_bias.stride(0) if dt_bias is not None else 0,

661 + A.stride(0),

662 + B.stride(0),

663 + B.stride(1),

664 + B.stride(2),

665 + C.stride(0),

666 + C.stride(1),

667 + C.stride(2),

668 + D.stride(0) if D is not None else 0,

669 + D.stride(1) if D is not None else 0,

670 + z_strides[0],

671 + z_strides[1],

672 + z_strides[2],

673 + out.stride(0),

674 + out.stride(1),

675 + out.stride(2),

676 + x_cache.stride(0),

677 + x_cache.stride(1),

678 + x_cache.stride(3),

679 + x_cache.stride(2),

680 + dt_cache.stride(0),

681 + dt_cache.stride(1),

682 + dt_cache.stride(2),

683 + B_cache.stride(0),

684 + B_cache.stride(1),

685 + B_cache.stride(2),

686 + B_cache.stride(3),

687 + bc_pre.stride(0),

688 + bc_pre.stride(1),

689 + bc_pre.stride(2),

690 + state_indices_strides[0],

691 + state_indices_strides[1],

692 + dt_softplus,

693 + max_cache_len,

694 + block_size_m,

695 + block_size_k_cache,

696 + block_size_k_dot,

697 + nf_dstate_tile,

698 + nf_nds,

699 + fl_dstate_tile,

700 + fl_nds,

701 + enable_stochastic_rounding,

702 + cache_philox_rounds,

703 + num_warps=num_warps,

704 + num_stages=num_stages,

705 + )

706 +

707 + if not has_heads:

708 + out = out.squeeze(1)

709 + return out

vllm/model_executor/models/interfaces.py

+25

@@ -984,6 +984,31 @@ def supports_mamba_prefix_caching(

984 return getattr(model, "supports_mamba_prefix_caching", False)

985

986

987 +@runtime_checkable

988 +class SupportsReplaySSM(Protocol):

989 + """The interface for models whose Mamba2 layers support ReplaySSM cached

990 + standard decode.

991 +

992 + This is currently experimental.

993 + """

994 +

995 + supports_replayssm: ClassVar[Literal[True]] = True

996 +

997 +

998 +@overload

999 +def supports_replayssm(model: object) -> TypeIs[SupportsReplaySSM]: ...

1000 +

1001 +

1002 +@overload

1003 +def supports_replayssm(model: type[object]) -> TypeIs[type[SupportsReplaySSM]]: ...

1004 +

1005 +

1006 +def supports_replayssm(

1007 + model: type[object] | object,

1008 +) -> TypeIs[type[SupportsReplaySSM]] | TypeIs[SupportsReplaySSM]:

1009 + return getattr(model, "supports_replayssm", False)

1010 +

1011 +

1012 @runtime_checkable

1013 class SupportsCrossEncoding(Protocol):

1014 """The interface required for all models that support cross encoding."""

vllm/model_executor/models/nemotron_h.py

+25 -6

@@ -69,6 +69,7 @@ from vllm.model_executor.models.interfaces import (

69 SupportsMambaPrefixCaching,

70 SupportsPP,

71 SupportsQuant,

72 + SupportsReplaySSM,

73 )

74 from vllm.model_executor.models.utils import (

75 AutoWeightsLoader,

@@ -707,6 +708,7 @@ class NemotronHForCausalLM(

708 SupportsQuant,

709 MixtureOfExperts,

710 SupportsMambaPrefixCaching,

711 + SupportsReplaySSM,

712 ):

713 # Relevant only if self.has_moe is True

714 is_non_gated_moe: bool = True

@@ -742,18 +744,24 @@ class NemotronHForCausalLM(

744 def get_mamba_state_dtype_from_config(

745 cls,

746 vllm_config: "VllmConfig",

745 - ) -> tuple[torch.dtype, torch.dtype]:

746 - return MambaStateDtypeCalculator.mamba2_state_dtype(

747 + ) -> tuple[torch.dtype, ...]:

748 + cache_config = vllm_config.cache_config

749 + base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype(

750 vllm_config.model_config.dtype,

748 - vllm_config.cache_config.mamba_cache_dtype,

749 - vllm_config.cache_config.mamba_ssm_cache_dtype,

751 + cache_config.mamba_cache_dtype,

752 + cache_config.mamba_ssm_cache_dtype,

753 )

754 + if cache_config.use_replayssm:

755 + return MambaStateDtypeCalculator.append_replayssm_ring(

756 + base_dtype, vllm_config.model_config.dtype

757 + )

758 + return base_dtype

759

760 @classmethod

761 def get_mamba_state_shape_from_config(

762 cls,

763 vllm_config: "VllmConfig",

756 - ) -> tuple[tuple[int, int], tuple[int, int, int]]:

764 + ) -> tuple[tuple[int, ...], ...]:

765 """Calculate shapes for Mamba's convolutional and state caches.

766

767 Args:

@@ -763,12 +771,15 @@ class NemotronHForCausalLM(

771 Tuple containing:

772 - conv_state_shape: Shape for convolutional state cache

773 - temporal_state_shape: Shape for state space model cache

774 + - (when use_replayssm is enabled) the x_cache/dt_cache/B_cache

775 + ring-buffer shapes

776 """

777 parallel_config = vllm_config.parallel_config

778 + cache_config = vllm_config.cache_config

779 hf_config = vllm_config.model_config.hf_config

780 intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim

781

771 - return MambaStateShapeCalculator.mamba2_state_shape(

782 + base_shape = MambaStateShapeCalculator.mamba2_state_shape(

783 intermediate_size=intermediate_size,

784 tp_world_size=parallel_config.tensor_parallel_size,

785 n_groups=hf_config.n_groups,

@@ -778,6 +789,14 @@ class NemotronHForCausalLM(

789 conv_kernel=hf_config.conv_kernel,

790 num_spec=vllm_config.num_speculative_tokens,

791 )

792 + if cache_config.use_replayssm:

793 + return MambaStateShapeCalculator.append_replayssm_ring(

794 + base_shape,

795 + hf_config.n_groups,

796 + parallel_config.tensor_parallel_size,

797 + cache_config.replayssm_buffer_len,

798 + )

799 + return base_shape

800

801 @classmethod

802 def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:

vllm/model_executor/models/registry.py

@@ -55,6 +55,7 @@ from .interfaces import (

55 supports_multimodal_encoder_tp_data,

56 supports_multimodal_raw_input_only,

57 supports_pp,

58 + supports_replayssm,

59 supports_transcription,

60 )

61 from .interfaces_base import (

@@ -789,6 +790,7 @@ class _ModelInfo:

790 is_hybrid: bool

791 has_noops: bool

792 supports_mamba_prefix_caching: bool

793 + supports_replayssm: bool

794 supports_transcription: bool

795 supports_transcription_only: bool

796

@@ -815,6 +817,7 @@ class _ModelInfo:

817 is_attention_free=is_attention_free(model),

818 is_hybrid=is_hybrid(model),

819 supports_mamba_prefix_caching=supports_mamba_prefix_caching(model),

820 + supports_replayssm=supports_replayssm(model),

821 supports_transcription=supports_transcription(model),

822 supports_transcription_only=(

823 supports_transcription(model) and model.supports_transcription_only

vllm/v1/attention/backend.py

@@ -480,6 +480,11 @@ class CommonAttentionMetadata:

480 fixed sliding window. None disables R-SWA. The attention backend copies this

481 into its own persistent buffer and reads ``rswa_window`` from model config."""

482

483 + replayssm_decode_base_cpu: torch.Tensor | None = None

484 + """(batch_size,) CPU ring origin for Mamba2 ReplaySSM decode: num_computed

485 + at the current decode run's last full-state write. write_pos counts from

486 + here, so a preemption-resumed request re-anchors past the prompt boundary."""

487 +

488 # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)

489 _seq_lens_cpu: torch.Tensor | None = None

490 _num_computed_tokens_cpu: torch.Tensor | None = None

@@ -591,6 +596,7 @@ class CommonAttentionMetadata:

596 dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),

597 is_prefilling=maybe_slice_reqs(self.is_prefilling),

598 rswa_prefix_lens=maybe_slice_reqs(self.rswa_prefix_lens),

599 + replayssm_decode_base_cpu=maybe_slice_reqs(self.replayssm_decode_base_cpu),

600 )

601

602

vllm/v1/attention/backends/mamba_attn.py

+144

@@ -74,6 +74,12 @@ class BaseMambaAttentionMetadata:

74 nums_dict: dict | None = None

75 batch_ptr: torch.Tensor | None = None

76 token_chunk_offset_ptr: torch.Tensor | None = None

77 + # ReplaySSM standard decode: per-row ring cursor and flush flag, plus the

78 + # per-step (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for the

79 + # precomputed k^T q products. All None when use_replayssm is disabled.

80 + write_pos_d: torch.Tensor | None = None

81 + is_flush_d: torch.Tensor | None = None

82 + bc_pre_scratch: torch.Tensor | None = None

85 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

@@ -98,6 +104,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

104 self.compilation_config = vllm_config.compilation_config

105 self.num_spec_tokens: int = vllm_config.num_speculative_tokens

106 self.use_spec_decode = self.num_spec_tokens > 0

107 + self.use_replayssm = vllm_config.cache_config.use_replayssm

108 + self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len

109

110 assert isinstance(kv_cache_spec, MambaSpec)

111 scheduler_config = vllm_config.scheduler_config

@@ -158,6 +166,36 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

166 dtype=torch.int32,

167 device=device,

168 )

169 + # ReplaySSM standard-decode CUDA-graph buffers: per-row ring cursor,

170 + # flush flag, and the k^T q precompute scratch.

171 + if self.use_replayssm:

172 + self.decode_write_pos_d: torch.Tensor = torch.empty(

173 + (self.decode_cudagraph_max_bs,),

174 + dtype=torch.int32,

175 + device=device,

176 + )

177 + self.decode_is_flush_d: torch.Tensor = torch.empty(

178 + (self.decode_cudagraph_max_bs,),

179 + dtype=torch.int8,

180 + device=device,

181 + )

182 + # B_cache shape = (ngroups, replayssm_buffer_len, dstate); the page

183 + # layout is (conv_state, ssm_state, x_cache, dt_cache, B_cache).

184 + bc_ngroups = kv_cache_spec.shapes[4][0]

185 + bc_scratch_bs = max(

186 + self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs

187 + )

188 + self.decode_bc_pre_scratch: torch.Tensor = torch.empty(

189 + (

190 + bc_scratch_bs,

191 + bc_ngroups,

192 + self.replayssm_buffer_len,

193 + ),

194 + dtype=torch.float32,

195 + device=device,

196 + )

197 + else:

198 + self.decode_bc_pre_scratch = None

199

200 self._init_reorder_batch_threshold(1, self.use_spec_decode)

201 if self.use_spec_decode:

@@ -414,6 +452,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

452 has_prior_state = seq_lens_cpu > 1

453 prefill_to_decode = single_token_prefill_rows & has_prior_state

454 if torch.any(prefill_to_decode).item():

455 + # ReplaySSM handles these rows as single-token flushes (see the

456 + # write-position derivation below), same as the baseline decode path.

457 is_prefilling = is_prefilling.clone()

458 is_prefilling[prefill_to_decode] = False

459 common_attn_metadata = common_attn_metadata.replace(

@@ -444,6 +484,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

484

485 # for causal_conv1d

486 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None

487 + write_pos_d = None

488 + is_flush_d = None

489

490 if self.vllm_config.cache_config.mamba_cache_mode == "all":

491 num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()

@@ -530,6 +572,79 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

572 num_reqs - num_prefills : num_reqs

573 ]

574

575 + if self.use_replayssm and num_decodes > 0:

576 + decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu

577 + num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu

578 + if decode_base_cpu is None or num_computed_tokens_cpu is None:

579 + raise ValueError(

580 + "--use-replayssm requires CPU decode-base and "

581 + "computed-token counts to derive decode write positions"

582 + )

583 + num_computed_d = num_computed_tokens_cpu[:num_decodes]

584 + decode_base_d = decode_base_cpu[:num_decodes]

585 + align_mode = self.vllm_config.cache_config.mamba_cache_mode == "align"

586 + block_size = self.kv_cache_spec.block_size

587 + if align_mode:

588 + # After a boundary the align copy leaves an exact checkpoint at

589 + # the block start and the new block's ring restarts empty, so

590 + # re-anchor there; max() keeps the prompt-end anchor for the

591 + # first (partial) block.

592 + effective_base = torch.maximum(

593 + decode_base_d, (num_computed_d // block_size) * block_size

594 + )

595 + else:

596 + effective_base = decode_base_d

597 + # write_pos counts decode steps since the ring's last full-state

598 + # write (the anchor), so a resumed request re-anchors correctly.

599 + decode_steps_cpu = num_computed_d - effective_base

600 + query_lens_cpu = (

601 + common_attn_metadata.query_start_loc_cpu[1 : num_decodes + 1]

602 + - common_attn_metadata.query_start_loc_cpu[:num_decodes]

603 + )

604 + valid_decode_rows = query_lens_cpu > 0

605 + # A single-token prefill row replayed as decode (query_len==1 with

606 + # prior state) has decode_steps < 0; force it to a one-token flush

607 + # (write_pos=0, is_flush=1). The flush branch reads an empty history

608 + # window, so it applies exactly one recurrence step off the checkpoint

609 + # -- identical to the baseline decode kernel for that row. The split

610 + # (treat_short_extends_as_decodes=False) admits only such rows here.

611 + leftover_prompt = valid_decode_rows & (decode_steps_cpu < 0)

612 + decode_steps_cpu = torch.where(

613 + valid_decode_rows & ~leftover_prompt,

614 + decode_steps_cpu,

615 + torch.zeros_like(decode_steps_cpu),

616 + )

617 + write_pos_cpu = torch.remainder(decode_steps_cpu, self.replayssm_buffer_len)

618 + is_flush_cpu = (

619 + write_pos_cpu == self.replayssm_buffer_len - 1

620 + ) | leftover_prompt

621 + if align_mode:

622 + # Force a flush on the step completing a mamba block so the exact

623 + # boundary state is materialized for prefix caching.

624 + is_flush_cpu = is_flush_cpu | (

625 + valid_decode_rows

626 + & ((num_computed_d + query_lens_cpu) % block_size == 0)

627 + )

628 + is_flush_cpu = is_flush_cpu.to(torch.int8)

629 + write_pos_d = async_tensor_h2d(

630 + write_pos_cpu.to(torch.int32).tolist(),

631 + dtype=torch.int32,

632 + device=common_attn_metadata.query_start_loc.device,

633 + )

634 + is_flush_d = async_tensor_h2d(

635 + is_flush_cpu.tolist(),

636 + dtype=torch.int8,

637 + device=common_attn_metadata.query_start_loc.device,

638 + )

639 +

640 + bc_pre_scratch = None

641 + if (

642 + self.use_replayssm

643 + and self.decode_bc_pre_scratch is not None

644 + and num_decodes > 0

645 + ):

646 + bc_pre_scratch = self.decode_bc_pre_scratch[:num_decodes]

647 +

648 metadata = self.metadata_cls(

649 num_prefills=num_prefills,

650 num_prefill_tokens=num_prefill_tokens,

@@ -539,6 +654,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

654 has_initial_states_p=has_initial_states_p,

655 state_indices_tensor_p=state_indices_tensor_p,

656 state_indices_tensor_d=state_indices_tensor_d,

657 + write_pos_d=write_pos_d,

658 + is_flush_d=is_flush_d,

659 + bc_pre_scratch=bc_pre_scratch,

660 num_accepted_tokens=num_accepted_tokens,

661 query_start_loc_d=query_start_loc_d,

662 block_idx_last_scheduled_token=block_idx_last_scheduled_token,

@@ -573,6 +691,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

691 block_idx_last_scheduled_token_prev_step = (

692 metadata.block_idx_last_scheduled_token_prev_step

693 )

694 + write_pos_d = metadata.write_pos_d

695 + is_flush_d = metadata.is_flush_d

696 + bc_pre_scratch = metadata.bc_pre_scratch

697 if (

698 metadata.num_prefills == 0

699 and metadata.num_decodes <= self.decode_cudagraph_max_bs

@@ -634,11 +755,34 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):

755 )

756 block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0

757

758 + if self.use_replayssm:

759 + assert write_pos_d is not None

760 + assert is_flush_d is not None

761 + self.decode_write_pos_d[: metadata.num_decodes].copy_(

762 + write_pos_d[: metadata.num_decodes],

763 + non_blocking=True,

764 + )

765 + write_pos_d = self.decode_write_pos_d[:padded_bs]

766 + write_pos_d[metadata.num_decodes :] = 0

767 +

768 + self.decode_is_flush_d[: metadata.num_decodes].copy_(

769 + is_flush_d[: metadata.num_decodes],

770 + non_blocking=True,

771 + )

772 + is_flush_d = self.decode_is_flush_d[:padded_bs]

773 + is_flush_d[metadata.num_decodes :] = 0

774 +

775 + if self.decode_bc_pre_scratch is not None:

776 + bc_pre_scratch = self.decode_bc_pre_scratch[:padded_bs]

777 +

778 return replace(

779 metadata,

780 state_indices_tensor_d=state_indices_tensor_d,

781 query_start_loc_d=query_start_loc_d,

782 num_accepted_tokens=num_accepted_tokens,

783 + write_pos_d=write_pos_d,

784 + is_flush_d=is_flush_d,

785 + bc_pre_scratch=bc_pre_scratch,

786 block_idx_last_scheduled_token=block_idx_last_scheduled_token,

787 block_idx_last_computed_token=block_idx_last_computed_token,

788 block_idx_last_scheduled_token_prev_step=(

vllm/v1/worker/gpu_input_batch.py

+26

@@ -106,6 +106,7 @@ class InputBatch:

106 is_pooling_model: bool = False,

107 cp_kv_cache_interleave_size: int = 1,

108 reasoning_config: ReasoningConfig | None = None,

109 + use_replayssm: bool = False,

110 slot_mapping_modes: list[SlotMappingMode] | None = None,

111 ):

112 self.thinking_budget_state_holder = maybe_create_thinking_budget_state_holder(

@@ -170,6 +171,17 @@ class InputBatch:

171 )

172 self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy()

173

174 + # Mamba2 ReplaySSM decode ring origin (num_computed at each request's

175 + # last full-state write); populated only when the feature is on.

176 + self.use_replayssm = use_replayssm

177 + self.replayssm_decode_base_cpu_tensor = torch.zeros(

178 + (max_num_reqs,),

179 + device="cpu",

180 + dtype=torch.int32,

181 + pin_memory=PIN_MEMORY,

182 + )

183 + self.replayssm_decode_base = self.replayssm_decode_base_cpu_tensor.numpy()

184 +

185 # Block table.

186 self.block_table = MultiGroupBlockTable(

187 max_num_reqs=max_num_reqs,

@@ -377,6 +389,11 @@ class InputBatch:

389 # Number of tokens without spec decode tokens.

390 self.num_tokens_no_spec[req_index] = request.num_tokens

391

392 + if self.use_replayssm:

393 + # Ring origin = full context at (re)admission (prompt + any resumed

394 + # output), so a resumed request re-anchors past the prompt.

395 + self.replayssm_decode_base[req_index] = request.num_tokens

396 +

397 self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens

398 self.block_table.add_row(request.block_ids, req_index)

399

@@ -597,6 +614,11 @@ class InputBatch:

614 self.num_prompt_tokens[i2],

615 self.num_prompt_tokens[i1],

616 )

617 + if self.use_replayssm:

618 + self.replayssm_decode_base[i1], self.replayssm_decode_base[i2] = (

619 + self.replayssm_decode_base[i2],

620 + self.replayssm_decode_base[i1],

621 + )

622 self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = (

623 self.num_computed_tokens_cpu[i2],

624 self.num_computed_tokens_cpu[i1],

@@ -754,6 +776,10 @@ class InputBatch:

776 last_req_index

777 ]

778 self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index]

779 + if self.use_replayssm:

780 + self.replayssm_decode_base[empty_index] = self.replayssm_decode_base[

781 + last_req_index

782 + ]

783 self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[

784 last_req_index

785 ]

vllm/v1/worker/gpu_model_runner.py

@@ -727,6 +727,7 @@ class GPUModelRunner(

727 is_pooling_model=self.is_pooling_model,

728 cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,

729 reasoning_config=self.vllm_config.reasoning_config,

730 + use_replayssm=self.cache_config.use_replayssm,

731 )

732

733 # Separate cuda stream for overlapping transfer of sampled token ids from

@@ -2413,6 +2414,12 @@ class GPUModelRunner(

2414 if self.model_config.rswa_window is not None:

2415 rswa_prefix_lens = num_prompt_tokens_cpu

2416

2417 + replayssm_decode_base_cpu = None

2418 + if self.cache_config.use_replayssm:

2419 + replayssm_decode_base_cpu = (

2420 + self.input_batch.replayssm_decode_base_cpu_tensor[:num_reqs_padded]

2421 + )

2422 +

2423 cm_base = CommonAttentionMetadata(

2424 query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],

2425 query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],

@@ -2420,6 +2427,7 @@ class GPUModelRunner(

2427 _seq_lens_cpu=seq_lens_cpu,

2428 _num_computed_tokens_cpu=num_computed_tokens_cpu,

2429 seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,

2430 + replayssm_decode_base_cpu=replayssm_decode_base_cpu,

2431 num_reqs=num_reqs_padded,

2432 num_actual_tokens=num_tokens_padded,

2433 max_query_len=max_query_len,

@@ -7262,6 +7270,7 @@ class GPUModelRunner(

7270 is_pooling_model=self.is_pooling_model,

7271 cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,

7272 reasoning_config=self.vllm_config.reasoning_config,

7273 + use_replayssm=self.cache_config.use_replayssm,

7274 slot_mapping_modes=slot_mapping_modes,

7275 )

7276