[Mamba] Unify FI ReplaySSM STP/MTP/Prefix Caching Lifecycle

.buildkite/test_areas/engine.yaml

@@ -95,6 +95,38 @@ steps:
       - tests/v1/e2e/general/
       - vllm/platforms/rocm.py
 
+- label: ":nvidia: (H100) ReplaySSM E2E"
+  device: h100
+  num_devices: 2
+  key: replayssm-e2e
+  timeout_in_minutes: 90
+  working_dir: "/vllm-workspace/tests"
+  source_file_dependencies:
+    - tests/v1/e2e/test_replayssm_decode.py
+    - tests/kernels/mamba/test_ssu_dispatch.py
+    - tests/model_executor/test_replayssm_warmup.py
+    - tests/v1/worker/test_kv_cache_allocation_scope.py
+    - vllm/v1/worker/utils.py
+    - vllm/v1/core/kv_cache_coordinator.py
+    - vllm/v1/core/single_type_kv_cache_manager.py
+    - vllm/config/cache.py
+    - vllm/config/vllm.py
+    - vllm/model_executor/layers/mamba/
+    - vllm/model_executor/models/nemotron_h.py
+    - vllm/model_executor/warmup/replayssm_warmup.py
+    - vllm/v1/attention/backends/mamba_attn.py
+    - vllm/v1/worker/mamba_utils.py
+    - vllm/v1/spec_decode/
+    - vllm/v1/worker/gpu/
+    - vllm/v1/worker/gpu_model_runner.py
+    - vllm/v1/worker/gpu_worker.py
+  commands:
+    - uv pip install --system flashinfer-python==0.7.0 --index-url https://pypi.org/simple
+    - uv pip install --system flashinfer-cubin==0.7.0 --index-url https://flashinfer.ai/whl/
+    - VLLM_TEST_REQUIRE_REPLAYSSM=1 pytest -v -s v1/e2e/test_replayssm_decode.py
+    - pytest -v kernels/mamba/test_ssu_dispatch.py -k replayssm
+    - pytest -v model_executor/test_replayssm_warmup.py v1/worker/test_kv_cache_allocation_scope.py
+
 - label: ":nvidia: (L4) V1 E2E"
   key: v1-e2e-2-gpus
   timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability

tests/kernels/mamba/test_precopy_mamba_align.py

@@ -260,6 +260,7 @@ def __init__(self, n):
         self.mamba_state_idx_buf = _FakeCpuGpuBuffer(n)
         self.precopy_src_col_buf = _FakeCpuGpuBuffer(n)
         self.precopy_token_bias_buf = _FakeCpuGpuBuffer(n)
+        self.replayssm = None
         self.calls = []
 
     def initialize_from_forward_context(self, *args, **kwargs):

tests/kernels/mamba/test_ssu_dispatch.py

@@ -1,6 +1,7 @@
 # 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
@@ -11,14 +12,15 @@
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     FlashInferSSUBackend,
     TritonSSUBackend,
