[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>

tests/kernels/mamba/test_ssu_dispatch.py

@@ -7,12 +7,15 @@ import pytest
 import torch
 
 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
+from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     FlashInferSSUBackend,
     TritonSSUBackend,
     get_mamba_ssu_backend,
     initialize_mamba_ssu_backend,
+    reset_replayssm_ring_trackers,
     selective_state_update,
+    update_replayssm_ring_trackers,
 )
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
@@ -30,6 +33,47 @@ except ImportError:
     HAS_FLASHINFER = False
 
 
+@pytest.fixture(autouse=True)
+def restore_backend_state():
+    import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
+
+    old_backend = mod._mamba_ssu_backend
+    old_replayssm_kernel = mod._flashinfer_replayssm_kernel
+    yield
+    mod._mamba_ssu_backend = old_backend
+    mod._flashinfer_replayssm_kernel = old_replayssm_kernel
+
+
+def test_flashinfer_replayssm_ring_tracker_lifecycle():
+    ring_start = torch.zeros(2, dtype=torch.int32, device="cuda")
+    prev_num_accepted = torch.zeros(2, dtype=torch.int32, device="cuda")
+    state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda")
+
+    observed = []
+    for _ in range(33):
+        update_replayssm_ring_trackers(
+            ring_start,
+            prev_num_accepted,
+            state_batch_indices,
+            logical_window=16,
+            ring_buffer_len=17,
+        )
+        observed.append((int(ring_start[1]), int(prev_num_accepted[1])))
+
+    assert observed[4] == (0, 5)
+    assert observed[15] == (0, 16)
+    assert observed[16] == (16, 1)
+    assert observed[31] == (16, 16)
+    assert observed[32] == (15, 1)
+
+    reset_replayssm_ring_trackers(
+        ring_start,
+        prev_num_accepted,
+        state_batch_indices,
+    )
+    assert (ring_start[1].item(), prev_num_accepted[1].item()) == (0, 0)
+
+
 def _kv_cache_config_with_ssu(
     mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2,
 ) -> KVCacheConfig:
@@ -116,11 +160,10 @@ def test_flashinfer_forwards_ssu_algorithm(
 def test_uninitialized_backend_raises():
     import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
 
-    old = mod._mamba_ssu_backend
+    # restore_backend_state (autouse) puts the global back afterwards.
     mod._mamba_ssu_backend = None
     with pytest.raises(RuntimeError, match="not been initialized"):
         get_mamba_ssu_backend()
-    mod._mamba_ssu_backend = old
 
 
 @pytest.mark.parametrize(
@@ -186,3 +229,28 @@ def test_triton_basic_call():
         out=out,
     )
     assert not torch.isnan(out).any()
+
+
+@pytest.mark.parametrize(
+    ("backend", "expected_ring_len"),
+    [
+        (MambaBackendEnum.TRITON, 16),
+        (MambaBackendEnum.FLASHINFER, 17),
+    ],
+)
+def test_replayssm_physical_ring_shape(backend, expected_ring_len):
+    base_shapes = ((64, 3), (8, 4, 16))
+
+    shapes = MambaStateShapeCalculator.append_replayssm_ring(
+        base_shapes,
+        n_groups=4,
+        tp_world_size=2,
+        logical_window=16,
+        backend=backend,
+    )
+
+    assert shapes[2:] == (
+        (8, expected_ring_len, 4),
+        (8, expected_ring_len),
+        (2, expected_ring_len, 16),
+    )

tests/model_executor/test_replayssm_warmup.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import Mock, patch
+
+import numpy as np
+import pytest
+import torch
+
+from vllm.config.mamba import MambaBackendEnum
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
+from vllm.model_executor.warmup import replayssm_warmup as warmup
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer
+
+pytestmark = pytest.mark.skipif(
+    not current_platform.is_cuda() or not has_flashinfer(),
+    reason="FlashInfer ReplaySSM warmup tests require CUDA and FlashInfer",
+)
+
+PREFILL_KWARGS = {
+    "num_tokens": 128,
+    "skip_eplb": True,
+    "is_profile": True,
+    "randomize_inputs": True,
+}
+
+
+def _autotune_runner(
+    *,
+    use_v2_model_runner: bool = False,
+    query_len: int = 6,
+    max_num_seqs: int = 32,
+    num_blocks: int = 17,
+    max_num_tokens: int = 100,
+    use_replayssm: bool = True,
+    backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER,
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        vllm_config=SimpleNamespace(
+            cache_config=SimpleNamespace(use_replayssm=use_replayssm),
+            mamba_config=SimpleNamespace(backend=backend),
+            use_v2_model_runner=use_v2_model_runner,
+        ),
+        uniform_decode_query_len=query_len,
+        decode_query_len=query_len,
+        max_num_tokens=max_num_tokens,
+        scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
+        kv_cache_config=SimpleNamespace(num_blocks=num_blocks),
+    )
+
+
+@pytest.mark.parametrize(
+    ("runner_kwargs", "expected_num_reqs"),
+    [
+        # max_num_seqs (32) vs max_num_tokens // query_len (16) vs blocks-1 (16).
+        (dict(query_len=6, use_v2_model_runner=False), 16),
+        (dict(query_len=6, use_v2_model_runner=True), 16),
+        # num_blocks - 1 is the binding constraint.
+        (dict(query_len=1, max_num_tokens=128, max_num_seqs=64, num_blocks=5), 4),
+    ],
+    ids=["v1", "v2", "clamped_to_state_capacity"],
+)
+def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
+    query_len = runner_kwargs["query_len"]
+    with patch.object(
+        warmup, "flashinfer_replayssm_autotune_supported", return_value=True
+    ):
+        result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs))
+
+    expected_kwargs = {
+        **PREFILL_KWARGS,
+        "num_tokens": expected_num_reqs * query_len,
+        "uniform_decode": True,
+    }
+    if runner_kwargs.get("use_v2_model_runner"):
+        expected_kwargs["valid_dummy_state_slots"] = True
+    else:
+        expected_kwargs.update(
+            allow_microbatching=False,
+            force_attention=True,
+            profile_seq_lens=query_len + 1,
+        )
+    assert result == (expected_num_reqs, expected_kwargs)
+
+
+@pytest.mark.parametrize(
+    ("runner_kwargs", "flashinfer_supported"),
+    [
+        (dict(use_replayssm=False), True),
+        (dict(backend=MambaBackendEnum.TRITON), True),
+        ({}, False),
+    ],
+    ids=["replayssm_disabled", "non_flashinfer_backend", "kernel_unavailable"],
+)
+def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported):
+    with patch.object(
+        warmup,
+        "flashinfer_replayssm_autotune_supported",
+        return_value=flashinfer_supported,
+    ):
+        result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs))
+    assert result is None
+
+
+def test_replayssm_autotune_slots_restore_state_and_trackers():
+    mixer = MambaMixer2.__new__(MambaMixer2)
+    torch.nn.Module.__init__(mixer)
+    mixer.use_replayssm = True
+    mixer.replayssm_buffer_len = 16
+    mixer.kv_cache = (
+        torch.full((4, 2), 3.0),
+        torch.full((4, 2), 3.0),
+        *(torch.full((4, 2, 17), 3.0) for _ in range(3)),
+    )
+    mixer._replayssm_ring_start = torch.full((4,), 3, dtype=torch.int32)
+    mixer._replayssm_prev_num_accepted = torch.full((4,), 3, dtype=torch.int32)
+    tracked = (
+        *mixer.kv_cache,
+        mixer._replayssm_ring_start,
+        mixer._replayssm_prev_num_accepted,
+    )
+
+    block_ids = np.arange(10, 14, dtype=np.int32).reshape(4, 1)
+    original_block_ids = block_ids.copy()
+    block_table = SimpleNamespace(block_table=SimpleNamespace(np=block_ids))
+    multi_group_block_table = SimpleNamespace(
+        block_tables=[block_table], commit_block_table=Mock()
+    )
+    runner = SimpleNamespace(
+        vllm_config=SimpleNamespace(use_v2_model_runner=False),
+        input_batch=SimpleNamespace(block_table=multi_group_block_table),
+        get_model=lambda: SimpleNamespace(modules=lambda: (mixer,)),
+    )
+
+    with warmup._temporary_replayssm_autotune_state(runner, 2):
+        assert block_ids[:2, 0].tolist() == [1, 2]
+        for tensor in tracked:
+            tensor[1:3].fill_(9)
+
+    assert np.array_equal(block_ids, original_block_ids)
+    assert multi_group_block_table.commit_block_table.call_args_list == [
+        ((2,), {}),
+        ((2,), {}),
+    ]
+    for tensor in tracked:
+        assert torch.count_nonzero(tensor[1:3]) == 0
+        assert torch.all(tensor[0] == 3)
+        assert torch.all(tensor[3] == 3)

