@tduncan / Nvidia-Optimization-Stack-VLLM / commits / adebc41b7e

[Mamba] Add FlashInfer ReplaySSM backend (#52506)

Signed-off-by: Andrii Skliar <askliar@nvidia.com> Signed-off-by: Andrii Skliar <andreyws96@gmail.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>

Andrii Skliar committed Sep 1, 2026 at 16:26 UTC adebc41b7e9f1085d3f73434e23beb76883b9eb4
19 files changed +1029 -105
tests/kernels/mamba/test_ssu_dispatch.py
+70 -2
@@ -7,12 +7,15 @@ import pytest
7 import torch
8
9 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
10 +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
11 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
12 FlashInferSSUBackend,
13 TritonSSUBackend,
14 get_mamba_ssu_backend,
15 initialize_mamba_ssu_backend,
16 + reset_replayssm_ring_trackers,
17 selective_state_update,
18 + update_replayssm_ring_trackers,
19 )
20 from vllm.utils.torch_utils import set_random_seed
21 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
@@ -30,6 +33,47 @@ except ImportError:
33 HAS_FLASHINFER = False
34
35
36 +@pytest.fixture(autouse=True)
37 +def restore_backend_state():
38 + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
39 +
40 + old_backend = mod._mamba_ssu_backend
41 + old_replayssm_kernel = mod._flashinfer_replayssm_kernel
42 + yield
43 + mod._mamba_ssu_backend = old_backend
44 + mod._flashinfer_replayssm_kernel = old_replayssm_kernel
45 +
46 +
47 +def test_flashinfer_replayssm_ring_tracker_lifecycle():
48 + ring_start = torch.zeros(2, dtype=torch.int32, device="cuda")
49 + prev_num_accepted = torch.zeros(2, dtype=torch.int32, device="cuda")
50 + state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda")
51 +
52 + observed = []
53 + for _ in range(33):
54 + update_replayssm_ring_trackers(
55 + ring_start,
56 + prev_num_accepted,
57 + state_batch_indices,
58 + logical_window=16,
59 + ring_buffer_len=17,
60 + )
61 + observed.append((int(ring_start[1]), int(prev_num_accepted[1])))
62 +
63 + assert observed[4] == (0, 5)
64 + assert observed[15] == (0, 16)
65 + assert observed[16] == (16, 1)
66 + assert observed[31] == (16, 16)
67 + assert observed[32] == (15, 1)
68 +
69 + reset_replayssm_ring_trackers(
70 + ring_start,
71 + prev_num_accepted,
72 + state_batch_indices,
73 + )
74 + assert (ring_start[1].item(), prev_num_accepted[1].item()) == (0, 0)
75 +
76 +
77 def _kv_cache_config_with_ssu(
78 mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2,
79 ) -> KVCacheConfig:
@@ -116,11 +160,10 @@ def test_flashinfer_forwards_ssu_algorithm(
160 def test_uninitialized_backend_raises():
161 import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
162
119 - old = mod._mamba_ssu_backend
163 + # restore_backend_state (autouse) puts the global back afterwards.
164 mod._mamba_ssu_backend = None
165 with pytest.raises(RuntimeError, match="not been initialized"):
166 get_mamba_ssu_backend()
123 - mod._mamba_ssu_backend = old
167
168
169 @pytest.mark.parametrize(
@@ -186,3 +229,28 @@ def test_triton_basic_call():
229 out=out,
230 )
231 assert not torch.isnan(out).any()
232 +
233 +
234 +@pytest.mark.parametrize(
235 + ("backend", "expected_ring_len"),
236 + [
237 + (MambaBackendEnum.TRITON, 16),
238 + (MambaBackendEnum.FLASHINFER, 17),
239 + ],
240 +)
241 +def test_replayssm_physical_ring_shape(backend, expected_ring_len):
242 + base_shapes = ((64, 3), (8, 4, 16))
243 +
244 + shapes = MambaStateShapeCalculator.append_replayssm_ring(
245 + base_shapes,
246 + n_groups=4,
247 + tp_world_size=2,
248 + logical_window=16,
249 + backend=backend,
250 + )
251 +
252 + assert shapes[2:] == (
253 + (8, expected_ring_len, 4),
254 + (8, expected_ring_len),
255 + (2, expected_ring_len, 16),
256 + )
tests/model_executor/test_replayssm_warmup.py new
+150
@@ -0,0 +1,150 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from types import SimpleNamespace
5 +from unittest.mock import Mock, patch
6 +
7 +import numpy as np
8 +import pytest
9 +import torch
10 +
11 +from vllm.config.mamba import MambaBackendEnum
12 +from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
13 +from vllm.model_executor.warmup import replayssm_warmup as warmup
14 +from vllm.platforms import current_platform
15 +from vllm.utils.flashinfer import has_flashinfer
16 +
17 +pytestmark = pytest.mark.skipif(
18 + not current_platform.is_cuda() or not has_flashinfer(),
19 + reason="FlashInfer ReplaySSM warmup tests require CUDA and FlashInfer",
20 +)
21 +
22 +PREFILL_KWARGS = {
23 + "num_tokens": 128,
24 + "skip_eplb": True,
25 + "is_profile": True,
26 + "randomize_inputs": True,
27 +}
28 +
29 +
30 +def _autotune_runner(
31 + *,
32 + use_v2_model_runner: bool = False,
33 + query_len: int = 6,
34 + max_num_seqs: int = 32,
35 + num_blocks: int = 17,
36 + max_num_tokens: int = 100,
37 + use_replayssm: bool = True,
38 + backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER,
39 +) -> SimpleNamespace:
40 + return SimpleNamespace(
41 + vllm_config=SimpleNamespace(
42 + cache_config=SimpleNamespace(use_replayssm=use_replayssm),
43 + mamba_config=SimpleNamespace(backend=backend),
44 + use_v2_model_runner=use_v2_model_runner,
45 + ),
46 + uniform_decode_query_len=query_len,
47 + decode_query_len=query_len,
48 + max_num_tokens=max_num_tokens,
49 + scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
50 + kv_cache_config=SimpleNamespace(num_blocks=num_blocks),
51 + )
52 +
53 +
54 +@pytest.mark.parametrize(
55 + ("runner_kwargs", "expected_num_reqs"),
56 + [
57 + # max_num_seqs (32) vs max_num_tokens // query_len (16) vs blocks-1 (16).
58 + (dict(query_len=6, use_v2_model_runner=False), 16),
59 + (dict(query_len=6, use_v2_model_runner=True), 16),
60 + # num_blocks - 1 is the binding constraint.
61 + (dict(query_len=1, max_num_tokens=128, max_num_seqs=64, num_blocks=5), 4),
62 + ],
63 + ids=["v1", "v2", "clamped_to_state_capacity"],
64 +)
65 +def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
66 + query_len = runner_kwargs["query_len"]
67 + with patch.object(
68 + warmup, "flashinfer_replayssm_autotune_supported", return_value=True
69 + ):
70 + result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs))
71 +
72 + expected_kwargs = {
73 + **PREFILL_KWARGS,
74 + "num_tokens": expected_num_reqs * query_len,
75 + "uniform_decode": True,
76 + }
77 + if runner_kwargs.get("use_v2_model_runner"):
78 + expected_kwargs["valid_dummy_state_slots"] = True
79 + else:
80 + expected_kwargs.update(
81 + allow_microbatching=False,
82 + force_attention=True,
83 + profile_seq_lens=query_len + 1,
84 + )
85 + assert result == (expected_num_reqs, expected_kwargs)
86 +
87 +
88 +@pytest.mark.parametrize(
89 + ("runner_kwargs", "flashinfer_supported"),
90 + [
91 + (dict(use_replayssm=False), True),
92 + (dict(backend=MambaBackendEnum.TRITON), True),
93 + ({}, False),
94 + ],
95 + ids=["replayssm_disabled", "non_flashinfer_backend", "kernel_unavailable"],
96 +)
97 +def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported):
98 + with patch.object(
99 + warmup,
100 + "flashinfer_replayssm_autotune_supported",
101 + return_value=flashinfer_supported,
102 + ):
103 + result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs))
104 + assert result is None
105 +
106 +
107 +def test_replayssm_autotune_slots_restore_state_and_trackers():
108 + mixer = MambaMixer2.__new__(MambaMixer2)
109 + torch.nn.Module.__init__(mixer)
110 + mixer.use_replayssm = True
111 + mixer.replayssm_buffer_len = 16
112 + mixer.kv_cache = (
113 + torch.full((4, 2), 3.0),
114 + torch.full((4, 2), 3.0),
115 + *(torch.full((4, 2, 17), 3.0) for _ in range(3)),
116 + )
117 + mixer._replayssm_ring_start = torch.full((4,), 3, dtype=torch.int32)
118 + mixer._replayssm_prev_num_accepted = torch.full((4,), 3, dtype=torch.int32)
119 + tracked = (
120 + *mixer.kv_cache,
121 + mixer._replayssm_ring_start,
122 + mixer._replayssm_prev_num_accepted,
123 + )
124 +
125 + block_ids = np.arange(10, 14, dtype=np.int32).reshape(4, 1)
126 + original_block_ids = block_ids.copy()
127 + block_table = SimpleNamespace(block_table=SimpleNamespace(np=block_ids))
128 + multi_group_block_table = SimpleNamespace(
129 + block_tables=[block_table], commit_block_table=Mock()
130 + )
131 + runner = SimpleNamespace(
132 + vllm_config=SimpleNamespace(use_v2_model_runner=False),
133 + input_batch=SimpleNamespace(block_table=multi_group_block_table),
134 + get_model=lambda: SimpleNamespace(modules=lambda: (mixer,)),
135 + )
136 +
137 + with warmup._temporary_replayssm_autotune_state(runner, 2):
138 + assert block_ids[:2, 0].tolist() == [1, 2]
139 + for tensor in tracked:
140 + tensor[1:3].fill_(9)
141 +
142 + assert np.array_equal(block_ids, original_block_ids)
143 + assert multi_group_block_table.commit_block_table.call_args_list == [
144 + ((2,), {}),
145 + ((2,), {}),
146 + ]
147 + for tensor in tracked:
148 + assert torch.count_nonzero(tensor[1:3]) == 0
149 + assert torch.all(tensor[0] == 3)
150 + assert torch.all(tensor[3] == 3)
tests/v1/attention/test_replayssm_metadata_builder.py
+46 -12
@@ -16,6 +16,7 @@ from tests.v1.attention.utils import (
16 create_common_attn_metadata,
17 create_vllm_config,
18 )
19 +from vllm.config.mamba import MambaBackendEnum
20 from vllm.v1.kv_cache_interface import MambaSpec
21
22 BLOCK_SIZE = 16
@@ -187,34 +188,47 @@ REPLAYSSM_BUILD_CASES = {
188 }
189
190
190 -def _make_mamba_spec(buffer_len: int) -> MambaSpec:
191 - # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups).
191 +def _make_mamba_spec(
192 + buffer_len: int,
193 + mamba_backend: MambaBackendEnum,
194 +) -> MambaSpec:
195 + ring_buffer_len = buffer_len + (
196 + 1 if mamba_backend == MambaBackendEnum.FLASHINFER else 0
197 + )
198 + shapes = (
199 + (1, 1),
200 + (1, 1, 1),
201 + (1, ring_buffer_len, 1),
202 + (1, ring_buffer_len),
203 + (1, ring_buffer_len, 1),
204 + )
205 return MambaSpec(
206 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 - ),
207 + shapes=shapes,
208 dtypes=(torch.float32,),
209 )
210
211
212 def _create_replayssm_builder(
206 - buffer_len: int, mamba_cache_mode: str = "none"
213 + buffer_len: int,
214 + mamba_cache_mode: str = "none",
215 + *,
216 + mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON,
217 ) -> MockMambaBuilder:
218 vllm_config = create_vllm_config(
219 model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE
220 )
221 # Set the flags after construction to skip validate_mamba_cached_kernel
212 - # (it requires a Triton backend) on the mock model.
222 + # (it requires a real SupportsReplaySSM model) on the mock model.
223 vllm_config.cache_config.use_replayssm = True
224 vllm_config.cache_config.replayssm_buffer_len = buffer_len
225 vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
226 + vllm_config.mamba_config.backend = mamba_backend
227 return MockMambaBuilder(
217 - _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE
228 + _make_mamba_spec(buffer_len, mamba_backend),
229 + ["layer0"],
230 + vllm_config,
231 + DEVICE,
232 )
233
234
@@ -254,3 +268,23 @@ def test_resumed_request_differs_from_fresh():
268
269 assert meta.write_pos_d.tolist()[:2] == [5, 0]
270 assert meta.is_flush_d.tolist()[:2] == [0, 0]
271 +
272 +
273 +def test_flashinfer_replayssm_scratch_metadata_fresh_decode():
274 + checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
275 + if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"):
276 + pytest.skip("FlashInfer does not expose ReplaySSM scratch allocation")
277 +
278 + builder = _create_replayssm_builder(16, mamba_backend=MambaBackendEnum.FLASHINFER)
279 + case = REPLAYSSM_BUILD_CASES["fresh_decode"]
280 + meta = _build(builder, case)
281 +
282 + assert meta.write_pos_d is None
283 + assert meta.is_flush_d is None
284 + assert meta.bc_pre_scratch is None
285 + assert meta.replayssm_scratch is not None
286 + assert [tensor.shape for tensor in meta.replayssm_scratch] == [
287 + (1, 1, 32, 8),
288 + (1, 1, 16),
289 + (1, 1, 32, 8),
290 + ]
tests/v1/e2e/test_replayssm_decode.py
+37 -2
@@ -4,6 +4,7 @@
4
5 import pytest
6
7 +import vllm.envs as envs
8 from vllm.v1.metrics.reader import Counter
9
10 from ...models.utils import check_logprobs_close
@@ -21,29 +22,48 @@ PROMPTS = [
22 ]
23
24
24 -def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1):
25 +def _check_replayssm_parity(
26 + vllm_runner,
27 + model_name,
28 + *,
29 + tensor_parallel_size=1,
30 + mamba_backend: str = "triton",
31 + name_1: str = "replayssm",
32 + require_v2: bool = False,
33 + monkeypatch: pytest.MonkeyPatch | None = None,
34 +):
35 # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a
36 # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are
37 # common-mode and only ReplaySSM varies.
38 + if require_v2:
39 + assert monkeypatch is not None
40 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
41 + envs.disable_envs_cache()
42 +
43 common = dict(
44 max_model_len=1024,
45 trust_remote_code=True,
46 enable_prefix_caching=False,
47 mamba_cache_mode="none",
48 tensor_parallel_size=tensor_parallel_size,
49 + mamba_backend=mamba_backend,
50 )
51 with vllm_runner(model_name, **common) as llm:
52 + if require_v2:
53 + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner
54 baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
55 with vllm_runner(
56 model_name, use_replayssm=True, replayssm_buffer_len=16, **common
57 ) as llm:
58 + if require_v2:
59 + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner
60 replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
61
62 check_logprobs_close(
63 outputs_0_lst=baseline,
64 outputs_1_lst=replay,
65 name_0="baseline",
46 - name_1="replayssm",
66 + name_1=name_1,
67 )
68
69
@@ -60,6 +80,21 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name):
80 _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2)
81
82
83 +@pytest.mark.parametrize("model_name", MODELS)
84 +def test_replayssm_flashinfer_decode_matches_baseline_v2(
85 + vllm_runner, model_name, monkeypatch
86 +):
87 + pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
88 + _check_replayssm_parity(
89 + vllm_runner,
90 + model_name,
91 + mamba_backend="flashinfer",
92 + name_1="replayssm_flashinfer_v2",
93 + require_v2=True,
94 + monkeypatch=monkeypatch,
95 + )
96 +
97 +
98 # Prefix spans several mamba blocks; prefix caching only reuses full blocks.
99 _PC_SENTENCE = (
100 "In a detailed survey of state space models, the authors compared many "
tests/v1/worker/test_kv_cache_allocation_scope.py
+2 -2
@@ -48,7 +48,7 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
48 result = attn_utils.init_kv_cache(
49 [],
50 {},
51 - object(),
51 + SimpleNamespace(kv_cache_groups=[]),
52 torch.device("cpu"),
53 [],
54 config,
@@ -83,7 +83,7 @@ def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
83 )
84 result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors(
85 runner,
86 - object(),
86 + SimpleNamespace(kv_cache_groups=[]),
87 [],
88 kv_cache_allocation_context=scope,
89 )
tests/v1/worker/test_utils.py
+57
@@ -1,11 +1,68 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from types import SimpleNamespace
5 +
6 import torch
7
8 +from vllm.config.mamba import MambaBackendEnum, MambaConfig
9 +from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
10 from vllm.v1.worker.utils import bind_kv_cache
11
12
13 +class _TestReplaySSMMixer(MambaMixer2):
14 + def __init__(self):
15 + torch.nn.Module.__init__(self)
16 + self.use_replayssm = True
17 + self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER)
18 + self._replayssm_ring_start = torch.empty(0, dtype=torch.int32)
19 + self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32)
20 + self._updates_replayssm_trackers = True
21 +
22 + def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
23 + return ((2,), (3,), (4,), (5,), (6,))
24 +
25 + def get_state_dtype(self) -> tuple[torch.dtype, ...]:
26 + return (torch.float32,) * 5
27 +
28 +
29 +def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor:
30 + return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8)
31 +
32 +
33 +def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
34 + mixers = [_TestReplaySSMMixer() for _ in range(3)]
35 + layer_names = [f"layers.{i}.mixer" for i in range(3)]
36 + ctx = dict(zip(layer_names, mixers))
37 + # Reverse insertion order: updater must follow layer index, not dict order.
38 + kv_cache = {
39 + layer_names[2]: _packed_replayssm_cache(4),
40 + layer_names[1]: _packed_replayssm_cache(4),
41 + layer_names[0]: _packed_replayssm_cache(4),
42 + }
43 + kv_cache_groups = [
44 + SimpleNamespace(layer_names=[layer_names[0], layer_names[2]]),
45 + SimpleNamespace(layer_names=[layer_names[1]]),
46 + ]
47 +
48 + bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
49 +
50 + assert (
51 + mixers[0]._replayssm_ring_start.data_ptr()
52 + == mixers[2]._replayssm_ring_start.data_ptr()
53 + )
54 + assert (
55 + mixers[0]._replayssm_prev_num_accepted.data_ptr()
56 + == mixers[2]._replayssm_prev_num_accepted.data_ptr()
57 + )
58 + assert (
59 + mixers[1]._replayssm_ring_start.data_ptr()
60 + != mixers[0]._replayssm_ring_start.data_ptr()
61 + )
62 + # Group {0, 2} shares trackers; layer 2 (not 0) updates after both run.
63 + assert [m._updates_replayssm_trackers for m in mixers] == [False, True, True]
64 +
65 +
66 def test_bind_kv_cache(default_vllm_config):
67 from vllm.model_executor.layers.attention import Attention
68
vllm/config/cache.py
+6 -5
@@ -196,15 +196,16 @@ class CacheConfig:
196 caching is enabled.
197 """
198 replayssm_buffer_len: int = Field(default=16, gt=0)
199 - """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3
200 - speculative decoding does not use B. Default 16."""
199 + """ReplaySSM logical history length B for Mamba2. Triton uses B physical
200 + rows and FlashInfer uses B+1. Kimi-K3 speculative decode does not use B.
201 + Default 16."""
202 use_replayssm: bool = False
203 """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip
204 the per-step full-state store, writing the checkpoint back only on flush.
205 Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton
205 - mamba backend; standard (non-speculative) decode only. In align mode flushes
206 - are most efficient when mamba_block_size is a multiple of replayssm_buffer_len,
207 - but this is not required."""
206 + or FlashInfer mamba backend; standard (non-speculative) decode only. In align
207 + mode flushes are most efficient when mamba_block_size is a multiple of
208 + replayssm_buffer_len, but this is not required."""
209 use_kda_recoverssm: bool = field(default=False, init=False)
210 """Whether Kimi-K3 KDA uses RecoverSSM speculative decode."""
211
vllm/config/vllm.py
+17 -2
@@ -2889,13 +2889,28 @@ class VllmConfig:
2889 raise ValueError(
2890 "RecoverSSM currently requires pipeline_parallel_size=1"
2891 )
2892 + if self.mamba_config.backend != MambaBackendEnum.TRITON:
2893 + raise ValueError("RecoverSSM requires --mamba-backend triton")
2894 elif self.cache_config.mamba_cache_mode == "all":
2895 raise ValueError(
2896 "--use-replayssm supports prefix caching only in align mode; "
2897 "pass --mamba-cache-mode align"
2898 )
2897 - if self.mamba_config.backend != MambaBackendEnum.TRITON:
2898 - raise ValueError("--use-replayssm requires --mamba-backend triton")
2899 + elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
2900 + if self.cache_config.mamba_cache_mode == "align":
2901 + raise ValueError(
2902 + "FlashInfer ReplaySSM does not support "
2903 + "--mamba-cache-mode align yet; use none"
2904 + )
2905 + elif self.mamba_config.backend != MambaBackendEnum.TRITON:
2906 + raise ValueError(
2907 + "--use-replayssm requires --mamba-backend triton or flashinfer"
2908 + )
2909 + elif self.use_v2_model_runner:
2910 + raise ValueError(
2911 + "Triton ReplaySSM requires Model Runner V1; use "
2912 + "--mamba-backend flashinfer or Model Runner V1"
2913 + )
2914 if (
2915 self.kv_transfer_config is not None
2916 and self.kv_transfer_config.is_kv_transfer_instance
vllm/model_executor/layers/mamba/mamba_mixer2.py
+148 -37
@@ -2,10 +2,13 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 +from collections.abc import Sequence
6 +
7 import torch
8 from torch import nn
9
10 from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
11 +from vllm.config.mamba import MambaBackendEnum
12 from vllm.distributed import (
13 divide,
14 get_tensor_model_parallel_rank,
@@ -16,6 +19,7 @@ from vllm.distributed import (
19 from vllm.forward_context import ForwardContext, get_forward_context
20 from vllm.logger import init_logger
21 from vllm.model_executor.custom_op import CustomOp, PluggableLayer
22 +from vllm.model_executor.layers.attention import Attention
23 from vllm.model_executor.layers.linear import (
24 ColumnParallelLinear,
25 MergedColumnParallelLinear,
@@ -38,7 +42,11 @@ from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_outpu
42 from vllm.model_executor.layers.mamba.ops.ssd_combined import (
43 mamba_chunk_scan_combined_varlen,
44 )
41 -from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update
45 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
46 + reset_replayssm_ring_trackers,
47 + selective_state_update,
48 + selective_state_update_replayssm_flashinfer,
49 +)
50 from vllm.model_executor.layers.quantization import QuantizationConfig
51 from vllm.model_executor.model_loader.weight_utils import (
52 LoaderFunction,
@@ -58,6 +66,7 @@ from vllm.utils.torch_utils import (
66 from vllm.v1.attention.backend import AttentionMetadata
67 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
68 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
69 +from vllm.v1.kv_cache_interface import KVCacheGroupSpec
70
71 logger = init_logger(__name__)
72
@@ -515,10 +524,12 @@ class MambaMixer2(MambaBase, PluggableLayer):
524 raise ValueError(
525 "--use-replayssm requires tensor-parallel heads to divide evenly"
526 )
518 - # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode
519 - # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache).
527 + # ReplaySSM appends x/dt/B rings to (conv_state, ssm_state).
528 _n_state = 5 if self.use_replayssm else 2
529 self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state))
530 + self._replayssm_ring_start = torch.empty(0, dtype=torch.int32)
531 + self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32)
532 + self._updates_replayssm_trackers = True
533
534 self.num_spec = vllm_config.num_speculative_tokens
535 if self.num_spec > 0:
@@ -706,6 +717,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
717 assert self.cache_config is not None
718 mamba_block_size = self.cache_config.mamba_block_size
719 is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
720 + ring_start = prev_num_accepted = None
721
722 attn_metadata: AttentionMetadata | None = None
723 if attn_metadata_raw is not None:
@@ -723,7 +735,10 @@ class MambaMixer2(MambaBase, PluggableLayer):
735 )
736 ssm_state = self.kv_cache[1]
737 if self.use_replayssm:
726 - x_cache, dt_cache, B_cache = self.kv_cache[2:]
738 + x_cache, dt_cache, B_cache = self.kv_cache[2:5]
739 + if self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
740 + ring_start = self._replayssm_ring_start
741 + prev_num_accepted = self._replayssm_prev_num_accepted
742 else:
743 x_cache = dt_cache = B_cache = None
744 has_initial_states_p = attn_metadata.has_initial_states_p
@@ -982,6 +997,13 @@ class MambaMixer2(MambaBase, PluggableLayer):
997 # tensor
998 assert state_indices_tensor_p is not None
999 ssm_state[state_indices_tensor_p] = varlen_states
1000 + if ring_start is not None and self._updates_replayssm_trackers:
1001 + assert prev_num_accepted is not None
1002 + reset_replayssm_ring_trackers(
1003 + ring_start,
1004 + prev_num_accepted,
1005 + state_indices_tensor_p,
1006 + )
1007
1008 # Process decode requests
1009 if has_decode:
@@ -1060,35 +1082,64 @@ class MambaMixer2(MambaBase, PluggableLayer):
1082 )
1083 if self.use_replayssm:
1084 assert self.replayssm_buffer_len is not None
1063 - selective_state_update_replayssm_output_only(
1064 - ssm_state,
1065 - hidden_states_d,
1066 - dt_d,
1067 - A_d,
1068 - B_d,
1069 - C_d,
1070 - D_d,
1071 - dt_bias,
1072 - dt_softplus=True,
1073 - x_cache=x_cache,
1074 - dt_cache=dt_cache,
1075 - B_cache=B_cache,
1076 - bc_pre=attn_metadata.bc_pre_scratch,
1077 - write_pos=attn_metadata.write_pos_d,
1078 - is_flush=attn_metadata.is_flush_d,
1079 - max_cache_len=self.replayssm_buffer_len,
1080 - state_batch_indices=state_indices_tensor_d_input,
1081 - out=preallocated_ssm_out_d,
1082 - # Stochastic Rounding for the vanilla decode path is read
1083 - # from mamba_config inside ssu_dispatch; the replay kernel
1084 - # isn't a dispatch backend, so pass it here.
1085 - enable_stochastic_rounding=(
1086 - self.mamba_config.enable_stochastic_rounding
1087 - ),
1088 - cache_philox_rounds=(
1089 - self.mamba_config.stochastic_rounding_philox_rounds
1090 - ),
1091 - )
1085 + if self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
1086 + assert ring_start is not None
1087 + assert prev_num_accepted is not None
1088 + assert attn_metadata.replayssm_scratch is not None
1089 + selective_state_update_replayssm_flashinfer(
1090 + ssm_state,
1091 + hidden_states_d,
1092 + dt_d,
1093 + A_d,
1094 + B_d,
1095 + C_d,
1096 + preallocated_ssm_out_d,
1097 + x_cache,
1098 + B_cache,
1099 + dt_cache,
1100 + ring_start,
1101 + prev_num_accepted,
1102 + logical_window=self.replayssm_buffer_len,
1103 + D=D_d,
1104 + dt_bias=dt_bias,
1105 + dt_softplus=True,
1106 + state_batch_indices=state_indices_tensor_d_input,
1107 + scratch=attn_metadata.replayssm_scratch,
1108 + update_trackers=self._updates_replayssm_trackers,
1109 + enable_stochastic_rounding=(
1110 + self.mamba_config.enable_stochastic_rounding
1111 + ),
1112 + stochastic_rounding_philox_rounds=(
1113 + self.mamba_config.stochastic_rounding_philox_rounds
1114 + ),
1115 + )
1116 + else:
1117 + selective_state_update_replayssm_output_only(
1118 + ssm_state,
1119 + hidden_states_d,
1120 + dt_d,
1121 + A_d,
1122 + B_d,
1123 + C_d,
1124 + D_d,
1125 + dt_bias,
1126 + dt_softplus=True,
1127 + x_cache=x_cache,
1128 + dt_cache=dt_cache,
1129 + B_cache=B_cache,
1130 + bc_pre=attn_metadata.bc_pre_scratch,
1131 + write_pos=attn_metadata.write_pos_d,
1132 + is_flush=attn_metadata.is_flush_d,
1133 + max_cache_len=self.replayssm_buffer_len,
1134 + state_batch_indices=state_indices_tensor_d_input,
1135 + out=preallocated_ssm_out_d,
1136 + enable_stochastic_rounding=(
1137 + self.mamba_config.enable_stochastic_rounding
1138 + ),
1139 + cache_philox_rounds=(
1140 + self.mamba_config.stochastic_rounding_philox_rounds
1141 + ),
1142 + )
1143 else:
1144 selective_state_update(
1145 ssm_state,
@@ -1137,10 +1188,11 @@ class MambaMixer2(MambaBase, PluggableLayer):
1188 if self.use_replayssm:
1189 assert self.replayssm_buffer_len is not None
1190 return MambaStateShapeCalculator.append_replayssm_ring(
1140 - base_shape,
1141 - self.n_groups,
1142 - tp_world_size,
1143 - self.replayssm_buffer_len,
1191 + base_shapes=base_shape,
1192 + n_groups=self.n_groups,
1193 + tp_world_size=tp_world_size,
1194 + logical_window=self.replayssm_buffer_len,
1195 + backend=self.mamba_config.backend,
1196 )
1197 return base_shape
1198
@@ -1149,6 +1201,65 @@ class MambaMixer2(MambaBase, PluggableLayer):
1201 return MambaAttentionBackendEnum.MAMBA2
1202
1203
1204 +def share_replayssm_ring_trackers(
1205 + ordered_layer_names: list[str],
1206 + forward_context: dict[str, Attention],
1207 + kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
1208 +) -> None:
1209 + """Share ring cursors within each cache-slot index namespace.
1210 +
1211 + Layers backed by one KV-cache group use the same physical block indices and
1212 + can therefore share cursors. Different KV-cache groups may assign different
1213 + block indices to the same request and must keep separate cursor tensors.
1214 + The final local layer in each group advances its cursors after every layer
1215 + in that group has consumed the previous values.
1216 + """
1217 +
1218 + replayssm_mixers: dict[str, MambaMixer2] = {}
1219 + for layer_name in ordered_layer_names:
1220 + layer = forward_context[layer_name]
1221 + if (
1222 + isinstance(layer, MambaMixer2)
1223 + and layer.use_replayssm
1224 + and layer.mamba_config.backend == MambaBackendEnum.FLASHINFER
1225 + ):
1226 + replayssm_mixers[layer_name] = layer
1227 +
1228 + layer_to_group: dict[str, int] = {}
1229 + if kv_cache_groups:
1230 + for group_idx, group in enumerate(kv_cache_groups):
1231 + for layer_name in group.layer_names:
1232 + layer_to_group[layer_name] = group_idx
1233 +
1234 + groups_by_namespace: dict[int | str, list[str]] = {}
1235 + for layer_name in ordered_layer_names:
1236 + if layer_name not in replayssm_mixers:
1237 + continue
1238 + namespace = layer_to_group.get(layer_name, layer_name)
1239 + groups_by_namespace.setdefault(namespace, []).append(layer_name)
1240 +
1241 + for group_layer_names in groups_by_namespace.values():
1242 + last_layer_name = group_layer_names[-1]
1243 +
1244 + first_mixer = replayssm_mixers[group_layer_names[0]]
1245 + first_state = first_mixer.kv_cache[1]
1246 + num_blocks, device = first_state.shape[0], first_state.device
1247 + for layer_name in group_layer_names:
1248 + state = replayssm_mixers[layer_name].kv_cache[1]
1249 + if (state.shape[0], state.device) != (num_blocks, device):
1250 + raise ValueError(
1251 + "ReplaySSM layers in one cache group must share cache capacity"
1252 + )
1253 +
1254 + ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device)
1255 + prev_num_accepted = torch.zeros_like(ring_start)
1256 + for layer_name in group_layer_names:
1257 + mixer = replayssm_mixers[layer_name]
1258 + mixer._replayssm_ring_start = ring_start
1259 + mixer._replayssm_prev_num_accepted = prev_num_accepted
1260 + mixer._updates_replayssm_trackers = layer_name == last_layer_name
1261 +
1262 +
1263 def mamba_mixer2(
1264 projected_states: torch.Tensor,
1265 output: torch.Tensor,
vllm/model_executor/layers/mamba/mamba_utils.py
+14 -8
@@ -10,6 +10,7 @@ import torch
10
11 import vllm.envs as envs
12 from vllm.config.cache import MambaDType
13 +from vllm.config.mamba import MambaBackendEnum
14 from vllm.config.model import ModelDType
15 from vllm.distributed import divide
16 from vllm.logger import init_logger
@@ -214,20 +215,25 @@ class MambaStateShapeCalculator:
215 base_shapes: tuple[tuple[int, ...], ...],
216 n_groups: int,
217 tp_world_size: int,
217 - replayssm_buffer_len: int,
218 + logical_window: int,
219 + backend: MambaBackendEnum,
220 ) -> tuple[tuple[int, ...], ...]:
219 - """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a
220 - base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape
221 - ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended
222 - ``n_groups``.
221 + """Append the physical ReplaySSM ring shapes.
222 +
223 + ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;
224 + B_cache uses the un-extended ``n_groups``.
225 """
226 + ring_buffer_len = logical_window
227 + if backend == MambaBackendEnum.FLASHINFER:
228 + # FlashInfer keeps the live window and appended token together.
229 + ring_buffer_len += 1
230 local_nheads, head_dim, state_size = base_shapes[1]
231 local_ngroups = divide(n_groups, tp_world_size)
232 return (
233 *base_shapes,
228 - (local_nheads, replayssm_buffer_len, head_dim),
229 - (local_nheads, replayssm_buffer_len),
230 - (local_ngroups, replayssm_buffer_len, state_size),
234 + (local_nheads, ring_buffer_len, head_dim),
235 + (local_nheads, ring_buffer_len),
236 + (local_ngroups, ring_buffer_len, state_size),
237 )
238
239 @classmethod
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+221 -14
@@ -10,11 +10,14 @@ the backend defaults to 'cpu'.
10 """
11
12 from abc import ABC, abstractmethod
13 +from collections.abc import Callable
14 +from functools import cache
15
16 import torch
17
18 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
19 from vllm.logger import init_logger
20 +from vllm.triton_utils import tl, triton
21 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
22 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
23 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
@@ -22,6 +25,99 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
25 logger = init_logger(__name__)
26
27
28 +@triton.jit(
29 + do_not_specialize=["n_slots", "state_batch_indices_stride"],
30 + do_not_specialize_on_alignment=["state_batch_indices"],
31 +)
32 +def _update_replayssm_ring_trackers_kernel(
33 + ring_start,
34 + prev_num_accepted,
35 + state_batch_indices,
36 + state_batch_indices_stride,
37 + n_slots,
38 + num_states,
39 + logical_window: tl.constexpr,
40 + ring_buffer_len: tl.constexpr,
41 + pad_slot_id: tl.constexpr,
42 + RESET: tl.constexpr,
43 + BLOCK: tl.constexpr,
44 +) -> None:
45 + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
46 + mask = offsets < n_slots
47 + slots = tl.load(
48 + state_batch_indices + offsets * state_batch_indices_stride,
49 + mask=mask,
50 + other=pad_slot_id,
51 + )
52 + valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states)
53 + if RESET:
54 + tl.store(ring_start + slots, 0, mask=valid)
55 + tl.store(prev_num_accepted + slots, 0, mask=valid)
56 + else:
57 + prev = tl.load(prev_num_accepted + slots, mask=valid, other=0)
58 + start = tl.load(ring_start + slots, mask=valid, other=0)
59 + must_checkpoint = prev + 1 > logical_window
60 + next_start = tl.where(
61 + must_checkpoint,
62 + (start + prev) % ring_buffer_len,
63 + start,
64 + )
65 + next_prev = tl.where(must_checkpoint, 1, prev + 1)
66 + tl.store(ring_start + slots, next_start, mask=valid)
67 + tl.store(prev_num_accepted + slots, next_prev, mask=valid)
68 +
69 +
70 +def update_replayssm_ring_trackers(
71 + ring_start: torch.Tensor,
72 + prev_num_accepted: torch.Tensor,
73 + state_batch_indices: torch.Tensor,
74 + logical_window: int | None = None,
75 + ring_buffer_len: int | None = None,
76 + pad_slot_id: int = NULL_BLOCK_ID,
77 +) -> None:
78 + """Reset selected trackers, or advance them when a window is provided."""
79 + if state_batch_indices.dim() > 1:
80 + state_batch_indices = state_batch_indices[:, 0]
81 + n_slots = state_batch_indices.numel()
82 + if n_slots == 0:
83 + return
84 + reset = logical_window is None
85 + if reset:
86 + logical_window = 0
87 + ring_buffer_len = 1
88 + else:
89 + assert ring_buffer_len is not None
90 + block = 128
91 + _update_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)](
92 + ring_start,
93 + prev_num_accepted,
94 + state_batch_indices,
95 + state_batch_indices.stride(0),
96 + n_slots,
97 + min(ring_start.numel(), prev_num_accepted.numel()),
98 + logical_window,
99 + ring_buffer_len,
100 + pad_slot_id,
101 + RESET=reset,
102 + BLOCK=block,
103 + )
104 +
105 +
106 +def reset_replayssm_ring_trackers(
107 + ring_start: torch.Tensor,
108 + prev_num_accepted: torch.Tensor,
109 + state_batch_indices: torch.Tensor,
110 + pad_slot_id: int = NULL_BLOCK_ID,
111 +) -> None:
112 + """Reset selected ReplaySSM ring trackers."""
113 + update_replayssm_ring_trackers(
114 + ring_start,
115 + prev_num_accepted,
116 + state_batch_indices,
117 + pad_slot_id=pad_slot_id,
118 + )
119 +
120 +
121 class MambaSSUBackend(ABC):
122 """Abstract base class for Mamba SSU backends."""
123
@@ -262,15 +358,115 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
358 _mamba_ssu_backend: MambaSSUBackend | None = None
359
360
361 +_flashinfer_replayssm_kernel: Callable[..., torch.Tensor] | None = None
362 +
363 +
364 +@cache
365 +def flashinfer_replayssm_autotune_supported() -> bool:
366 + """Return True when FlashInfer exposes ReplaySSM autotuning."""
367 + try:
368 + from flashinfer.mamba.checkpointing_ssu import ( # noqa: F401
369 + CheckpointingSSURunner,
370 + )
371 + except ImportError:
372 + return False
373 + return True
374 +
375 +
376 +def selective_state_update_replayssm_flashinfer(
377 + state: torch.Tensor,
378 + x: torch.Tensor,
379 + dt: torch.Tensor,
380 + A: torch.Tensor,
381 + B: torch.Tensor,
382 + C: torch.Tensor,
383 + out: torch.Tensor,
384 + x_cache: torch.Tensor,
385 + B_cache: torch.Tensor,
386 + dt_cache: torch.Tensor,
387 + ring_start: torch.Tensor,
388 + prev_num_accepted_tokens: torch.Tensor,
389 + logical_window: int,
390 + D: torch.Tensor | None = None,
391 + dt_bias: torch.Tensor | None = None,
392 + dt_softplus: bool = False,
393 + state_batch_indices: torch.Tensor | None = None,
394 + null_block_id: int = NULL_BLOCK_ID,
395 + scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
396 + update_trackers: bool = True,
397 + enable_stochastic_rounding: bool = False,
398 + stochastic_rounding_philox_rounds: int = 0,
399 +) -> torch.Tensor:
400 + """Run FlashInfer checkpointing SSU and optionally advance shared trackers."""
401 + if _flashinfer_replayssm_kernel is None:
402 + raise RuntimeError(
403 + "FlashInfer ReplaySSM has not been initialized. "
404 + "Call initialize_mamba_ssu_backend() with use_replayssm=True."
405 + )
406 +
407 + if x.dim() == 3:
408 + x = x.unsqueeze(1)
409 + dt = dt.unsqueeze(1)
410 + B = B.unsqueeze(1)
411 + C = C.unsqueeze(1)
412 + out = out.unsqueeze(1)
413 +
414 + indices = state_batch_indices
415 + if indices is not None and indices.dim() > 1:
416 + indices = indices[:, 0]
417 +
418 + cb_scaled = cumAdt_vec = cb_old = None
419 + if scratch is not None:
420 + cb_scaled, cumAdt_vec, cb_old = scratch
421 +
422 + rand_seed = (
423 + torch.randint(0, 2**32, (1,), device=state.device, dtype=torch.int64)
424 + if enable_stochastic_rounding
425 + else None
426 + )
427 + result = _flashinfer_replayssm_kernel(
428 + state,
429 + x_cache,
430 + B_cache,
431 + dt_cache,
432 + ring_start,
433 + prev_num_accepted_tokens,
434 + x,
435 + dt,
436 + A,
437 + B,
438 + C,
439 + out,
440 + D=D,
441 + dt_bias=dt_bias,
442 + dt_softplus=dt_softplus,
443 + state_batch_indices=indices,
444 + pad_slot_id=null_block_id,
445 + rand_seed=rand_seed,
446 + philox_rounds=stochastic_rounding_philox_rounds or 10,
447 + cb_scaled=cb_scaled,
448 + cumAdt_vec=cumAdt_vec,
449 + cb_old=cb_old,
450 + )
451 + if update_trackers and indices is not None:
452 + update_replayssm_ring_trackers(
453 + ring_start,
454 + prev_num_accepted_tokens,
455 + indices,
456 + logical_window=logical_window,
457 + ring_buffer_len=x_cache.size(2),
458 + pad_slot_id=null_block_id,
459 + )
460 + return result
461 +
462 +
463 def initialize_mamba_ssu_backend(
464 mamba_config: MambaConfig,
465 kv_cache_config: KVCacheConfig,
466 + *,
467 + use_replayssm: bool = False,
468 ) -> None:
269 - """Initialize the global Mamba SSU backend.
270 -
271 - No-op if `kv_cache_config` contains no specs that call
272 - selective_state_update.
273 - """
469 + """Initialize the Mamba SSU backend and optional FlashInfer ReplaySSM."""
470 if not any(
471 isinstance(g.kv_cache_spec, MambaSpec)
472 and g.kv_cache_spec.mamba_type
@@ -279,13 +475,9 @@ def initialize_mamba_ssu_backend(
475 ):
476 return
477
282 - global _mamba_ssu_backend
283 -
478 + global _flashinfer_replayssm_kernel, _mamba_ssu_backend
479 backend = mamba_config.backend
480
286 - # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is
287 - # unstable or unavailable. Silently fall back to the CPU
288 - # backend unless the user explicitly chose something other than "triton".
481 if backend == MambaBackendEnum.TRITON:
482 from vllm.platforms import current_platform
483
@@ -301,13 +493,28 @@ def initialize_mamba_ssu_backend(
493 f"Unknown Mamba SSU backend: {backend}. "
494 f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
495 )
496 + if use_replayssm and backend not in (
497 + MambaBackendEnum.TRITON,
498 + MambaBackendEnum.FLASHINFER,
499 + ):
500 + raise ValueError(f"ReplaySSM does not support mamba backend {backend.value!r}")
501
502 backend_cls = _BACKEND_REGISTRY[backend]
306 - if isinstance(_mamba_ssu_backend, backend_cls):
307 - return
503 + if not isinstance(_mamba_ssu_backend, backend_cls):
504 + _mamba_ssu_backend = backend_cls(mamba_config)
505 + logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
506
309 - _mamba_ssu_backend = backend_cls(mamba_config)
310 - logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
507 + _flashinfer_replayssm_kernel = None
508 + if use_replayssm and backend == MambaBackendEnum.FLASHINFER:
509 + try:
510 + from flashinfer.mamba.checkpointing_ssu import checkpointing_ssu
511 + except ImportError as e:
512 + raise ImportError(
513 + "FlashInfer ReplaySSM requires a compatible flashinfer-python package"
514 + ) from e
515 + _flashinfer_replayssm_kernel = checkpointing_ssu
516 + if use_replayssm:
517 + logger.info("Using %s ReplaySSM backend.", backend.value)
518
519
520 def get_mamba_ssu_backend() -> MambaSSUBackend:
vllm/model_executor/models/nemotron_h.py
+7 -5
@@ -747,7 +747,8 @@ class NemotronHForCausalLM(
747 )
748 if cache_config.use_replayssm:
749 return MambaStateDtypeCalculator.append_replayssm_ring(
750 - base_dtype, vllm_config.model_config.dtype
750 + base_dtype,
751 + vllm_config.model_config.dtype,
752 )
753 return base_dtype
754
@@ -784,10 +785,11 @@ class NemotronHForCausalLM(
785 )
786 if cache_config.use_replayssm:
787 return MambaStateShapeCalculator.append_replayssm_ring(
787 - base_shape,
788 - hf_config.n_groups,
789 - parallel_config.tensor_parallel_size,
790 - cache_config.replayssm_buffer_len,
788 + base_shapes=base_shape,
789 + n_groups=hf_config.n_groups,
790 + tp_world_size=parallel_config.tensor_parallel_size,
791 + logical_window=cache_config.replayssm_buffer_len,
792 + backend=vllm_config.mamba_config.backend,
793 )
794 return base_shape
795
vllm/model_executor/warmup/kernel_warmup.py
+4
@@ -38,6 +38,9 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
38 kimi_k3_triton_warmup,
39 )
40 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
41 +from vllm.model_executor.warmup.replayssm_warmup import (
42 + replayssm_autotune_warmup,
43 +)
44 from vllm.platforms import current_platform
45 from vllm.utils.deep_gemm import is_deep_gemm_supported
46 from vllm.utils.flashinfer import has_flashinfer
@@ -335,6 +338,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
338 fi_utils.autotune(tune_mode=True, **autotune_kwargs),
339 ):
340 _run_flashinfer_autotune_dummy_runs(runner)
341 + replayssm_autotune_warmup(runner)
342 finally:
343 set_autotune_process_group(None)
344
vllm/model_executor/warmup/replayssm_warmup.py new
+159
@@ -0,0 +1,159 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from collections.abc import Iterator
5 +from contextlib import contextmanager
6 +from typing import TYPE_CHECKING, Any
7 +
8 +import torch
9 +
10 +from vllm.config.mamba import MambaBackendEnum
11 +from vllm.logger import init_logger
12 +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
13 + flashinfer_replayssm_autotune_supported,
14 +)
15 +
16 +if TYPE_CHECKING:
17 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +def _replayssm_autotune_kwargs(
23 + runner: "GPUModelRunner",
24 +) -> tuple[int, dict[str, Any]] | None:
25 + config = runner.vllm_config
26 + if not (
27 + config.cache_config.use_replayssm
28 + and config.mamba_config.backend == MambaBackendEnum.FLASHINFER
29 + ):
30 + return None
31 + if not flashinfer_replayssm_autotune_supported():
32 + logger.info_once(
33 + "Skipping FlashInfer ReplaySSM autotuning because "
34 + "flashinfer.mamba.checkpointing_ssu.CheckpointingSSURunner "
35 + "is unavailable."
36 + )
37 + return None
38 + v2_runner: Any = runner
39 + query_len = (
40 + v2_runner.decode_query_len
41 + if config.use_v2_model_runner
42 + else runner.uniform_decode_query_len
43 + )
44 + max_num_reqs = min(
45 + runner.scheduler_config.max_num_seqs,
46 + runner.max_num_tokens // query_len,
47 + runner.kv_cache_config.num_blocks - 1,
48 + )
49 + decode_kwargs = {
50 + "num_tokens": max_num_reqs * query_len,
51 + "skip_eplb": True,
52 + "is_profile": True,
53 + "randomize_inputs": True,
54 + "uniform_decode": True,
55 + }
56 + if config.use_v2_model_runner:
57 + decode_kwargs["valid_dummy_state_slots"] = True
58 + else:
59 + decode_kwargs.update(
60 + allow_microbatching=False,
61 + force_attention=True,
62 + profile_seq_lens=query_len + 1,
63 + )
64 + return max_num_reqs, decode_kwargs
65 +
66 +
67 +@contextmanager
68 +def _temporary_replayssm_autotune_state(
69 + runner: "GPUModelRunner", max_num_reqs: int
70 +) -> Iterator[None]:
71 + from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
72 + from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
73 + reset_replayssm_ring_trackers,
74 + update_replayssm_ring_trackers,
75 + )
76 +
77 + reset_tensors: dict[int, torch.Tensor] = {}
78 + tracker_specs: dict[int, tuple[torch.Tensor, torch.Tensor, int, int]] = {}
79 + for module in runner.get_model().modules():
80 + if not isinstance(module, MambaMixer2) or not module.use_replayssm:
81 + continue
82 + assert module.replayssm_buffer_len is not None
83 + ring_start = module._replayssm_ring_start
84 + prev_num_accepted = module._replayssm_prev_num_accepted
85 + tracker_specs.setdefault(
86 + ring_start.data_ptr(),
87 + (
88 + ring_start,
89 + prev_num_accepted,
90 + module.replayssm_buffer_len,
91 + module.kv_cache[2].size(2),
92 + ),
93 + )
94 + tensors = (
95 + *module.kv_cache,
96 + ring_start,
97 + prev_num_accepted,
98 + )
99 + for tensor in tensors:
100 + if tensor.numel():
101 + reset_tensors.setdefault(tensor.data_ptr(), tensor)
102 +
103 + v2_runner: Any = runner
104 + block_tables = saved_block_ids = None
105 + if not runner.vllm_config.use_v2_model_runner:
106 + block_tables = runner.input_batch.block_table.block_tables
107 + saved_block_ids = tuple(
108 + block_table.block_table.np[:max_num_reqs, 0].copy()
109 + for block_table in block_tables
110 + )
111 + dummy_block_ids = range(1, max_num_reqs + 1)
112 + for block_table in block_tables:
113 + block_table.block_table.np[:max_num_reqs, 0] = dummy_block_ids
114 + runner.input_batch.block_table.commit_block_table(max_num_reqs)
115 +
116 + first_tracker = next(iter(tracker_specs.values()), None)
117 + if first_tracker is not None and first_tracker[0].is_cuda:
118 + state_slots = torch.arange(
119 + 1, max_num_reqs + 1, dtype=torch.int32, device=first_tracker[0].device
120 + )
121 + for (
122 + ring_start,
123 + prev_num_accepted,
124 + logical_window,
125 + ring_buffer_len,
126 + ) in tracker_specs.values():
127 + # Compile reset (prefill) and advance (decode) before inference.
128 + # The final reset leaves the decode tuning run in a clean state.
129 + reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots)
130 + update_replayssm_ring_trackers(
131 + ring_start,
132 + prev_num_accepted,
133 + state_slots,
134 + logical_window,
135 + ring_buffer_len,
136 + )
137 + reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots)
138 +
139 + try:
140 + yield
141 + finally:
142 + if runner.vllm_config.use_v2_model_runner:
143 + v2_runner.block_tables.get_dummy_block_tables(max_num_reqs)
144 + else:
145 + assert block_tables is not None and saved_block_ids is not None
146 + for block_table, block_ids in zip(block_tables, saved_block_ids):
147 + block_table.block_table.np[:max_num_reqs, 0] = block_ids
148 + runner.input_batch.block_table.commit_block_table(max_num_reqs)
149 + for tensor in reset_tensors.values():
150 + tensor[1 : max_num_reqs + 1].zero_()
151 +
152 +
153 +def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None:
154 + autotune = _replayssm_autotune_kwargs(runner)
155 + if autotune is None:
156 + return
157 + max_num_reqs, decode_kwargs = autotune
158 + with _temporary_replayssm_autotune_state(runner, max_num_reqs):
159 + runner._dummy_run(**decode_kwargs)
vllm/v1/attention/backends/mamba_attn.py
+55 -11
@@ -8,6 +8,7 @@ from typing import Any, ClassVar, TypeVar
8 import torch
9
10 from vllm.config import VllmConfig
11 +from vllm.config.mamba import MambaBackendEnum
12 from vllm.utils.math_utils import cdiv
13 from vllm.utils.torch_utils import async_tensor_h2d
14 from vllm.v1.attention.backend import (
@@ -74,12 +75,15 @@ class BaseMambaAttentionMetadata:
75 nums_dict: dict | None = None
76 batch_ptr: torch.Tensor | None = None
77 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.
78 + # ReplaySSM standard decode — Triton: per-row ring cursor and flush flag,
79 + # plus (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for
80 + # precomputed B·C products. All None when use_replayssm is disabled or
81 + # the FlashInfer ReplaySSM backend is selected.
82 write_pos_d: torch.Tensor | None = None
83 is_flush_d: torch.Tensor | None = None
84 bc_pre_scratch: torch.Tensor | None = None
85 + # ReplaySSM — FlashInfer checkpointing_ssu two-kernel scratch.
86 + replayssm_scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
87
88
89 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
@@ -107,6 +111,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
111 self.use_spec_decode = self.num_spec_tokens > 0
112 self.use_replayssm = vllm_config.cache_config.use_replayssm
113 self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len
114 + self.use_flashinfer_replayssm = (
115 + self.use_replayssm
116 + and vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER
117 + )
118
119 scheduler_config = vllm_config.scheduler_config
120 self.decode_cudagraph_max_bs: int = scheduler_config.max_num_seqs
@@ -166,9 +174,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
174 dtype=torch.int32,
175 device=device,
176 )
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:
177 + self.decode_bc_pre_scratch: torch.Tensor | None = None
178 + self.decode_replayssm_scratch: (
179 + tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None
180 + ) = None
181 + # ReplaySSM CUDA-graph buffers for the selected backend.
182 + if self.use_replayssm and not self.use_flashinfer_replayssm:
183 self.decode_write_pos_d: torch.Tensor = torch.empty(
184 (self.decode_cudagraph_max_bs,),
185 dtype=torch.int32,
@@ -185,7 +196,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
196 bc_scratch_bs = max(
197 self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs
198 )
188 - self.decode_bc_pre_scratch: torch.Tensor = torch.empty(
199 + self.decode_bc_pre_scratch = torch.empty(
200 (
201 bc_scratch_bs,
202 bc_ngroups,
@@ -194,8 +205,20 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
205 dtype=torch.float32,
206 device=device,
207 )
197 - else:
198 - self.decode_bc_pre_scratch = None
208 + elif self.use_flashinfer_replayssm:
209 + from flashinfer.mamba.checkpointing_ssu import (
210 + allocate_checkpointing_ssu_scratch,
211 + )
212 +
213 + nheads = kv_cache_spec.shapes[2][0]
214 + self.decode_replayssm_scratch = allocate_checkpointing_ssu_scratch(
215 + batch_size=scheduler_config.max_num_seqs,
216 + num_heads=nheads,
217 + num_predicted_tokens=1,
218 + max_window=self.replayssm_buffer_len,
219 + dtype=vllm_config.model_config.dtype,
220 + device=device,
221 + )
222
223 self._init_reorder_batch_threshold(1, self.use_spec_decode)
224 if self.use_spec_decode:
@@ -500,6 +523,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
523 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
524 write_pos_d = None
525 is_flush_d = None
526 + replayssm_scratch = None
527
528 if self.vllm_config.cache_config.mamba_cache_mode == "all":
529 num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
@@ -586,7 +610,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
610 num_reqs - num_prefills : num_reqs
611 ]
612
589 - if self.use_replayssm and num_decodes > 0:
613 + if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0:
614 decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu
615 num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu
616 if decode_base_cpu is None or num_computed_tokens_cpu is None:
@@ -651,6 +675,15 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
675 device=common_attn_metadata.query_start_loc.device,
676 )
677
678 + if self.use_flashinfer_replayssm and num_decodes > 0:
679 + assert self.decode_replayssm_scratch is not None
680 + cb_scaled, cumAdt_vec, cb_old = self.decode_replayssm_scratch
681 + replayssm_scratch = (
682 + cb_scaled[:num_decodes],
683 + cumAdt_vec[:num_decodes],
684 + cb_old[:num_decodes],
685 + )
686 +
687 bc_pre_scratch = None
688 if (
689 self.use_replayssm
@@ -671,6 +704,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
704 write_pos_d=write_pos_d,
705 is_flush_d=is_flush_d,
706 bc_pre_scratch=bc_pre_scratch,
707 + replayssm_scratch=replayssm_scratch,
708 num_accepted_tokens=num_accepted_tokens,
709 query_start_loc_d=query_start_loc_d,
710 block_idx_last_scheduled_token=block_idx_last_scheduled_token,
@@ -708,6 +742,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
742 write_pos_d = metadata.write_pos_d
743 is_flush_d = metadata.is_flush_d
744 bc_pre_scratch = metadata.bc_pre_scratch
745 + replayssm_scratch = metadata.replayssm_scratch
746 if (
747 metadata.num_prefills == 0
748 and metadata.num_decodes <= self.decode_cudagraph_max_bs
@@ -769,7 +804,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
804 )
805 block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0
806
772 - if self.use_replayssm:
807 + if self.use_replayssm and not self.use_flashinfer_replayssm:
808 assert write_pos_d is not None
809 assert is_flush_d is not None
810 self.decode_write_pos_d[: metadata.num_decodes].copy_(
@@ -788,6 +823,14 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
823
824 if self.decode_bc_pre_scratch is not None:
825 bc_pre_scratch = self.decode_bc_pre_scratch[:padded_bs]
826 + elif self.use_flashinfer_replayssm:
827 + assert self.decode_replayssm_scratch is not None
828 + cb_scaled, cumAdt_vec, cb_old = self.decode_replayssm_scratch
829 + replayssm_scratch = (
830 + cb_scaled[:padded_bs],
831 + cumAdt_vec[:padded_bs],
832 + cb_old[:padded_bs],
833 + )
834
835 return replace(
836 metadata,
@@ -797,6 +840,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
840 write_pos_d=write_pos_d,
841 is_flush_d=is_flush_d,
842 bc_pre_scratch=bc_pre_scratch,
843 + replayssm_scratch=replayssm_scratch,
844 block_idx_last_scheduled_token=block_idx_last_scheduled_token,
845 block_idx_last_computed_token=block_idx_last_computed_token,
846 block_idx_last_scheduled_token_prev_step=(
vllm/v1/worker/gpu/attn_utils.py
+7 -1
@@ -229,7 +229,13 @@ def init_kv_cache(
229 in ("longcat_flash", "longcat_flash_ngram")
230 else 1
231 )
232 - bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module)
232 + bind_kv_cache(
233 + kv_caches,
234 + forward_context,
235 + runner_kv_caches,
236 + num_attn_module,
237 + kv_cache_groups=kv_cache_config.kv_cache_groups,
238 + )
239 return kv_caches
240
241
vllm/v1/worker/gpu/model_runner.py
+19 -3
@@ -633,7 +633,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
633 cls=self.pcp_manager_cls,
634 )
635 initialize_mamba_ssu_backend(
636 - self.vllm_config.mamba_config, self.kv_cache_config
636 + self.vllm_config.mamba_config,
637 + self.kv_cache_config,
638 + use_replayssm=self.vllm_config.cache_config.use_replayssm,
639 )
640 if self.adaptive_verification is not None:
641 self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
@@ -707,6 +709,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
709 context_len: int = 0,
710 skip_eplb: bool = False,
711 is_profile: bool = False,
712 + valid_dummy_state_slots: bool = False,
713 **kwargs,
714 ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
715 if skip_attn and not is_profile:
@@ -764,6 +767,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
767 skip_attn_for_dummy_run=skip_attn,
768 is_profile=is_profile,
769 context_len=context_len,
770 + valid_dummy_state_slots=valid_dummy_state_slots,
771 )
772 self.kv_connector.set_disabled(False)
773
@@ -1390,9 +1394,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1394 return block_tables, slot_mappings
1395
1396 def prepare_dummy_attn(
1393 - self, input_batch: InputBatch
1397 + self, input_batch: InputBatch, valid_state_slots: bool = False
1398 ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
1399 block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
1400 + if valid_state_slots:
1401 + state_slots = torch.arange(
1402 + 1,
1403 + input_batch.num_reqs + 1,
1404 + dtype=torch.int32,
1405 + device=self.device,
1406 + )
1407 + for block_table in block_tables:
1408 + block_table[:, 0].copy_(state_slots)
1409 slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
1410 self.pcp_manager, self.block_tables, input_batch.num_tokens
1411 )
@@ -1525,6 +1538,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1538 skip_attn_for_dummy_run: bool = False,
1539 is_profile: bool = False,
1540 context_len: int = 0,
1541 + valid_dummy_state_slots: bool = False,
1542 ) -> ModelRunnerOutput | IntermediateTensors | None:
1543 if not dummy_run:
1544 # Update the request states.
@@ -1624,7 +1638,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1638 max_query_len=batch_desc.max_query_len,
1639 )
1640 if not skip_attn_for_dummy_run:
1627 - block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
1641 + block_tables, slot_mappings = self.prepare_dummy_attn(
1642 + input_batch, valid_dummy_state_slots
1643 + )
1644 if context_len:
1645 set_dummy_context(
1646 input_batch,
vllm/v1/worker/gpu_model_runner.py
+4 -1
@@ -7434,6 +7434,7 @@ class GPUModelRunner(
7434 self.compilation_config.static_forward_context,
7435 self.kv_caches,
7436 num_attn_module,
7437 + kv_cache_groups=kv_cache_config.kv_cache_groups,
7438 )
7439 return kv_caches
7440
@@ -7485,7 +7486,9 @@ class GPUModelRunner(
7486 self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
7487 self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
7488 initialize_mamba_ssu_backend(
7488 - self.vllm_config.mamba_config, self.kv_cache_config
7489 + self.vllm_config.mamba_config,
7490 + self.kv_cache_config,
7491 + use_replayssm=self.vllm_config.cache_config.use_replayssm,
7492 )
7493 # The kernel block size for all KV cache groups. For example, if
7494 # kv_cache_manager uses block_size 256 for a given group, but the attention
vllm/v1/worker/utils.py
+6
@@ -13,6 +13,7 @@ import torch
13 from vllm.config import CacheConfig, VllmConfig
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.attention import Attention
16 +from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers
17 from vllm.model_executor.layers.utils import warmup_rocm_skinny_gemm_workspaces
18 from vllm.model_executor.models.interfaces import MultiModalEmbeddings
19 from vllm.model_executor.models.utils import extract_layer_index
@@ -575,6 +576,7 @@ def bind_kv_cache(
576 forward_context: dict[str, Attention],
577 runner_kv_caches: list[torch.Tensor],
578 num_attn_module: int = 1,
579 + kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
580 ) -> None:
581 """
582 Bind the allocated KV cache to both ModelRunner and forward context so
@@ -600,6 +602,7 @@ def bind_kv_cache(
602 for layer_name in kv_caches:
603 index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)
604
605 + ordered_layer_names: list[str] = []
606 for layer_index in sorted(index2name.keys()):
607 layer_names = index2name[layer_index]
608 if len(layer_names) > 1:
@@ -613,6 +616,7 @@ def bind_kv_cache(
616 current_platform.check_runner_kv_caches_multi_layer()
617 for layer_name in layer_names:
618 runner_kv_caches.append(kv_caches[layer_name])
619 + ordered_layer_names.append(layer_name)
620
621 # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
622 # its raw allocation into the per-layer view(s) it needs (e.g. Mamba
@@ -621,6 +625,8 @@ def bind_kv_cache(
625 for layer_name, kv_cache in kv_caches.items():
626 forward_context[layer_name].bind_kv_cache(kv_cache)
627
628 + share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups)
629 +
630
631 def clear_layer_kv_caches(layers: Iterable[Any]) -> None:
632 """Detach the KV/state cache tensors installed by bind_kv_cache().