[K3] support recoverssm for K3 (#51855)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/models/kimi_k3/test_kda.py

@@ -6,6 +6,8 @@ Compares chunk_kda against a naive recurrent reference (float32).
 Uses torch.rand for q/k/v to match FLA's test pattern.
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.nn.functional as F
@@ -22,6 +24,12 @@ from vllm.models.kimi_k3.nvidia.kda import (
     is_flashkda_supported,
     is_fused_kda_decode_supported,
 )
+from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM
+from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops
+from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
+    KDARecoverSSMCommitContext,
+    kda_recoverssm_verify,
+)
 from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
     chunk_kda,
     chunk_kda_with_fused_gate,
@@ -48,6 +56,38 @@ PACKED_DECODE_IMPLS = {
 }
 
 
+def test_kda_recoverssm_config_state_layout():
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            dtype=torch.bfloat16,
+            hf_config=SimpleNamespace(
+                linear_attn_config={
+                    "num_heads": 4,
+                    "head_dim": 32,
+                    "short_conv_kernel_size": 4,
+                }
+            ),
+        ),
+        cache_config=SimpleNamespace(
+            mamba_cache_dtype="auto",
+            use_kda_recoverssm=True,
+        ),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        speculative_config=SimpleNamespace(num_speculative_tokens=2),
+    )
+
+    assert KimiLinearForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == (
+        torch.bfloat16,
+        torch.float32,
+        torch.float32,
+        torch.bfloat16,
+    )
+    assert KimiLinearForCausalLM.get_mamba_state_shape_from_config(vllm_config)[2:] == (
+        (4, 3, 32),
+        (4, 3, 64),
+    )
+
+
 @torch.inference_mode()
 def test_gather_initial_states_correctness():
     row_size = 8 * 128 * 128