tests/v1/attention/test_replayssm_metadata_builder.py

@@ -16,6 +16,7 @@ from tests.v1.attention.utils import (
     create_common_attn_metadata,
     create_vllm_config,
 )
+from vllm.config.mamba import MambaBackendEnum
 from vllm.v1.kv_cache_interface import MambaSpec
 
 BLOCK_SIZE = 16
@@ -187,34 +188,47 @@ REPLAYSSM_BUILD_CASES = {
 }
 
 
-def _make_mamba_spec(buffer_len: int) -> MambaSpec:
-    # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups).
+def _make_mamba_spec(
+    buffer_len: int,
+    mamba_backend: MambaBackendEnum,
+) -> MambaSpec:
+    ring_buffer_len = buffer_len + (
+        1 if mamba_backend == MambaBackendEnum.FLASHINFER else 0
+    )
+    shapes = (
+        (1, 1),
+        (1, 1, 1),
+        (1, ring_buffer_len, 1),
+        (1, ring_buffer_len),
+        (1, ring_buffer_len, 1),
+    )
     return MambaSpec(
         block_size=BLOCK_SIZE,
-        shapes=(
-            (1, 1),
-            (1, 1, 1),
-            (1, buffer_len, 1),
-            (1, buffer_len),
-            (1, buffer_len, 1),
-        ),
+        shapes=shapes,
         dtypes=(torch.float32,),
     )
 
 
 def _create_replayssm_builder(
-    buffer_len: int, mamba_cache_mode: str = "none"
+    buffer_len: int,
+    mamba_cache_mode: str = "none",
+    *,
+    mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON,
 ) -> MockMambaBuilder:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE
     )
     # Set the flags after construction to skip validate_mamba_cached_kernel
-    # (it requires a Triton backend) on the mock model.
+    # (it requires a real SupportsReplaySSM model) on the mock model.
     vllm_config.cache_config.use_replayssm = True
     vllm_config.cache_config.replayssm_buffer_len = buffer_len
     vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
+    vllm_config.mamba_config.backend = mamba_backend
     return MockMambaBuilder(
-        _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE
+        _make_mamba_spec(buffer_len, mamba_backend),
+        ["layer0"],
+        vllm_config,
+        DEVICE,
     )
 
 
@@ -254,3 +268,23 @@ def test_resumed_request_differs_from_fresh():
 
     assert meta.write_pos_d.tolist()[:2] == [5, 0]
     assert meta.is_flush_d.tolist()[:2] == [0, 0]
+
+
+def test_flashinfer_replayssm_scratch_metadata_fresh_decode():
+    checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
+    if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"):
+        pytest.skip("FlashInfer does not expose ReplaySSM scratch allocation")
+
+    builder = _create_replayssm_builder(16, mamba_backend=MambaBackendEnum.FLASHINFER)
+    case = REPLAYSSM_BUILD_CASES["fresh_decode"]
+    meta = _build(builder, case)
+
+    assert meta.write_pos_d is None
+    assert meta.is_flush_d is None
+    assert meta.bc_pre_scratch is None
+    assert meta.replayssm_scratch is not None
+    assert [tensor.shape for tensor in meta.replayssm_scratch] == [
+        (1, 1, 32, 8),
+        (1, 1, 16),
+        (1, 1, 32, 8),
+    ]

tests/v1/e2e/test_replayssm_decode.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+import vllm.envs as envs
 from vllm.v1.metrics.reader import Counter
 
 from ...models.utils import check_logprobs_close
@@ -21,29 +22,48 @@ PROMPTS = [
 ]
 
 
-def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1):
+def _check_replayssm_parity(
+    vllm_runner,
+    model_name,
+    *,
+    tensor_parallel_size=1,
+    mamba_backend: str = "triton",
+    name_1: str = "replayssm",
+    require_v2: bool = False,
+    monkeypatch: pytest.MonkeyPatch | None = None,
+):
     # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a
     # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are
     # common-mode and only ReplaySSM varies.
+    if require_v2:
+        assert monkeypatch is not None
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
+        envs.disable_envs_cache()
+
     common = dict(
         max_model_len=1024,
         trust_remote_code=True,
         enable_prefix_caching=False,
         mamba_cache_mode="none",
         tensor_parallel_size=tensor_parallel_size,
+        mamba_backend=mamba_backend,
     )
     with vllm_runner(model_name, **common) as llm:
+        if require_v2:
+            assert llm.llm.llm_engine.vllm_config.use_v2_model_runner
         baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
     with vllm_runner(
         model_name, use_replayssm=True, replayssm_buffer_len=16, **common
     ) as llm:
+        if require_v2:
+            assert llm.llm.llm_engine.vllm_config.use_v2_model_runner
         replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
 
     check_logprobs_close(
         outputs_0_lst=baseline,
         outputs_1_lst=replay,
         name_0="baseline",
-        name_1="replayssm",
+        name_1=name_1,
     )
 
 
@@ -60,6 +80,21 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name):
     _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2)
 
 
