Enable FlashInfer checkpointing SSU STP no-fallback path

tests/v1/worker/test_mamba_utils.py

@@ -2,6 +2,13 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from unittest.mock import MagicMock, patch
 
+import torch
+
+from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFuncCalculator,
+    get_stable_temporal_copy_spec,
+    get_temporal_copy_spec,
+)
 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.worker.mamba_utils import preprocess_mamba
 
@@ -67,3 +74,29 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx():
         )
 
     assert mamba_state_idx == {"keep": 99}
+
+
+def test_checkpointing_temporal_copy_uses_stable_slot():
+    state = torch.arange(10 * 2, dtype=torch.float32).view(10, 2)
+    block_ids = [3, 4, 5, 6]
+
+    old_fi_spec = get_temporal_copy_spec(
+        state, block_ids, cur_block_idx=1, num_accepted_tokens=3
+    )
+    ckpt_spec = get_stable_temporal_copy_spec(
+        state, block_ids, cur_block_idx=1, num_accepted_tokens=3
+    )
+
+    assert old_fi_spec.start_addr == state[block_ids[3]].data_ptr()
+    assert ckpt_spec.start_addr == state[block_ids[1]].data_ptr()
+    assert old_fi_spec.num_elements == ckpt_spec.num_elements == 2
+
+
+def test_mamba2_checkpointing_copy_funcs_keep_ssm_replay_state_stable():
+    funcs = MambaStateCopyFuncCalculator.mamba2_checkpointing_state_copy_func()
+
+    assert len(funcs) == 8
+    assert funcs[1:] == (get_stable_temporal_copy_spec,) * 7
+    assert MambaStateCopyFuncCalculator.mamba2_state_copy_func()[1] is (
+        get_temporal_copy_spec
+    )

tests/v1/worker/test_utils.py

@@ -1,9 +1,15 @@
 # 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.v1.worker.utils import bind_kv_cache
+from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
+from vllm.v1.kv_cache_interface import MambaSpec
+from vllm.v1.worker.utils import AttentionGroup, KVBlockZeroer, bind_kv_cache
 
 
 def test_bind_kv_cache(default_vllm_config):
@@ -90,3 +96,94 @@ def test_bind_kv_cache_draft_model(default_vllm_config):
     assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"]
     assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"]
     assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"]
+
+
+def test_kv_block_zeroer_zeros_mamba_cache_blocks():
+    num_blocks = 4
+    state_storage = torch.full((num_blocks * 16,), 1.0)
+    state = torch.as_strided(
+        state_storage,
+        size=(num_blocks, 2, 3),
+        stride=(16, 3, 1),
+    )
+    tracker = torch.full((num_blocks,), 9, dtype=torch.int32)
+    spec = MambaSpec(
+        block_size=1,
+        shapes=((2, 3), ()),
+        dtypes=(torch.float32, torch.int32),
+    )
+    group = AttentionGroup(
+        backend=Mock(),
+        layer_names=["layer.0"],
+        kv_cache_spec=spec,
+        kv_cache_group_id=0,
+    )
+    zeroer = KVBlockZeroer(torch.device("cpu"), pin_memory=False)
+    zeroer.init_meta(
+        [group],
+        [1],
+        "auto",
+        set(),
+        {"layer.0": SimpleNamespace(kv_cache=(state, tracker))},
+    )
+
+    zeroer.zero_block_ids([1, 3])
+
+    assert torch.all(state[1] == 0)
+    assert torch.all(state[3] == 0)
+    assert torch.all(tracker[torch.tensor([1, 3])] == 0)
+    assert torch.all(state[0] == 1)
+    assert torch.all(state[2] == 1)
+    assert torch.all(tracker[torch.tensor([0, 2])] == 9)
+    assert torch.all(state_storage[22:32] == 1)
+
+
+@pytest.mark.parametrize("checkpoint_interval", [1, 6])
+def test_mamba2_state_shape_honors_checkpoint_interval(checkpoint_interval):
+    shapes = MambaStateShapeCalculator.mamba2_state_shape(
+        tp_world_size=1,
+        intermediate_size=8192,
+        n_groups=8,
+        num_heads=128,
+        head_dim=64,
+        state_size=128,
+        conv_kernel=4,
+        checkpoint_interval=checkpoint_interval,
+    )
+
+    old_x_shape = shapes[2]
+    old_B_shape = shapes[3]
+    old_dt_shape = shapes[4]
+    old_cumAdt_shape = shapes[5]
+
+    assert old_x_shape[0] == checkpoint_interval
+    assert old_B_shape[1] == checkpoint_interval
+    assert old_dt_shape[-1] == checkpoint_interval
+    assert old_cumAdt_shape[-1] == checkpoint_interval
+
+
+@pytest.mark.parametrize("checkpoint_interval", [1, 6])
+def test_nemotron_h_state_shape_honors_checkpoint_interval(checkpoint_interval):
+    from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+    vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        model_config=SimpleNamespace(
+            hf_config=SimpleNamespace(
+                mamba_num_heads=128,
+                mamba_head_dim=64,
+                n_groups=8,
+                ssm_state_size=128,
+                conv_kernel=4,
+            ),
+        ),
+        num_speculative_tokens=0,
+        mamba_config=SimpleNamespace(checkpoint_interval=checkpoint_interval),
+    )
+
+    shapes = NemotronHForCausalLM.get_mamba_state_shape_from_config(vllm_config)
+
+    assert shapes[2][0] == checkpoint_interval
+    assert shapes[3][1] == checkpoint_interval
+    assert shapes[4][-1] == checkpoint_interval
+    assert shapes[5][-1] == checkpoint_interval