@@ -535,6 +575,308 @@ def test_kda_spec_decode_correctness(
     assert torch.isnan(output_storage[..., H * D :]).all()
 
 
+@pytest.mark.parametrize(
+    (
+        "conv_state_dim_first",
+        "use_request_indices",
+        "lower_bound",
+        "align_mode",
+    ),
+    [
+        pytest.param(False, False, None, False, id="baseline"),
+        pytest.param(True, True, -5.0, True, id="all-features"),
+        pytest.param(False, True, -5.0, False, id="request-indexed"),
+        pytest.param(True, False, None, True, id="aligned"),
+    ],
+)
+@torch.inference_mode()
+def test_kda_recoverssm_verify_and_group_commit(
+    monkeypatch: pytest.MonkeyPatch,
+    lower_bound: float | None,
+    use_request_indices: bool,
+    conv_state_dim_first: bool,
+    align_mode: bool,
+):
+    monkeypatch.setattr(
+        recoverssm_ops,
+        "is_conv_state_dim_first",
+        lambda: conv_state_dim_first,
+    )
+    num_layers, num_seqs, query_len = 2, 2, 8
+    num_blocks, num_heads, dim = (7 if align_mode else 3), 4, 128
+    total_tokens = num_seqs * query_len
+    torch.manual_seed(20260808)
+
+    q, k, v, raw_g = [
+        torch.randn(
+            1,
+            total_tokens,
+            num_heads,
+            dim,
+            dtype=torch.bfloat16,
+            device=DEVICE,
+        )
+        for _ in range(4)
+    ]
+    raw_beta = torch.randn(
+        1,
+        total_tokens,
+        num_heads,
+        dtype=torch.bfloat16,
+        device=DEVICE,
+    )
+    query_start_loc = torch.arange(
+        0,
+        total_tokens + 1,
+        query_len,
+        dtype=torch.int32,
+        device=DEVICE,
+    )
+    state_indices = torch.tensor(
+        [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE
+    )
+    accepted = [2, 8]
+    if use_request_indices:
+        global_num_accepted = torch.tensor(
+            [0, accepted[0], 0, accepted[1]],
+            dtype=torch.int32,
+            device=DEVICE,
+        )
+        request_indices = torch.tensor([1, 3], dtype=torch.int32, device=DEVICE)
+    else:
+        global_num_accepted = torch.tensor(accepted, dtype=torch.int32, device=DEVICE)
+        request_indices = None
+
+    block_table = None
+    num_computed_tokens = None
+    mamba_block_size = None
+    if align_mode:
+        batch_size = 4 if use_request_indices else num_seqs
+        block_table = torch.full((batch_size, 2), -1, dtype=torch.int32, device=DEVICE)
+        rows = (
+            request_indices
+            if request_indices is not None
+            else torch.arange(num_seqs, device=DEVICE)
+        )
+        block_table[rows] = torch.tensor(
+            [[1, 5], [2, 6]],
+            dtype=torch.int32,
+            device=DEVICE,
+        )
+        num_computed_tokens = torch.zeros(batch_size, dtype=torch.int32, device=DEVICE)
+        num_computed_tokens[rows] = 4
+        mamba_block_size = 8
+
+    layers = []
+    expected_outputs = []
+    expected_states = []
+    initial_states = []
+    initial_conv_states = []
+    history_len, conv_dim = 3, 12
+    for layer_idx in range(num_layers):
+        A_log = (
+            0.2 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE)
+            + layer_idx * 0.03
+        ).contiguous()
+        dt_bias = (
+            0.1 * torch.randn(num_heads, dim, dtype=torch.float32, device=DEVICE)
+        ).contiguous()
+        checkpoint = 0.01 * torch.randn(
+            num_blocks,
+            num_heads,
+            dim,
+            dim,
+            dtype=torch.float32,
+            device=DEVICE,
+        )
+        conv_shape = (
+            (num_blocks, conv_dim, history_len + query_len - 1)
+            if conv_state_dim_first
+            else (num_blocks, history_len + query_len - 1, conv_dim)
+        )
+        conv_state = torch.randn(conv_shape, dtype=torch.bfloat16, device=DEVICE)
+        correction_cache = torch.empty(
+            num_blocks,
+            num_heads,
+            query_len,
+            dim,
+            dtype=torch.float32,
+            device=DEVICE,
+        )
+        kg_cache = torch.empty(
+            num_blocks,
+            num_heads,
+            query_len,
+            2 * dim,
+            dtype=torch.bfloat16,
+            device=DEVICE,
+        )
+        layer = SimpleNamespace(
+            kv_cache=(
+                conv_state,
+                checkpoint,
+                correction_cache,
+                kg_cache,
+            ),
+            A_log=A_log,
+            dt_bias=dt_bias,
+            local_num_heads=num_heads,
+            head_dim=dim,
+            gate_lower_bound=lower_bound,
+        )
+        layers.append(layer)
+        initial_states.append(checkpoint.clone())
+        initial_conv_states.append(conv_state.clone())
+
+        actual_output = kda_recoverssm_verify(
+            q=q,
+            k=k,
+            v=v,
+            raw_g=raw_g,
+            raw_beta=raw_beta,
+            A_log=A_log,
+            dt_bias=dt_bias,
+            lower_bound=lower_bound,
+            checkpoint_state=checkpoint,
+            correction_cache=correction_cache,
+            kg_cache=kg_cache,
+            query_start_loc=query_start_loc,
+            state_indices=state_indices,
+            spec_query_len=query_len,
+        )
+
+        normalized_q = q.float() * torch.rsqrt(
+            q.float().square().sum(dim=-1, keepdim=True) + 1e-6
+        )
+        normalized_k = k.float() * torch.rsqrt(
+            k.float().square().sum(dim=-1, keepdim=True) + 1e-6
+        )
+        gate_input = raw_g.float() + dt_bias.view(1, 1, num_heads, dim)
+        if lower_bound is None:
+            gate = -A_log.exp().view(1, 1, num_heads, 1) * F.softplus(gate_input)
+        else:
+            gate = lower_bound * torch.sigmoid(
+                A_log.exp().view(1, 1, num_heads, 1) * gate_input
+            )
+        beta = raw_beta.float().sigmoid()
+
+        reference_output = []
+        committed_states = checkpoint.clone()
+        for seq_idx, commit_len in enumerate(accepted):
+            start = seq_idx * query_len
+            end = start + query_len
+            output, _ = naive_recurrent_kda(
+                normalized_q[:, start:end],
+                normalized_k[:, start:end],
+                v[:, start:end],
+                gate[:, start:end],
+                beta[:, start:end],
+                initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2),
+            )
+            reference_output.append(output)
+            _, committed_state = naive_recurrent_kda(
+                normalized_q[:, start : start + commit_len],
+                normalized_k[:, start : start + commit_len],
+                v[:, start : start + commit_len],
+                gate[:, start : start + commit_len],
+                beta[:, start : start + commit_len],
+                initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2),
+                output_final_state=True,
+            )
+            assert committed_state is not None
+            final_block = state_indices[seq_idx]
+            if align_mode:
+                assert block_table is not None
+                row = request_indices[seq_idx] if use_request_indices else seq_idx
+                final_block = block_table[row, (4 + commit_len) // 8]
+            committed_states[final_block] = committed_state.transpose(-1, -2)
+            if align_mode and 4 + commit_len >= 8:
+                _, boundary_state = naive_recurrent_kda(
+                    normalized_q[:, start : start + 4],
+                    normalized_k[:, start : start + 4],
+                    v[:, start : start + 4],
+                    gate[:, start : start + 4],
+                    beta[:, start : start + 4],
+                    initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2),
+                    output_final_state=True,
+                )
+                assert boundary_state is not None
+                assert block_table is not None
+                row = request_indices[seq_idx] if use_request_indices else seq_idx
+                committed_states[block_table[row, 0]] = boundary_state.transpose(-1, -2)
+        expected_outputs.append(torch.cat(reference_output, dim=1))
+        expected_states.append(committed_states)
+        torch.testing.assert_close(checkpoint, initial_states[-1])
+        torch.testing.assert_close(
+            actual_output,
+            expected_outputs[-1],
+            atol=3e-2,
+            rtol=3e-2,
+        )
+
+    context = KDARecoverSSMCommitContext.create(
+        layers,
+        spec_query_len=query_len,
+        max_num_reqs=global_num_accepted.shape[0],
+    )
+    context.commit(
+        global_num_accepted,
+        state_indices,
+        query_start_loc,
+        request_indices=request_indices,
+        block_table=block_table,
+        num_computed_tokens=num_computed_tokens,
+        mamba_block_size=mamba_block_size,
+    )
+
+    for layer_idx, layer in enumerate(layers):
+        torch.testing.assert_close(
+            layer.kv_cache[1],
+            expected_states[layer_idx],
+            atol=3e-3,
+            rtol=3e-3,
+        )
+        for seq_idx, commit_len in enumerate(accepted):
+            block = state_indices[seq_idx]
+            if align_mode:
+                assert block_table is not None
+                row = request_indices[seq_idx] if use_request_indices else seq_idx
+                block = block_table[row, (4 + commit_len) // 8]
+            source_block = state_indices[seq_idx] if align_mode else block
+            if conv_state_dim_first:
+                actual_conv = layer.kv_cache[0][block, :, :history_len]
+                expected_conv = initial_conv_states[layer_idx][
+                    source_block,
+                    :,
+                    commit_len - 1 : commit_len - 1 + history_len,
+                ]
+            else:
+                actual_conv = layer.kv_cache[0][block, :history_len]
+                expected_conv = initial_conv_states[layer_idx][
+                    source_block,
+                    commit_len - 1 : commit_len - 1 + history_len,
+                ]
+            torch.testing.assert_close(actual_conv, expected_conv)
+            if align_mode and 4 + commit_len >= 8:
+                assert block_table is not None
+                boundary_block = block_table[row, 0]
+                if conv_state_dim_first:
+                    actual_boundary_conv = layer.kv_cache[0][
+                        boundary_block, :, :history_len
+                    ]
+                    expected_boundary_conv = initial_conv_states[layer_idx][
+                        state_indices[seq_idx], :, 3 : 3 + history_len
+                    ]
+                else:
+                    actual_boundary_conv = layer.kv_cache[0][
+                        boundary_block, :history_len
+                    ]
+                    expected_boundary_conv = initial_conv_states[layer_idx][
+                        state_indices[seq_idx], 3 : 3 + history_len
+                    ]
+                torch.testing.assert_close(actual_boundary_conv, expected_boundary_conv)
+
+
 @pytest.mark.parametrize(
     ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"),
     [

tests/models/kimi_k3/test_kda_metadata.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from dataclasses import fields
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -26,6 +27,9 @@ from vllm.v1.attention.backends.gdn_attn import (
     GDNAttentionMetadata,
     GDNAttentionMetadataBuilder,
 )
+from vllm.v1.attention.backends.recoverssm_metadata import (
+    RecoverSSMPostprocessMetadata,
+)
 from vllm.v1.attention.backends.utils import (
     NULL_BLOCK_ID,
     mamba_get_block_table_tensor,
@@ -44,8 +48,12 @@ PRUNED_METADATA_FIELDS = {
 }
 
 
-def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata):
-    for field in fields(KimiK3KDAMetadata):
+def _assert_matches_shared_gdn(
+    reference: GDNAttentionMetadata, actual: KimiK3KDAMetadata
+):
+    assert actual.recoverssm_commit is None
+    assert actual.recoverssm_context is None
+    for field in fields(GDNAttentionMetadata):
         actual_value = getattr(actual, field.name)
         expected_value = getattr(reference, field.name)
         if field.name in PRUNED_METADATA_FIELDS:
@@ -78,6 +86,7 @@ def _make_builder(
     full_cuda_graph: bool,
     device: torch.device = DEVICE,
     mamba_cache_mode: str = "none",
+    use_recoverssm: bool = False,
 ) -> AttentionMetadataBuilder:
     vllm_config = create_vllm_config(
         model_name="Qwen/Qwen3.5-0.8B",
@@ -92,17 +101,24 @@ def _make_builder(
         CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE
     )
     vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
-    return builder_cls(
+    vllm_config.cache_config.use_replayssm = use_recoverssm
+    vllm_config.cache_config.use_kda_recoverssm = use_recoverssm
+    builder = builder_cls(
         kv_cache_spec=MambaSpec(
             block_size=BLOCK_SIZE,
             shapes=((16, 64),),
             dtypes=(torch.float16,),
-            num_speculative_blocks=num_speculative_tokens,
+            mamba_cache_mode=mamba_cache_mode,
+            num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens),
         ),
         layer_names=["layer.0"],
         vllm_config=vllm_config,
         device=device,
     )
+    if use_recoverssm:
+        assert isinstance(builder, KimiK3KDAMetadataBuilder)
+        builder.recoverssm_context = Mock()
+    return builder
 
 
 @pytest.mark.parametrize(
@@ -244,6 +260,99 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding():
     torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3]))
 
 
+@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
+def test_recoverssm_spec_uses_one_state_slot_and_current_window(
+    mamba_cache_mode: str,
+):
+    if mamba_cache_mode == "align" and not torch.cuda.is_available():
+        pytest.skip("align metadata construction requires CUDA")
+    device = torch.device("cuda") if mamba_cache_mode == "align" else DEVICE
+    batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, device
+    ).replace(is_prefilling=torch.tensor([True, True, False]))
+    builder = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=False,
+        device=device,
+        mamba_cache_mode=mamba_cache_mode,
+        use_recoverssm=True,
+    )
+    assert isinstance(builder, KimiK3KDAMetadataBuilder)
+    context = builder.recoverssm_context
+    assert context is not None
+    actual = builder.build(
+        0,
+        common_attn_metadata,
+        num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32),
+        num_accepted_tokens=torch.tensor([3, 2, 2], dtype=torch.int32, device=device),
+    )
+
+    assert actual.spec_state_indices_tensor is not None
+    assert actual.spec_state_indices_tensor.shape == (1, 1)
+    torch.testing.assert_close(
+        actual.num_accepted_tokens,
+        torch.ones(1, dtype=torch.int32, device=device),
+    )
+    commit_metadata = actual.recoverssm_commit
+    assert commit_metadata is not None
+    torch.testing.assert_close(
+        commit_metadata.request_indices,
+        torch.tensor([2], dtype=torch.int32, device=device),
+    )
+    assert actual.recoverssm_context is context
+    num_accepted_tokens = torch.tensor([3, 2, 1], dtype=torch.int32, device=device)
+
+    postprocess = actual.commit_recoverssm_state(num_accepted_tokens)
+
+    if mamba_cache_mode == "none":
+        assert commit_metadata.align is None
+        assert postprocess is None
+    else:
+        assert isinstance(postprocess, RecoverSSMPostprocessMetadata)
+        assert postprocess.num_spec_decodes == 1
+        assert postprocess.request_indices is commit_metadata.request_indices
+        assert postprocess.block_table is common_attn_metadata.block_table_tensor
+        assert (
+            postprocess.num_computed_tokens
+            is common_attn_metadata.compute_num_computed_tokens()
+        )
+        assert postprocess.block_size == BLOCK_SIZE
+    args = context.commit.call_args.args
+    assert args[0] is num_accepted_tokens
+    torch.testing.assert_close(args[1], commit_metadata.state_indices[:, 0])
+    torch.testing.assert_close(args[2], commit_metadata.query_start_loc)
+
+
+def test_recoverssm_distinguishes_draftless_decode_from_one_token_prefill():
+    batch = BatchSpec(seq_lens=[40, 30], query_lens=[1, 1])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, DEVICE
+    ).replace(is_prefilling=torch.tensor([False, True]))
+    actual = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=False,
+        use_recoverssm=True,
+    ).build(
+        0,
+        common_attn_metadata,
+        num_decode_draft_tokens_cpu=torch.full((2,), -1, dtype=torch.int32),
+        num_accepted_tokens=torch.ones(2, dtype=torch.int32),
+    )
+
+    assert actual.num_spec_decodes == 1
+    assert actual.num_decodes == 0
+    assert actual.num_prefills == 1
+    assert actual.spec_state_indices_tensor is not None
+    assert actual.spec_state_indices_tensor.shape == (1, 1)
+    torch.testing.assert_close(
+        actual.spec_query_start_loc,
+        torch.tensor([0, 1], dtype=torch.int32),
+    )
+
+
 @pytest.mark.parametrize(
     ("seq_len", "expected_has_initial_state"),
     [
@@ -307,6 +416,38 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn():
     _assert_matches_shared_gdn(reference, actual)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+def test_recoverssm_spec_cudagraph_stages_one_checkpoint_per_request():
+    device = torch.device("cuda")
+    batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3])
+    common_attn_metadata = create_common_attn_metadata(
+        batch, BLOCK_SIZE, device
+    ).replace(is_prefilling=torch.tensor([False, False]))
+    builder = _make_builder(
+        KimiK3KDAMetadataBuilder,
+        num_speculative_tokens=2,
+        full_cuda_graph=True,
+        device=device,
+        use_recoverssm=True,
+    )
+    assert isinstance(builder, KimiK3KDAMetadataBuilder)
+    assert builder.spec_state_indices_tensor.shape == (
+        builder.vllm_config.scheduler_config.max_num_seqs,
+        1,
+    )
+    actual = builder.build_for_cudagraph_capture(common_attn_metadata)
+
+    assert actual.spec_state_indices_tensor is not None
+    assert actual.spec_state_indices_tensor.shape == (batch.batch_size, 1)
+    assert actual.num_accepted_tokens is not None
+    torch.testing.assert_close(
+        actual.num_accepted_tokens,
+        torch.ones(batch.batch_size, dtype=torch.int32, device=device),
+    )
+    assert actual.recoverssm_commit is not None
+    assert actual.recoverssm_commit.request_indices is None
+
+
 def test_kimi_k3_kda_backend_uses_private_metadata_builder():
     assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder
     assert KimiK3KDAAttentionBackend.is_ssm()