+@pytest.mark.parametrize("model_name", MODELS)
+def test_replayssm_flashinfer_decode_matches_baseline_v2(
+    vllm_runner, model_name, monkeypatch
+):
+    pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
+    _check_replayssm_parity(
+        vllm_runner,
+        model_name,
+        mamba_backend="flashinfer",
+        name_1="replayssm_flashinfer_v2",
+        require_v2=True,
+        monkeypatch=monkeypatch,
+    )
+
+
 # Prefix spans several mamba blocks; prefix caching only reuses full blocks.
 _PC_SENTENCE = (
     "In a detailed survey of state space models, the authors compared many "

tests/v1/worker/test_kv_cache_allocation_scope.py

@@ -48,7 +48,7 @@ def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
     result = attn_utils.init_kv_cache(
         [],
         {},
-        object(),
+        SimpleNamespace(kv_cache_groups=[]),
         torch.device("cpu"),
         [],
         config,
@@ -83,7 +83,7 @@ def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None:
     )
     result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors(
         runner,
-        object(),
+        SimpleNamespace(kv_cache_groups=[]),
         [],
         kv_cache_allocation_context=scope,
     )

tests/v1/worker/test_utils.py

@@ -1,11 +1,68 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+
 import torch
 
+from vllm.config.mamba import MambaBackendEnum, MambaConfig
+from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
 from vllm.v1.worker.utils import bind_kv_cache
 
 
+class _TestReplaySSMMixer(MambaMixer2):
+    def __init__(self):
+        torch.nn.Module.__init__(self)
+        self.use_replayssm = True
+        self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER)
+        self._replayssm_ring_start = torch.empty(0, dtype=torch.int32)
+        self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32)
+        self._updates_replayssm_trackers = True
+
+    def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
+        return ((2,), (3,), (4,), (5,), (6,))
+
+    def get_state_dtype(self) -> tuple[torch.dtype, ...]:
+        return (torch.float32,) * 5
+
+
+def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor:
+    return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8)
+
+
+def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group():
+    mixers = [_TestReplaySSMMixer() for _ in range(3)]
+    layer_names = [f"layers.{i}.mixer" for i in range(3)]
+    ctx = dict(zip(layer_names, mixers))
+    # Reverse insertion order: updater must follow layer index, not dict order.
+    kv_cache = {
+        layer_names[2]: _packed_replayssm_cache(4),
+        layer_names[1]: _packed_replayssm_cache(4),
+        layer_names[0]: _packed_replayssm_cache(4),
+    }
+    kv_cache_groups = [
+        SimpleNamespace(layer_names=[layer_names[0], layer_names[2]]),
+        SimpleNamespace(layer_names=[layer_names[1]]),
+    ]
+
+    bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
+
+    assert (
+        mixers[0]._replayssm_ring_start.data_ptr()
+        == mixers[2]._replayssm_ring_start.data_ptr()
+    )
+    assert (
+        mixers[0]._replayssm_prev_num_accepted.data_ptr()
+        == mixers[2]._replayssm_prev_num_accepted.data_ptr()
+    )
+    assert (
+        mixers[1]._replayssm_ring_start.data_ptr()
+        != mixers[0]._replayssm_ring_start.data_ptr()
+    )
+    # Group {0, 2} shares trackers; layer 2 (not 0) updates after both run.
+    assert [m._updates_replayssm_trackers for m in mixers] == [False, True, True]
+
+
 def test_bind_kv_cache(default_vllm_config):
     from vllm.model_executor.layers.attention import Attention
 

vllm/config/cache.py

@@ -196,15 +196,16 @@ class CacheConfig:
       caching is enabled.
     """
     replayssm_buffer_len: int = Field(default=16, gt=0)
-    """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3
-    speculative decoding does not use B. Default 16."""
+    """ReplaySSM logical history length B for Mamba2. Triton uses B physical
+    rows and FlashInfer uses B+1. Kimi-K3 speculative decode does not use B.
+    Default 16."""
     use_replayssm: bool = False
     """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip
     the per-step full-state store, writing the checkpoint back only on flush.
     Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton
-    mamba backend; standard (non-speculative) decode only. In align mode flushes
-    are most efficient when mamba_block_size is a multiple of replayssm_buffer_len,
-    but this is not required."""
+    or FlashInfer mamba backend; standard (non-speculative) decode only. In align
+    mode flushes are most efficient when mamba_block_size is a multiple of
+    replayssm_buffer_len, but this is not required."""
     use_kda_recoverssm: bool = field(default=False, init=False)
     """Whether Kimi-K3 KDA uses RecoverSSM speculative decode."""
 

vllm/config/vllm.py

@@ -2889,13 +2889,28 @@ class VllmConfig:
                 raise ValueError(
                     "RecoverSSM currently requires pipeline_parallel_size=1"
                 )
+            if self.mamba_config.backend != MambaBackendEnum.TRITON:
+                raise ValueError("RecoverSSM requires --mamba-backend triton")
         elif self.cache_config.mamba_cache_mode == "all":
             raise ValueError(
                 "--use-replayssm supports prefix caching only in align mode; "
                 "pass --mamba-cache-mode align"
             )
-        if self.mamba_config.backend != MambaBackendEnum.TRITON:
-            raise ValueError("--use-replayssm requires --mamba-backend triton")
+        elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
+            if self.cache_config.mamba_cache_mode == "align":
+                raise ValueError(
+                    "FlashInfer ReplaySSM does not support "
+                    "--mamba-cache-mode align yet; use none"
+                )
+        elif self.mamba_config.backend != MambaBackendEnum.TRITON:
+            raise ValueError(
+                "--use-replayssm requires --mamba-backend triton or flashinfer"
+            )
+        elif self.use_v2_model_runner:
+            raise ValueError(
+                "Triton ReplaySSM requires Model Runner V1; use "
+                "--mamba-backend flashinfer or Model Runner V1"
+            )
         if (
             self.kv_transfer_config is not None
             and self.kv_transfer_config.is_kv_transfer_instance

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -2,10 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
+from collections.abc import Sequence
+
 import torch
 from torch import nn
 
 from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config
+from vllm.config.mamba import MambaBackendEnum
 from vllm.distributed import (
     divide,
     get_tensor_model_parallel_rank,
@@ -16,6 +19,7 @@ from vllm.distributed import (
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp, PluggableLayer
+from vllm.model_executor.layers.attention import Attention
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
@@ -38,7 +42,11 @@ from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_outpu
 from vllm.model_executor.layers.mamba.ops.ssd_combined import (
     mamba_chunk_scan_combined_varlen,
 )
-from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    reset_replayssm_ring_trackers,
+    selective_state_update,
+    selective_state_update_replayssm_flashinfer,
+)
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.model_loader.weight_utils import (
     LoaderFunction,
@@ -58,6 +66,7 @@ from vllm.utils.torch_utils import (
 from vllm.v1.attention.backend import AttentionMetadata
 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
+from vllm.v1.kv_cache_interface import KVCacheGroupSpec
 
 logger = init_logger(__name__)
 
@@ -515,10 +524,12 @@ class MambaMixer2(MambaBase, PluggableLayer):
             raise ValueError(
                 "--use-replayssm requires tensor-parallel heads to divide evenly"
             )
-        # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode
-        # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache).
+        # ReplaySSM appends x/dt/B rings to (conv_state, ssm_state).
         _n_state = 5 if self.use_replayssm else 2
         self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state))
+        self._replayssm_ring_start = torch.empty(0, dtype=torch.int32)
+        self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32)
+        self._updates_replayssm_trackers = True
 
         self.num_spec = vllm_config.num_speculative_tokens
         if self.num_spec > 0:
@@ -706,6 +717,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
         assert self.cache_config is not None
         mamba_block_size = self.cache_config.mamba_block_size
         is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
+        ring_start = prev_num_accepted = None
 
         attn_metadata: AttentionMetadata | None = None
         if attn_metadata_raw is not None:
@@ -723,7 +735,10 @@ class MambaMixer2(MambaBase, PluggableLayer):
             )
             ssm_state = self.kv_cache[1]
             if self.use_replayssm:
-                x_cache, dt_cache, B_cache = self.kv_cache[2:]
+                x_cache, dt_cache, B_cache = self.kv_cache[2:5]
+                if self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
+                    ring_start = self._replayssm_ring_start
+                    prev_num_accepted = self._replayssm_prev_num_accepted
             else:
                 x_cache = dt_cache = B_cache = None
             has_initial_states_p = attn_metadata.has_initial_states_p