vllm/config/mamba.py

@@ -45,6 +45,10 @@ class MambaConfig:
     generation. 0 uses the Triton default. Higher values improve randomness
     quality at the cost of compute."""
 
+    checkpoint_interval: int = 1
+    """Number of decode tokens to keep in the FlashInfer checkpointing SSU
+    replay window before materializing a new SSM state checkpoint."""
+
     @field_validator("backend", mode="before")
     @classmethod
     def validate_backend_before(cls, value: Any) -> Any:
@@ -53,6 +57,15 @@ def validate_backend_before(cls, value: Any) -> Any:
             return MambaBackendEnum[value.upper()]
         return value
 
+    @field_validator("checkpoint_interval")
+    @classmethod
+    def validate_checkpoint_interval(cls, value: int) -> int:
+        if value < 1:
+            raise ValueError("Mamba checkpoint interval must be >= 1")
+        if value > 16:
+            raise ValueError("FlashInfer checkpointing SSU supports interval <= 16")
+        return value
+
     def __post_init__(self):
         if self.enable_stochastic_rounding:
             from vllm.platforms import current_platform

vllm/engine/arg_utils.py

@@ -676,6 +676,7 @@ class EngineArgs:
         MambaConfig.enable_stochastic_rounding
     )
     mamba_cache_philox_rounds: int = MambaConfig.stochastic_rounding_philox_rounds
+    mamba_checkpoint_interval: int = MambaConfig.checkpoint_interval
 
     additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config")
 
@@ -907,6 +908,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
             "--mamba-cache-philox-rounds",
             **mamba_kwargs["stochastic_rounding_philox_rounds"],
         )
+        mamba_group.add_argument(
+            "--mamba-checkpoint-interval",
+            **mamba_kwargs["checkpoint_interval"],
+        )
 
         # Structured outputs arguments
         structured_outputs_kwargs = get_kwargs(StructuredOutputsConfig)
@@ -2106,6 +2111,8 @@ def create_engine_config(
             mamba_config.stochastic_rounding_philox_rounds = (
                 self.mamba_cache_philox_rounds
             )
+        if self.mamba_checkpoint_interval:
+            mamba_config.checkpoint_interval = self.mamba_checkpoint_interval
 
         # Kernel config overrides
         kernel_config = copy.deepcopy(self.kernel_config)