+    _postprocess_replayssm_kernel,
+    _ReplaySSMGroupContext,
     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
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.kv_cache_interface import (
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -44,36 +46,6 @@ def restore_backend_state():
     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:
@@ -90,6 +62,84 @@ def _kv_cache_config_with_ssu(
     )
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("post_step", [False, True], ids=["v1", "v2"])
+def test_replayssm_materializes_only_accepted_boundary_with_compacted_rows(post_step):
+    """Publish a boundary without modifying its live or cached-prefix sources."""
+    from flashinfer.mamba.replayssm_materialize import replayssm_materialize
+
+    def ints(values):
+        return torch.tensor(values, dtype=torch.int32, device="cuda")
+
+    set_random_seed(42)
+    slots, heads, dim, dstate, ring_len = 6, 8, 64, 128, 20
+    state = torch.randn(slots, heads, dim, dstate, device="cuda") * 0.1
+    # Slot 1 is an immutable cached prefix, copied to private live slot 3.
+    state[3].copy_(state[1])
+    original = state.clone()
+    x = torch.randn(slots, heads, ring_len, dim, device="cuda", dtype=torch.bfloat16)
+    dt = torch.full((slots, heads, ring_len), 0.01, device="cuda")
+    B = torch.randn(slots, 1, ring_len, dstate, device="cuda", dtype=torch.bfloat16)
+    A = -torch.ones(heads, device="cuda")
+    mixer = SimpleNamespace(
+        kv_cache=(torch.empty(0, device="cuda"), state),
+        replayssm_cache=(x, dt, B),
+        _replayssm_ring_start=ints([0, 0, 0, 7, 0, 2]),
+        _replayssm_prev_num_accepted=ints([0, 0, 0, 2, 0, 1]),
+        replayssm_buffer_len=16,
+        A=A,
+        mamba_config=MambaConfig(backend=MambaBackendEnum.FLASHINFER),
+    )
+    group = _ReplaySSMGroupContext.create(
+        [mixer],
+        ints(
+            [
+                [2, 3],
+                [NULL_BLOCK_ID, 5],
+                [NULL_BLOCK_ID, NULL_BLOCK_ID],
+                [NULL_BLOCK_ID, NULL_BLOCK_ID],
+            ]
+        ),
+        "all",
+        16,
+        4,
+    )
+    # Batch row 0 crosses token 16; row 1 rejects all drafts and stops at 15.
+    # Rows 2 and 3 exercise padded physical slots and padded request indices.
+    mapping = ints([2, 0, 1, -1])
+    group.postprocess(
+        idx_mapping=mapping,
+        query_metadata=ints([0, 4, 8, 12, 16]) if post_step else ints([4] * 4),
+        query_metadata_is_cumulative=post_step,
+        num_computed_tokens=ints([15, 17, 17]) if post_step else ints([14] * 3),
+        num_computed_is_post_step=post_step,
+        num_accepted_tokens=ints([1, 3, 3]),
+        is_prefilling=torch.zeros(4, dtype=torch.bool, device="cuda"),
+        live_cols=ints([1, 1, 1]),
+        num_reqs=4,
+    )
+    assert group.active_request_indices.tolist() == [0, -1, -1, -1]
+    assert group.plan_flush_count.tolist() == [4, -1, -1, -1]
+    group.materialize(replayssm_materialize)
+
+    expected = original[3].clone()
+    for offset in range(4):
+        pos = (7 + offset) % ring_len
+        delta = dt[3, :, pos]
+        expected *= torch.exp(delta * A)[:, None, None]
+        expected += (
+            delta[:, None, None]
+            * x[3, :, pos].float()[:, :, None]
+            * B[3, 0, pos].float()[None, None, :]
+        )
+    torch.testing.assert_close(state[2], expected, atol=2e-3, rtol=2e-3)
+    # The snapshot excludes the accepted tail past 16 and all rejected drafts.
+    for slot in (0, 1, 3, 4, 5):
+        torch.testing.assert_close(state[slot], original[slot], atol=0, rtol=0)
+    assert mixer._replayssm_prev_num_accepted[3].item() == 5
+    assert mixer._replayssm_prev_num_accepted[2].item() == 0
+
+
 def test_default_backend_is_triton():
     initialize_mamba_ssu_backend(MambaConfig(), _kv_cache_config_with_ssu())
     backend = get_mamba_ssu_backend()
@@ -196,6 +246,7 @@ def test_flashinfer_import_error():
         FlashInferSSUBackend(MambaConfig())
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_triton_basic_call():
     set_random_seed(0)
     initialize_mamba_ssu_backend(
@@ -232,25 +283,88 @@ def test_triton_basic_call():
 
 
 @pytest.mark.parametrize(
-    ("backend", "expected_ring_len"),
+    ("backend", "num_speculative_tokens", "expected_ring_len"),
     [
-        (MambaBackendEnum.TRITON, 16),
-        (MambaBackendEnum.FLASHINFER, 17),
+        (MambaBackendEnum.TRITON, 0, 16),
+        (MambaBackendEnum.FLASHINFER, 0, 17),
+        (MambaBackendEnum.FLASHINFER, 3, 20),
     ],
 )
-def test_replayssm_physical_ring_shape(backend, expected_ring_len):
-    base_shapes = ((64, 3), (8, 4, 16))
-
-    shapes = MambaStateShapeCalculator.append_replayssm_ring(
-        base_shapes,
+def test_replayssm_physical_ring_shape(
+    backend, num_speculative_tokens, expected_ring_len
+):
+    shapes = MambaStateShapeCalculator.replayssm_ring_shapes(
+        num_heads=16,
+        head_dim=4,
+        state_size=16,
         n_groups=4,
         tp_world_size=2,
         logical_window=16,
         backend=backend,
+        num_speculative_tokens=num_speculative_tokens,
     )
 
-    assert shapes[2:] == (
+    assert shapes == (
         (8, expected_ring_len, 4),
         (8, expected_ring_len),
         (2, expected_ring_len, 16),
     )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
+@pytest.mark.parametrize("post_step", [False, True])
+@pytest.mark.parametrize(
+    ("computed_before", "query_len", "prefilling", "accepted", "expected"),
+    [
+        (256, 4, False, 1, 3),  # Padded cached prompt tail commits its real token.
+        (1, 4, False, 1, 3),  # Rejected placeholders exceed the cached prefix.
+        (256, 1, False, 1, 3),
+        (0, 4, True, 1, 0),  # Initial prefill clears stale trackers.
+        (256, 4, True, 1, 0),  # Multi-token prefill also clears trackers.
+        (256, 4, False, 3, 5),  # Ordinary speculative decode commits acceptance.
+    ],
+)
+def test_replayssm_postprocess_commits_staged_transition(
+    post_step, computed_before, query_len, prefilling, accepted, expected
+):
+    """The staged kernel path must preserve accepted history on prompt tails."""
+
+    def tensor(values):
+        return torch.tensor(values, dtype=torch.int32, device="cuda")
+
+    computed = computed_before
+    if post_step:
+        computed += query_len if prefilling else accepted
+    ring_start = tensor([3])
+    committed = tensor([2])
+    plan_start, plan_flush = tensor([0]), tensor([-1])
+    slots = tensor([[0]])
+    _postprocess_replayssm_kernel[(1,)](
+        tensor([0]),
+        tensor([0, query_len]) if post_step else tensor([query_len]),
+        tensor([computed]),
+        tensor([accepted]),
+        torch.tensor([prefilling], device="cuda"),
+        None,
+        tensor([[0, 0]]),
+        ring_start,
+        committed,
+        slots,
+        slots,
+        plan_start,
+        plan_flush,
+        2,
+        1,
+        MAMBA_BLOCK_SIZE=256,
+        LOGICAL_WINDOW=16,
+        RING_BUFFER_LEN=20,
+        NUM_LAYERS=1,
+        PAD_SLOT_ID=-1,
+        QUERY_METADATA_IS_CUMULATIVE=post_step,
+        NUM_COMPUTED_IS_POST_STEP=post_step,
+        HAS_IDX_MAPPING=post_step,
+        MATERIALIZE_PREFIXES=False,
+        LIVE_COL_IS_ZERO=True,
+    )
+    assert committed.item() == expected
+    assert ring_start.item() == (0 if prefilling else 3)

tests/model_executor/test_nemotron_h_mtp.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch.nn as nn
+
+from vllm.config import CompilationMode
+from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig
+
+
+class _StubModule(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+
+def test_native_mtp_attention_registers_after_target_attention():
+    from vllm.model_executor.models import nemotron_h, nemotron_h_mtp
+
+    static_forward_context = {}
+
+    class StaticContextAttention(_StubModule):
+        def __init__(self, *args, prefix: str = "", **kwargs):
+            super().__init__()
+            static_forward_context[prefix] = self
+
+    config = NemotronHConfig(
+        vocab_size=8,
+        hidden_size=4,
+        num_hidden_layers=1,
+        hybrid_override_pattern="*",
+        mtp_hybrid_override_pattern="*",
+        num_attention_heads=1,
+        num_key_value_heads=1,
+        head_dim=4,
+        num_nextn_predict_layers=1,
+    )
+    model_config = SimpleNamespace(hf_config=config)
+    vllm_config = SimpleNamespace(
+        model_config=model_config,
+        cache_config=None,
+        quant_config=None,
+        parallel_config=None,
+        compilation_config=SimpleNamespace(mode=CompilationMode.NONE),
+    )
+
+    with (
+        patch.object(
+            nemotron_h, "get_tensor_model_parallel_world_size", return_value=1
+        ),
+        patch.object(nemotron_h, "Attention", StaticContextAttention),
+        patch.object(nemotron_h, "QKVParallelLinear", _StubModule),
+        patch.object(nemotron_h, "RowParallelLinear", _StubModule),
+        patch.object(nemotron_h, "RMSNorm", _StubModule),
+        patch.object(nemotron_h_mtp, "VocabParallelEmbedding", _StubModule),
+        patch.object(nemotron_h_mtp, "ColumnParallelLinear", _StubModule),
+        patch.object(nemotron_h_mtp, "ParallelLMHead", _StubModule),
+        patch.object(nemotron_h_mtp, "LogitsProcessor", _StubModule),
+        patch.object(nemotron_h_mtp, "RMSNorm", _StubModule),
+    ):
+        nemotron_h.NemotronHAttentionDecoderLayer(
+            config=config,
+            layer_idx=0,
+            model_config=model_config,
+            prefix="model.layers.0",
+        )
+        nemotron_h_mtp.NemotronHMTP(
+            vllm_config=vllm_config,
+            prefix="draft_model",
+        )
+
+    assert list(static_forward_context) == [
+        "model.layers.0.mixer.attn",
+        "draft_model.mtp.layers.0.mixer.attn",
+    ]