@@ -982,6 +997,13 @@ class MambaMixer2(MambaBase, PluggableLayer):
                 #   tensor
                 assert state_indices_tensor_p is not None
                 ssm_state[state_indices_tensor_p] = varlen_states
+                if ring_start is not None and self._updates_replayssm_trackers:
+                    assert prev_num_accepted is not None
+                    reset_replayssm_ring_trackers(
+                        ring_start,
+                        prev_num_accepted,
+                        state_indices_tensor_p,
+                    )
 
         # Process decode requests
         if has_decode:
@@ -1060,35 +1082,64 @@ class MambaMixer2(MambaBase, PluggableLayer):
             )
             if self.use_replayssm:
                 assert self.replayssm_buffer_len is not None
-                selective_state_update_replayssm_output_only(
-                    ssm_state,
-                    hidden_states_d,
-                    dt_d,
-                    A_d,
-                    B_d,
-                    C_d,
-                    D_d,
-                    dt_bias,
-                    dt_softplus=True,
-                    x_cache=x_cache,
-                    dt_cache=dt_cache,
-                    B_cache=B_cache,
-                    bc_pre=attn_metadata.bc_pre_scratch,
-                    write_pos=attn_metadata.write_pos_d,
-                    is_flush=attn_metadata.is_flush_d,
-                    max_cache_len=self.replayssm_buffer_len,
-                    state_batch_indices=state_indices_tensor_d_input,
-                    out=preallocated_ssm_out_d,
-                    # Stochastic Rounding for the vanilla decode path is read
-                    # from mamba_config inside ssu_dispatch; the replay kernel
-                    # isn't a dispatch backend, so pass it here.
-                    enable_stochastic_rounding=(
-                        self.mamba_config.enable_stochastic_rounding
-                    ),
-                    cache_philox_rounds=(
-                        self.mamba_config.stochastic_rounding_philox_rounds
-                    ),
-                )
+                if self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
+                    assert ring_start is not None
+                    assert prev_num_accepted is not None
+                    assert attn_metadata.replayssm_scratch is not None
+                    selective_state_update_replayssm_flashinfer(
+                        ssm_state,
+                        hidden_states_d,
+                        dt_d,
+                        A_d,
+                        B_d,
+                        C_d,
+                        preallocated_ssm_out_d,
+                        x_cache,
+                        B_cache,
+                        dt_cache,
+                        ring_start,
+                        prev_num_accepted,
+                        logical_window=self.replayssm_buffer_len,
+                        D=D_d,
+                        dt_bias=dt_bias,
+                        dt_softplus=True,
+                        state_batch_indices=state_indices_tensor_d_input,
+                        scratch=attn_metadata.replayssm_scratch,
+                        update_trackers=self._updates_replayssm_trackers,
+                        enable_stochastic_rounding=(
+                            self.mamba_config.enable_stochastic_rounding
+                        ),
+                        stochastic_rounding_philox_rounds=(
+                            self.mamba_config.stochastic_rounding_philox_rounds
+                        ),
+                    )
+                else:
+                    selective_state_update_replayssm_output_only(
+                        ssm_state,
+                        hidden_states_d,
+                        dt_d,
+                        A_d,
+                        B_d,
+                        C_d,
+                        D_d,
+                        dt_bias,
+                        dt_softplus=True,
+                        x_cache=x_cache,
+                        dt_cache=dt_cache,
+                        B_cache=B_cache,
+                        bc_pre=attn_metadata.bc_pre_scratch,
+                        write_pos=attn_metadata.write_pos_d,
+                        is_flush=attn_metadata.is_flush_d,
+                        max_cache_len=self.replayssm_buffer_len,
+                        state_batch_indices=state_indices_tensor_d_input,
+                        out=preallocated_ssm_out_d,
+                        enable_stochastic_rounding=(
+                            self.mamba_config.enable_stochastic_rounding
+                        ),
+                        cache_philox_rounds=(
+                            self.mamba_config.stochastic_rounding_philox_rounds
+                        ),
+                    )
             else:
                 selective_state_update(
                     ssm_state,
@@ -1137,10 +1188,11 @@ class MambaMixer2(MambaBase, PluggableLayer):
         if self.use_replayssm:
             assert self.replayssm_buffer_len is not None
             return MambaStateShapeCalculator.append_replayssm_ring(
-                base_shape,
-                self.n_groups,
-                tp_world_size,
-                self.replayssm_buffer_len,
+                base_shapes=base_shape,
+                n_groups=self.n_groups,
+                tp_world_size=tp_world_size,
+                logical_window=self.replayssm_buffer_len,
+                backend=self.mamba_config.backend,
             )
         return base_shape
 
@@ -1149,6 +1201,65 @@ class MambaMixer2(MambaBase, PluggableLayer):
         return MambaAttentionBackendEnum.MAMBA2
 
 
+def share_replayssm_ring_trackers(
+    ordered_layer_names: list[str],
+    forward_context: dict[str, Attention],
+    kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
+) -> None:
+    """Share ring cursors within each cache-slot index namespace.
+
+    Layers backed by one KV-cache group use the same physical block indices and
+    can therefore share cursors. Different KV-cache groups may assign different
+    block indices to the same request and must keep separate cursor tensors.
+    The final local layer in each group advances its cursors after every layer
+    in that group has consumed the previous values.
+    """
+
+    replayssm_mixers: dict[str, MambaMixer2] = {}
+    for layer_name in ordered_layer_names:
+        layer = forward_context[layer_name]
+        if (
+            isinstance(layer, MambaMixer2)
+            and layer.use_replayssm
+            and layer.mamba_config.backend == MambaBackendEnum.FLASHINFER
+        ):
+            replayssm_mixers[layer_name] = layer
+
+    layer_to_group: dict[str, int] = {}
+    if kv_cache_groups:
+        for group_idx, group in enumerate(kv_cache_groups):
+            for layer_name in group.layer_names:
+                layer_to_group[layer_name] = group_idx
+
+    groups_by_namespace: dict[int | str, list[str]] = {}
+    for layer_name in ordered_layer_names:
+        if layer_name not in replayssm_mixers:
+            continue
+        namespace = layer_to_group.get(layer_name, layer_name)
+        groups_by_namespace.setdefault(namespace, []).append(layer_name)
+
+    for group_layer_names in groups_by_namespace.values():
+        last_layer_name = group_layer_names[-1]
+
+        first_mixer = replayssm_mixers[group_layer_names[0]]
+        first_state = first_mixer.kv_cache[1]
+        num_blocks, device = first_state.shape[0], first_state.device
+        for layer_name in group_layer_names:
+            state = replayssm_mixers[layer_name].kv_cache[1]
+            if (state.shape[0], state.device) != (num_blocks, device):
+                raise ValueError(
+                    "ReplaySSM layers in one cache group must share cache capacity"
+                )
+
+        ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device)
+        prev_num_accepted = torch.zeros_like(ring_start)
+        for layer_name in group_layer_names:
+            mixer = replayssm_mixers[layer_name]
+            mixer._replayssm_ring_start = ring_start
+            mixer._replayssm_prev_num_accepted = prev_num_accepted
+            mixer._updates_replayssm_trackers = layer_name == last_layer_name
+
+
 def mamba_mixer2(
     projected_states: torch.Tensor,
     output: torch.Tensor,

vllm/model_executor/layers/mamba/mamba_utils.py

@@ -10,6 +10,7 @@ import torch
 
 import vllm.envs as envs
 from vllm.config.cache import MambaDType
+from vllm.config.mamba import MambaBackendEnum
 from vllm.config.model import ModelDType
 from vllm.distributed import divide
 from vllm.logger import init_logger
@@ -214,20 +215,25 @@ class MambaStateShapeCalculator:
         base_shapes: tuple[tuple[int, ...], ...],
         n_groups: int,
         tp_world_size: int,
-        replayssm_buffer_len: int,
+        logical_window: int,
+        backend: MambaBackendEnum,
     ) -> tuple[tuple[int, ...], ...]:
