From deb148e730225c64427f8ccf1f313cb4b33325b0 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 30 Sep 2026 23:41:17 +0200 Subject: [PATCH] [Mamba] Add FlashInfer ReplaySSM support for MTP (#52928) Signed-off-by: Andrii Skliar Co-authored-by: Andrii Skliar --- .buildkite/test_areas/engine.yaml | 21 + tests/kernels/mamba/test_ssu_dispatch.py | 398 +++++++++++++++++- tests/model_executor/test_replayssm_warmup.py | 23 +- tests/test_config.py | 6 +- .../test_attention_backends_selection.py | 28 ++ .../test_replayssm_metadata_builder.py | 94 ++++- tests/v1/e2e/test_replayssm_decode.py | 146 +++++-- tests/v1/worker/test_utils.py | 26 +- tools/pre_commit/test_tethering_allowlist.txt | 1 - vllm/config/cache.py | 11 +- vllm/config/vllm.py | 47 ++- vllm/model_executor/layers/mamba/abstract.py | 4 +- .../layers/mamba/mamba_mixer2.py | 109 ++++- .../layers/mamba/mamba_utils.py | 5 +- .../layers/mamba/ops/ssu_dispatch.py | 143 ++++++- vllm/model_executor/models/nemotron_h.py | 1 + .../model_executor/warmup/replayssm_warmup.py | 37 +- vllm/v1/attention/backends/mamba_attn.py | 42 +- 18 files changed, 1029 insertions(+), 113 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index da02f47827..fa71c9f737 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -133,6 +133,27 @@ steps: - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async +- 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 + - 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/spec_decode/ + - vllm/v1/worker/gpu/ + - vllm/v1/worker/gpu_model_runner.py + - vllm/v1/worker/gpu_worker.py + commands: + - pytest -v -s v1/e2e/test_replayssm_decode.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 diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 6226340a99..082b61c43d 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -1,24 +1,30 @@ # 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.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm +from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( FlashInferSSUBackend, TritonSSUBackend, + commit_replayssm_ring_trackers, get_mamba_ssu_backend, initialize_mamba_ssu_backend, reset_replayssm_ring_trackers, selective_state_update, + selective_state_update_replayssm_flashinfer, update_replayssm_ring_trackers, ) from vllm.utils.torch_utils import set_random_seed +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata 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, @@ -32,6 +38,13 @@ try: except ImportError: HAS_FLASHINFER = False +try: + from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner + + HAS_FLASHINFER_CHECKPOINTING_SSU = CheckpointingSSURunner is not None +except ImportError: + HAS_FLASHINFER_CHECKPOINTING_SSU = False + @pytest.fixture(autouse=True) def restore_backend_state(): @@ -47,6 +60,7 @@ def restore_backend_state(): 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") + prev_query_len = torch.zeros(2, dtype=torch.int32, device="cuda") state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda") observed = [] @@ -54,6 +68,7 @@ def test_flashinfer_replayssm_ring_tracker_lifecycle(): update_replayssm_ring_trackers( ring_start, prev_num_accepted, + prev_query_len, state_batch_indices, logical_window=16, ring_buffer_len=17, @@ -69,9 +84,205 @@ def test_flashinfer_replayssm_ring_tracker_lifecycle(): reset_replayssm_ring_trackers( ring_start, prev_num_accepted, + prev_query_len, state_batch_indices, ) - assert (ring_start[1].item(), prev_num_accepted[1].item()) == (0, 0) + assert ( + ring_start[1].item(), + prev_num_accepted[1].item(), + prev_query_len[1].item(), + ) == (0, 0, 0) + + +@pytest.mark.parametrize( + ("accepted_sequence", "expected"), + [ + pytest.param( + [4] * 22, + [ + (0, 0, 4), + (0, 4, 4), + (0, 8, 4), + (0, 12, 4), + (0, 16, 4), + (16, 4, 4), + (16, 8, 4), + (16, 12, 4), + (16, 16, 4), + (12, 4, 4), + (12, 8, 4), + (12, 12, 4), + (12, 16, 4), + (8, 4, 4), + (8, 8, 4), + (8, 12, 4), + (8, 16, 4), + (4, 4, 4), + (4, 8, 4), + (4, 12, 4), + (4, 16, 4), + (0, 4, 4), + ], + id="all-accepted", + ), + pytest.param( + [4, 4, 0, 3, 4, 2, 4, 1], + [ + (0, 0, 4), + (0, 4, 4), + (0, 4, 4), + (0, 7, 4), + (0, 11, 4), + (0, 13, 4), + (13, 4, 4), + (13, 5, 4), + ], + id="mixed", + ), + ], +) +def test_replayssm_commit_tracker_acceptance_sequence(accepted_sequence, expected): + logical_window = 16 + num_speculative_tokens = 3 + query_len = 1 + num_speculative_tokens + ring_buffer_len = logical_window + 1 + num_speculative_tokens + ring_start = torch.zeros(2, dtype=torch.int32, device="cuda") + prev_num_accepted = torch.zeros(2, dtype=torch.int32, device="cuda") + prev_query_len = torch.zeros(2, dtype=torch.int32, device="cuda") + state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda") + query_start_loc = torch.tensor([0, query_len], dtype=torch.int32, device="cuda") + + observed = [] + for accepted in accepted_sequence: + commit_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + torch.tensor([accepted], dtype=torch.int32, device="cuda"), + query_start_loc, + logical_window, + ring_buffer_len, + ) + snapshot = ( + ring_start[1].item(), + prev_num_accepted[1].item(), + prev_query_len[1].item(), + ) + observed.append(snapshot) + assert snapshot[1] + snapshot[2] <= ring_buffer_len + + assert observed == expected + + +def test_replayssm_resume_resets_commit_history(): + ring_start = torch.tensor([0, 13], dtype=torch.int32, device="cuda") + prev_num_accepted = torch.tensor([0, 13], dtype=torch.int32, device="cuda") + prev_query_len = torch.tensor([0, 4], dtype=torch.int32, device="cuda") + state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda") + + reset_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + ) + assert ( + ring_start[1].item(), + prev_num_accepted[1].item(), + prev_query_len[1].item(), + ) == (0, 0, 0) + + commit_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + torch.tensor([3], dtype=torch.int32, device="cuda"), + torch.tensor([0, 4], dtype=torch.int32, device="cuda"), + logical_window=16, + ring_buffer_len=20, + ) + assert ( + ring_start[1].item(), + prev_num_accepted[1].item(), + prev_query_len[1].item(), + ) == (0, 0, 4) + + +def test_replayssm_commit_tracker_ragged_query_lengths(): + ring_start = torch.zeros(3, dtype=torch.int32, device="cuda") + prev_num_accepted = torch.zeros(3, dtype=torch.int32, device="cuda") + prev_query_len = torch.zeros(3, dtype=torch.int32, device="cuda") + state_batch_indices = torch.tensor([1, 2], dtype=torch.int32, device="cuda") + query_start_loc = torch.tensor([0, 4, 6], dtype=torch.int32, device="cuda") + + observed = [] + for accepted in ([4, 2], [3, 1]): + commit_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + torch.tensor(accepted, dtype=torch.int32, device="cuda"), + query_start_loc, + logical_window=16, + ring_buffer_len=20, + ) + observed.append( + [ + ( + ring_start[slot].item(), + prev_num_accepted[slot].item(), + prev_query_len[slot].item(), + ) + for slot in (1, 2) + ] + ) + + assert observed == [[(0, 0, 4), (0, 0, 2)], [(0, 3, 4), (0, 1, 2)]] + + +@pytest.mark.parametrize("operation", ["commit", "reset"]) +def test_replayssm_tracker_kernels_mask_invalid_slots(operation): + num_states = 3 + ring_start = torch.tensor([11, 2, 33], dtype=torch.int32, device="cuda") + prev_num_accepted = torch.tensor([11, 3, 33], dtype=torch.int32, device="cuda") + prev_query_len = torch.tensor([11, 4, 33], dtype=torch.int32, device="cuda") + state_batch_indices = torch.tensor( + [-1, num_states, NULL_BLOCK_ID, 1], dtype=torch.int32, device="cuda" + ) + + if operation == "commit": + commit_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + torch.tensor([4, 4, 4, 2], dtype=torch.int32, device="cuda"), + torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32, device="cuda"), + logical_window=16, + ring_buffer_len=20, + ) + expected_valid = (2, 5, 4) + else: + reset_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + ) + expected_valid = (0, 0, 0) + + assert ( + ring_start.tolist(), + prev_num_accepted.tolist(), + prev_query_len.tolist(), + ) == ( + [11, expected_valid[0], 33], + [11, expected_valid[1], 33], + [11, expected_valid[2], 33], + ) def _kv_cache_config_with_ssu( @@ -231,14 +442,190 @@ def test_triton_basic_call(): assert not torch.isnan(out).any() +@pytest.mark.parametrize("layout", ["packed", "dense"]) +def test_replayssm_flashinfer_call_forwards_mtp_layout(monkeypatch, layout): + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod + + kernel = Mock(return_value=torch.empty(0)) + monkeypatch.setattr(mod, "_flashinfer_replayssm_kernel", kernel) + + batch, max_seqlen, nheads, dim, dstate, ngroups = 2, 4, 2, 4, 8, 1 + state = torch.empty(2, nheads, dim, dstate) + x_shape: tuple[int, ...] + B_shape: tuple[int, ...] + expected_x_shape: tuple[int, ...] + expected_B_shape: tuple[int, ...] + if layout == "packed": + x_shape = (6, nheads, dim) + B_shape = (6, ngroups, dstate) + expected_x_shape = (1, 6, nheads, dim) + expected_B_shape = (1, 6, ngroups, dstate) + cu_seqlens = torch.tensor([0, 4, 6], dtype=torch.int32) + kernel_max_seqlen = max_seqlen + else: + x_shape = (batch, max_seqlen, nheads, dim) + B_shape = (batch, max_seqlen, ngroups, dstate) + expected_x_shape = x_shape + expected_B_shape = B_shape + cu_seqlens = None + kernel_max_seqlen = None + x = torch.empty(x_shape) + dt = torch.empty_like(x) + A = torch.empty(nheads, dim, dstate) + B = torch.empty(B_shape) + C = torch.empty_like(B) + out = torch.empty_like(x) + x_cache = torch.empty(2, nheads, 20, dim) + dt_cache = torch.empty(2, nheads, 20) + B_cache = torch.empty(2, ngroups, 20, dstate) + ring_start = torch.zeros(2, dtype=torch.int32) + prev_num_accepted = torch.zeros(2, dtype=torch.int32) + prev_query_len = torch.zeros(2, dtype=torch.int32) + selective_state_update_replayssm_flashinfer( + state, + x, + dt, + A, + B, + C, + out, + x_cache, + B_cache, + dt_cache, + ring_start, + prev_num_accepted, + prev_query_len, + logical_window=16, + state_batch_indices=torch.tensor([0, 1], dtype=torch.int32), + cu_seqlens=cu_seqlens, + max_seqlen=kernel_max_seqlen, + update_trackers=False, + ) + + args = kernel.call_args.args + assert args[6].shape == expected_x_shape + assert args[7].shape == expected_x_shape + assert args[9].shape == expected_B_shape + assert args[10].shape == expected_B_shape + assert args[11].shape == expected_x_shape + assert kernel.call_args.kwargs["cu_seqlens"] is cu_seqlens + assert kernel.call_args.kwargs["max_seqlen"] == kernel_max_seqlen + + @pytest.mark.parametrize( - ("backend", "expected_ring_len"), + ("query_start_loc", "expected_shape", "expected_max_seqlen"), [ - (MambaBackendEnum.TRITON, 16), - (MambaBackendEnum.FLASHINFER, 17), + pytest.param([0, 4, 8], (2, 4, 2, 4), None, id="dense"), + pytest.param([0, 4, 6], (6, 2, 4), 4, id="packed"), ], ) -def test_replayssm_physical_ring_shape(backend, expected_ring_len): +def test_replayssm_mixer_selects_mtp_layout( + monkeypatch, query_start_loc, expected_shape, expected_max_seqlen +): + import vllm.model_executor.layers.mamba.mamba_mixer2 as mod + + mixer = MambaMixer2.__new__(MambaMixer2) + torch.nn.Module.__init__(mixer) + mixer.prefix = "mixer" + mixer.tped_intermediate_size = 0 + mixer.tped_conv_size = 1 + mixer.tped_dt_size = 2 + mixer.num_heads = 2 + mixer.head_dim = 4 + mixer.n_groups = mixer.tp_size = 1 + mixer.ssm_state_size = 8 + mixer.num_spec = 3 + mixer.use_replayssm = True + mixer.replayssm_buffer_len = 16 + mixer._commits_replayssm_trackers = True + mixer._updates_replayssm_trackers = False + mixer.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER) + mixer.cache_config = SimpleNamespace(mamba_block_size=16, mamba_cache_mode="none") + mixer.conv_weights = torch.empty(0) + mixer.conv1d = SimpleNamespace(bias=None) + mixer.activation = "silu" + mixer.A = torch.empty(2) + mixer.dt_bias = torch.empty(2) + mixer.D = torch.empty(2) + mixer._replayssm_ring_start = torch.zeros(3, dtype=torch.int32) + mixer._replayssm_prev_num_accepted = torch.zeros(3, dtype=torch.int32) + mixer._replayssm_prev_query_len = torch.zeros(3, dtype=torch.int32) + mixer.kv_cache = ( + torch.empty(3, 1), + torch.empty(3, 2, 4, 8), + torch.empty(3, 2, 20, 4), + torch.empty(3, 2, 20), + torch.empty(3, 1, 20, 8), + ) + + num_decode_tokens = query_start_loc[-1] + query_start_loc_d = torch.tensor(query_start_loc, dtype=torch.int32) + metadata = Mamba2AttentionMetadata( + num_prefills=0, + num_prefill_tokens=0, + num_decodes=2, + num_decode_tokens=num_decode_tokens, + num_reqs=2, + has_initial_states_p=None, + query_start_loc_p=None, + num_computed_tokens_p=None, + state_indices_tensor_p=None, + state_indices_tensor_d=torch.tensor([[1], [2]], dtype=torch.int32), + query_start_loc_d=query_start_loc_d, + num_accepted_tokens=torch.tensor([4, 2], dtype=torch.int32), + block_idx_last_scheduled_token=None, + block_idx_first_scheduled_token_p=None, + block_idx_last_computed_token=None, + block_idx_last_scheduled_token_prev_step=None, + seq_lens=torch.tensor([104, 102], dtype=torch.int32), + replayssm_scratch=(torch.empty(0), torch.empty(0), torch.empty(0)), + replayssm_state_indices_d=torch.tensor([1, 2], dtype=torch.int32), + ) + + def split_hidden_states_B_C(values): + tokens = values.size(0) + return ( + torch.empty(tokens, 8), + torch.empty(tokens, 8), + torch.empty(tokens, 8), + ) + + mixer.split_hidden_states_B_C_fn = split_hidden_states_B_C + kernel = Mock() + monkeypatch.setattr( + mod, + "get_forward_context", + lambda: SimpleNamespace(attn_metadata={mixer.prefix: metadata}), + ) + monkeypatch.setattr(mod, "commit_replayssm_ring_trackers", Mock()) + monkeypatch.setattr( + mod, "causal_conv1d_update", lambda values, *args, **kwargs: values + ) + monkeypatch.setattr(mod, "selective_state_update_replayssm_flashinfer", kernel) + + mixer.conv_ssm_forward( + torch.empty(num_decode_tokens, 3), torch.empty(num_decode_tokens, 8) + ) + + assert kernel.call_args.args[1].shape == expected_shape + if expected_max_seqlen is None: + assert kernel.call_args.kwargs["cu_seqlens"] is None + else: + assert kernel.call_args.kwargs["cu_seqlens"] is query_start_loc_d + assert kernel.call_args.kwargs["max_seqlen"] == expected_max_seqlen + + +@pytest.mark.parametrize( + ("backend", "num_speculative_tokens", "expected_ring_len"), + [ + (MambaBackendEnum.TRITON, 0, 16), + (MambaBackendEnum.FLASHINFER, 0, 17), + (MambaBackendEnum.FLASHINFER, 3, 20), + ], +) +def test_replayssm_physical_ring_shape( + backend, num_speculative_tokens, expected_ring_len +): base_shapes = ((64, 3), (8, 4, 16)) shapes = MambaStateShapeCalculator.append_replayssm_ring( @@ -247,6 +634,7 @@ def test_replayssm_physical_ring_shape(backend, expected_ring_len): tp_world_size=2, logical_window=16, backend=backend, + num_speculative_tokens=num_speculative_tokens, ) assert shapes[2:] == ( diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index 272fa2cf85..5db0f9a50d 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -67,8 +67,13 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): with patch.object( warmup, "flashinfer_replayssm_autotune_supported", return_value=True ): - result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs)) + result = warmup._replayssm_autotune_kwargs( + _autotune_runner(**runner_kwargs), PREFILL_KWARGS + ) + assert result is not None + max_num_reqs, decode_kwargs = result + assert max_num_reqs == expected_num_reqs expected_kwargs = { **PREFILL_KWARGS, "num_tokens": expected_num_reqs * query_len, @@ -82,7 +87,7 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): force_attention=True, profile_seq_lens=query_len + 1, ) - assert result == (expected_num_reqs, expected_kwargs) + assert decode_kwargs == expected_kwargs @pytest.mark.parametrize( @@ -91,8 +96,14 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): (dict(use_replayssm=False), True), (dict(backend=MambaBackendEnum.TRITON), True), ({}, False), + (dict(num_blocks=1), True), + ], + ids=[ + "replayssm_disabled", + "non_flashinfer_backend", + "kernel_unavailable", + "zero_non_padding_slots", ], - ids=["replayssm_disabled", "non_flashinfer_backend", "kernel_unavailable"], ) def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): with patch.object( @@ -100,7 +111,9 @@ def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): "flashinfer_replayssm_autotune_supported", return_value=flashinfer_supported, ): - result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs)) + result = warmup._replayssm_autotune_kwargs( + _autotune_runner(**runner_kwargs), PREFILL_KWARGS + ) assert result is None @@ -116,10 +129,12 @@ def test_replayssm_autotune_slots_restore_state_and_trackers(): ) mixer._replayssm_ring_start = torch.full((4,), 3, dtype=torch.int32) mixer._replayssm_prev_num_accepted = torch.full((4,), 3, dtype=torch.int32) + mixer._replayssm_prev_query_len = torch.full((4,), 3, dtype=torch.int32) tracked = ( *mixer.kv_cache, mixer._replayssm_ring_start, mixer._replayssm_prev_num_accepted, + mixer._replayssm_prev_query_len, ) block_ids = np.arange(10, 14, dtype=np.int32).reshape(4, 1) diff --git a/tests/test_config.py b/tests/test_config.py index 93578ab918..712d01fd33 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -153,6 +153,7 @@ def test_kda_recoverssm_derivation_is_revalidated(): use_replayssm=True, use_kda_recoverssm=False, mamba_cache_mode="none", + replayssm_buffer_len=16, ), num_speculative_tokens=3, model_config=SimpleNamespace( @@ -184,8 +185,9 @@ def test_kda_recoverssm_derivation_is_revalidated(): 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.mamba_config.backend = MambaBackendEnum.FLASHINFER + VllmConfig.validate_mamba_cached_kernel(config) + assert not config.cache_config.use_kda_recoverssm config.model_config.architecture = "KimiLinearForCausalLM" config.parallel_config.pipeline_parallel_size = 2 diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index 8486d216a1..2695517e14 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/tests/v1/attention/test_attention_backends_selection.py @@ -5,7 +5,9 @@ from types import SimpleNamespace import pytest +import torch +from vllm.model_executor.layers.mamba.abstract import MambaBase from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import ( MiniMaxText01LinearAttention, ) @@ -19,6 +21,32 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend +@pytest.mark.parametrize(("use_replayssm", "expected_blocks"), [(False, 3), (True, 0)]) +def test_replayssm_does_not_reserve_speculative_state_blocks( + use_replayssm, expected_blocks +): + layer = SimpleNamespace( + get_state_shape=lambda: ((2,),), + get_state_dtype=lambda: (torch.float32,), + mamba_type=MambaAttentionBackendEnum.MAMBA2, + is_kv_cache_tp_replicated=False, + ) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace( + mamba_block_size=1, + mamba_page_size_padded=None, + mamba_cache_mode="none", + use_replayssm=use_replayssm, + ), + num_speculative_tokens=3, + ) + + spec = MambaBase.get_kv_cache_spec(layer, vllm_config) + + assert spec is not None + assert spec.num_speculative_blocks == expected_blocks + + @pytest.mark.parametrize( "layer_class, init_kwargs, expected_backend, expected_mamba_type", [ diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index 89cb0f5664..58679ecac9 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -16,6 +16,8 @@ from tests.v1.attention.utils import ( create_common_attn_metadata, create_vllm_config, ) +from vllm.config import SpeculativeConfig +from vllm.config.compilation import CUDAGraphMode from vllm.config.mamba import MambaBackendEnum from vllm.v1.kv_cache_interface import MambaSpec @@ -191,9 +193,12 @@ REPLAYSSM_BUILD_CASES = { def _make_mamba_spec( buffer_len: int, mamba_backend: MambaBackendEnum, + num_speculative_tokens: int = 0, ) -> MambaSpec: ring_buffer_len = buffer_len + ( - 1 if mamba_backend == MambaBackendEnum.FLASHINFER else 0 + 1 + num_speculative_tokens + if mamba_backend == MambaBackendEnum.FLASHINFER + else 0 ) shapes = ( (1, 1), @@ -214,6 +219,7 @@ def _create_replayssm_builder( mamba_cache_mode: str = "none", *, mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON, + num_speculative_tokens: int = 0, ) -> MockMambaBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE @@ -224,21 +230,32 @@ def _create_replayssm_builder( 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 + if num_speculative_tokens > 0: + # Triton ReplaySSM requires exact synchronous CPU metadata. + vllm_config.scheduler_config.async_scheduling = False + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=num_speculative_tokens, + ) return MockMambaBuilder( - _make_mamba_spec(buffer_len, mamba_backend), + _make_mamba_spec(buffer_len, mamba_backend, num_speculative_tokens), ["layer0"], vllm_config, DEVICE, ) -def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase): +def _build( + builder: MockMambaBuilder, + case: ReplaySSMBuildCase, + num_accepted_tokens: torch.Tensor | None = None, +): batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens) common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace( is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool), replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32), ) - return builder.build(0, common) + return builder.build(0, common, num_accepted_tokens=num_accepted_tokens) @pytest.mark.parametrize( @@ -270,21 +287,60 @@ def test_resumed_request_differs_from_fresh(): 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") +def test_spec_decode_single_token_chunk_synthesizes_acceptance_metadata(): + builder = _create_replayssm_builder(16, num_speculative_tokens=3) + case = REPLAYSSM_BUILD_CASES["leftover_prompt_one_token_flush"] - 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), - ] + assert meta.query_start_loc_d is not None + assert meta.query_start_loc_d.tolist() == [0, 1] + assert meta.num_accepted_tokens is not None + assert meta.num_accepted_tokens.tolist() == [1] + + +def test_flashinfer_replayssm_state_indices_are_stable_for_full_cudagraph(): + checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu") + if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"): + pytest.skip("requires FlashInfer ReplaySSM autotuning support") + + builder = _create_replayssm_builder( + 16, + mamba_backend=MambaBackendEnum.FLASHINFER, + num_speculative_tokens=3, + ) + builder.compilation_config.cudagraph_mode = CUDAGraphMode.FULL + + first = _build( + builder, + ReplaySSMBuildCase( + seq_lens=[106, 106], + query_lens=[1, 1], + is_prefilling=[False, False], + decode_base=[100, 100], + buffer_len=16, + expected_write_pos=[], + expected_is_flush=[], + ), + ) + first_indices = first.replayssm_state_indices_d + assert first_indices is not None + assert first_indices.is_contiguous() + first_ptr = first_indices.data_ptr() + + second = _build( + builder, + ReplaySSMBuildCase( + seq_lens=[122, 122], + query_lens=[1, 1], + is_prefilling=[False, False], + decode_base=[116, 116], + buffer_len=16, + expected_write_pos=[], + expected_is_flush=[], + ), + ) + second_indices = second.replayssm_state_indices_d + assert second_indices is not None + assert second_indices.data_ptr() == first_ptr + assert torch.equal(second_indices, second.state_indices_tensor_d[:, 0]) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 2e66a46689..1a369c781f 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -10,11 +10,16 @@ from vllm.v1.metrics.reader import Counter from ...models.utils import check_logprobs_close from ...utils import large_gpu_mark, multi_gpu_test -# Mamba2 (Nemotron-3) hybrid. +# Mamba2 (Nemotron-3) hybrids. MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +# FlashInfer 0.7.0's standard SSU cannot compile Nano's DIM=80 configuration. +MAMBA2_FLASHINFER_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4" MODELS = [ pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), ] +FLASHINFER_MODELS = [ + pytest.param(MAMBA2_FLASHINFER_MODEL, marks=large_gpu_mark(min_gb=40)), +] PROMPTS = [ "The capital of France is", @@ -22,6 +27,16 @@ PROMPTS = [ ] +@pytest.fixture(autouse=True) +def _use_v1_model_runner_by_default(monkeypatch): + # Triton ReplaySSM is V1-only. FlashInfer V2 tests override this locally. + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + envs.disable_envs_cache() + yield + envs.disable_envs_cache() + + def _check_replayssm_parity( vllm_runner, model_name, @@ -29,17 +44,11 @@ def _check_replayssm_parity( tensor_parallel_size=1, mamba_backend: str = "triton", name_1: str = "replayssm", - require_v2: bool = False, - monkeypatch: pytest.MonkeyPatch | None = None, + expected_v2: bool | 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() - + # near-tie. Baseline and ReplaySSM use the same backend so only ReplaySSM + # behavior varies. common = dict( max_model_len=1024, trust_remote_code=True, @@ -49,14 +58,17 @@ def _check_replayssm_parity( 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 + if expected_v2 is not None: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2 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 + 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 + if expected_v2 is not None: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2 replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) check_logprobs_close( @@ -80,19 +92,101 @@ 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.mark.parametrize("model_name", FLASHINFER_MODELS) +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_replayssm_flashinfer_decode_matches_baseline( + vllm_runner, model_name, monkeypatch, use_v2_model_runner ): - 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, + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) + envs.disable_envs_cache() + _check_replayssm_parity( + vllm_runner, + model_name, + mamba_backend="flashinfer", + name_1="replayssm_flashinfer", + expected_v2=use_v2_model_runner, + ) + finally: + # The context restores the environment before the final cache reset. + envs.disable_envs_cache() + + +@pytest.mark.parametrize("model_name", FLASHINFER_MODELS) +def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_name): + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + speculative_config={ + "method": "ngram", + "num_speculative_tokens": 3, + "prompt_lookup_max": 3, + }, ) + with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: + baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + with vllm_runner( + model_name, + use_replayssm=True, + replayssm_buffer_len=16, + mamba_backend="flashinfer", + **common, + ) as llm: + 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_spec", + name_1="replayssm_flashinfer_spec", + ) + + +@multi_gpu_test(num_gpus=2) +@large_gpu_mark(min_gb=40) +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) +def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner): + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + mamba_backend="flashinfer", + tensor_parallel_size=2, + disable_log_stats=False, + speculative_config={"method": "mtp", "num_speculative_tokens": 3}, + ) + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) + envs.disable_envs_cache() + with vllm_runner( + MAMBA2_FLASHINFER_MODEL, + use_replayssm=True, + replayssm_buffer_len=16, + **common, + ) as llm: + assert ( + llm.llm.llm_engine.vllm_config.use_v2_model_runner + is use_v2_model_runner + ) + outputs = llm.generate_greedy(PROMPTS, max_tokens=32) + draft_count = sum( + metric.value + for metric in llm.llm.get_metrics() + if isinstance(metric, Counter) + and metric.name == "vllm:spec_decode_num_drafts" + ) + finally: + envs.disable_envs_cache() + + # At least one request must run past the 16-token replay window; another + # may legitimately stop early on EOS. + assert any(len(token_ids) > 16 for token_ids, _ in outputs) + assert draft_count > 0 # Prefix spans several mamba blocks; prefix caching only reuses full blocks. diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 1833bca8ab..43ac710f35 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1260,6 +1260,8 @@ class _TestReplaySSMMixer(MambaMixer2): 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._replayssm_prev_query_len = torch.empty(0, dtype=torch.int32) + self._commits_replayssm_trackers = True self._updates_replayssm_trackers = True def get_state_shape(self) -> tuple[tuple[int, ...], ...]: @@ -1294,19 +1296,19 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only): else: 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() + tracker_names = ( + "_replayssm_ring_start", + "_replayssm_prev_num_accepted", + "_replayssm_prev_query_len", ) - 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. + for tracker_name in tracker_names: + group_tracker = getattr(mixers[0], tracker_name) + assert group_tracker.data_ptr() == getattr(mixers[2], tracker_name).data_ptr() + assert group_tracker.data_ptr() != getattr(mixers[1], tracker_name).data_ptr() + assert group_tracker.shape == (4,) + assert torch.count_nonzero(group_tracker) == 0 + + assert [m._commits_replayssm_trackers for m in mixers] == [True, True, False] assert [m._updates_replayssm_trackers for m in mixers] == [False, True, True] diff --git a/tools/pre_commit/test_tethering_allowlist.txt b/tools/pre_commit/test_tethering_allowlist.txt index cc5294cf26..d7f19ec79a 100644 --- a/tools/pre_commit/test_tethering_allowlist.txt +++ b/tools/pre_commit/test_tethering_allowlist.txt @@ -60,7 +60,6 @@ tests/models/test_language_model_cache_is_weak.py # no models_*.yaml job collec tests/models/transformers/test_layer_registry.py # no models_*.yaml job collects this path # tests/v1/ - the v1 subtree is otherwise covered by many `pytest v1/...` jobs. -tests/v1/e2e/test_replayssm_decode.py # ReplaySSM decode test (#52506) wired into no job tests/v1/ec_connector/integration/test_epd_correctness.py # run_epd_correctness_test.sh exists but no Buildkite job invokes it tests/v1/ec_connector/integration/test_nixl_failure.py # 2-GPU + NIXL E/PD integration test (#55290); no Buildkite job collects ec_connector/integration/ tests/v1/kv_connector/nixl_integration/test_disagg_accuracy.py # run by intel_jobs/misc_intel.yaml's run_xpu_disagg_accuracy_test.sh, which names it via $EXP_ROOT (the script's own dir) - not statically resolvable diff --git a/vllm/config/cache.py b/vllm/config/cache.py index eeff06933e..13665f3cff 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -204,15 +204,16 @@ class CacheConfig: Mamba block size.""" replayssm_buffer_len: int = Field(default=16, gt=0) """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.""" + rows and FlashInfer uses B+T, where T is the target verification length. + 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 - 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.""" + or FlashInfer mamba backend. Mamba2 speculative decode requires FlashInfer + and mamba_cache_mode 'none'. 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.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2da17707e5..0fd536ff87 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3458,19 +3458,35 @@ class VllmConfig: if not self.cache_config.use_replayssm: self.cache_config.use_kda_recoverssm = False return self - self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0 + + kda_architectures = ( + "KimiLinearForCausalLM", + "KimiK3ForConditionalGeneration", + ) + is_kda_model = ( + self.model_config is not None + and self.model_config.architecture in kda_architectures + ) + self.cache_config.use_kda_recoverssm = ( + self.num_speculative_tokens > 0 and is_kda_model + ) + use_mamba_replayssm_spec = ( + self.num_speculative_tokens > 0 and not self.cache_config.use_kda_recoverssm + ) if self.model_config is not None and not self.model_config.supports_replayssm: raise ValueError( "--use-replayssm is not supported for architecture " f"{self.model_config.architecture!r}" ) + if ( + self.mamba_config.backend == MambaBackendEnum.FLASHINFER + and self.cache_config.replayssm_buffer_len > 16 + ): + raise ValueError( + "FlashInfer ReplaySSM requires --replayssm-buffer-len <= 16" + ) 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 " @@ -3494,6 +3510,25 @@ class VllmConfig: ) if self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError("RecoverSSM requires --mamba-backend triton") + elif use_mamba_replayssm_spec: + if self.cache_config.mamba_cache_mode != "none": + raise ValueError( + "FlashInfer ReplaySSM speculative decoding requires " + "--mamba-cache-mode none" + ) + query_len = 1 + self.num_speculative_tokens + if self.cache_config.replayssm_buffer_len < query_len: + raise ValueError( + "FlashInfer ReplaySSM speculative decoding requires " + "--replayssm-buffer-len >= 1 + num_speculative_tokens " + f"({query_len}); got " + f"{self.cache_config.replayssm_buffer_len}" + ) + if self.mamba_config.backend != MambaBackendEnum.FLASHINFER: + raise ValueError( + "Mamba2 ReplaySSM speculative decoding requires " + "--mamba-backend flashinfer" + ) elif self.cache_config.mamba_cache_mode == "all": raise ValueError( "--use-replayssm supports prefix caching only in align mode; " diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index fa0eb25758..5322f899dd 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -74,11 +74,11 @@ class MambaBase(AttentionLayerBase): mamba_type=self.mamba_type, tp_replicated=self.is_kv_cache_tp_replicated, mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode, - # RecoverSSM verifies the whole window off one checkpoint, so it + # ReplaySSM verifies the whole window off one checkpoint, so it # never writes the baseline's per-draft-token state slots. num_speculative_blocks=( 0 - if vllm_config.cache_config.use_kda_recoverssm + if vllm_config.cache_config.use_replayssm else vllm_config.num_speculative_tokens ), ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 812cab0212..fc58aed91c 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -43,6 +43,7 @@ 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 ( + commit_replayssm_ring_trackers, reset_replayssm_ring_trackers, selective_state_update, selective_state_update_replayssm_flashinfer, @@ -69,6 +70,13 @@ from vllm.v1.kv_cache_interface import KVCacheGroupSpec logger = init_logger(__name__) + +def _view_mtp_decode_tensor( + tensor: torch.Tensor, decode_batch: int, spec_query_len: int +) -> torch.Tensor: + return tensor.view(decode_batch, spec_query_len, *tensor.shape[1:]) + + # Added by the IBM Team, 2024 @@ -536,6 +544,8 @@ class MambaMixer2(MambaBase, PluggableLayer): 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._replayssm_prev_query_len = torch.empty(0, dtype=torch.int32) + self._commits_replayssm_trackers = True self._updates_replayssm_trackers = True self.num_spec = vllm_config.num_speculative_tokens @@ -724,7 +734,8 @@ 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 + use_spec_decode = self.num_spec > 0 + ring_start = prev_num_accepted = prev_query_len = None attn_metadata: AttentionMetadata | None = None if attn_metadata_raw is not None: @@ -746,6 +757,7 @@ class MambaMixer2(MambaBase, PluggableLayer): if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: ring_start = self._replayssm_ring_start prev_num_accepted = self._replayssm_prev_num_accepted + prev_query_len = self._replayssm_prev_query_len else: x_cache = dt_cache = B_cache = None has_initial_states_p = attn_metadata.has_initial_states_p @@ -757,6 +769,7 @@ class MambaMixer2(MambaBase, PluggableLayer): last_chunk_indices_p = attn_metadata.last_chunk_indices_p state_indices_tensor_p = attn_metadata.state_indices_tensor_p state_indices_tensor_d = attn_metadata.state_indices_tensor_d + replayssm_state_indices_d = attn_metadata.replayssm_state_indices_d num_accepted_tokens = attn_metadata.num_accepted_tokens query_start_loc_d = attn_metadata.query_start_loc_d num_decodes = attn_metadata.num_decodes @@ -996,6 +1009,7 @@ class MambaMixer2(MambaBase, PluggableLayer): reset_replayssm_ring_trackers( ring_start, prev_num_accepted, + prev_query_len, state_indices_tensor_p, ) @@ -1003,7 +1017,7 @@ class MambaMixer2(MambaBase, PluggableLayer): if has_decode: assert state_indices_tensor_d is not None if is_mamba_cache_all: - if self.num_spec > 0: + if use_spec_decode: assert block_idx_last_scheduled_token_prev_step_d is not None input_indices = ( block_idx_last_scheduled_token_prev_step_d.unsqueeze(1) @@ -1031,6 +1045,34 @@ class MambaMixer2(MambaBase, PluggableLayer): state_indices_tensor_d_input = state_indices_tensor_d state_indices_tensor_d_output = state_indices_tensor_d + if ( + self.use_replayssm + and self.mamba_config.backend == MambaBackendEnum.FLASHINFER + and use_spec_decode + and self._commits_replayssm_trackers + ): + assert ring_start is not None + assert prev_num_accepted is not None + assert prev_query_len is not None + assert replayssm_state_indices_d is not None + assert num_accepted_tokens is not None + assert query_start_loc_d is not None + assert x_cache is not None + assert self.replayssm_buffer_len is not None + # The previous MTP query's accepted prefix is known only now. + # Commit it once, before this cache group's layers evaluate the + # current speculative query with the shared ring positions. + commit_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + prev_query_len, + replayssm_state_indices_d, + num_accepted_tokens, + query_start_loc_d, + logical_window=self.replayssm_buffer_len, + ring_buffer_len=x_cache.size(2), + ) + # 2. Convolution sequence transformation hidden_states_B_C_d = causal_conv1d_update( hidden_states_B_C_d, @@ -1043,7 +1085,13 @@ class MambaMixer2(MambaBase, PluggableLayer): initial_state_idx=block_idx_last_computed_token_d, num_accepted_tokens=num_accepted_tokens, query_start_loc=query_start_loc_d, - max_query_len=state_indices_tensor_d.size(-1), + # ReplaySSM keeps one physical state block while a speculative + # decode call still processes the full target + draft window. + max_query_len=( + 1 + self.num_spec + if self.use_replayssm and use_spec_decode + else state_indices_tensor_d.size(-1) + ), ) hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn( @@ -1079,7 +1127,35 @@ class MambaMixer2(MambaBase, PluggableLayer): if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: assert ring_start is not None assert prev_num_accepted is not None + assert prev_query_len is not None assert attn_metadata.replayssm_scratch is not None + fi_cu_seqlens = query_start_loc_d + fi_max_seqlen = None + if use_spec_decode: + spec_query_len = 1 + self.num_spec + fi_max_seqlen = spec_query_len + assert replayssm_state_indices_d is not None + decode_batch = replayssm_state_indices_d.size(0) + if num_decode_tokens == decode_batch * spec_query_len: + hidden_states_d = _view_mtp_decode_tensor( + hidden_states_d, decode_batch, spec_query_len + ) + dt_d = _view_mtp_decode_tensor( + dt_d, decode_batch, spec_query_len + ) + B_d = _view_mtp_decode_tensor( + B_d, decode_batch, spec_query_len + ) + C_d = _view_mtp_decode_tensor( + C_d, decode_batch, spec_query_len + ) + preallocated_ssm_out_d = _view_mtp_decode_tensor( + preallocated_ssm_out_d, + decode_batch, + spec_query_len, + ) + fi_cu_seqlens = None + fi_max_seqlen = None selective_state_update_replayssm_flashinfer( ssm_state, hidden_states_d, @@ -1093,19 +1169,25 @@ class MambaMixer2(MambaBase, PluggableLayer): dt_cache, ring_start, prev_num_accepted, + prev_query_len, logical_window=self.replayssm_buffer_len, D=D_d, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=state_indices_tensor_d_input, + state_batch_indices=replayssm_state_indices_d, scratch=attn_metadata.replayssm_scratch, - update_trackers=self._updates_replayssm_trackers, + update_trackers=( + self._updates_replayssm_trackers and not use_spec_decode + ), enable_stochastic_rounding=( self.mamba_config.enable_stochastic_rounding ), stochastic_rounding_philox_rounds=( self.mamba_config.stochastic_rounding_philox_rounds ), + cu_seqlens=fi_cu_seqlens, + max_seqlen=fi_max_seqlen, + enable_pdl=False, ) else: selective_state_update_replayssm_output_only( @@ -1187,6 +1269,7 @@ class MambaMixer2(MambaBase, PluggableLayer): tp_world_size=tp_world_size, logical_window=self.replayssm_buffer_len, backend=self.mamba_config.backend, + num_speculative_tokens=self.num_spec, ) return base_shape @@ -1204,9 +1287,10 @@ def share_replayssm_ring_trackers( 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. + block indices to the same request and must keep separate cursor tensors. For + speculative decode, the first local layer commits the preceding acceptance; + for standard decode, the final local layer advances after all layers consume + the previous values. """ replayssm_mixers: dict[str, MambaMixer2] = {} for layer_name in ordered_layer_names: @@ -1232,6 +1316,7 @@ def share_replayssm_ring_trackers( groups_by_namespace.setdefault(namespace, []).append(layer_name) for group_layer_names in groups_by_namespace.values(): + first_layer_name = group_layer_names[0] last_layer_name = group_layer_names[-1] first_mixer = replayssm_mixers[group_layer_names[0]] @@ -1246,11 +1331,17 @@ def share_replayssm_ring_trackers( ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device) prev_num_accepted = torch.zeros_like(ring_start) + prev_query_len = 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 + mixer._replayssm_prev_query_len = prev_query_len + mixer._commits_replayssm_trackers = False + mixer._updates_replayssm_trackers = False + + replayssm_mixers[first_layer_name]._commits_replayssm_trackers = True + replayssm_mixers[last_layer_name]._updates_replayssm_trackers = True def mamba_mixer2( diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 8c89488d2d..31e89ffe55 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -227,6 +227,7 @@ class MambaStateShapeCalculator: tp_world_size: int, logical_window: int, backend: MambaBackendEnum, + num_speculative_tokens: int = 0, ) -> tuple[ tuple[int, int], tuple[int, int, int], @@ -241,8 +242,8 @@ class MambaStateShapeCalculator: """ ring_buffer_len = logical_window if backend == MambaBackendEnum.FLASHINFER: - # FlashInfer keeps the live window and appended token together. - ring_buffer_len += 1 + # FlashInfer keeps the live window and current verify window together. + ring_buffer_len += 1 + num_speculative_tokens local_nheads, head_dim, state_size = base_shapes[1] local_ngroups = divide(n_groups, tp_world_size) return ( diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index b622924327..b6e9108a94 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -31,6 +31,7 @@ logger = init_logger(__name__) def _update_replayssm_ring_trackers_kernel( ring_start, prev_num_accepted, + prev_query_len, state_batch_indices, state_batch_indices_stride, n_slots, @@ -52,6 +53,9 @@ def _update_replayssm_ring_trackers_kernel( if RESET: tl.store(ring_start + slots, 0, mask=valid) tl.store(prev_num_accepted + slots, 0, mask=valid) + # MTP alone reads prev_query_len. Standard decode deliberately leaves + # it unchanged; reset clears the pending MTP window after prefill. + tl.store(prev_query_len + 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) @@ -66,9 +70,70 @@ def _update_replayssm_ring_trackers_kernel( tl.store(prev_num_accepted + slots, next_prev, mask=valid) +@triton.jit( + do_not_specialize=["n_slots", "state_batch_indices_stride"], + do_not_specialize_on_alignment=["state_batch_indices"], +) +def _commit_replayssm_ring_trackers_kernel( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + num_accepted_tokens, + query_start_loc, + state_batch_indices_stride, + n_slots, + num_states, + logical_window: tl.constexpr, + ring_buffer_len: tl.constexpr, + pad_slot_id: 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) + accepted_since_checkpoint = tl.load(prev_num_accepted + slots, mask=valid, other=0) + current_ring_start = tl.load(ring_start + slots, mask=valid, other=0) + previous_speculative_query_len = tl.load( + prev_query_len + slots, mask=valid, other=0 + ) + accepted_from_previous_query = tl.load( + num_accepted_tokens + offsets, mask=mask, other=0 + ) + must_checkpoint = (previous_speculative_query_len > 0) & ( + accepted_since_checkpoint + previous_speculative_query_len > logical_window + ) + next_start = tl.where( + must_checkpoint, + (current_ring_start + accepted_since_checkpoint) % ring_buffer_len, + current_ring_start, + ) + next_prev = tl.where( + previous_speculative_query_len == 0, + 0, + tl.where( + must_checkpoint, + accepted_from_previous_query, + accepted_since_checkpoint + accepted_from_previous_query, + ), + ) + current_speculative_query_len = tl.load( + query_start_loc + offsets + 1, mask=mask, other=0 + ) - tl.load(query_start_loc + offsets, mask=mask, other=0) + tl.store(ring_start + slots, next_start, mask=valid) + tl.store(prev_num_accepted + slots, next_prev, mask=valid) + tl.store(prev_query_len + slots, current_speculative_query_len, mask=valid) + + def update_replayssm_ring_trackers( ring_start: torch.Tensor, prev_num_accepted: torch.Tensor, + prev_query_len: torch.Tensor, state_batch_indices: torch.Tensor, logical_window: int | None = None, ring_buffer_len: int | None = None, @@ -90,10 +155,15 @@ def update_replayssm_ring_trackers( _update_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)]( ring_start, prev_num_accepted, + prev_query_len, state_batch_indices, state_batch_indices.stride(0), n_slots, - min(ring_start.numel(), prev_num_accepted.numel()), + min( + ring_start.numel(), + prev_num_accepted.numel(), + prev_query_len.numel(), + ), logical_window, ring_buffer_len, pad_slot_id, @@ -105,6 +175,7 @@ def update_replayssm_ring_trackers( def reset_replayssm_ring_trackers( ring_start: torch.Tensor, prev_num_accepted: torch.Tensor, + prev_query_len: torch.Tensor, state_batch_indices: torch.Tensor, pad_slot_id: int = NULL_BLOCK_ID, ) -> None: @@ -112,11 +183,62 @@ def reset_replayssm_ring_trackers( update_replayssm_ring_trackers( ring_start, prev_num_accepted, + prev_query_len, state_batch_indices, pad_slot_id=pad_slot_id, ) +def commit_replayssm_ring_trackers( + ring_start: torch.Tensor, + prev_num_accepted: torch.Tensor, + prev_query_len: torch.Tensor, + state_batch_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + query_start_loc: torch.Tensor, + logical_window: int, + ring_buffer_len: int, + pad_slot_id: int = NULL_BLOCK_ID, +) -> None: + """Commit the preceding speculative window and record the current one. + + MTP evaluates a target token and its draft tokens together, but the number + accepted from that query is available only on the next forward pass. This + function then advances the per-request ring by the accepted prefix and + records the current query length for the following pass. A zero + ``prev_query_len`` means that reset/prefill left no prior MTP query to + commit; slot validity is handled independently by the kernel mask. + + Standard single-token decode needs no delayed commit: its ReplaySSM kernel + advances the shared ring trackers directly after every token. + """ + if state_batch_indices.dim() > 1: + state_batch_indices = state_batch_indices[:, 0] + n_slots = state_batch_indices.numel() + if n_slots == 0: + return + block = 128 + _commit_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)]( + ring_start, + prev_num_accepted, + prev_query_len, + state_batch_indices, + num_accepted_tokens, + query_start_loc, + state_batch_indices.stride(0), + n_slots, + min( + ring_start.numel(), + prev_num_accepted.numel(), + prev_query_len.numel(), + ), + logical_window, + ring_buffer_len, + pad_slot_id, + BLOCK=block, + ) + + class MambaSSUBackend(ABC): """Abstract base class for Mamba SSU backends.""" @@ -385,6 +507,7 @@ def selective_state_update_replayssm_flashinfer( dt_cache: torch.Tensor, ring_start: torch.Tensor, prev_num_accepted_tokens: torch.Tensor, + prev_query_len: torch.Tensor, logical_window: int, D: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, @@ -395,6 +518,9 @@ def selective_state_update_replayssm_flashinfer( update_trackers: bool = True, enable_stochastic_rounding: bool = False, stochastic_rounding_philox_rounds: int = 0, + cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, + enable_pdl: bool = False, ) -> torch.Tensor: """Run FlashInfer checkpointing SSU and optionally advance shared trackers.""" if _flashinfer_replayssm_kernel is None: @@ -404,11 +530,12 @@ def selective_state_update_replayssm_flashinfer( ) if x.dim() == 3: - x = x.unsqueeze(1) - dt = dt.unsqueeze(1) - B = B.unsqueeze(1) - C = C.unsqueeze(1) - out = out.unsqueeze(1) + dim = 0 if cu_seqlens is not None else 1 + x = x.unsqueeze(dim) + dt = dt.unsqueeze(dim) + B = B.unsqueeze(dim) + C = C.unsqueeze(dim) + out = out.unsqueeze(dim) indices = state_batch_indices if indices is not None and indices.dim() > 1: @@ -443,6 +570,9 @@ def selective_state_update_replayssm_flashinfer( pad_slot_id=null_block_id, rand_seed=rand_seed, philox_rounds=stochastic_rounding_philox_rounds or 10, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + enable_pdl=enable_pdl, cb_scaled=cb_scaled, cumAdt_vec=cumAdt_vec, cb_old=cb_old, @@ -451,6 +581,7 @@ def selective_state_update_replayssm_flashinfer( update_replayssm_ring_trackers( ring_start, prev_num_accepted_tokens, + prev_query_len, indices, logical_window=logical_window, ring_buffer_len=x_cache.size(2), diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 3f97deebe4..bef3da5cc5 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -815,6 +815,7 @@ class NemotronHForCausalLM( tp_world_size=parallel_config.tensor_parallel_size, logical_window=cache_config.replayssm_buffer_len, backend=vllm_config.mamba_config.backend, + num_speculative_tokens=vllm_config.num_speculative_tokens, ) return base_shape diff --git a/vllm/model_executor/warmup/replayssm_warmup.py b/vllm/model_executor/warmup/replayssm_warmup.py index 60b3eced20..28c45804ca 100644 --- a/vllm/model_executor/warmup/replayssm_warmup.py +++ b/vllm/model_executor/warmup/replayssm_warmup.py @@ -21,6 +21,7 @@ logger = init_logger(__name__) def _replayssm_autotune_kwargs( runner: "GPUModelRunner", + max_token_prefill_kwargs: dict[str, Any], ) -> tuple[int, dict[str, Any]] | None: config = runner.vllm_config if not ( @@ -46,11 +47,16 @@ def _replayssm_autotune_kwargs( runner.max_num_tokens // query_len, runner.kv_cache_config.num_blocks - 1, ) + if max_num_reqs <= 0: + logger.warning_once( + "Skipping FlashInfer ReplaySSM autotuning because no non-padding " + "state slot is available." + ) + return None + decode_kwargs = { + **max_token_prefill_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: @@ -75,18 +81,22 @@ def _temporary_replayssm_autotune_state( ) reset_tensors: dict[int, torch.Tensor] = {} - tracker_specs: dict[int, tuple[torch.Tensor, torch.Tensor, int, int]] = {} + tracker_specs: dict[ + int, tuple[torch.Tensor, 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 + prev_query_len = module._replayssm_prev_query_len tracker_specs.setdefault( ring_start.data_ptr(), ( ring_start, prev_num_accepted, + prev_query_len, module.replayssm_buffer_len, module.kv_cache[2].size(2), ), @@ -95,6 +105,7 @@ def _temporary_replayssm_autotune_state( *module.kv_cache, ring_start, prev_num_accepted, + prev_query_len, ) for tensor in tensors: if tensor.numel(): @@ -121,20 +132,26 @@ def _temporary_replayssm_autotune_state( for ( ring_start, prev_num_accepted, + prev_query_len, 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) + reset_replayssm_ring_trackers( + ring_start, prev_num_accepted, prev_query_len, state_slots + ) update_replayssm_ring_trackers( ring_start, prev_num_accepted, + prev_query_len, state_slots, logical_window, ring_buffer_len, ) - reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots) + reset_replayssm_ring_trackers( + ring_start, prev_num_accepted, prev_query_len, state_slots + ) try: yield @@ -151,7 +168,13 @@ def _temporary_replayssm_autotune_state( def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None: - autotune = _replayssm_autotune_kwargs(runner) + max_token_prefill_kwargs = { + "num_tokens": runner.scheduler_config.max_num_batched_tokens, + "skip_eplb": True, + "is_profile": True, + "randomize_inputs": True, + } + autotune = _replayssm_autotune_kwargs(runner, max_token_prefill_kwargs) if autotune is None: return max_num_reqs, decode_kwargs = autotune diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index c831266109..7c543570aa 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -87,6 +87,8 @@ class BaseMambaAttentionMetadata: 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 + # Contiguous cache-slot indices shared by all FlashInfer ReplaySSM layers. + replayssm_state_indices_d: torch.Tensor | None = None class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): @@ -182,6 +184,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): self.decode_replayssm_scratch: ( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None ) = None + self.decode_replayssm_state_indices_d: 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( @@ -218,11 +221,17 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): self.decode_replayssm_scratch = allocate_checkpointing_ssu_scratch( batch_size=scheduler_config.max_num_seqs, num_heads=nheads, - num_predicted_tokens=1, + num_predicted_tokens=1 + self.num_spec_tokens, max_window=self.replayssm_buffer_len, dtype=vllm_config.model_config.dtype, device=device, ) + # Full CUDA graphs retain capture-time tensor addresses. Keep the + # contiguous first-column view used by FlashInfer in a persistent + # buffer and refresh its contents before each replay. + self.decode_replayssm_state_indices_d = torch.empty( + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device + ) self._init_reorder_batch_threshold(1, self.use_spec_decode) if self.use_spec_decode: @@ -592,13 +601,16 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): ] state_indices_tensor_p = state_indices_tensor_p[:, 0] - # Sometimes even with specdec enabled we get single-token prefill chunks that - # should be treated as decodes but don't have num_accepted_tokens set. - # These should be fine to process as non-spec decodes since there's only - # one token, so no risk of placing accepted tokens in the wrong slot. - if num_decodes > 0 and self.use_spec_decode and num_accepted_tokens is not None: + if num_decodes > 0 and self.use_spec_decode: query_start_loc_d = common_attn_metadata.query_start_loc[: num_decodes + 1] - num_accepted_tokens = num_accepted_tokens[:num_decodes] + if num_accepted_tokens is None: + # Single-token prefill chunks can be reclassified as decodes before + # speculative decoding has produced acceptance counts. Treat each + # token as accepted so recurrent state and ReplaySSM trackers follow + # the normal speculative-decode path. + num_accepted_tokens = torch.diff(query_start_loc_d) + else: + num_accepted_tokens = num_accepted_tokens[:num_decodes] if num_prefills > 0: if num_computed_tokens is None: @@ -788,6 +800,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): is_flush_d = metadata.is_flush_d bc_pre_scratch = metadata.bc_pre_scratch replayssm_scratch = metadata.replayssm_scratch + replayssm_state_indices_d = None if ( metadata.num_prefills == 0 and metadata.num_decodes <= self.decode_cudagraph_max_bs @@ -876,6 +889,20 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): cumAdt_vec[:padded_bs], cb_old[:padded_bs], ) + assert self.decode_replayssm_state_indices_d is not None + self.decode_replayssm_state_indices_d[:padded_bs].copy_( + state_indices_tensor_d[:, 0], non_blocking=True + ) + replayssm_state_indices_d = self.decode_replayssm_state_indices_d[ + :padded_bs + ] + + if ( + self.use_flashinfer_replayssm + and state_indices_tensor_d is not None + and replayssm_state_indices_d is None + ): + replayssm_state_indices_d = state_indices_tensor_d[:, 0].contiguous() return replace( metadata, @@ -886,6 +913,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): is_flush_d=is_flush_d, bc_pre_scratch=bc_pre_scratch, replayssm_scratch=replayssm_scratch, + replayssm_state_indices_d=replayssm_state_indices_d, 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=(