tests/models/test_registry.py

@@ -143,8 +143,10 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda):
 @pytest.mark.parametrize(
     "model_arch,supported",
     [
-        # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today.
+        # ReplaySSM is opt-in per model.
         ("NemotronHForCausalLM", True),
+        ("KimiLinearForCausalLM", not current_platform.is_rocm()),
+        ("KimiK3ForConditionalGeneration", not current_platform.is_rocm()),
         ("Mamba2ForCausalLM", False),
         ("Zamba2ForCausalLM", False),
     ],

tests/test_config.py

@@ -29,6 +29,7 @@ from vllm.config import (
 from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.config.kernel import IrOpPriorityConfig
 from vllm.config.load import LoadConfig
+from vllm.config.mamba import MambaBackendEnum
 from vllm.config.utils import get_field
 from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
 from vllm.platforms import current_platform
@@ -37,6 +38,52 @@ from vllm.v1.attention.backend import AttentionCGSupport
 DEVICE_TYPE = current_platform.device_type
 
 
+def test_kda_recoverssm_derivation_is_revalidated():
+    config = SimpleNamespace(
+        cache_config=SimpleNamespace(
+            use_replayssm=True,
+            use_kda_recoverssm=False,
+            mamba_cache_mode="none",
+        ),
+        num_speculative_tokens=3,
+        model_config=SimpleNamespace(
+            supports_replayssm=True,
+            architecture="KimiLinearForCausalLM",
+        ),
+        mamba_config=SimpleNamespace(
+            backend=MambaBackendEnum.TRITON,
+            enable_stochastic_rounding=False,
+        ),
+        parallel_config=SimpleNamespace(pipeline_parallel_size=1),
+        kv_transfer_config=None,
+        use_v2_model_runner=True,
+    )
+
+    VllmConfig.validate_mamba_cached_kernel(config)
+    assert config.cache_config.use_replayssm
+    assert config.cache_config.use_kda_recoverssm
+
+    config.cache_config.mamba_cache_mode = "align"
+    VllmConfig.validate_mamba_cached_kernel(config)
+    config.use_v2_model_runner = False
+    with pytest.raises(ValueError, match="VLLM_USE_V2_MODEL_RUNNER=1"):
+        VllmConfig.validate_mamba_cached_kernel(config)
+    config.use_v2_model_runner = True
+    config.cache_config.mamba_cache_mode = "all"
+    with pytest.raises(ValueError, match="only none and align"):
+        VllmConfig.validate_mamba_cached_kernel(config)
+    config.cache_config.mamba_cache_mode = "none"
+
+    config.model_config.architecture = "NemotronHForCausalLM"
+    with pytest.raises(ValueError, match="only supported for Kimi-K3 KDA"):
+        VllmConfig.validate_mamba_cached_kernel(config)
+
+    config.model_config.architecture = "KimiLinearForCausalLM"
+    config.parallel_config.pipeline_parallel_size = 2
+    with pytest.raises(ValueError, match="pipeline_parallel_size=1"):
+        VllmConfig.validate_mamba_cached_kernel(config)
+
+
 def test_compile_config_repr_succeeds():
     # setup: VllmBackend mutates the config object
     config = VllmConfig()

tests/v1/worker/test_mamba_hybrid_model_state.py

@@ -1,11 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from types import SimpleNamespace
+from unittest.mock import Mock
+
 import pytest
 import torch
 
 from vllm.platforms import current_platform
+from vllm.v1.attention.backends.recoverssm_metadata import (
+    RecoverSSMMetadata,
+    RecoverSSMPostprocessMetadata,
+)
 from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
+from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
 
 
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
@@ -18,6 +26,7 @@ def test_postprocess_state_scalar_with_int32_mapping(
         (4,), 9, dtype=torch.int32, device="cuda"
     )
     state._align_mode = False
+    state.recoverssm = None
     state._mamba_ctx = None
     idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda")
 
@@ -27,3 +36,61 @@ def test_postprocess_state_scalar_with_int32_mapping(
         [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda"
     )
     torch.testing.assert_close(state.num_accepted_tokens_gpu, expected)
+
+
+def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None:
+    state = RecoverSSMState()
+    metadata = Mock(spec=RecoverSSMMetadata)
+    metadata.commit_recoverssm_state.return_value = None
+    num_sampled = torch.tensor([3, 1], dtype=torch.int32)
+    idx_mapping = torch.tensor([0, 1], dtype=torch.int32)
+    num_accepted_tokens = torch.ones(2, dtype=torch.int32)
+    group = SimpleNamespace(layer_names=["layer"])
+
+    state.record_step({"layer": metadata}, [[group]], for_capture=False)
+    state.commit_step(
+        num_sampled,
+        idx_mapping,
+        state_indices=None,
+        num_accepted_tokens=num_accepted_tokens,
+    )
+    state.commit_step(
+        num_sampled,
+        idx_mapping,
+        state_indices=None,
+        num_accepted_tokens=num_accepted_tokens,
+    )
+
+    metadata.commit_recoverssm_state.assert_called_once_with(num_sampled)
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
+def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None:
+    state = object.__new__(MambaHybridModelState)
+    state._align_mode = True
+    state._mamba_ctx = None
+    state._mamba_state_idx_gpu = torch.full((5,), -1, dtype=torch.int32, device="cuda")
+    state.recoverssm = RecoverSSMState()
+    state.num_accepted_tokens_gpu = torch.full(
+        (5,), 9, dtype=torch.int32, device="cuda"
+    )
+    metadata = Mock(spec=RecoverSSMMetadata)
+    metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata(
+        num_spec_decodes=1,
+        request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"),
+        num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"),
+        block_size=8,
+        block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"),
+    )
+    num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda")
+    idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda")
+    group = SimpleNamespace(layer_names=["layer"])
+
+    state.recoverssm.record_step({"layer": metadata}, [[group]], for_capture=False)
+
+    state.postprocess_state(idx_mapping, num_sampled)
+
+    expected_state_indices = [-1, 1, -1, -1, -1]
+    assert state._mamba_state_idx_gpu.tolist() == expected_state_indices
+    expected_accepted = [9, 1, 9, 2, 9]
+    assert state.num_accepted_tokens_gpu.tolist() == expected_accepted

tests/v1/worker/test_mamba_utils.py

@@ -490,6 +490,42 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx():
         )
 
 
+def test_gpu_context_ignores_auxiliary_cache_tensors() -> None:
+    device = torch.device("cpu")
+    config = _TestConfig(num_layers=1)
+    layer_names = ["layer_0"]
+    kv_cache_config = _make_kv_cache_config(config, layer_names)
+    conv_state = torch.empty(
+        config.num_blocks,
+        config.conv_width,
+        config.conv_inner_dim,
+        dtype=config.dtype,
+    )
+    temporal_state = torch.empty(
+        config.num_blocks, config.temporal_state_dim, dtype=config.dtype
+    )
+    attention = MagicMock()
+    attention.kv_cache = [
+        conv_state,
+        temporal_state,
+        *(torch.empty(config.num_blocks, 1) for _ in range(4)),
+    ]
+    context = _make_gpu_ctx(config, kv_cache_config, device)
+
+    context.initialize_from_forward_context(
+        kv_cache_config,
+        {"layer_0": attention},
+        _COPY_FUNCS,
+        [torch.zeros(1, 1, dtype=torch.int32)],
+    )
+
+    assert context.is_initialized
+    assert context.state_base_addrs.tolist() == [
+        conv_state.data_ptr(),
+        temporal_state.data_ptr(),
+    ]
+
+
 def _run_gpu_postprocess(
     gpu_ctx: MambaSpecDecodeGPUContext,
     *,

vllm/config/cache.py

@@ -144,9 +144,8 @@ class CacheConfig:
       caching is enabled.
     """
     replayssm_buffer_len: int = Field(default=16, gt=0)
-    """ReplaySSM history buffer length B: with use_replayssm, standard decode
-    caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint
-    state to HBM every B steps. Default 16."""
+    """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3
+    speculative decoding 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.
@@ -154,6 +153,8 @@ class CacheConfig:
     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."""
 
     # Will be set after profiling.
     num_gpu_blocks: int | None = field(default=None, init=False)

vllm/config/vllm.py

@@ -2540,22 +2540,47 @@ class VllmConfig:
     @model_validator(mode="after")
     def validate_mamba_cached_kernel(self) -> "VllmConfig":
         if not self.cache_config.use_replayssm:
+            self.cache_config.use_kda_recoverssm = False
             return self
-        # ReplaySSM adds a 3-tensor ring to the mamba state; only models that
-        # opt in (supports_replayssm) build a consistent shape on both the layer
-        # and config paths. Reject others so the mamba page size cannot desync.
+        self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0
+
         if self.model_config is not None and not self.model_config.supports_replayssm:
             raise ValueError(
-                "--use-replayssm is only supported for Nemotron-H models "
-                f"(got architecture {self.model_config.architecture!r})"
+                "--use-replayssm is not supported for architecture "
+                f"{self.model_config.architecture!r}"
             )
-        if self.cache_config.mamba_cache_mode == "all":
+        if self.cache_config.use_kda_recoverssm:
+            if self.model_config is not None and self.model_config.architecture not in (
+                "KimiLinearForCausalLM",
+                "KimiK3ForConditionalGeneration",
+            ):
+                raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA")
+            if self.mamba_config.enable_stochastic_rounding:
+                raise ValueError(
+                    "RecoverSSM supports bfloat16/float32 "
+                    "SSM state caches, not --enable-mamba-cache-stochastic-"
+                    "rounding, which requires an explicit float16 cache"
+                )
+            if self.cache_config.mamba_cache_mode not in ("none", "align"):
+                raise ValueError(
+                    "RecoverSSM supports only none and align Mamba cache modes"
+                )
+            if (
+                self.cache_config.mamba_cache_mode == "align"
+                and not self.use_v2_model_runner
+            ):
+                raise ValueError(
+                    "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1"
+                )
+            if self.parallel_config.pipeline_parallel_size > 1:
+                raise ValueError(
+                    "RecoverSSM currently requires pipeline_parallel_size=1"
+                )
+        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.num_speculative_tokens > 0:
-            raise ValueError("--use-replayssm does not support speculative decoding")
         if self.mamba_config.backend != MambaBackendEnum.TRITON:
             raise ValueError("--use-replayssm requires --mamba-backend triton")
         if (

vllm/model_executor/layers/mamba/abstract.py

@@ -71,10 +71,12 @@ class MambaBase(AttentionLayerBase):
             page_size_padded=page_size_padded,
             mamba_type=self.mamba_type,
             mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
+            # RecoverSSM verifies the whole window off one checkpoint, so it
+            # never writes the baseline's per-draft-token state slots.
             num_speculative_blocks=(
-                vllm_config.speculative_config.num_speculative_tokens
-                if vllm_config.speculative_config
-                else 0
+                0
+                if vllm_config.cache_config.use_kda_recoverssm
+                else vllm_config.num_speculative_tokens
             ),
         )
 

vllm/model_executor/layers/mamba/mamba_utils.py

@@ -136,6 +136,15 @@ class MambaStateDtypeCalculator:
         state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
         return (state_dtype, torch.float32)
 
+    @classmethod
+    def append_kda_recoverssm_record(
+        cls,
+        base_dtypes: tuple[torch.dtype, ...],
+        model_dtype: ModelDType | torch.dtype,
+    ) -> tuple[torch.dtype, ...]:
+        activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype)
+        return (*base_dtypes, torch.float32, activation_dtype)
+
 
 class MambaStateShapeCalculator:
     @classmethod
@@ -293,6 +302,27 @@ class MambaStateShapeCalculator:
         recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
         return (conv_state_shape, recurrent_state_shape)
 
+    @classmethod
+    def append_kda_recoverssm_record(
+        cls,
+        base_shapes: tuple[tuple[int, int], tuple[int, int, int]],
+        num_heads: int,
+        head_dim: int,
+        tp_world_size: int,
+        spec_query_len: int,
+    ) -> tuple[
+        tuple[int, int],
+        tuple[int, int, int],
+        tuple[int, int, int],
+        tuple[int, int, int],
+    ]:
+        local_num_heads = divide(num_heads, tp_world_size)
+        return (
+            *base_shapes,
+            (local_num_heads, spec_query_len, head_dim),
+            (local_num_heads, spec_query_len, 2 * head_dim),
+        )
+
 
 @dataclass
 class MambaCopySpec:

vllm/model_executor/models/interfaces.py

@@ -75,6 +75,12 @@ MambaStateShapes: TypeAlias = (
     | tuple[tuple[int, int, int]]
     | tuple[tuple[int, int], tuple[int, int]]
     | tuple[tuple[int, int], tuple[int, int, int]]
+    | tuple[
+        tuple[int, int],
+        tuple[int, int, int],
+        tuple[int, int, int],
+        tuple[int, int, int],
+    ]
 )
 
 
@@ -1099,8 +1105,8 @@ def supports_mamba_prefix_caching(
 
 @runtime_checkable
 class SupportsReplaySSM(Protocol):
-    """The interface for models whose Mamba2 layers support ReplaySSM cached
-    standard decode.
+    """The interface for models whose recurrent layers support ReplaySSM
+    cached decode.
 
     This is currently experimental.
     """

vllm/models/kimi_k3/nvidia/kda.py

@@ -282,23 +282,37 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
 
     def get_state_dtype(
         self,
-    ) -> tuple[torch.dtype, torch.dtype]:
+    ) -> tuple[torch.dtype, ...]:
         if self.model_config is None or self.cache_config is None:
             raise ValueError("model_config and cache_config must be set")
-        return MambaStateDtypeCalculator.kda_state_dtype(
+        base_dtypes = MambaStateDtypeCalculator.kda_state_dtype(
             self.model_config.dtype, self.cache_config.mamba_cache_dtype
         )
+        if self.cache_config.use_kda_recoverssm:
+            return MambaStateDtypeCalculator.append_kda_recoverssm_record(
+                base_dtypes, self.model_config.dtype
+            )
+        return base_dtypes
 
     def get_state_shape(
         self,
-    ) -> tuple[tuple[int, ...], tuple[int, ...]]:
-        return MambaStateShapeCalculator.kda_state_shape(
+    ) -> tuple[tuple[int, ...], ...]:
+        base_shapes = MambaStateShapeCalculator.kda_state_shape(
             self.tp_size,
             self.num_heads,
             self.head_dim,
             conv_kernel_size=self.conv_size,
             num_spec=self.num_spec,
         )
+        if self.cache_config.use_kda_recoverssm:
+            return MambaStateShapeCalculator.append_kda_recoverssm_record(
+                base_shapes,
+                self.num_heads,
+                self.head_dim,
+                tp_world_size=self.tp_size,
+                spec_query_len=1 + self.num_spec,
+            )
+        return base_shapes
 
     def __init__(
         self,
@@ -308,6 +322,12 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         run_gemm_rs: bool = False,
     ) -> None:
         super().__init__(config, vllm_config, prefix)
+        self.use_recoverssm = self.cache_config.use_kda_recoverssm
+        if self.cache_config.use_replayssm and not self.use_recoverssm:
+            raise ValueError(
+                "Kimi-K3 supports --use-replayssm only with speculative decoding"
+            )
+        self.spec_query_len = 1 + self.num_spec
 
         kda_config = config.linear_attn_config  # type: ignore[attr-defined]
         assert kda_config is not None, "linear_attn_config must be set"
@@ -371,7 +391,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
         # Keep a width-major copy for fused decode without changing the layout
         # consumed by the prefill and fallback decode kernels.
-        conv_state_dtype, _ = self.get_state_dtype()
+        conv_state_dtype = self.get_state_dtype()[0]
         decode_conv1d_weight = None
         if is_fused_kda_decode_supported(
             self.local_num_heads,
@@ -567,7 +587,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
         g1 = g1[:, :num_actual_tokens]
         beta = beta[:, :num_actual_tokens]
 
-        conv_state, recurrent_state = self.kv_cache
+        conv_state, recurrent_state, *recoverssm_records = self.kv_cache
         # The convolution kernels consume (..., dim, width - 1).
         if not is_conv_state_dim_first():
             conv_state = conv_state.transpose(-1, -2)
@@ -634,7 +654,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
             assert spec_state_indices_tensor is not None
             assert spec_query_start_loc is not None
             spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
-            spec_max_query_len = spec_state_indices_tensor.size(-1)
+            spec_max_query_len = (
+                self.spec_query_len
+                if self.use_recoverssm
+                else spec_state_indices_tensor.size(-1)
+            )
             spec_conv_out = torch.empty_like(mixed_qkv_spec)
             mixed_qkv_spec = causal_conv1d_update(
                 mixed_qkv_spec,
@@ -659,21 +683,48 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
                 if m.num_prefills == 0 and m.num_decodes == 0
                 else None
             )
-            core_attn_out_spec, _ = fused_recurrent_kda(
-                q=q_spec,
-                k=k_spec,
-                v=v_spec,
-                raw_g=g1_spec,
-                raw_beta=beta_spec,
-                A_log=self.A_log,
-                dt_bias=self.dt_bias,
-                lower_bound=self.gate_lower_bound,
-                initial_state=recurrent_state,
-                cu_seqlens=spec_cu_seqlens,
-                ssm_state_indices=spec_state_indices_tensor,
-                num_accepted_tokens=num_accepted_tokens,
-                out=spec_out,
-            )
+            if self.use_recoverssm:
+                from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
+                    kda_recoverssm_verify,
+                )
+
+                if len(recoverssm_records) != 2:
+                    raise ValueError(
+                        "KDA RecoverSSM requires correction and key/gate buffers"
+                    )
+                core_attn_out_spec = kda_recoverssm_verify(
+                    q=q_spec,
+                    k=k_spec,
+                    v=v_spec,
+                    raw_g=g1_spec,
+                    raw_beta=beta_spec,
+                    A_log=self.A_log,
+                    dt_bias=self.dt_bias,
+                    lower_bound=self.gate_lower_bound,
+                    checkpoint_state=recurrent_state,
+                    correction_cache=recoverssm_records[0],
+                    kg_cache=recoverssm_records[1],
+                    query_start_loc=spec_cu_seqlens,
+                    state_indices=spec_state_indices_tensor[: m.num_spec_decodes, 0],
+                    spec_query_len=self.spec_query_len,
+                    out=spec_out,
+                )
+            else:
+                core_attn_out_spec, _ = fused_recurrent_kda(
+                    q=q_spec,
+                    k=k_spec,
+                    v=v_spec,
+                    raw_g=g1_spec,
+                    raw_beta=beta_spec,
+                    A_log=self.A_log,
+                    dt_bias=self.dt_bias,
+                    lower_bound=self.gate_lower_bound,
+                    initial_state=recurrent_state,
+                    cu_seqlens=spec_cu_seqlens,
+                    ssm_state_indices=spec_state_indices_tensor,
+                    num_accepted_tokens=num_accepted_tokens,
+                    out=spec_out,
+                )
 
         # Prefill or plain-decode path.
         core_attn_out_non_spec = None

vllm/models/kimi_k3/nvidia/kda_metadata.py

@@ -6,13 +6,18 @@ The request classification and cudagraph staging intentionally mirror
 ``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its
 prefill KDA kernel internally, so this builder omits the shared FLA chunk
 metadata construction.
+
+For Kimi-K3 speculative decoding, ``--use-replayssm`` selects the simplified
+RecoverSSM path implemented here instead of the Mamba2 ReplaySSM kernel.
 """
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from functools import cache
+from typing import TYPE_CHECKING
 
 import torch
 
+from vllm.config import VllmConfig
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import async_tensor_h2d
@@ -22,6 +27,10 @@ from vllm.v1.attention.backends.gdn_attn import (
     GDNAttentionMetadata,
     GDNAttentionMetadataBuilder,
 )
+from vllm.v1.attention.backends.recoverssm_metadata import (
+    RecoverSSMMetadata,
+    RecoverSSMPostprocessMetadata,
+)
 from vllm.v1.attention.backends.utils import (
     NULL_BLOCK_ID,
     compute_causal_conv1d_metadata,
@@ -29,6 +38,11 @@ from vllm.v1.attention.backends.utils import (
 )
 from vllm.v1.kv_cache_interface import MambaSpec
 
+if TYPE_CHECKING:
+    from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
+        KDARecoverSSMCommitContext,
+    )
+
 
 @cache
 def _metadata_launch_pdl() -> bool:
@@ -229,11 +243,103 @@ def stage_spec_decode_metadata(
 
 
 @dataclass
-class KimiK3KDAMetadata(GDNAttentionMetadata):
-    pass
+class KDARecoverSSMAlignMetadata:
+    block_table: torch.Tensor
+    num_computed_tokens: torch.Tensor
+    block_size: int
+
+
+@dataclass
+class KDARecoverSSMCommitMetadata:
+    state_indices: torch.Tensor
+    query_start_loc: torch.Tensor
+    request_indices: torch.Tensor | None
+    align: KDARecoverSSMAlignMetadata | None
+
+
+@dataclass
+class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
+    recoverssm_commit: KDARecoverSSMCommitMetadata | None = None
+    recoverssm_context: "KDARecoverSSMCommitContext | None" = field(
+        default=None, repr=False, compare=False
+    )
+
+    def commit_recoverssm_state(
+        self, num_accepted_tokens: torch.Tensor
+    ) -> RecoverSSMPostprocessMetadata | None:
+        commit = self.recoverssm_commit
+        if commit is None:
+            return None
+        context = self.recoverssm_context
+        assert context is not None
+        align = commit.align
+        context.commit(
+            num_accepted_tokens,
+            commit.state_indices[: self.num_spec_decodes, 0],
+            commit.query_start_loc[: self.num_spec_decodes + 1],
+            request_indices=commit.request_indices,
+            block_table=align.block_table if align is not None else None,
+            num_computed_tokens=(
+                align.num_computed_tokens if align is not None else None
+            ),
+            mamba_block_size=align.block_size if align is not None else None,
+        )
+        if align is None:
+            return None
+        return RecoverSSMPostprocessMetadata(
+            num_spec_decodes=self.num_spec_decodes,
+            request_indices=commit.request_indices,
+            block_table=align.block_table,
+            num_computed_tokens=align.num_computed_tokens,
+            block_size=align.block_size,
+        )
 
 
 class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
+    def __init__(
+        self,
+        kv_cache_spec: MambaSpec,
+        layer_names: list[str],
+        vllm_config: VllmConfig,
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm
+        self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1
+        self.recoverssm_num_accepted_tokens: torch.Tensor | None = None
+        self.recoverssm_context: KDARecoverSSMCommitContext | None = None
+        if self.use_recoverssm:
+            max_num_reqs = vllm_config.scheduler_config.max_num_seqs
+            self.spec_state_indices_tensor = torch.empty(
+                (max_num_reqs, 1),
+                dtype=torch.int32,
+                device=device,
+            )
+            self.recoverssm_num_accepted_tokens = torch.ones(
+                max_num_reqs,
+                dtype=torch.int32,
+                device=device,
+            )
+
+    def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext":
+        context = self.recoverssm_context
+        if context is not None:
+            return context
+
+        from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
+            KDARecoverSSMCommitContext,
+        )
+
+        forward_context = self.vllm_config.compilation_config.static_forward_context
+        layers = [forward_context[layer_name] for layer_name in self.layer_names]
+        context = KDARecoverSSMCommitContext.create(
+            layers,
+            spec_query_len=1 + self.vllm_config.num_speculative_tokens,
+            max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs,
+        )
+        self.recoverssm_context = context
+        return context
+
     def build(  # type: ignore[override]
         self,
         common_prefix_len: int,
@@ -263,14 +369,26 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
             num_spec_decodes = 0
         else:
             spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
-            # A nonnegative entry identifies a spec request. If no draft token
-            # was scheduled, process the whole batch as non-spec instead.
-            if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0:
+            if self.use_recoverssm:
+                assert m.is_prefilling is not None
+                assert m.is_prefilling.device.type == "cpu"
+                active_decode_mask_cpu = (~m.is_prefilling) & (
+                    query_start_loc_cpu.diff() > 0
+                )
+                spec_sequence_masks_cpu |= active_decode_mask_cpu
+            # Native KDA can use its regular decode path when no draft token
+            # was scheduled. RecoverSSM must preserve its extended conv window.
+            if (
+                not self.use_recoverssm
+                and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
+                == 0
+            ):
                 spec_sequence_masks_cpu = None
                 num_spec_decodes = 0
             else:
                 num_spec_decodes = spec_sequence_masks_cpu.sum().item()
 
+        spec_request_indices = None
         if num_spec_decodes == 0:
             # The runner orders ordinary decodes before prefills.
             num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
@@ -289,6 +407,16 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
             assert spec_sequence_masks_cpu is not None
             assert num_accepted_tokens is not None
             query_lens_cpu = query_start_loc_cpu.diff()
+            if (
+                self.use_recoverssm
+                and torch.any(
+                    query_lens_cpu[spec_sequence_masks_cpu] > self.num_spec + 1
+                ).item()
+            ):
+                raise ValueError(
+                    "KDA RecoverSSM speculative decode query length exceeds "
+                    f"its activation capacity ({self.num_spec + 1})"
+                )
             num_query_tokens = query_start_loc_cpu[-1].item()
 
             # Exclude zero-length cudagraph padding from request-indexed
@@ -325,7 +453,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 non_spec_token_indx = None
                 # Real requests precede trailing cudagraph padding.
                 spec_state_indices_tensor = block_table_tensor[
-                    :num_spec_decodes, : self.num_spec + 1
+                    :num_spec_decodes, : self.spec_state_slots
                 ]
                 non_spec_state_indices_tensor = None
                 # Padding trails real requests, so this prefix already contains
@@ -339,6 +467,12 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 spec_sequence_masks_gpu = async_tensor_h2d(
                     spec_sequence_masks_cpu, device=query_start_loc.device
                 )
+                if self.use_recoverssm:
+                    spec_request_indices = async_tensor_h2d(
+                        spec_sequence_masks_cpu.nonzero(as_tuple=True)[0],
+                        dtype=torch.int32,
+                        device=query_start_loc.device,
+                    )
                 spec_token_masks = torch.repeat_interleave(
                     spec_sequence_masks_gpu,
                     query_lens,
@@ -351,10 +485,10 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 non_spec_token_indx = index[:num_non_spec_tokens]
                 spec_token_indx = index[num_non_spec_tokens:]
 
-                # Spec requests carry one state slot per speculative step;
-                # non-spec requests use only their current state slot.
+                # Native spec uses one state slot per step. RecoverSSM keeps
+                # only the current checkpoint slot.
                 spec_state_indices_tensor = block_table_tensor[
-                    spec_sequence_masks_cpu, : self.num_spec + 1
+                    spec_sequence_masks_cpu, : self.spec_state_slots
                 ]
                 non_spec_state_indices_tensor = block_table_tensor[
                     active_non_spec_mask_cpu, 0
@@ -400,12 +534,18 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
 
                 num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
 
+            if self.use_recoverssm:
+                assert self.recoverssm_num_accepted_tokens is not None
+                num_accepted_tokens = self.recoverssm_num_accepted_tokens[
+                    :num_spec_decodes
+                ]
+
         # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares
         # its own chunk indices. Only causal-convolution metadata is needed here.
         nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
         if num_prefills > 0:
             has_initial_state = m.compute_num_computed_tokens() > 0
-            if spec_sequence_masks_cpu is not None:
+            if num_spec_decodes > 0:
                 has_initial_state = has_initial_state[active_non_spec_mask_cpu]
                 assert non_spec_query_start_loc_cpu is not None
             nums_dict, batch_ptr, token_chunk_offset_ptr = (
@@ -426,7 +566,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
             and num_spec_decodes > 0
             and num_prefills == 0
             and num_decodes == 0
-            and num_spec_decodes <= self.decode_cudagraph_max_bs
+            and batch_size <= self.spec_state_indices_tensor.shape[0]
             and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
         ):
             # Equivalent PyTorch staging:
@@ -463,6 +603,24 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
                 :batch_size
             ]
 
+        recoverssm_commit = None
+        if self.use_recoverssm and num_spec_decodes > 0:
+            assert spec_state_indices_tensor is not None
+            assert spec_query_start_loc is not None
+            align = None
+            if self.kv_cache_spec.mamba_cache_mode == "align":
+                align = KDARecoverSSMAlignMetadata(
+                    block_table=m.block_table_tensor,
+                    num_computed_tokens=m.compute_num_computed_tokens(),
+                    block_size=self.kv_cache_spec.block_size,
+                )
+            recoverssm_commit = KDARecoverSSMCommitMetadata(
+                state_indices=spec_state_indices_tensor,
+                query_start_loc=spec_query_start_loc,
+                request_indices=spec_request_indices,
+                align=align,
+            )
+
         return KimiK3KDAMetadata(
             num_prefills=num_prefills,
             num_prefill_tokens=num_prefill_tokens,
@@ -480,6 +638,12 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
             spec_token_indx=spec_token_indx,
             non_spec_token_indx=non_spec_token_indx,
             num_accepted_tokens=num_accepted_tokens,
+            recoverssm_commit=recoverssm_commit,
+            recoverssm_context=(
+                self._get_recoverssm_context()
+                if recoverssm_commit is not None
+                else None
+            ),
             nums_dict=nums_dict,
             batch_ptr=batch_ptr,
             token_chunk_offset_ptr=token_chunk_offset_ptr,

vllm/models/kimi_k3/nvidia/model.py

@@ -62,12 +62,14 @@ from vllm.model_executor.models.interfaces import (
     EagleModelMixin,
     HasInnerState,
     IsHybrid,
+    MambaStateShapes,
     MixtureOfExperts,
     SupportsEagle3,
     SupportsEncoderCudaGraph,
     SupportsMultiModal,
     SupportsPP,
     SupportsQuant,
+    SupportsReplaySSM,
 )
 from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs
 from vllm.model_executor.models.kimi_k25_vit import (
@@ -1559,7 +1561,13 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
 
 
 class KimiLinearForCausalLM(
-    nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid, SupportsEagle3
+    nn.Module,
+    HasInnerState,
+    SupportsPP,
+    MixtureOfExperts,
+    IsHybrid,
+    SupportsEagle3,
+    SupportsReplaySSM,
 ):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -1613,15 +1621,20 @@ class KimiLinearForCausalLM(
     def get_mamba_state_dtype_from_config(
         cls,
         vllm_config: "VllmConfig",
-    ) -> tuple[torch.dtype, torch.dtype]:
-        return MambaStateDtypeCalculator.kda_state_dtype(
+    ) -> tuple[torch.dtype, ...]:
+        dtypes = MambaStateDtypeCalculator.kda_state_dtype(
             vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
         )
+        if vllm_config.cache_config.use_kda_recoverssm:
+            dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record(
+                dtypes, vllm_config.model_config.dtype
+            )
+        return dtypes
 
     @classmethod
     def get_mamba_state_shape_from_config(
         cls, vllm_config: "VllmConfig"
-    ) -> tuple[tuple[int, int], tuple[int, int, int]]:
+    ) -> MambaStateShapes:
         parallel_config = vllm_config.parallel_config
         hf_config = vllm_config.model_config.hf_config
         tp_size = parallel_config.tensor_parallel_size
@@ -1630,13 +1643,22 @@ class KimiLinearForCausalLM(
             if vllm_config.speculative_config
             else 0
         )
-        return MambaStateShapeCalculator.kda_state_shape(
+        shapes = MambaStateShapeCalculator.kda_state_shape(
             tp_size,
             hf_config.linear_attn_config["num_heads"],
             hf_config.linear_attn_config["head_dim"],
             conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
             num_spec=num_spec,
         )
+        if vllm_config.cache_config.use_kda_recoverssm:
+            return MambaStateShapeCalculator.append_kda_recoverssm_record(
+                shapes,
+                hf_config.linear_attn_config["num_heads"],
+                hf_config.linear_attn_config["head_dim"],
+                tp_world_size=tp_size,
+                spec_query_len=1 + num_spec,
+            )
+        return shapes
 
     @classmethod
     def get_mamba_state_copy_func(
@@ -1697,6 +1719,7 @@ class KimiK3ForConditionalGeneration(
     SupportsEagle3,
     HasInnerState,
     IsHybrid,
+    SupportsReplaySSM,
 ):
     """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text."""
 

vllm/v1/attention/backends/recoverssm_metadata.py

@@ -0,0 +1,26 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import abc
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class RecoverSSMPostprocessMetadata:
+    """Metadata used during postprocessing for align-mode prefix caching."""
+
+    num_spec_decodes: int
+    request_indices: torch.Tensor | None
+    block_table: torch.Tensor
+    num_computed_tokens: torch.Tensor
+    block_size: int
+
+
+class RecoverSSMMetadata(abc.ABC):
+    @abc.abstractmethod
+    def commit_recoverssm_state(
+        self, num_accepted_tokens: torch.Tensor
+    ) -> RecoverSSMPostprocessMetadata | None:
+        raise NotImplementedError

vllm/v1/worker/gpu/model_states/mamba_hybrid.py

@@ -20,6 +20,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
+from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
 from vllm.v1.worker.mamba_utils import (
     MambaSpecDecodeGPUContext,
     preprocess_mamba_align_fused_kernel,
@@ -80,6 +81,9 @@ class MambaHybridModelState(DefaultModelState):
         # kernel reusing the postprocess copy machinery, so the per-step src
         # columns and the running state_idx are kept GPU-resident.
         self._align_mode = self.cache_config.mamba_cache_mode == "align"
+        self.recoverssm = (
+            RecoverSSMState() if self.cache_config.use_kda_recoverssm else None
+        )
         if self._align_mode:
             self._mamba_state_idx_gpu = torch.zeros(
                 self.max_num_reqs, dtype=torch.int32, device=self.device
@@ -267,7 +271,7 @@ class MambaHybridModelState(DefaultModelState):
             num_accepted_tokens=num_accepted_tokens,
             num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
         )
-        return build_attn_metadata(
+        attn_metadata = build_attn_metadata(
             attn_groups=attn_groups,
             num_reqs=num_reqs,
             num_tokens=num_tokens,
@@ -285,6 +289,13 @@ class MambaHybridModelState(DefaultModelState):
             for_cudagraph_capture=for_capture,
             rswa_prefix_lens=input_batch.prompt_lens,
         )
+        if self.recoverssm is not None:
+            self.recoverssm.record_step(
+                attn_metadata,
+                attn_groups,
+                for_capture=for_capture,
+            )
+        return attn_metadata
 
     def postprocess_state(
         self,
@@ -295,21 +306,34 @@ class MambaHybridModelState(DefaultModelState):
         # Chunked prefill does not sample a token, so num_sampled can be 0.
         # Mamba treats num_accepted_tokens=1 as the neutral non-spec value.
         num_reqs = idx_mapping.shape[0]
+        if num_reqs:
+            if not isinstance(num_sampled, int):
+                # idx_mapping may contain -1 sentinels (filtered rows) under PP; the
+                # kernel skips them rather than scattering with a host-side gather.
+                _scatter_num_accepted_kernel[(num_reqs,)](
+                    idx_mapping,
+                    num_sampled,
+                    self.num_accepted_tokens_gpu,
+                )
+            else:
+                # Fill with single value.
+                _fill_num_accepted_kernel[(num_reqs,)](
+                    idx_mapping,
+                    self.num_accepted_tokens_gpu,
+                    max(num_sampled, 1),
+                )
+
+        if self.recoverssm is not None:
+            self.recoverssm.commit_step(
+                num_sampled,
+                idx_mapping,
+                state_indices=(self._mamba_state_idx_gpu if self._align_mode else None),
+                num_accepted_tokens=self.num_accepted_tokens_gpu,
+            )
+
         if not num_reqs:
             return
 
-        if not isinstance(num_sampled, int):
-            # idx_mapping may contain -1 sentinels (filtered rows) under PP; the
-            # kernel skips them rather than scattering with a host-side gather.
-            _scatter_num_accepted_kernel[(num_reqs,)](
-                idx_mapping, num_sampled, self.num_accepted_tokens_gpu
-            )
-        else:
-            # Fill with single value.
-            _fill_num_accepted_kernel[(num_reqs,)](
-                idx_mapping, self.num_accepted_tokens_gpu, max(num_sampled, 1)
-            )
-
         # Align: save the running state to the block-aligned position when
         # spec-decode acceptance leaves the sequence non-block-aligned (mirrors
         # the V1 align postprocess). num_computed_tokens already holds the

vllm/v1/worker/gpu/model_states/recoverssm.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+
+from vllm.triton_utils import tl, triton
+from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata
+from vllm.v1.worker.utils import AttentionGroup
+
+
+class RecoverSSMState:
+    """Coordinates RecoverSSM metadata between attention and postprocessing."""
+
+    def __init__(self) -> None:
+        self._step: tuple[RecoverSSMMetadata, ...] | None = None
+
+    def record_step(
+        self,
+        attn_metadata: dict[str, Any],
+        attn_groups: list[list[AttentionGroup]],
+        *,
+        for_capture: bool,
+    ) -> None:
+        if for_capture:
+            self._step = None
+            return
+
+        step: list[RecoverSSMMetadata] = []
+        for group_list in attn_groups:
+            for group in group_list:
+                metadata = attn_metadata[group.layer_names[0]]
+                if isinstance(metadata, RecoverSSMMetadata):
+                    step.append(metadata)
+        self._step = tuple(step)
+
+    def commit_step(
+        self,
+        num_sampled: torch.Tensor | int,
+        idx_mapping: torch.Tensor,
+        *,
+        state_indices: torch.Tensor | None,
+        num_accepted_tokens: torch.Tensor,
+    ) -> None:
+        step = self._step
+        self._step = None
+        if isinstance(num_sampled, int) or step is None:
+            return
+
+        for metadata in step:
+            postprocess_meta = metadata.commit_recoverssm_state(num_sampled)
+            if postprocess_meta is None:
+                continue
+            assert state_indices is not None
+            # RecoverSSM already restored the accepted state. Update its running
+            # column and reset the next-step copy bias to the neutral value.
+            _postprocess_recoverssm_align_kernel[(postprocess_meta.num_spec_decodes,)](
+                idx_mapping,
+                num_sampled,
+                postprocess_meta.request_indices,
+                postprocess_meta.num_computed_tokens,
+                state_indices,
+                num_accepted_tokens,
+                MAMBA_BLOCK_SIZE=postprocess_meta.block_size,
+                BLOCK_TABLE_WIDTH=postprocess_meta.block_table.shape[1],
+            )
+
+
+@triton.heuristics(
+    {"HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None}
+)
+@triton.jit
+def _postprocess_recoverssm_align_kernel(
+    idx_mapping_ptr,
+    num_sampled_ptr,
+    request_indices_ptr,
+    num_computed_ptr,
+    state_idx_ptr,
+    num_accepted_ptr,
+    HAS_REQUEST_INDICES: tl.constexpr,
+    MAMBA_BLOCK_SIZE: tl.constexpr,
+    BLOCK_TABLE_WIDTH: tl.constexpr,
+):
+    spec_idx = tl.program_id(0)
+    batch_idx = spec_idx
+    if HAS_REQUEST_INDICES:
+        batch_idx = tl.load(request_indices_ptr + spec_idx)
+    req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
+    if req_state_idx < 0:
+        return
+    num_sampled = tl.load(num_sampled_ptr + batch_idx)
+    num_computed = tl.load(num_computed_ptr + batch_idx)
+    tl.store(
+        state_idx_ptr + req_state_idx,
+        tl.minimum(
+            (num_computed + num_sampled) // MAMBA_BLOCK_SIZE,
+            BLOCK_TABLE_WIDTH - 1,
+        ),
+    )
+    tl.store(num_accepted_ptr + req_state_idx, 1)

vllm/v1/worker/mamba_utils.py

@@ -828,7 +828,13 @@ class MambaSpecDecodeGPUContext:
                 attention = forward_context[layer_name]
                 kv_caches: list[torch.Tensor] = attention.kv_cache
 
-                for state_type_idx, state in enumerate(kv_caches):
+                if len(kv_caches) < self.num_state_types:
+                    raise ValueError(
+                        f"Expected at least {self.num_state_types} Mamba state "
+                        f"tensors, got {len(kv_caches)}"
+                    )
+                for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs):
+                    state = kv_caches[state_type_idx]
                     # Base address
                     self.state_base_addrs[idx] = state.data_ptr()
 
@@ -845,7 +851,6 @@ class MambaSpecDecodeGPUContext:
                     # Element size
                     self.state_elem_sizes[idx] = state.element_size()
 
-                    copy_func = mamba_state_copy_funcs[state_type_idx]
                     assert (
                         copy_func is get_conv_copy_spec
                         or copy_func is get_temporal_copy_spec