-        """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a
-        base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape
-        ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended
-        ``n_groups``.
+        """Append the physical ReplaySSM ring shapes.
+
+        ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``;
+        B_cache uses the un-extended ``n_groups``.
         """
+        ring_buffer_len = logical_window
+        if backend == MambaBackendEnum.FLASHINFER:
+            # FlashInfer keeps the live window and appended token together.
+            ring_buffer_len += 1
         local_nheads, head_dim, state_size = base_shapes[1]
         local_ngroups = divide(n_groups, tp_world_size)
         return (
             *base_shapes,
-            (local_nheads, replayssm_buffer_len, head_dim),
-            (local_nheads, replayssm_buffer_len),
-            (local_ngroups, replayssm_buffer_len, state_size),
+            (local_nheads, ring_buffer_len, head_dim),
+            (local_nheads, ring_buffer_len),
+            (local_ngroups, ring_buffer_len, state_size),
         )
 
     @classmethod

vllm/model_executor/layers/mamba/ops/ssu_dispatch.py

@@ -10,11 +10,14 @@ the backend defaults to 'cpu'.
 """
 
 from abc import ABC, abstractmethod
+from collections.abc import Callable
+from functools import cache
 
 import torch
 
 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
 from vllm.logger import init_logger
+from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
@@ -22,6 +25,99 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
 logger = init_logger(__name__)
 
 
+@triton.jit(
+    do_not_specialize=["n_slots", "state_batch_indices_stride"],
+    do_not_specialize_on_alignment=["state_batch_indices"],
+)
+def _update_replayssm_ring_trackers_kernel(
+    ring_start,
+    prev_num_accepted,
+    state_batch_indices,
+    state_batch_indices_stride,
+    n_slots,
+    num_states,
+    logical_window: tl.constexpr,
+    ring_buffer_len: tl.constexpr,
+    pad_slot_id: tl.constexpr,
+    RESET: tl.constexpr,
+    BLOCK: tl.constexpr,
+) -> None:
+    offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < n_slots
+    slots = tl.load(
+        state_batch_indices + offsets * state_batch_indices_stride,
+        mask=mask,
+        other=pad_slot_id,
+    )
+    valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states)
+    if RESET:
+        tl.store(ring_start + slots, 0, mask=valid)
+        tl.store(prev_num_accepted + slots, 0, mask=valid)
+    else:
+        prev = tl.load(prev_num_accepted + slots, mask=valid, other=0)
+        start = tl.load(ring_start + slots, mask=valid, other=0)
+        must_checkpoint = prev + 1 > logical_window
+        next_start = tl.where(
+            must_checkpoint,
+            (start + prev) % ring_buffer_len,
+            start,
+        )
+        next_prev = tl.where(must_checkpoint, 1, prev + 1)
+        tl.store(ring_start + slots, next_start, mask=valid)
+        tl.store(prev_num_accepted + slots, next_prev, mask=valid)
+
+
+def update_replayssm_ring_trackers(
+    ring_start: torch.Tensor,
+    prev_num_accepted: torch.Tensor,
+    state_batch_indices: torch.Tensor,
+    logical_window: int | None = None,
+    ring_buffer_len: int | None = None,
+    pad_slot_id: int = NULL_BLOCK_ID,
+) -> None:
+    """Reset selected trackers, or advance them when a window is provided."""
+    if state_batch_indices.dim() > 1:
+        state_batch_indices = state_batch_indices[:, 0]
+    n_slots = state_batch_indices.numel()
+    if n_slots == 0:
+        return
+    reset = logical_window is None
+    if reset:
+        logical_window = 0
+        ring_buffer_len = 1
+    else:
+        assert ring_buffer_len is not None
+    block = 128
+    _update_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)](
+        ring_start,
+        prev_num_accepted,
+        state_batch_indices,
+        state_batch_indices.stride(0),
+        n_slots,
+        min(ring_start.numel(), prev_num_accepted.numel()),
+        logical_window,
+        ring_buffer_len,
+        pad_slot_id,
+        RESET=reset,
+        BLOCK=block,
+    )
+
+
+def reset_replayssm_ring_trackers(
+    ring_start: torch.Tensor,
+    prev_num_accepted: torch.Tensor,
+    state_batch_indices: torch.Tensor,
+    pad_slot_id: int = NULL_BLOCK_ID,
+) -> None:
+    """Reset selected ReplaySSM ring trackers."""
+    update_replayssm_ring_trackers(
+        ring_start,
+        prev_num_accepted,
+        state_batch_indices,
+        pad_slot_id=pad_slot_id,
+    )
+
+
 class MambaSSUBackend(ABC):
     """Abstract base class for Mamba SSU backends."""
 
@@ -262,15 +358,115 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = {
 _mamba_ssu_backend: MambaSSUBackend | None = None
 
 
+_flashinfer_replayssm_kernel: Callable[..., torch.Tensor] | None = None
+
+
+@cache
+def flashinfer_replayssm_autotune_supported() -> bool:
+    """Return True when FlashInfer exposes ReplaySSM autotuning."""
+    try:
+        from flashinfer.mamba.checkpointing_ssu import (  # noqa: F401
+            CheckpointingSSURunner,
+        )
+    except ImportError:
+        return False
+    return True
+
+
+def selective_state_update_replayssm_flashinfer(
+    state: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    A: torch.Tensor,
+    B: torch.Tensor,
+    C: torch.Tensor,
+    out: torch.Tensor,
+    x_cache: torch.Tensor,
+    B_cache: torch.Tensor,
+    dt_cache: torch.Tensor,
+    ring_start: torch.Tensor,
+    prev_num_accepted_tokens: torch.Tensor,
+    logical_window: int,
+    D: torch.Tensor | None = None,
+    dt_bias: torch.Tensor | None = None,
+    dt_softplus: bool = False,
+    state_batch_indices: torch.Tensor | None = None,
+    null_block_id: int = NULL_BLOCK_ID,
+    scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+    update_trackers: bool = True,
+    enable_stochastic_rounding: bool = False,
+    stochastic_rounding_philox_rounds: int = 0,
+) -> torch.Tensor:
+    """Run FlashInfer checkpointing SSU and optionally advance shared trackers."""
+    if _flashinfer_replayssm_kernel is None:
+        raise RuntimeError(
+            "FlashInfer ReplaySSM has not been initialized. "
+            "Call initialize_mamba_ssu_backend() with use_replayssm=True."
+        )
+
+    if x.dim() == 3:
+        x = x.unsqueeze(1)
+        dt = dt.unsqueeze(1)
+        B = B.unsqueeze(1)
+        C = C.unsqueeze(1)
+        out = out.unsqueeze(1)
+
+    indices = state_batch_indices
+    if indices is not None and indices.dim() > 1:
+        indices = indices[:, 0]
+
+    cb_scaled = cumAdt_vec = cb_old = None
+    if scratch is not None:
+        cb_scaled, cumAdt_vec, cb_old = scratch
+
+    rand_seed = (
+        torch.randint(0, 2**32, (1,), device=state.device, dtype=torch.int64)
+        if enable_stochastic_rounding
+        else None
+    )
+    result = _flashinfer_replayssm_kernel(
+        state,
+        x_cache,
+        B_cache,
+        dt_cache,
+        ring_start,
+        prev_num_accepted_tokens,
+        x,
+        dt,
+        A,
+        B,
+        C,
+        out,
+        D=D,
+        dt_bias=dt_bias,
+        dt_softplus=dt_softplus,
+        state_batch_indices=indices,
+        pad_slot_id=null_block_id,
+        rand_seed=rand_seed,
+        philox_rounds=stochastic_rounding_philox_rounds or 10,
+        cb_scaled=cb_scaled,
+        cumAdt_vec=cumAdt_vec,
+        cb_old=cb_old,
+    )
+    if update_trackers and indices is not None:
+        update_replayssm_ring_trackers(
+            ring_start,
+            prev_num_accepted_tokens,
+            indices,
+            logical_window=logical_window,
+            ring_buffer_len=x_cache.size(2),
+            pad_slot_id=null_block_id,
+        )
+    return result
+
+
 def initialize_mamba_ssu_backend(
     mamba_config: MambaConfig,
     kv_cache_config: KVCacheConfig,
+    *,
+    use_replayssm: bool = False,
 ) -> None:
-    """Initialize the global Mamba SSU backend.
-
-    No-op if `kv_cache_config` contains no specs that call
-    selective_state_update.
-    """
+    """Initialize the Mamba SSU backend and optional FlashInfer ReplaySSM."""
     if not any(
         isinstance(g.kv_cache_spec, MambaSpec)
         and g.kv_cache_spec.mamba_type
@@ -279,13 +475,9 @@ def initialize_mamba_ssu_backend(
     ):
         return
 
-    global _mamba_ssu_backend
-
+    global _flashinfer_replayssm_kernel, _mamba_ssu_backend
     backend = mamba_config.backend
 
-    # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is
-    # unstable or unavailable.  Silently fall back to the CPU
-    # backend unless the user explicitly chose something other than "triton".
     if backend == MambaBackendEnum.TRITON:
         from vllm.platforms import current_platform
 
@@ -301,13 +493,28 @@ def initialize_mamba_ssu_backend(
             f"Unknown Mamba SSU backend: {backend}. "
             f"Valid options: {list(_BACKEND_REGISTRY.keys())}"
         )
+    if use_replayssm and backend not in (
+        MambaBackendEnum.TRITON,
+        MambaBackendEnum.FLASHINFER,
+    ):
+        raise ValueError(f"ReplaySSM does not support mamba backend {backend.value!r}")
 
     backend_cls = _BACKEND_REGISTRY[backend]
-    if isinstance(_mamba_ssu_backend, backend_cls):
-        return
+    if not isinstance(_mamba_ssu_backend, backend_cls):
+        _mamba_ssu_backend = backend_cls(mamba_config)
+        logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
 
-    _mamba_ssu_backend = backend_cls(mamba_config)
-    logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name)
+    _flashinfer_replayssm_kernel = None
+    if use_replayssm and backend == MambaBackendEnum.FLASHINFER:
+        try:
+            from flashinfer.mamba.checkpointing_ssu import checkpointing_ssu
+        except ImportError as e:
+            raise ImportError(
+                "FlashInfer ReplaySSM requires a compatible flashinfer-python package"
+            ) from e
+        _flashinfer_replayssm_kernel = checkpointing_ssu
+    if use_replayssm:
+        logger.info("Using %s ReplaySSM backend.", backend.value)
 
 
 def get_mamba_ssu_backend() -> MambaSSUBackend:

vllm/model_executor/models/nemotron_h.py

@@ -747,7 +747,8 @@ class NemotronHForCausalLM(
         )
         if cache_config.use_replayssm:
             return MambaStateDtypeCalculator.append_replayssm_ring(
-                base_dtype, vllm_config.model_config.dtype
+                base_dtype,
+                vllm_config.model_config.dtype,
             )
         return base_dtype
 
@@ -784,10 +785,11 @@ class NemotronHForCausalLM(
         )
         if cache_config.use_replayssm:
             return MambaStateShapeCalculator.append_replayssm_ring(
-                base_shape,
-                hf_config.n_groups,
-                parallel_config.tensor_parallel_size,
-                cache_config.replayssm_buffer_len,
+                base_shapes=base_shape,
+                n_groups=hf_config.n_groups,
+                tp_world_size=parallel_config.tensor_parallel_size,
+                logical_window=cache_config.replayssm_buffer_len,
+                backend=vllm_config.mamba_config.backend,
             )
         return base_shape
 

vllm/model_executor/warmup/kernel_warmup.py

@@ -38,6 +38,9 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import (
     kimi_k3_triton_warmup,
 )
 from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup
+from vllm.model_executor.warmup.replayssm_warmup import (
+    replayssm_autotune_warmup,
+)
 from vllm.platforms import current_platform
 from vllm.utils.deep_gemm import is_deep_gemm_supported
 from vllm.utils.flashinfer import has_flashinfer
@@ -335,6 +338,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
             fi_utils.autotune(tune_mode=True, **autotune_kwargs),
         ):
             _run_flashinfer_autotune_dummy_runs(runner)
+            replayssm_autotune_warmup(runner)
     finally:
         set_autotune_process_group(None)
 

vllm/model_executor/warmup/replayssm_warmup.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.config.mamba import MambaBackendEnum
+from vllm.logger import init_logger
+from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+    flashinfer_replayssm_autotune_supported,
+)
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+logger = init_logger(__name__)
+
+
+def _replayssm_autotune_kwargs(
+    runner: "GPUModelRunner",
+) -> tuple[int, dict[str, Any]] | None:
+    config = runner.vllm_config
+    if not (
+        config.cache_config.use_replayssm
+        and config.mamba_config.backend == MambaBackendEnum.FLASHINFER
+    ):
+        return None
+    if not flashinfer_replayssm_autotune_supported():
+        logger.info_once(
+            "Skipping FlashInfer ReplaySSM autotuning because "
+            "flashinfer.mamba.checkpointing_ssu.CheckpointingSSURunner "
+            "is unavailable."
+        )
+        return None
+    v2_runner: Any = runner
+    query_len = (
+        v2_runner.decode_query_len
+        if config.use_v2_model_runner
+        else runner.uniform_decode_query_len
+    )
+    max_num_reqs = min(
+        runner.scheduler_config.max_num_seqs,
+        runner.max_num_tokens // query_len,
+        runner.kv_cache_config.num_blocks - 1,
+    )
+    decode_kwargs = {
+        "num_tokens": max_num_reqs * query_len,
+        "skip_eplb": True,
+        "is_profile": True,
+        "randomize_inputs": True,
+        "uniform_decode": True,
+    }
+    if config.use_v2_model_runner:
+        decode_kwargs["valid_dummy_state_slots"] = True
+    else:
+        decode_kwargs.update(
+            allow_microbatching=False,
+            force_attention=True,
+            profile_seq_lens=query_len + 1,
+        )
+    return max_num_reqs, decode_kwargs
+
+
+@contextmanager
+def _temporary_replayssm_autotune_state(
+    runner: "GPUModelRunner", max_num_reqs: int
+) -> Iterator[None]:
+    from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
+    from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
+        reset_replayssm_ring_trackers,
+        update_replayssm_ring_trackers,
+    )
+
+    reset_tensors: dict[int, torch.Tensor] = {}
+    tracker_specs: dict[int, tuple[torch.Tensor, torch.Tensor, int, int]] = {}
+    for module in runner.get_model().modules():
+        if not isinstance(module, MambaMixer2) or not module.use_replayssm:
+            continue
+        assert module.replayssm_buffer_len is not None
+        ring_start = module._replayssm_ring_start
+        prev_num_accepted = module._replayssm_prev_num_accepted
+        tracker_specs.setdefault(
+            ring_start.data_ptr(),
+            (
+                ring_start,
+                prev_num_accepted,
+                module.replayssm_buffer_len,
+                module.kv_cache[2].size(2),
+            ),
+        )
+        tensors = (
+            *module.kv_cache,
+            ring_start,
+            prev_num_accepted,
+        )
+        for tensor in tensors:
+            if tensor.numel():
+                reset_tensors.setdefault(tensor.data_ptr(), tensor)
+
+    v2_runner: Any = runner
+    block_tables = saved_block_ids = None
+    if not runner.vllm_config.use_v2_model_runner:
+        block_tables = runner.input_batch.block_table.block_tables
+        saved_block_ids = tuple(
+            block_table.block_table.np[:max_num_reqs, 0].copy()
+            for block_table in block_tables
+        )
+        dummy_block_ids = range(1, max_num_reqs + 1)
+        for block_table in block_tables:
+            block_table.block_table.np[:max_num_reqs, 0] = dummy_block_ids
+        runner.input_batch.block_table.commit_block_table(max_num_reqs)
+
+    first_tracker = next(iter(tracker_specs.values()), None)
+    if first_tracker is not None and first_tracker[0].is_cuda:
+        state_slots = torch.arange(
+            1, max_num_reqs + 1, dtype=torch.int32, device=first_tracker[0].device
+        )
+        for (
+            ring_start,
+            prev_num_accepted,
+            logical_window,
+            ring_buffer_len,
+        ) in tracker_specs.values():
+            # Compile reset (prefill) and advance (decode) before inference.
+            # The final reset leaves the decode tuning run in a clean state.
+            reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots)
+            update_replayssm_ring_trackers(
+                ring_start,
+                prev_num_accepted,
+                state_slots,
+                logical_window,
+                ring_buffer_len,
+            )
+            reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots)
+
+    try:
+        yield
+    finally:
+        if runner.vllm_config.use_v2_model_runner:
+            v2_runner.block_tables.get_dummy_block_tables(max_num_reqs)
+        else:
+            assert block_tables is not None and saved_block_ids is not None
+            for block_table, block_ids in zip(block_tables, saved_block_ids):
+                block_table.block_table.np[:max_num_reqs, 0] = block_ids
+            runner.input_batch.block_table.commit_block_table(max_num_reqs)
+        for tensor in reset_tensors.values():
+            tensor[1 : max_num_reqs + 1].zero_()
+
+
+def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None:
+    autotune = _replayssm_autotune_kwargs(runner)
+    if autotune is None:
+        return
+    max_num_reqs, decode_kwargs = autotune
+    with _temporary_replayssm_autotune_state(runner, max_num_reqs):
+        runner._dummy_run(**decode_kwargs)

vllm/v1/attention/backends/mamba_attn.py

@@ -8,6 +8,7 @@ from typing import Any, ClassVar, TypeVar
 import torch
 
 from vllm.config import VllmConfig
+from vllm.config.mamba import MambaBackendEnum
 from vllm.utils.math_utils import cdiv
 from vllm.utils.torch_utils import async_tensor_h2d
 from vllm.v1.attention.backend import (
@@ -74,12 +75,15 @@ class BaseMambaAttentionMetadata:
     nums_dict: dict | None = None
     batch_ptr: torch.Tensor | None = None
     token_chunk_offset_ptr: torch.Tensor | None = None
-    # ReplaySSM standard decode: per-row ring cursor and flush flag, plus the
-    # per-step (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for the
-    # precomputed k^T q products. All None when use_replayssm is disabled.
+    # ReplaySSM standard decode — Triton: per-row ring cursor and flush flag,
+    # plus (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for
+    # precomputed B·C products. All None when use_replayssm is disabled or
+    # the FlashInfer ReplaySSM backend is selected.
     write_pos_d: torch.Tensor | None = None
     is_flush_d: torch.Tensor | None = None
     bc_pre_scratch: torch.Tensor | None = None
+    # ReplaySSM — FlashInfer checkpointing_ssu two-kernel scratch.
+    replayssm_scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
 
 
 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
@@ -107,6 +111,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
         self.use_spec_decode = self.num_spec_tokens > 0
         self.use_replayssm = vllm_config.cache_config.use_replayssm
         self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len
+        self.use_flashinfer_replayssm = (
+            self.use_replayssm
+            and vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER
+        )
 
         scheduler_config = vllm_config.scheduler_config
         self.decode_cudagraph_max_bs: int = scheduler_config.max_num_seqs
@@ -166,9 +174,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
                 dtype=torch.int32,
                 device=device,
             )
-        # ReplaySSM standard-decode CUDA-graph buffers: per-row ring cursor,
-        # flush flag, and the k^T q precompute scratch.
-        if self.use_replayssm:
+        self.decode_bc_pre_scratch: torch.Tensor | None = None
+        self.decode_replayssm_scratch: (
+            tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None
+        ) = None
+        # ReplaySSM CUDA-graph buffers for the selected backend.
+        if self.use_replayssm and not self.use_flashinfer_replayssm:
             self.decode_write_pos_d: torch.Tensor = torch.empty(
                 (self.decode_cudagraph_max_bs,),
                 dtype=torch.int32,
@@ -185,7 +196,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
             bc_scratch_bs = max(
                 self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs
             )
-            self.decode_bc_pre_scratch: torch.Tensor = torch.empty(
+            self.decode_bc_pre_scratch = torch.empty(
                 (
                     bc_scratch_bs,
                     bc_ngroups,
@@ -194,8 +205,20 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
                 dtype=torch.float32,
                 device=device,
             )
-        else:
-            self.decode_bc_pre_scratch = None
+        elif self.use_flashinfer_replayssm:
+            from flashinfer.mamba.checkpointing_ssu import (
+                allocate_checkpointing_ssu_scratch,
+            )
+
+            nheads = kv_cache_spec.shapes[2][0]
+            self.decode_replayssm_scratch = allocate_checkpointing_ssu_scratch(
+                batch_size=scheduler_config.max_num_seqs,
+                num_heads=nheads,
+                num_predicted_tokens=1,
+                max_window=self.replayssm_buffer_len,
+                dtype=vllm_config.model_config.dtype,
+                device=device,
+            )
 
         self._init_reorder_batch_threshold(1, self.use_spec_decode)
         if self.use_spec_decode:
@@ -500,6 +523,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
         nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
         write_pos_d = None
         is_flush_d = None
+        replayssm_scratch = None
 
         if self.vllm_config.cache_config.mamba_cache_mode == "all":
             num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
@@ -586,7 +610,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
                     num_reqs - num_prefills : num_reqs
                 ]
 
-        if self.use_replayssm and num_decodes > 0:
+        if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0:
             decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu
             num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu
             if decode_base_cpu is None or num_computed_tokens_cpu is None:
@@ -651,6 +675,15 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
                 device=common_attn_metadata.query_start_loc.device,
             )
 
+        if self.use_flashinfer_replayssm and num_decodes > 0:
+            assert self.decode_replayssm_scratch is not None
+            cb_scaled, cumAdt_vec, cb_old = self.decode_replayssm_scratch
+            replayssm_scratch = (
+                cb_scaled[:num_decodes],
+                cumAdt_vec[:num_decodes],
+                cb_old[:num_decodes],
+            )
+
         bc_pre_scratch = None
         if (
             self.use_replayssm
@@ -671,6 +704,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
             write_pos_d=write_pos_d,
             is_flush_d=is_flush_d,
             bc_pre_scratch=bc_pre_scratch,
+            replayssm_scratch=replayssm_scratch,
             num_accepted_tokens=num_accepted_tokens,
             query_start_loc_d=query_start_loc_d,
             block_idx_last_scheduled_token=block_idx_last_scheduled_token,
@@ -708,6 +742,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
         write_pos_d = metadata.write_pos_d
         is_flush_d = metadata.is_flush_d
         bc_pre_scratch = metadata.bc_pre_scratch
+        replayssm_scratch = metadata.replayssm_scratch
         if (
             metadata.num_prefills == 0
             and metadata.num_decodes <= self.decode_cudagraph_max_bs
@@ -769,7 +804,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
                     )
                     block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0
 
-            if self.use_replayssm:
+            if self.use_replayssm and not self.use_flashinfer_replayssm:
                 assert write_pos_d is not None
                 assert is_flush_d is not None
                 self.decode_write_pos_d[: metadata.num_decodes].copy_(
@@ -788,6 +823,14 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
 
                 if self.decode_bc_pre_scratch is not None:
                     bc_pre_scratch = self.decode_bc_pre_scratch[:padded_bs]
+            elif self.use_flashinfer_replayssm:
+                assert self.decode_replayssm_scratch is not None
+                cb_scaled, cumAdt_vec, cb_old = self.decode_replayssm_scratch
+                replayssm_scratch = (
+                    cb_scaled[:padded_bs],
+                    cumAdt_vec[:padded_bs],
+                    cb_old[:padded_bs],
+                )
 
         return replace(
             metadata,
@@ -797,6 +840,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
             write_pos_d=write_pos_d,
             is_flush_d=is_flush_d,
             bc_pre_scratch=bc_pre_scratch,
+            replayssm_scratch=replayssm_scratch,
             block_idx_last_scheduled_token=block_idx_last_scheduled_token,
             block_idx_last_computed_token=block_idx_last_computed_token,
             block_idx_last_scheduled_token_prev_step=(

vllm/v1/worker/gpu/attn_utils.py

@@ -229,7 +229,13 @@ def init_kv_cache(
         in ("longcat_flash", "longcat_flash_ngram")
         else 1
     )
-    bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module)
+    bind_kv_cache(
+        kv_caches,
+        forward_context,
+        runner_kv_caches,
+        num_attn_module,
+        kv_cache_groups=kv_cache_config.kv_cache_groups,
+    )
     return kv_caches
 
 

vllm/v1/worker/gpu/model_runner.py

@@ -633,7 +633,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             cls=self.pcp_manager_cls,
         )
         initialize_mamba_ssu_backend(
-            self.vllm_config.mamba_config, self.kv_cache_config
+            self.vllm_config.mamba_config,
+            self.kv_cache_config,
+            use_replayssm=self.vllm_config.cache_config.use_replayssm,
         )
         if self.adaptive_verification is not None:
             self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE
@@ -707,6 +709,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         context_len: int = 0,
         skip_eplb: bool = False,
         is_profile: bool = False,
+        valid_dummy_state_slots: bool = False,
         **kwargs,
     ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
         if skip_attn and not is_profile:
@@ -764,6 +767,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 skip_attn_for_dummy_run=skip_attn,
                 is_profile=is_profile,
                 context_len=context_len,
+                valid_dummy_state_slots=valid_dummy_state_slots,
             )
         self.kv_connector.set_disabled(False)
 
@@ -1390,9 +1394,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         return block_tables, slot_mappings
 
     def prepare_dummy_attn(
-        self, input_batch: InputBatch
+        self, input_batch: InputBatch, valid_state_slots: bool = False
     ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
         block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
+        if valid_state_slots:
+            state_slots = torch.arange(
+                1,
+                input_batch.num_reqs + 1,
+                dtype=torch.int32,
+                device=self.device,
+            )
+            for block_table in block_tables:
+                block_table[:, 0].copy_(state_slots)
         slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
             self.pcp_manager, self.block_tables, input_batch.num_tokens
         )
@@ -1525,6 +1538,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         skip_attn_for_dummy_run: bool = False,
         is_profile: bool = False,
         context_len: int = 0,
+        valid_dummy_state_slots: bool = False,
     ) -> ModelRunnerOutput | IntermediateTensors | None:
         if not dummy_run:
             # Update the request states.
@@ -1624,7 +1638,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 max_query_len=batch_desc.max_query_len,
             )
             if not skip_attn_for_dummy_run:
-                block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
+                block_tables, slot_mappings = self.prepare_dummy_attn(
+                    input_batch, valid_dummy_state_slots
+                )
                 if context_len:
                     set_dummy_context(
                         input_batch,

vllm/v1/worker/gpu_model_runner.py

@@ -7434,6 +7434,7 @@ class GPUModelRunner(
             self.compilation_config.static_forward_context,
             self.kv_caches,
             num_attn_module,
+            kv_cache_groups=kv_cache_config.kv_cache_groups,
         )
         return kv_caches
 
@@ -7485,7 +7486,9 @@ class GPUModelRunner(
         self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config)
         self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling)
         initialize_mamba_ssu_backend(
-            self.vllm_config.mamba_config, self.kv_cache_config
+            self.vllm_config.mamba_config,
+            self.kv_cache_config,
+            use_replayssm=self.vllm_config.cache_config.use_replayssm,
         )
         # The kernel block size for all KV cache groups. For example, if
         # kv_cache_manager uses block_size 256 for a given group, but the attention

vllm/v1/worker/utils.py

@@ -13,6 +13,7 @@ import torch
 from vllm.config import CacheConfig, VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention import Attention
+from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers
 from vllm.model_executor.layers.utils import warmup_rocm_skinny_gemm_workspaces
 from vllm.model_executor.models.interfaces import MultiModalEmbeddings
 from vllm.model_executor.models.utils import extract_layer_index
@@ -575,6 +576,7 @@ def bind_kv_cache(
     forward_context: dict[str, Attention],
     runner_kv_caches: list[torch.Tensor],
     num_attn_module: int = 1,
+    kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None,
 ) -> None:
     """
     Bind the allocated KV cache to both ModelRunner and forward context so
@@ -600,6 +602,7 @@ def bind_kv_cache(
     for layer_name in kv_caches:
         index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name)
 
+    ordered_layer_names: list[str] = []
     for layer_index in sorted(index2name.keys()):
         layer_names = index2name[layer_index]
         if len(layer_names) > 1:
@@ -613,6 +616,7 @@ def bind_kv_cache(
             current_platform.check_runner_kv_caches_multi_layer()
         for layer_name in layer_names:
             runner_kv_caches.append(kv_caches[layer_name])
+            ordered_layer_names.append(layer_name)
 
     # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks
     # its raw allocation into the per-layer view(s) it needs (e.g. Mamba
@@ -621,6 +625,8 @@ def bind_kv_cache(
     for layer_name, kv_cache in kv_caches.items():
         forward_context[layer_name].bind_kv_cache(kv_cache)
 
+    share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups)
+
 
 def clear_layer_kv_caches(layers: Iterable[Any]) -> None:
     """Detach the KV/state cache tensors installed by bind_kv_cache().