From 069ba91903bdd58cf05b6a21c9674f586d3182c7 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Tue, 1 Sep 2026 16:56:57 +0200 Subject: [PATCH 01/53] [Mamba] Add FlashInfer ReplaySSM support for MTP Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 420 +++++++++++++++++- tests/model_executor/test_replayssm_warmup.py | 55 ++- tests/test_config.py | 40 +- .../test_attention_backends_selection.py | 25 ++ .../test_replayssm_metadata_builder.py | 88 +++- tests/v1/e2e/test_replayssm_decode.py | 112 ++++- .../worker/test_kv_cache_allocation_scope.py | 4 +- tests/v1/worker/test_utils.py | 26 +- vllm/config/cache.py | 11 +- vllm/config/vllm.py | 47 +- vllm/model_executor/layers/mamba/abstract.py | 4 +- .../layers/mamba/mamba_mixer2.py | 100 ++++- .../layers/mamba/mamba_utils.py | 5 +- .../layers/mamba/ops/ssu_dispatch.py | 145 +++++- vllm/model_executor/models/nemotron_h.py | 1 + .../model_executor/warmup/replayssm_warmup.py | 37 +- vllm/v1/attention/backends/mamba_attn.py | 42 +- 17 files changed, 1000 insertions(+), 162 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 6226340a9936..7e5d3aa1bb53 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -1,24 +1,31 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from importlib import import_module +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 +39,13 @@ 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 +61,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 +69,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 +85,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 +443,211 @@ 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( + ("query_start_loc", "expected_shape", "expected_max_seqlen"), + [ + 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_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.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="compatible flashinfer checkpointing_ssu not available", +) +def test_replayssm_flashinfer_backend_rejects_missing_mtp_api(monkeypatch): + checkpointing_ssu_module = import_module("flashinfer.mamba.checkpointing_ssu") + + def legacy_checkpointing_ssu(): + pass + + monkeypatch.setattr( + checkpointing_ssu_module, "checkpointing_ssu", legacy_checkpointing_ssu + ) + with pytest.raises(ImportError, match="native MTP and PDL support"): + initialize_mamba_ssu_backend( + MambaConfig(backend=MambaBackendEnum.FLASHINFER), + _kv_cache_config_with_ssu(), + use_replayssm=True, + ) + + @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): +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 +656,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 272fa2cf8528..f26f2c1b160c 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -11,13 +11,6 @@ from vllm.config.mamba import MambaBackendEnum from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.model_executor.warmup import replayssm_warmup as warmup -from vllm.platforms import current_platform -from vllm.utils.flashinfer import has_flashinfer - -pytestmark = pytest.mark.skipif( - not current_platform.is_cuda() or not has_flashinfer(), - reason="FlashInfer ReplaySSM warmup tests require CUDA and FlashInfer", -) PREFILL_KWARGS = { "num_tokens": 128, @@ -67,22 +60,22 @@ 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 + ) - expected_kwargs = { - **PREFILL_KWARGS, - "num_tokens": expected_num_reqs * query_len, - "uniform_decode": True, - } + assert result is not None + max_num_reqs, decode_kwargs = result + assert max_num_reqs == expected_num_reqs + assert decode_kwargs["num_tokens"] == expected_num_reqs * query_len + assert decode_kwargs["uniform_decode"] is True + assert decode_kwargs["is_profile"] is True if runner_kwargs.get("use_v2_model_runner"): - expected_kwargs["valid_dummy_state_slots"] = True + assert decode_kwargs["valid_dummy_state_slots"] is True + assert "profile_seq_lens" not in decode_kwargs else: - expected_kwargs.update( - allow_microbatching=False, - force_attention=True, - profile_seq_lens=query_len + 1, - ) - assert result == (expected_num_reqs, expected_kwargs) + assert decode_kwargs["profile_seq_lens"] == query_len + 1 + assert decode_kwargs["force_attention"] is True @pytest.mark.parametrize( @@ -92,7 +85,11 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): (dict(backend=MambaBackendEnum.TRITON), True), ({}, False), ], - ids=["replayssm_disabled", "non_flashinfer_backend", "kernel_unavailable"], + ids=[ + "replayssm_disabled", + "non_flashinfer_backend", + "kernel_unavailable", + ], ) def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): with patch.object( @@ -100,7 +97,19 @@ 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 + + +def test_replayssm_autotune_kwargs_skipped_without_non_padding_slot(): + with patch.object( + warmup, "flashinfer_replayssm_autotune_supported", return_value=True + ): + result = warmup._replayssm_autotune_kwargs( + _autotune_runner(num_blocks=1), PREFILL_KWARGS + ) assert result is None @@ -116,10 +125,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 aa0c838e3f87..2d9898e85a09 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -50,6 +50,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( @@ -66,28 +67,12 @@ def test_kda_recoverssm_derivation_is_revalidated(): ) 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) + config.mamba_config.backend = MambaBackendEnum.FLASHINFER + VllmConfig.validate_mamba_cached_kernel(config) + assert not config.cache_config.use_kda_recoverssm def test_per_request_spec_decode_metrics_requires_spec_decode(): @@ -2152,23 +2137,6 @@ def test_draft_sample_method_probabilistic_is_accepted(): assert speculative_config.draft_sample_method == "probabilistic" -@pytest.mark.parametrize("disable_eagle_block_drop", [False, True]) -def test_eagle_block_drop_can_be_disabled_without_disabling_eagle( - disable_eagle_block_drop: bool, -): - # Start from an ngram config to avoid loading model metadata: these predicates - # depend only on the speculative method and the new switch. - speculative_config = SpeculativeConfig( - method="ngram", - num_speculative_tokens=3, - disable_eagle_block_drop=disable_eagle_block_drop, - ) - speculative_config.method = "eagle3" - - assert speculative_config.use_eagle() - assert speculative_config.use_eagle_block_drop() is not disable_eagle_block_drop - - def test_draft_sample_method_gumbel_is_rejected(): with pytest.raises(ValidationError): SpeculativeConfig( diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index 8486d216a125..420fd5725964 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,29 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend +def test_replayssm_does_not_reserve_speculative_state_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=True, + ), + num_speculative_tokens=3, + ) + + spec = MambaBase.get_kv_cache_spec(layer, vllm_config) + + assert spec is not None + assert spec.num_speculative_blocks == 0 + + @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 89cb0f566467..e1e2875e2539 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -16,6 +16,8 @@ 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 @@ class ReplaySSMBuildCase: 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,30 @@ 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: + 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 +285,56 @@ 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(): + 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 2e66a4668924..1b61fd1f878e 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -12,6 +12,7 @@ # Mamba2 (Nemotron-3) hybrid. MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" +MAMBA2_MTP_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4" MODELS = [ pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), ] @@ -21,6 +22,13 @@ "Once upon a time, in a small village,", ] +try: + from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner + + HAS_FLASHINFER_CHECKPOINTING_SSU = CheckpointingSSURunner is not None +except ImportError: + HAS_FLASHINFER_CHECKPOINTING_SSU = False + def _check_replayssm_parity( vllm_runner, @@ -30,16 +38,10 @@ def _check_replayssm_parity( mamba_backend: str = "triton", name_1: str = "replayssm", require_v2: bool = False, - monkeypatch: pytest.MonkeyPatch | None = None, ): # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are # common-mode and only ReplaySSM varies. - if require_v2: - assert monkeypatch is not None - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - envs.disable_envs_cache() - common = dict( max_model_len=1024, trust_remote_code=True, @@ -80,19 +82,103 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2) +@pytest.mark.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="flashinfer.mamba.checkpointing_ssu not available", +) @pytest.mark.parametrize("model_name", MODELS) def test_replayssm_flashinfer_decode_matches_baseline_v2( vllm_runner, model_name, monkeypatch ): - pytest.importorskip("flashinfer.mamba.checkpointing_ssu") - _check_replayssm_parity( - vllm_runner, - model_name, + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + envs.disable_envs_cache() + _check_replayssm_parity( + vllm_runner, + model_name, + mamba_backend="flashinfer", + name_1="replayssm_flashinfer_v2", + require_v2=True, + ) + finally: + # The context restores the environment before the final cache reset. + envs.disable_envs_cache() + + +@pytest.mark.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="flashinfer.mamba.checkpointing_ssu not available", +) +@pytest.mark.parametrize("model_name", 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", + mamba_backend="flashinfer", + speculative_config={ + "method": "ngram", + "num_speculative_tokens": 3, + "prompt_lookup_max": 3, + }, + ) + with vllm_runner(model_name, **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, **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", + ) + + +@pytest.mark.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="flashinfer.mamba.checkpointing_ssu not available", +) +@large_gpu_mark(min_gb=40) +def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", mamba_backend="flashinfer", - name_1="replayssm_flashinfer_v2", - require_v2=True, - monkeypatch=monkeypatch, + 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", "1") + envs.disable_envs_cache() + with vllm_runner( + MAMBA2_MTP_MODEL, + use_replayssm=True, + replayssm_buffer_len=16, + **common, + ) as llm: + assert llm.llm.llm_engine.vllm_config.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_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index a9579f1af69b..1818f173d978 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -48,7 +48,7 @@ def bind(*args, **kwargs): result = attn_utils.init_kv_cache( [], {}, - SimpleNamespace(kv_cache_groups=[]), + object(), torch.device("cpu"), [], config, @@ -83,7 +83,7 @@ def bind(*args, **kwargs): ) result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors( runner, - SimpleNamespace(kv_cache_groups=[]), + object(), [], kv_cache_allocation_context=scope, ) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index a4474c8a5843..a5589c5048ba 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -17,6 +17,8 @@ def __init__(self): 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, ...], ...]: @@ -47,19 +49,19 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) - assert ( - mixers[0]._replayssm_ring_start.data_ptr() - == mixers[2]._replayssm_ring_start.data_ptr() - ) - assert ( - mixers[0]._replayssm_prev_num_accepted.data_ptr() - == mixers[2]._replayssm_prev_num_accepted.data_ptr() - ) - assert ( - mixers[1]._replayssm_ring_start.data_ptr() - != mixers[0]._replayssm_ring_start.data_ptr() + tracker_names = ( + "_replayssm_ring_start", + "_replayssm_prev_num_accepted", + "_replayssm_prev_query_len", ) - # 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/vllm/config/cache.py b/vllm/config/cache.py index 0ffd27c28157..6c6e8f4fe636 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -197,15 +197,16 @@ class CacheConfig: """ 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 7fb53cdc24e6..e48a9178fe02 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2855,19 +2855,35 @@ def validate_mamba_cached_kernel(self) -> "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 " @@ -2891,6 +2907,25 @@ def validate_mamba_cached_kernel(self) -> "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 10ed7d05fc8f..1dd942221b87 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -76,11 +76,11 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: 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 0de22ead8339..29edb4751cd1 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -43,6 +43,7 @@ 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, @@ -529,6 +530,8 @@ def __init__( 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 @@ -717,7 +720,7 @@ def conv_ssm_forward( 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 + ring_start = prev_num_accepted = prev_query_len = None attn_metadata: AttentionMetadata | None = None if attn_metadata_raw is not None: @@ -739,6 +742,7 @@ def conv_ssm_forward( 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 @@ -750,6 +754,7 @@ def conv_ssm_forward( 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 @@ -1002,6 +1007,7 @@ def conv_ssm_forward( reset_replayssm_ring_trackers( ring_start, prev_num_accepted, + prev_query_len, state_indices_tensor_p, ) @@ -1037,6 +1043,31 @@ def conv_ssm_forward( 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 self.num_spec > 0 + 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 + 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, @@ -1049,7 +1080,13 @@ def conv_ssm_forward( 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 self.num_spec > 0 + else state_indices_tensor_d.size(-1) + ), ) hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn( @@ -1085,33 +1122,61 @@ def conv_ssm_forward( 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_x = hidden_states_d + fi_dt = dt_d + fi_B = B_d + fi_C = C_d + fi_out = preallocated_ssm_out_d + fi_cu_seqlens = query_start_loc_d + fi_max_seqlen = None + if self.num_spec > 0: + 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: + fi_shape = (decode_batch, spec_query_len) + fi_x = fi_x.view(*fi_shape, *fi_x.shape[1:]) + fi_dt = fi_dt.view(*fi_shape, *fi_dt.shape[1:]) + fi_B = fi_B.view(*fi_shape, *fi_B.shape[1:]) + fi_C = fi_C.view(*fi_shape, *fi_C.shape[1:]) + fi_out = fi_out.view(*fi_shape, *fi_out.shape[1:]) + fi_cu_seqlens = None + fi_max_seqlen = None selective_state_update_replayssm_flashinfer( ssm_state, - hidden_states_d, - dt_d, + fi_x, + fi_dt, A_d, - B_d, - C_d, - preallocated_ssm_out_d, + fi_B, + fi_C, + fi_out, x_cache, B_cache, 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 self.num_spec == 0 + ), 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( @@ -1193,6 +1258,7 @@ def get_state_shape(self) -> tuple[tuple[int, ...], ...]: 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 @@ -1210,9 +1276,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] = {} @@ -1239,6 +1306,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]] @@ -1253,11 +1321,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 73b1d934e825..9a91173abd84 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -217,6 +217,7 @@ def append_replayssm_ring( tp_world_size: int, logical_window: int, backend: MambaBackendEnum, + num_speculative_tokens: int = 0, ) -> tuple[tuple[int, ...], ...]: """Append the physical ReplaySSM ring shapes. @@ -225,8 +226,8 @@ def append_replayssm_ring( """ 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 62e4307b37b7..722d385d7162 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -12,6 +12,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable from functools import cache +from inspect import signature import torch @@ -32,6 +33,7 @@ def _update_replayssm_ring_trackers_kernel( ring_start, prev_num_accepted, + prev_query_len, state_batch_indices, state_batch_indices_stride, n_slots, @@ -53,6 +55,7 @@ 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) + 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) @@ -67,9 +70,62 @@ 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) + prev = tl.load(prev_num_accepted + slots, mask=valid, other=0) + start = tl.load(ring_start + slots, mask=valid, other=0) + previous_query_len = tl.load(prev_query_len + slots, mask=valid, other=0) + accepted = tl.load(num_accepted_tokens + offsets, mask=mask, other=0) + must_checkpoint = (previous_query_len > 0) & ( + prev + previous_query_len > logical_window + ) + next_start = tl.where( + must_checkpoint, + (start + prev) % ring_buffer_len, + start, + ) + next_prev = tl.where( + previous_query_len == 0, + 0, + tl.where(must_checkpoint, accepted, prev + accepted), + ) + current_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_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, @@ -91,10 +147,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, @@ -106,6 +167,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: @@ -113,11 +175,51 @@ 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.""" + 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.""" @@ -365,12 +467,10 @@ def __call__( def flashinfer_replayssm_autotune_supported() -> bool: """Return True when FlashInfer exposes ReplaySSM autotuning.""" try: - from flashinfer.mamba.checkpointing_ssu import ( # noqa: F401 - CheckpointingSSURunner, - ) + from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner except ImportError: return False - return True + return callable(CheckpointingSSURunner) def selective_state_update_replayssm_flashinfer( @@ -386,6 +486,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, @@ -396,6 +497,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: @@ -405,11 +509,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: @@ -444,6 +549,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, @@ -452,6 +560,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), @@ -507,11 +616,25 @@ def initialize_mamba_ssu_backend( _flashinfer_replayssm_kernel = None if use_replayssm and backend == MambaBackendEnum.FLASHINFER: try: - from flashinfer.mamba.checkpointing_ssu import checkpointing_ssu + from flashinfer.mamba.checkpointing_ssu import ( + CheckpointingSSURunner, + checkpointing_ssu, + ) except ImportError as e: raise ImportError( "FlashInfer ReplaySSM requires a compatible flashinfer-python package" ) from e + if not callable(CheckpointingSSURunner): + raise ImportError("FlashInfer ReplaySSM requires native autotuning support") + required_parameters = {"cu_seqlens", "max_seqlen", "enable_pdl"} + missing_parameters = ( + required_parameters - signature(checkpointing_ssu).parameters.keys() + ) + if missing_parameters: + raise ImportError( + "FlashInfer ReplaySSM requires native MTP and PDL support; missing " + + ", ".join(sorted(missing_parameters)) + ) _flashinfer_replayssm_kernel = checkpointing_ssu if use_replayssm: logger.info("Using %s ReplaySSM backend.", backend.value) diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index e65e51f3bb9d..a3401dfcf522 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -790,6 +790,7 @@ def get_mamba_state_shape_from_config( 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 60b3eced2016..28c45804ca7e 100644 --- a/vllm/model_executor/warmup/replayssm_warmup.py +++ b/vllm/model_executor/warmup/replayssm_warmup.py @@ -21,6 +21,7 @@ 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 b1a0c2d0efbf..f7591c557eef 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -84,6 +84,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): @@ -178,6 +180,7 @@ def __init__( 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( @@ -214,11 +217,17 @@ def __init__( 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: @@ -569,13 +578,16 @@ def _compute_common_metadata( ] 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: @@ -743,6 +755,7 @@ def _update_metadata_for_cudagraph_capture( 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 @@ -831,6 +844,20 @@ def _update_metadata_for_cudagraph_capture( 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, @@ -841,6 +868,7 @@ def _update_metadata_for_cudagraph_capture( 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=( From f86b0fd6212988a610dac747c96ed556e9ae3d1f Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Tue, 1 Sep 2026 17:10:12 +0200 Subject: [PATCH 02/53] fix merge inconsistencies Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 22 ------------ tests/model_executor/test_replayssm_warmup.py | 12 ++----- tests/test_config.py | 34 +++++++++++++++++++ .../test_attention_backends_selection.py | 9 +++-- .../worker/test_kv_cache_allocation_scope.py | 4 +-- 5 files changed, 44 insertions(+), 37 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 7e5d3aa1bb53..082b61c43d2f 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from importlib import import_module from types import SimpleNamespace from unittest.mock import Mock @@ -616,27 +615,6 @@ def split_hidden_states_B_C(values): assert kernel.call_args.kwargs["max_seqlen"] == expected_max_seqlen -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="compatible flashinfer checkpointing_ssu not available", -) -def test_replayssm_flashinfer_backend_rejects_missing_mtp_api(monkeypatch): - checkpointing_ssu_module = import_module("flashinfer.mamba.checkpointing_ssu") - - def legacy_checkpointing_ssu(): - pass - - monkeypatch.setattr( - checkpointing_ssu_module, "checkpointing_ssu", legacy_checkpointing_ssu - ) - with pytest.raises(ImportError, match="native MTP and PDL support"): - initialize_mamba_ssu_backend( - MambaConfig(backend=MambaBackendEnum.FLASHINFER), - _kv_cache_config_with_ssu(), - use_replayssm=True, - ) - - @pytest.mark.parametrize( ("backend", "num_speculative_tokens", "expected_ring_len"), [ diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index f26f2c1b160c..a1a3daf167e4 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -84,11 +84,13 @@ 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", ], ) def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): @@ -103,16 +105,6 @@ def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): assert result is None -def test_replayssm_autotune_kwargs_skipped_without_non_padding_slot(): - with patch.object( - warmup, "flashinfer_replayssm_autotune_supported", return_value=True - ): - result = warmup._replayssm_autotune_kwargs( - _autotune_runner(num_blocks=1), PREFILL_KWARGS - ) - assert result is None - - def test_replayssm_autotune_slots_restore_state_and_trackers(): mixer = MambaMixer2.__new__(MambaMixer2) torch.nn.Module.__init__(mixer) diff --git a/tests/test_config.py b/tests/test_config.py index 2d9898e85a09..0fe3c365d1c7 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -67,13 +67,30 @@ def test_kda_recoverssm_derivation_is_revalidated(): ) 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" 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 + with pytest.raises(ValueError, match="pipeline_parallel_size=1"): + VllmConfig.validate_mamba_cached_kernel(config) + def test_per_request_spec_decode_metrics_requires_spec_decode(): # The flag only makes sense with speculative decoding configured; enabling @@ -2137,6 +2154,23 @@ def test_draft_sample_method_probabilistic_is_accepted(): assert speculative_config.draft_sample_method == "probabilistic" +@pytest.mark.parametrize("disable_eagle_block_drop", [False, True]) +def test_eagle_block_drop_can_be_disabled_without_disabling_eagle( + disable_eagle_block_drop: bool, +): + # Start from an ngram config to avoid loading model metadata: these predicates + # depend only on the speculative method and the new switch. + speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=3, + disable_eagle_block_drop=disable_eagle_block_drop, + ) + speculative_config.method = "eagle3" + + assert speculative_config.use_eagle() + assert speculative_config.use_eagle_block_drop() is not disable_eagle_block_drop + + def test_draft_sample_method_gumbel_is_rejected(): with pytest.raises(ValidationError): SpeculativeConfig( diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index 420fd5725964..2695517e1421 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/tests/v1/attention/test_attention_backends_selection.py @@ -21,7 +21,10 @@ from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend -def test_replayssm_does_not_reserve_speculative_state_blocks(): +@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,), @@ -33,7 +36,7 @@ def test_replayssm_does_not_reserve_speculative_state_blocks(): mamba_block_size=1, mamba_page_size_padded=None, mamba_cache_mode="none", - use_replayssm=True, + use_replayssm=use_replayssm, ), num_speculative_tokens=3, ) @@ -41,7 +44,7 @@ def test_replayssm_does_not_reserve_speculative_state_blocks(): spec = MambaBase.get_kv_cache_spec(layer, vllm_config) assert spec is not None - assert spec.num_speculative_blocks == 0 + assert spec.num_speculative_blocks == expected_blocks @pytest.mark.parametrize( diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 1818f173d978..a9579f1af69b 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -48,7 +48,7 @@ def bind(*args, **kwargs): result = attn_utils.init_kv_cache( [], {}, - object(), + SimpleNamespace(kv_cache_groups=[]), torch.device("cpu"), [], config, @@ -83,7 +83,7 @@ def bind(*args, **kwargs): ) result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors( runner, - object(), + SimpleNamespace(kv_cache_groups=[]), [], kv_cache_allocation_context=scope, ) From 829531ca79173b9a0be3f31929c73d0918532a81 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Tue, 1 Sep 2026 17:58:08 +0200 Subject: [PATCH 03/53] restore test skips Signed-off-by: Andrii Skliar --- tests/model_executor/test_replayssm_warmup.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index a1a3daf167e4..5db0f9a50d0f 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -11,6 +11,13 @@ from vllm.config.mamba import MambaBackendEnum from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.model_executor.warmup import replayssm_warmup as warmup +from vllm.platforms import current_platform +from vllm.utils.flashinfer import has_flashinfer + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda() or not has_flashinfer(), + reason="FlashInfer ReplaySSM warmup tests require CUDA and FlashInfer", +) PREFILL_KWARGS = { "num_tokens": 128, @@ -67,15 +74,20 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): assert result is not None max_num_reqs, decode_kwargs = result assert max_num_reqs == expected_num_reqs - assert decode_kwargs["num_tokens"] == expected_num_reqs * query_len - assert decode_kwargs["uniform_decode"] is True - assert decode_kwargs["is_profile"] is True + expected_kwargs = { + **PREFILL_KWARGS, + "num_tokens": expected_num_reqs * query_len, + "uniform_decode": True, + } if runner_kwargs.get("use_v2_model_runner"): - assert decode_kwargs["valid_dummy_state_slots"] is True - assert "profile_seq_lens" not in decode_kwargs + expected_kwargs["valid_dummy_state_slots"] = True else: - assert decode_kwargs["profile_seq_lens"] == query_len + 1 - assert decode_kwargs["force_attention"] is True + expected_kwargs.update( + allow_microbatching=False, + force_attention=True, + profile_seq_lens=query_len + 1, + ) + assert decode_kwargs == expected_kwargs @pytest.mark.parametrize( From 2a20d7903e7f71b44ae2acc92d044a0e5641eaec Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Tue, 1 Sep 2026 19:27:35 +0200 Subject: [PATCH 04/53] remove unnecessary code Signed-off-by: Andrii Skliar --- .../layers/mamba/ops/ssu_dispatch.py | 23 ++++--------------- 1 file changed, 5 insertions(+), 18 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 722d385d7162..e64653afcbaf 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -12,7 +12,6 @@ from abc import ABC, abstractmethod from collections.abc import Callable from functools import cache -from inspect import signature import torch @@ -467,10 +466,12 @@ def __call__( def flashinfer_replayssm_autotune_supported() -> bool: """Return True when FlashInfer exposes ReplaySSM autotuning.""" try: - from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner + from flashinfer.mamba.checkpointing_ssu import ( # noqa: F401 + CheckpointingSSURunner, + ) except ImportError: return False - return callable(CheckpointingSSURunner) + return True def selective_state_update_replayssm_flashinfer( @@ -616,25 +617,11 @@ def initialize_mamba_ssu_backend( _flashinfer_replayssm_kernel = None if use_replayssm and backend == MambaBackendEnum.FLASHINFER: try: - from flashinfer.mamba.checkpointing_ssu import ( - CheckpointingSSURunner, - checkpointing_ssu, - ) + from flashinfer.mamba.checkpointing_ssu import checkpointing_ssu except ImportError as e: raise ImportError( "FlashInfer ReplaySSM requires a compatible flashinfer-python package" ) from e - if not callable(CheckpointingSSURunner): - raise ImportError("FlashInfer ReplaySSM requires native autotuning support") - required_parameters = {"cu_seqlens", "max_seqlen", "enable_pdl"} - missing_parameters = ( - required_parameters - signature(checkpointing_ssu).parameters.keys() - ) - if missing_parameters: - raise ImportError( - "FlashInfer ReplaySSM requires native MTP and PDL support; missing " - + ", ".join(sorted(missing_parameters)) - ) _flashinfer_replayssm_kernel = checkpointing_ssu if use_replayssm: logger.info("Using %s ReplaySSM backend.", backend.value) From 0f7abb732d4cb62490815df073c8ed3de18363e4 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 15:02:47 +0200 Subject: [PATCH 05/53] Add ReplaySSM prefix caching on the MTP branch Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 766 ++++++++------- tests/model_executor/test_replayssm_warmup.py | 33 +- tests/test_config.py | 107 ++ .../test_mamba_update_block_table.py | 48 +- .../test_replayssm_metadata_builder.py | 80 +- tests/v1/core/test_contiguous_kv_packing.py | 50 +- .../v1/e2e/general/test_mamba_prefix_cache.py | 16 +- tests/v1/e2e/test_replayssm_decode.py | 247 ++++- .../worker/test_kv_cache_allocation_scope.py | 16 +- .../worker/test_mamba_hybrid_model_state.py | 2 + tests/v1/worker/test_mamba_utils.py | 205 +++- tests/v1/worker/test_utils.py | 75 +- vllm/config/cache.py | 9 +- vllm/config/vllm.py | 29 +- vllm/model_executor/layers/mamba/abstract.py | 18 + .../layers/mamba/mamba_mixer2.py | 121 +-- .../layers/mamba/ops/ssu_dispatch.py | 910 ++++++++++++++---- vllm/model_executor/models/nemotron_h.py | 16 - .../model_executor/warmup/replayssm_warmup.py | 48 - vllm/v1/attention/backends/mamba_attn.py | 50 +- vllm/v1/core/kv_cache_utils.py | 43 +- vllm/v1/kv_cache_interface.py | 15 + vllm/v1/worker/gpu/attn_utils.py | 3 + vllm/v1/worker/gpu/model_runner.py | 9 +- vllm/v1/worker/gpu/model_states/interface.py | 2 + .../worker/gpu/model_states/mamba_hybrid.py | 95 +- vllm/v1/worker/gpu/pp_utils.py | 15 + vllm/v1/worker/gpu_model_runner.py | 79 +- vllm/v1/worker/mamba_utils.py | 180 +++- vllm/v1/worker/utils.py | 36 + 30 files changed, 2434 insertions(+), 889 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 082b61c43d2f..da169d89d13a 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -7,22 +7,19 @@ import pytest import torch +import vllm.model_executor.layers.mamba.ops.ssu_dispatch as ssu_dispatch 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, + ReplaySSMModelContext, 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 ( @@ -40,8 +37,11 @@ try: from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner + from flashinfer.mamba.checkpointing_ssu import ( + checkpointing_ssu as checkpointing_ssu_kernel, + ) - HAS_FLASHINFER_CHECKPOINTING_SSU = CheckpointingSSURunner is not None + HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) except ImportError: HAS_FLASHINFER_CHECKPOINTING_SSU = False @@ -57,234 +57,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") - prev_query_len = 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, - prev_query_len, - 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, - prev_query_len, - state_batch_indices, - ) - 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( mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2, ) -> KVCacheConfig: @@ -407,6 +179,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( @@ -442,37 +215,18 @@ 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): +def test_replayssm_flashinfer_call_forwards_packed_mtp(monkeypatch): import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod - kernel = Mock(return_value=torch.empty(0)) + kernel = Mock(return_value=torch.empty(1, 6, 2, 4)) monkeypatch.setattr(mod, "_flashinfer_replayssm_kernel", kernel) - batch, max_seqlen, nheads, dim, dstate, ngroups = 2, 4, 2, 4, 8, 1 + tokens, nheads, dim, dstate, ngroups = 6, 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) + x = torch.empty(tokens, nheads, dim) dt = torch.empty_like(x) A = torch.empty(nheads, dim, dstate) - B = torch.empty(B_shape) + B = torch.empty(tokens, ngroups, dstate) C = torch.empty_like(B) out = torch.empty_like(x) x_cache = torch.empty(2, nheads, 20, dim) @@ -480,7 +234,8 @@ def test_replayssm_flashinfer_call_forwards_mtp_layout(monkeypatch, layout): 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) + cu_seqlens = torch.tensor([0, 4, 6], dtype=torch.int32) + selective_state_update_replayssm_flashinfer( state, x, @@ -494,125 +249,36 @@ def test_replayssm_flashinfer_call_forwards_mtp_layout(monkeypatch, layout): 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, + max_seqlen=4, ) 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( - ("query_start_loc", "expected_shape", "expected_max_seqlen"), - [ - pytest.param([0, 4, 8], (2, 4, 2, 4), None, id="dense"), - pytest.param([0, 4, 6], (6, 2, 4), 4, id="packed"), - ], + kwargs = kernel.call_args.kwargs + assert args[6].shape == (1, tokens, nheads, dim) + assert args[7].shape == (1, tokens, nheads, dim) + assert args[9].shape == (1, tokens, ngroups, dstate) + assert args[10].shape == (1, tokens, ngroups, dstate) + assert args[11].shape == (1, tokens, nheads, dim) + assert kwargs["cu_seqlens"] is cu_seqlens + assert kwargs["max_seqlen"] == 4 + + +@pytest.mark.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="compatible flashinfer checkpointing_ssu not available", ) -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) +def test_replayssm_flashinfer_backend_init(): + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod - mixer.conv_ssm_forward( - torch.empty(num_decode_tokens, 3), torch.empty(num_decode_tokens, 8) + initialize_mamba_ssu_backend( + MambaConfig(backend=MambaBackendEnum.FLASHINFER), + _kv_cache_config_with_ssu(), + use_replayssm=True, ) - - 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 + assert isinstance(get_mamba_ssu_backend(), FlashInferSSUBackend) + assert mod._flashinfer_replayssm_kernel is checkpointing_ssu_kernel @pytest.mark.parametrize( @@ -642,3 +308,363 @@ def test_replayssm_physical_ring_shape( (8, expected_ring_len), (2, expected_ring_len, 16), ) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (0, 0), + ((1 << 63) - 1, (1 << 63) - 1), + (1 << 63, -(1 << 63)), + ((1 << 64) - 1, -1), + ], +) +def test_reinterpret_u64_as_i64(value: int, expected: int): + assert ssu_dispatch._reinterpret_u64_as_i64(value) == expected + + +def _materialize_mixer(device: str = "cpu") -> Mock: + mixer = Mock() + mixer.kv_cache = [ + torch.empty(0, device=device), + torch.empty(8, 4, 3, 5, device=device), + torch.empty(8, 4, 20, 3, device=device), + torch.empty(8, 4, 20, device=device), + torch.empty(8, 2, 20, 5, device=device), + ] + mixer.A = torch.empty(4, 3, 5, device=device) + mixer._replayssm_ring_start = torch.arange(8, dtype=torch.int32, device=device) + mixer._replayssm_prev_num_accepted = torch.zeros( + 8, dtype=torch.int32, device=device + ) + mixer.replayssm_buffer_len = 16 + mixer.mamba_config = SimpleNamespace( + backend=MambaBackendEnum.FLASHINFER, + enable_stochastic_rounding=True, + stochastic_rounding_philox_rounds=6, + ) + return mixer + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_replayssm_materialize_ready_rejects_incomplete_cache(): + mixer = _materialize_mixer(device="cuda") + + mixer.kv_cache[1] = torch.empty(0, device="cuda") + assert not ssu_dispatch._replayssm_materialize_ready([mixer]) + + mixer = _materialize_mixer(device="cuda") + mixer.kv_cache[2] = torch.empty(0, device="cuda") + with pytest.raises(RuntimeError, match="replay ring buffers"): + ssu_dispatch._replayssm_materialize_ready([mixer]) + + mixer = _materialize_mixer(device="cuda") + mixer._replayssm_ring_start = torch.empty(0, dtype=torch.int32, device="cuda") + with pytest.raises(RuntimeError, match="ring trackers"): + ssu_dispatch._replayssm_materialize_ready([mixer]) + + mixer = _materialize_mixer(device="cuda") + mixer.replayssm_buffer_len = 0 + with pytest.raises(RuntimeError, match="buffer-len >= 1"): + ssu_dispatch._replayssm_materialize_ready([mixer]) + + +def test_replayssm_materialize_ready_requires_cuda_ssm_state(): + mixer = _materialize_mixer(device="cpu") + + with pytest.raises(RuntimeError, match="requires a CUDA SSM state cache"): + ssu_dispatch._replayssm_materialize_ready([mixer]) + + +def _modelwide_replayssm_fixture(): + groups = [ + [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], + [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], + ] + layer_names: list[list[str]] = [] + forward_context = {} + for group_idx, mixers in enumerate(groups): + # Layers in one cache group share the physical tracker namespace. + for mixer in mixers[1:]: + mixer._replayssm_ring_start = mixers[0]._replayssm_ring_start + mixer._replayssm_prev_num_accepted = mixers[0]._replayssm_prev_num_accepted + names = [f"group{group_idx}.layer{layer_idx}" for layer_idx in range(2)] + layer_names.append(names) + for name, mixer in zip(names, mixers): + mixer.use_replayssm = True + forward_context[name] = mixer + + config = Mock() + config.kv_cache_groups = [Mock(layer_names=names) for names in layer_names] + block_tables = [ + torch.tensor([[1, 2, 3], [0, 0, 0]], dtype=torch.int32, device="cuda"), + torch.tensor([[4, 5, 6], [0, 0, 0]], dtype=torch.int32, device="cuda"), + ] + return groups, config, forward_context, block_tables + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for mixers, source_slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[source_slot] = 2 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 + kernel = Mock() + monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.postprocess_and_materialize( + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), + query_is_cumulative=True, + num_computed_tokens=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + num_computed_is_after=True, + num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([False, False], device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), + mamba_block_size=4, + num_reqs=1, + ) + torch.cuda.synchronize() + + assert kernel.call_count == 1 + args = kernel.call_args.args + kwargs = kernel.call_args.kwargs + assert args[11] is ctx.src_slots + assert args[12] is ctx.dst_slots + assert args[13] is ctx.plan_ring_start + assert args[14] is ctx.plan_flush_count + assert kwargs["num_heads"] == 4 + assert kwargs["heads_per_group"] == 2 + assert kwargs["max_window"] == 16 + assert kwargs["ring_buffer_len"] == 20 + assert ctx.src_slots[:, 0].tolist() == [1, 1, 4, 4] + assert ctx.dst_slots[:, 0].tolist() == [2, 2, 5, 5] + assert ctx.src_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 4 + assert ctx.dst_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 4 + assert ctx.plan_ring_start.tolist() == [2, 0] + assert ctx.plan_flush_count.tolist() == [6, -1] + assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 6 + assert groups[0][0]._replayssm_prev_num_accepted[2].item() == 0 + assert groups[1][0]._replayssm_prev_num_accepted[4].item() == 6 + assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for mixers, source_slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[source_slot] = 2 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 13 + kernel = Mock() + monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.postprocess_and_materialize( + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), + query_is_cumulative=True, + num_computed_tokens=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + num_computed_is_after=True, + num_accepted_tokens=torch.tensor([3, 1], dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([False, False], device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + mamba_block_size=4, + num_reqs=1, + ) + torch.cuda.synchronize() + + assert kernel.call_count == 1 + assert ctx.plan_ring_start.tolist() == [15, 0] + assert ctx.plan_flush_count.tolist() == [1, -1] + for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): + assert mixers[0]._replayssm_ring_start[source_slot].item() == 15 + assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 3 + assert mixers[0]._replayssm_ring_start[destination_slot].item() == 0 + assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): + for source_slot in source_slots: + mixers[0]._replayssm_ring_start[source_slot] = 7 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 9 + kernel = Mock() + monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.postprocess_and_materialize( + idx_mapping=None, + query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), + query_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_after=False, + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([True, False], device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.tensor([-1, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), + mamba_block_size=4, + num_reqs=1, + materialize_possible=False, + ) + torch.cuda.synchronize() + + assert kernel.call_count == 0 + assert ctx.plan_flush_count.tolist() == [-1, -1] + for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): + for source_slot in source_slots: + assert mixers[0]._replayssm_ring_start[source_slot].item() == 0 + assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for mixers, source_slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[source_slot] = 2 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 + kernel = Mock() + monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.copy_reassigned_slots( + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + torch.cuda.synchronize() + + assert kernel.call_count == 1 + assert ctx.precopy_ring_start.tolist() == [2, 0] + assert ctx.precopy_flush_count.tolist() == [4, -1] + assert ctx.precopy_src_slots[:, 0].tolist() == [1, 1, 4, 4] + assert ctx.precopy_dst_slots[:, 0].tolist() == [2, 2, 5, 5] + for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): + assert mixers[0]._replayssm_ring_start[source_slot].item() == 2 + assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 4 + assert mixers[0]._replayssm_ring_start[destination_slot].item() == 0 + assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_does_not_copy_unchanged_physical_slots(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for block_table in block_tables: + block_table[0, 1] = block_table[0, 0] + for mixers, source_slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[source_slot] = 2 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 + kernel = Mock() + monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.copy_reassigned_slots( + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + torch.cuda.synchronize() + + assert kernel.call_count == 1 + assert ctx.precopy_flush_count.tolist() == [4, -1] + assert ctx.precopy_src_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 4 + assert ctx.precopy_dst_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 4 + for mixers, source_slot in zip(groups, (1, 4)): + assert mixers[0]._replayssm_ring_start[source_slot].item() == 2 + assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 4 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + block_tables[0][0, 1] = block_tables[0][0, 0] + for mixers, source_slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[source_slot] = 2 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 + kernel = Mock() + monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.copy_reassigned_slots( + idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + torch.cuda.synchronize() + + assert kernel.call_count == 1 + assert ctx.precopy_ring_start.tolist() == [2, 0] + assert ctx.precopy_flush_count.tolist() == [4, -1] + assert ctx.precopy_src_slots[:, 0].tolist() == [ + NULL_BLOCK_ID, + NULL_BLOCK_ID, + 4, + 4, + ] + assert ctx.precopy_dst_slots[:, 0].tolist() == [ + NULL_BLOCK_ID, + NULL_BLOCK_ID, + 5, + 5, + ] + assert groups[0][0]._replayssm_ring_start[1].item() == 2 + assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 4 + assert groups[1][0]._replayssm_ring_start[4].item() == 2 + assert groups[1][0]._replayssm_prev_num_accepted[4].item() == 4 + assert groups[1][0]._replayssm_ring_start[5].item() == 0 + assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index 5db0f9a50d0f..cb41a23f1336 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -74,20 +74,15 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): 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, - "uniform_decode": True, - } + assert decode_kwargs["num_tokens"] == expected_num_reqs * query_len + assert decode_kwargs["uniform_decode"] is True + assert decode_kwargs["is_profile"] is True if runner_kwargs.get("use_v2_model_runner"): - expected_kwargs["valid_dummy_state_slots"] = True + assert decode_kwargs["valid_dummy_state_slots"] is True + assert "profile_seq_lens" not in decode_kwargs else: - expected_kwargs.update( - allow_microbatching=False, - force_attention=True, - profile_seq_lens=query_len + 1, - ) - assert decode_kwargs == expected_kwargs + assert decode_kwargs["profile_seq_lens"] == query_len + 1 + assert decode_kwargs["force_attention"] is True @pytest.mark.parametrize( @@ -96,13 +91,11 @@ 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", ], ) def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): @@ -117,6 +110,16 @@ def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): assert result is None +def test_replayssm_autotune_kwargs_skipped_without_non_padding_slot(): + with patch.object( + warmup, "flashinfer_replayssm_autotune_supported", return_value=True + ): + result = warmup._replayssm_autotune_kwargs( + _autotune_runner(num_blocks=1), PREFILL_KWARGS + ) + assert result is None + + def test_replayssm_autotune_slots_restore_state_and_trackers(): mixer = MambaMixer2.__new__(MambaMixer2) torch.nn.Module.__init__(mixer) @@ -129,12 +132,10 @@ 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 0fe3c365d1c7..b2ae877ed9ef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -86,7 +86,18 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) assert not config.cache_config.use_kda_recoverssm + config.mamba_config.backend = MambaBackendEnum.TRITON + with pytest.raises(ValueError, match="requires --mamba-backend flashinfer"): + VllmConfig.validate_mamba_cached_kernel(config) + + config.mamba_config.backend = MambaBackendEnum.FLASHINFER + config.cache_config.replayssm_buffer_len = 3 + with pytest.raises(ValueError, match="replayssm-buffer-len"): + VllmConfig.validate_mamba_cached_kernel(config) + config.cache_config.replayssm_buffer_len = 16 + config.model_config.architecture = "KimiLinearForCausalLM" + config.mamba_config.backend = MambaBackendEnum.TRITON config.parallel_config.pipeline_parallel_size = 2 with pytest.raises(ValueError, match="pipeline_parallel_size=1"): VllmConfig.validate_mamba_cached_kernel(config) @@ -160,6 +171,102 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected +def _replayssm_config( + *, + backend: MambaBackendEnum, + use_v2_model_runner: bool = False, +) -> SimpleNamespace: + return SimpleNamespace( + cache_config=SimpleNamespace( + use_replayssm=True, + mamba_cache_mode="none", + replayssm_buffer_len=16, + ), + model_config=None, + num_speculative_tokens=0, + mamba_config=SimpleNamespace(backend=backend), + use_v2_model_runner=use_v2_model_runner, + kv_transfer_config=None, + ) + + +def test_v2_replayssm_requires_flashinfer(): + config = _replayssm_config( + backend=MambaBackendEnum.TRITON, + use_v2_model_runner=True, + ) + + with pytest.raises(ValueError, match="requires Model Runner V1"): + VllmConfig.validate_mamba_cached_kernel(config) + + +def test_v2_flashinfer_replayssm_is_supported(): + config = _replayssm_config( + backend=MambaBackendEnum.FLASHINFER, + use_v2_model_runner=True, + ) + + assert VllmConfig.validate_mamba_cached_kernel(config) is config + + +def test_flashinfer_replayssm_allows_align_prefix_caching(): + config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) + config.cache_config.mamba_cache_mode = "align" + + assert VllmConfig.validate_mamba_cached_kernel(config) is config + + +def test_flashinfer_replayssm_allows_all_prefix_caching(): + config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) + config.cache_config.mamba_cache_mode = "all" + + assert VllmConfig.validate_mamba_cached_kernel(config) is config + + +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) +def test_flashinfer_replayssm_spec_decode_allows_align_prefix_caching( + use_v2_model_runner, +): + config = _replayssm_config( + backend=MambaBackendEnum.FLASHINFER, + use_v2_model_runner=use_v2_model_runner, + ) + config.num_speculative_tokens = 3 + config.cache_config.mamba_cache_mode = "align" + + assert VllmConfig.validate_mamba_cached_kernel(config) is config + + +@pytest.mark.parametrize("use_v2_model_runner", [False, True]) +def test_flashinfer_replayssm_spec_decode_allows_all_prefix_caching( + use_v2_model_runner, +): + config = _replayssm_config( + backend=MambaBackendEnum.FLASHINFER, + use_v2_model_runner=use_v2_model_runner, + ) + config.num_speculative_tokens = 3 + config.cache_config.mamba_cache_mode = "all" + + assert VllmConfig.validate_mamba_cached_kernel(config) is config + + +def test_triton_replayssm_rejects_all_prefix_caching(): + config = _replayssm_config(backend=MambaBackendEnum.TRITON) + config.cache_config.mamba_cache_mode = "all" + + with pytest.raises(ValueError, match="all mode requires.*flashinfer"): + VllmConfig.validate_mamba_cached_kernel(config) + + +def test_flashinfer_replayssm_rejects_unsupported_buffer_length(): + config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) + config.cache_config.replayssm_buffer_len = 17 + + with pytest.raises(ValueError, match="replayssm-buffer-len <= 16"): + VllmConfig.validate_mamba_cached_kernel(config) + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 4ec138270203..9a909952ce0f 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -16,7 +16,11 @@ import torch -from tests.v1.attention.utils import MockMambaBuilder +from tests.v1.attention.utils import ( + BatchSpec, + MockMambaBuilder, + create_common_attn_metadata, +) from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadata from vllm.v1.kv_cache_interface import MambaSpec @@ -321,6 +325,48 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): assert builder.block_idx_last_scheduled_token_prev_step.dtype == torch.int32 +def test_all_spec_decode_without_drafts_uses_computed_state_anchor(): + """A step with no scheduled drafts still needs a valid input-state table. + + This occurs after a prefix hit when speculative decoding is configured but + the current step contains only one target token. The previous-step buffer + is not passed for that step, so the last computed block is the input anchor. + """ + block_size = 16 + seq_lens = [33, 49] + query_lens = [1, 1] + config = _make_vllm_config( + max_model_len=256, + max_num_seqs=len(seq_lens), + num_speculative_tokens=3, + ) + config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + spec = MambaSpec( + block_size=block_size, + shapes=((1,), (1,)), + dtypes=(torch.float32,), + mamba_cache_mode="all", + num_speculative_blocks=2, + ) + builder = MockMambaBuilder(spec, ["layer0"], config, torch.device("cpu")) + common = create_common_attn_metadata( + BatchSpec(seq_lens=seq_lens, query_lens=query_lens), + block_size, + torch.device("cpu"), + arange_block_indices=True, + ).replace(is_prefilling=torch.zeros(len(seq_lens), dtype=torch.bool)) + + metadata = builder.build(0, common) + + assert metadata.num_decodes == len(seq_lens) + assert metadata.block_idx_last_scheduled_token_prev_step is not None + expected = torch.tensor([1, 2], dtype=torch.int32) + torch.testing.assert_close( + metadata.block_idx_last_scheduled_token_prev_step, + expected, + ) + + def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode(): """Without spec decode, the prev-step buffer is unused and must not be allocated — the input anchor reduces to last_computed_token.""" diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index e1e2875e2539..6d4704136685 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -194,23 +194,25 @@ def _make_mamba_spec( buffer_len: int, mamba_backend: MambaBackendEnum, num_speculative_tokens: int = 0, + mamba_cache_mode: str = "none", ) -> MambaSpec: ring_buffer_len = buffer_len + ( 1 + num_speculative_tokens if mamba_backend == MambaBackendEnum.FLASHINFER else 0 ) - shapes = ( - (1, 1), - (1, 1, 1), + replayssm_shapes = ( (1, ring_buffer_len, 1), (1, ring_buffer_len), (1, ring_buffer_len, 1), ) return MambaSpec( block_size=BLOCK_SIZE, - shapes=shapes, - dtypes=(torch.float32,), + shapes=((1, 1), (1, 1, 1)), + dtypes=(torch.float32,) * 2, + replayssm_shapes=replayssm_shapes, + replayssm_dtypes=(torch.float32,) * 3, + mamba_cache_mode=mamba_cache_mode, ) @@ -236,24 +238,25 @@ def _create_replayssm_builder( num_speculative_tokens=num_speculative_tokens, ) return MockMambaBuilder( - _make_mamba_spec(buffer_len, mamba_backend, num_speculative_tokens), + _make_mamba_spec( + buffer_len, + mamba_backend, + num_speculative_tokens, + mamba_cache_mode, + ), ["layer0"], vllm_config, DEVICE, ) -def _build( - builder: MockMambaBuilder, - case: ReplaySSMBuildCase, - num_accepted_tokens: torch.Tensor | None = None, -): +def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase): 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, num_accepted_tokens=num_accepted_tokens) + return builder.build(0, common) @pytest.mark.parametrize( @@ -338,3 +341,56 @@ def test_flashinfer_replayssm_state_indices_are_stable_for_full_cudagraph(): assert second_indices is not None assert second_indices.data_ptr() == first_ptr assert torch.equal(second_indices, second.state_indices_tensor_d[:, 0]) + + +def test_flashinfer_replayssm_all_uses_last_scheduled_state_page(): + builder = _create_replayssm_builder( + 16, + mamba_cache_mode="all", + mamba_backend=MambaBackendEnum.FLASHINFER, + num_speculative_tokens=3, + ) + builder.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + + metadata = _build( + builder, + ReplaySSMBuildCase( + seq_lens=[34, 50], + query_lens=[1, 1], + is_prefilling=[False, False], + decode_base=[33, 49], + buffer_len=16, + expected_write_pos=[], + expected_is_flush=[], + mamba_cache_mode="all", + ), + ) + + assert metadata.replayssm_state_indices_d is not None + assert metadata.block_idx_last_scheduled_token is not None + live_cols = metadata.block_idx_last_scheduled_token[:2].to(torch.int64) + expected = metadata.state_indices_tensor_d.gather( + 1, live_cols.unsqueeze(1) + ).squeeze(1) + assert torch.equal(metadata.replayssm_state_indices_d, expected) + assert not torch.equal(expected, metadata.state_indices_tensor_d[:, 0]) + + +def test_flashinfer_replayssm_scratch_metadata_fresh_decode(): + checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu") + if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"): + pytest.skip("FlashInfer does not expose ReplaySSM scratch allocation") + + builder = _create_replayssm_builder(16, mamba_backend=MambaBackendEnum.FLASHINFER) + case = REPLAYSSM_BUILD_CASES["fresh_decode"] + meta = _build(builder, case) + + assert meta.write_pos_d is None + assert meta.is_flush_d is None + assert meta.bc_pre_scratch is None + assert meta.replayssm_scratch is not None + assert [tensor.shape for tensor in meta.replayssm_scratch] == [ + (1, 1, 32, 8), + (1, 1, 16), + (1, 1, 32, 8), + ] diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 5249de7f8c11..069ce9641d7b 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -38,7 +38,7 @@ UniformTypeKVCacheSpecs, iter_layer_specs, ) -from vllm.v1.worker.utils import allocate_kv_cache +from vllm.v1.worker.utils import allocate_kv_cache, allocate_replayssm_caches MEMORY = 8 * 1024 * 1024 @@ -97,6 +97,54 @@ def _bind(config, layout: str): return allocate_kv_cache(config, torch.device("cpu"), KVCacheLayout[layout], None) +def test_replayssm_rings_do_not_expand_canonical_mamba_page(): + full_spec = FullAttentionSpec( + block_size=2, + num_kv_heads=1, + head_size=8, + dtype=torch.float32, + ) + mamba_spec = MambaSpec( + block_size=2, + shapes=((16,), (16,)), + dtypes=(torch.float32, torch.float32), + replayssm_shapes=((2,), (1,), (1,)), + replayssm_dtypes=(torch.float32,) * 3, + ) + groups = [ + KVCacheGroupSpec(["full"], full_spec), + KVCacheGroupSpec(["mamba"], mamba_spec), + ] + + assert full_spec.page_size_bytes == mamba_spec.page_size_bytes == 128 + assert mamba_spec.replayssm_size_bytes == 16 + assert _get_kv_cache_bytes_per_block(groups) == 144 + + config = get_kv_cache_config_from_groups( + _mock_vllm_config("LBNHC"), groups, available_memory=4 * 144 + ) + assert config.num_blocks == 4 + caches = allocate_kv_cache( + config, + torch.device("cpu"), + KVCacheLayout.LBNHC, + ) + replayssm_caches = allocate_replayssm_caches(config, torch.device("cpu")) + assert ( + caches["mamba"].untyped_storage().data_ptr() + != replayssm_caches["mamba"][0].untyped_storage().data_ptr() + ) + assert caches["mamba"].untyped_storage().nbytes() == 4 * 128 + assert [tuple(state.shape) for state in replayssm_caches["mamba"]] == [ + (4, 2), + (4, 1), + (4, 1), + ] + replayssm_caches["mamba"][0].fill_(7) + assert torch.count_nonzero(caches["mamba"]) == 0 + assert torch.count_nonzero(replayssm_caches["mamba"][1]) == 0 + + MAIN_KV_PAGE_BYTES = 2_048 COMPRESSED_PAGE_BYTES = 128 NUM_CACHE_TUPLES = 3 diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 27331516cbd3..830bfccc9183 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -1037,6 +1037,8 @@ def wrapped_postprocess_state( idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int, num_computed_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, ) -> None: action = cur_step_action block_tables = captured.get("block_tables") @@ -1050,7 +1052,12 @@ def wrapped_postprocess_state( or action.postprocess_copy_idx == (-1, -1) ): return original_postprocess_state( - self, idx_mapping, num_sampled, num_computed_tokens + self, + idx_mapping, + num_sampled, + num_computed_tokens, + query_start_loc, + is_prefilling, ) expected = action.postprocess_copy_idx snapshots = [ @@ -1058,7 +1065,12 @@ def wrapped_postprocess_state( for temporal, bt in temporal_states(self, block_tables, kv_cache_config) ] ret = original_postprocess_state( - self, idx_mapping, num_sampled, num_computed_tokens + self, + idx_mapping, + num_sampled, + num_computed_tokens, + query_start_loc, + is_prefilling, ) # Comparing device tensors for the assertion is a deliberate D2H. with gpu_sync_allowed(): diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 1b61fd1f878e..690b57d8eb0e 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -10,9 +10,27 @@ from ...models.utils import check_logprobs_close from ...utils import large_gpu_mark, multi_gpu_test +try: + from flashinfer.mamba.checkpointing_ssu import ( + CheckpointingSSURunner, + allocate_checkpointing_ssu_scratch, # noqa: F401 + ) + + HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) +except ImportError: + HAS_FLASHINFER_CHECKPOINTING_SSU = False + +try: + from flashinfer.mamba.replayssm_materialize import replayssm_materialize + + HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = callable(replayssm_materialize) +except ImportError: + HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = False + # Mamba2 (Nemotron-3) hybrid. MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" MAMBA2_MTP_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4" +MAMBA2_PREFIX_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" MODELS = [ pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), ] @@ -25,7 +43,7 @@ try: from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner - HAS_FLASHINFER_CHECKPOINTING_SSU = CheckpointingSSURunner is not None + HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) except ImportError: HAS_FLASHINFER_CHECKPOINTING_SSU = False @@ -38,10 +56,16 @@ def _check_replayssm_parity( mamba_backend: str = "triton", name_1: str = "replayssm", require_v2: bool = False, + monkeypatch: pytest.MonkeyPatch | None = None, ): # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are # common-mode and only ReplaySSM varies. + if require_v2: + assert monkeypatch is not None + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + envs.disable_envs_cache() + common = dict( max_model_len=1024, trust_remote_code=True, @@ -82,6 +106,17 @@ 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(vllm_runner, model_name): + pytest.importorskip("flashinfer.mamba.checkpointing_ssu") + _check_replayssm_parity( + vllm_runner, + model_name, + mamba_backend="flashinfer", + name_1="replayssm_flashinfer", + ) + + @pytest.mark.skipif( not HAS_FLASHINFER_CHECKPOINTING_SSU, reason="flashinfer.mamba.checkpointing_ssu not available", @@ -90,20 +125,30 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): def test_replayssm_flashinfer_decode_matches_baseline_v2( vllm_runner, model_name, monkeypatch ): - try: - with monkeypatch.context() as patch: - patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - envs.disable_envs_cache() - _check_replayssm_parity( - vllm_runner, - model_name, - mamba_backend="flashinfer", - name_1="replayssm_flashinfer_v2", - require_v2=True, - ) - finally: - # The context restores the environment before the final cache reset. - envs.disable_envs_cache() + _check_replayssm_parity( + vllm_runner, + model_name, + mamba_backend="flashinfer", + name_1="replayssm_flashinfer_v2", + require_v2=True, + monkeypatch=monkeypatch, + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="flashinfer.mamba.checkpointing_ssu not available", +) +@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) +def test_replayssm_flashinfer_decode_matches_baseline_tp2(vllm_runner, model_name): + _check_replayssm_parity( + vllm_runner, + model_name, + tensor_parallel_size=2, + mamba_backend="flashinfer", + name_1="replayssm_flashinfer_tp2", + ) @pytest.mark.skipif( @@ -139,6 +184,37 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na ) +@pytest.mark.skipif( + not HAS_FLASHINFER_CHECKPOINTING_SSU, + reason="flashinfer.mamba.checkpointing_ssu not available", +) +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_flashinfer_matches_triton_replayssm(vllm_runner, model_name): + # Both backends implement ReplaySSM; compare them directly on V1 because + # Triton ReplaySSM is not supported on Model Runner V2. + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + use_replayssm=True, + replayssm_buffer_len=16, + ) + with vllm_runner(model_name, mamba_backend="triton", **common) as llm: + triton = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: + flashinfer = llm.generate_greedy_logprobs( + PROMPTS, max_tokens=32, num_logprobs=5 + ) + + check_logprobs_close( + outputs_0_lst=triton, + outputs_1_lst=flashinfer, + name_0="replayssm_triton", + name_1="replayssm_flashinfer", + ) + + @pytest.mark.skipif( not HAS_FLASHINFER_CHECKPOINTING_SSU, reason="flashinfer.mamba.checkpointing_ssu not available", @@ -158,6 +234,11 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): with monkeypatch.context() as patch: patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") envs.disable_envs_cache() + with vllm_runner(MAMBA2_MTP_MODEL, **common) as llm: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner + baseline = llm.generate_greedy_logprobs( + PROMPTS, max_tokens=32, num_logprobs=5 + ) with vllm_runner( MAMBA2_MTP_MODEL, use_replayssm=True, @@ -165,7 +246,9 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): **common, ) as llm: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner - outputs = llm.generate_greedy(PROMPTS, max_tokens=32) + replay = llm.generate_greedy_logprobs( + PROMPTS, max_tokens=32, num_logprobs=5 + ) draft_count = sum( metric.value for metric in llm.llm.get_metrics() @@ -175,10 +258,14 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): 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 any(len(token_ids) > 16 for token_ids, _ in replay) assert draft_count > 0 + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline_mtp_v2", + name_1="replayssm_flashinfer_mtp_v2", + ) # Prefix spans several mamba blocks; prefix caching only reuses full blocks. @@ -203,28 +290,58 @@ def _prefix_cache_hits(llm) -> int: ) -def _check_replayssm_prefix_caching_parity( - vllm_runner, model_name, *, tensor_parallel_size=1 +def _check_flashinfer_replayssm_prefix_caching( + vllm_runner, + model_name, + monkeypatch: pytest.MonkeyPatch, + *, + mamba_cache_mode: str, + moe_backend: str | None = None, + use_ngram: bool, + use_v2: bool, + tensor_parallel_size: int, ): - # align mode materializes the exact SSM state at each block boundary, so - # ReplaySSM's cached prefixes must match the always-materialized baseline. + # ReplaySSM materializes the exact SSM state at each cacheable block + # boundary, so cached prefixes must match the always-materialized baseline. + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2 else "0") + envs.disable_envs_cache() + common = dict( max_model_len=8192, trust_remote_code=True, enable_prefix_caching=True, enable_chunked_prefill=True, - mamba_cache_mode="align", + mamba_cache_mode=mamba_cache_mode, + mamba_backend="flashinfer", disable_log_stats=False, # required for llm.get_metrics() tensor_parallel_size=tensor_parallel_size, ) + if moe_backend is not None: + common["moe_backend"] = moe_backend + if use_ngram: + common["speculative_config"] = { + "method": "ngram", + "num_speculative_tokens": 3, + "prompt_lookup_max": 3, + } + with vllm_runner(model_name, **common) as llm: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 + baseline_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size + llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) baseline = llm.generate_greedy_logprobs( PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 ) + baseline_hits = _prefix_cache_hits(llm) + with vllm_runner( model_name, use_replayssm=True, replayssm_buffer_len=16, **common ) as llm: - # Prime the cache, then measure, so cache hits are deterministic. + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 + replay_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size + assert replay_block_size == baseline_block_size llm.generate_greedy_logprobs( PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 ) @@ -233,27 +350,87 @@ def _check_replayssm_prefix_caching_parity( ) replay_hits = _prefix_cache_hits(llm) - # Without real cache hits the cached path is never exercised. + assert baseline_hits > 0 assert replay_hits > 0, ( - "ReplaySSM align-mode run produced no prefix-cache hits; the shared " - "prefix may be shorter than one mamba block, so prefix caching is inert" + f"ReplaySSM {mamba_cache_mode}-mode run produced no prefix-cache hits; " + "the shared prefix may be shorter than one mamba block, so prefix " + "caching is inert" ) check_logprobs_close( outputs_0_lst=baseline, outputs_1_lst=replay, - name_0="baseline_align_pc", - name_1="replayssm_align_pc", + name_0=f"flashinfer_baseline_{mamba_cache_mode}_pc", + name_1=f"flashinfer_replayssm_{mamba_cache_mode}_pc", ) +@pytest.mark.skipif( + not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), + reason="FlashInfer ReplaySSM materialization APIs not available", +) @pytest.mark.parametrize("model_name", MODELS) -def test_replayssm_prefix_caching_matches_baseline(vllm_runner, model_name): - _check_replayssm_prefix_caching_parity(vllm_runner, model_name) +@pytest.mark.parametrize( + ("mamba_cache_mode", "use_v2", "use_ngram"), + [ + pytest.param("align", False, False, id="align-v1-stp"), + pytest.param("align", False, True, id="align-v1-ngram-t4"), + pytest.param("align", True, False, id="align-v2-stp"), + ], +) +def test_flashinfer_replayssm_prefix_cache_tp1( + vllm_runner, + model_name, + monkeypatch: pytest.MonkeyPatch, + mamba_cache_mode: str, + use_v2: bool, + use_ngram: bool, +): + _check_flashinfer_replayssm_prefix_caching( + vllm_runner, + model_name, + monkeypatch, + mamba_cache_mode=mamba_cache_mode, + use_ngram=use_ngram, + use_v2=use_v2, + tensor_parallel_size=1, + ) +@pytest.mark.skipif( + not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), + reason="FlashInfer ReplaySSM materialization APIs not available", +) +@large_gpu_mark(min_gb=40) +def test_flashinfer_replayssm_all_prefix_cache_v2(vllm_runner, monkeypatch): + _check_flashinfer_replayssm_prefix_caching( + vllm_runner, + MAMBA2_PREFIX_MODEL, + monkeypatch, + mamba_cache_mode="all", + moe_backend="triton", + use_ngram=False, + use_v2=True, + tensor_parallel_size=1, + ) + + +@pytest.mark.skipif( + not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), + reason="FlashInfer ReplaySSM materialization APIs not available", +) @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) -def test_replayssm_prefix_caching_matches_baseline_tp2(vllm_runner, model_name): - _check_replayssm_prefix_caching_parity( - vllm_runner, model_name, tensor_parallel_size=2 +def test_flashinfer_replayssm_prefix_cache_v2_tp2( + vllm_runner, + model_name, + monkeypatch: pytest.MonkeyPatch, +): + _check_flashinfer_replayssm_prefix_caching( + vllm_runner, + model_name, + monkeypatch, + mamba_cache_mode="align", + use_ngram=False, + use_v2=True, + tensor_parallel_size=2, ) diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index a9579f1af69b..39e4e8cd46e7 100644 --- a/tests/v1/worker/test_kv_cache_allocation_scope.py +++ b/tests/v1/worker/test_kv_cache_allocation_scope.py @@ -26,7 +26,7 @@ def __exit__(self, *args: Any) -> None: self.active = False -def test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: +def test_mrv2_kv_pool_wraps_all_cache_allocations(monkeypatch) -> None: scope = _AllocationScope() kv_caches = {"layer": torch.empty(0)} @@ -37,7 +37,12 @@ def allocate(*args, **kwargs): def bind(*args, **kwargs): assert not scope.active + def allocate_replayssm(*args, **kwargs): + assert scope.active + return {} + monkeypatch.setattr(attn_utils, "allocate_kv_cache", allocate) + monkeypatch.setattr(attn_utils, "allocate_replayssm_caches", allocate_replayssm) monkeypatch.setattr(attn_utils, "bind_kv_cache", bind) monkeypatch.setattr(attn_utils, "get_shared_kv_cache_layers", lambda config: {}) @@ -59,7 +64,7 @@ def bind(*args, **kwargs): assert not scope.active -def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: +def test_mrv1_kv_pool_wraps_all_cache_allocations(monkeypatch) -> None: scope = _AllocationScope() kv_caches = {"layer": torch.empty(0)} @@ -70,7 +75,14 @@ def allocate(*args, **kwargs): def bind(*args, **kwargs): assert not scope.active + def allocate_replayssm(*args, **kwargs): + assert scope.active + return {} + monkeypatch.setattr(gpu_model_runner, "allocate_kv_cache", allocate) + monkeypatch.setattr( + gpu_model_runner, "allocate_replayssm_caches", allocate_replayssm + ) monkeypatch.setattr(gpu_model_runner, "bind_kv_cache", bind) runner = SimpleNamespace( diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 749821274318..f02c7be84217 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -26,6 +26,7 @@ def test_postprocess_state_scalar_with_int32_mapping( (4,), 9, dtype=torch.int32, device="cuda" ) state._align_mode = False + state._mamba_lifecycle_mode = False state.recoverssm = None state._mamba_ctx = None idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda") @@ -68,6 +69,7 @@ def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None: state = object.__new__(MambaHybridModelState) state._align_mode = True + state._mamba_lifecycle_mode = True state._mamba_ctx = None state._mamba_state_idx_gpu = torch.full((5,), -1, dtype=torch.int32, device="cuda") state.recoverssm = RecoverSSMState() diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a162c39a8695..1608c0d1d762 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -8,6 +8,7 @@ import pytest import torch +from vllm.config.mamba import MambaBackendEnum from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncsByType, @@ -31,8 +32,10 @@ collect_mamba_copy_meta, do_mamba_copy_block, get_mamba_groups, + postprocess_mamba_align_gpu, preprocess_mamba, stage_postprocess_inputs_to_gpu, + validate_mamba_state_copy_funcs, ) # Conv + temporal copy specs, in the order the tests' MambaSpec shapes expect. @@ -166,6 +169,109 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): assert mamba_state_idx == {"keep": 99} +@pytest.mark.parametrize( + ("with_replayssm", "num_computed_tokens", "expected_order"), + [ + pytest.param(True, 4, ["materialize", "copy"], id="replayssm-boundary"), + pytest.param(True, 3, ["copy"], id="replayssm-no-boundary"), + pytest.param(False, 4, ["copy"], id="generic"), + ], +) +def test_preprocess_mamba_uses_modelwide_materializer_when_present( + with_replayssm: bool, + num_computed_tokens: int, + expected_order: list[str], +): + spec = MagicMock(block_size=4, num_speculative_blocks=0) + cache_config = MagicMock(enable_prefix_caching=True, use_replayssm=True) + input_batch = MagicMock() + input_batch.req_ids = ["r0"] + input_batch.num_accepted_tokens_cpu = np.array([1], dtype=np.int32) + copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec) + requests = {"r0": MagicMock(num_computed_tokens=num_computed_tokens)} + mamba_state_idx: dict[str, int] = {"r0": 0} + sched = _make_scheduler_output(set(), None, set()) + sched.num_scheduled_tokens = {"r0": 1} + + order: list[str] = [] + device = torch.device("cpu") + align_ctx = MagicMock(is_initialized=True) + align_ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(1, torch.int32, device) + align_ctx.precopy_src_col_buf = _MockCpuGpuBuffer(1, torch.int32, device) + align_ctx.precopy_token_bias_buf = _MockCpuGpuBuffer(1, torch.int32, device) + align_ctx.replayssm = MagicMock() if with_replayssm else None + if align_ctx.replayssm is not None: + align_ctx.replayssm.copy_reassigned_slots.side_effect = lambda **kwargs: ( + order.append("materialize") + ) + align_ctx.run_fused_precopy.side_effect = lambda **kwargs: order.append("copy") + + preprocess_mamba( + sched, + MagicMock(), + cache_config, + mamba_state_idx, + input_batch, + requests, + {}, + {}, + copy_bufs, + align_ctx=align_ctx, + ) + + assert order == expected_order + + +def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): + order: list[str] = [] + ctx = MagicMock() + ctx.is_initialized = True + ctx.mamba_group_ids = [0] + ctx.mamba_state_idx_buf = MagicMock() + ctx.num_scheduled_tokens_buf = MagicMock() + ctx.num_computed_tokens_buf = MagicMock() + ctx.num_draft_tokens_buf = MagicMock() + ctx.is_prefilling_buf = MagicMock() + ctx.num_accepted_tokens_out = torch.tensor([3], dtype=torch.int32) + ctx.materialize_src_cols = torch.tensor([2], dtype=torch.int32) + ctx.materialize_dst_cols = torch.tensor([1], dtype=torch.int32) + ctx.materialize_token_counts = torch.tensor([2], dtype=torch.int32) + ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") + ctx.replayssm = MagicMock() + ctx.replayssm.postprocess_and_materialize.side_effect = lambda **kwargs: ( + order.append("materialize") + ) + ctx.replayssm_materialize_possible = False + + block_table = MagicMock() + block_table.get_device_tensor.return_value = torch.zeros((1, 4), dtype=torch.int32) + input_batch = MagicMock() + input_batch.block_table = [block_table] + kv_cache_config = MagicMock() + kv_cache_config.kv_cache_groups = [MagicMock()] + accepted_cpu = torch.zeros(1, dtype=torch.int32) + + postprocess_mamba_align_gpu( + bufs=MagicMock(postprocess_align=ctx), + num_reqs=1, + num_accepted_tokens_gpu=torch.tensor([3], dtype=torch.int32), + num_accepted_tokens_cpu_tensor=accepted_cpu, + input_batch=input_batch, + kv_cache_config=kv_cache_config, + forward_context={}, + mamba_state_copy_funcs={}, + ) + + assert order == ["copy", "materialize"] + assert ( + ctx.replayssm.postprocess_and_materialize.call_args.kwargs[ + "materialize_possible" + ] + is False + ) + assert accepted_cpu.tolist() == [3] + + # ----------------------------------------------------------------------------- # Golden tests for postprocess_mamba_fused_kernel # ----------------------------------------------------------------------------- @@ -290,6 +396,28 @@ def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): ] +def test_gpu_context_marks_flashinfer_replayssm_temporal_state(): + cfg = _TestConfig(num_layers=1) + device = torch.device("cpu") + kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) + gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) + attention = _make_mock_attention( + torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), + torch.empty(cfg.num_blocks, cfg.temporal_state_dim), + ) + attention.use_replayssm = True + attention.mamba_config.backend = MambaBackendEnum.FLASHINFER + + gpu_ctx.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, + [torch.empty(1, 4, dtype=torch.int32)], + ) + + assert gpu_ctx.state_skip_postprocess.tolist() == [0, 1] + + def _make_postprocess_scheduler_output( req_ids: list[str], num_scheduled_tokens: dict[str, int], @@ -430,12 +558,18 @@ def _make_requests( req_ids: list[str], num_computed_tokens: list[int], block_ids_per_req: list[list[int]], + num_prompt_tokens: list[int] | None = None, ) -> dict[str, MagicMock]: """Create mock CachedRequestState objects.""" requests = {} for i, req_id in enumerate(req_ids): req = MagicMock() req.num_computed_tokens = num_computed_tokens[i] + req.num_prompt_tokens = ( + num_computed_tokens[i] + if num_prompt_tokens is None + else num_prompt_tokens[i] + ) req.block_ids = {0: block_ids_per_req[i]} # group_id=0 requests[req_id] = req return requests @@ -554,6 +688,31 @@ def test_mamba_groups_support_different_state_specs(): assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] +def test_mamba_copy_funcs_ignore_replayssm_state_tensors(): + replayssm_spec = MambaSpec( + block_size=16, + shapes=((4, 4), (2, 4, 4)), + dtypes=(torch.float16,) * 2, + replayssm_shapes=((2, 8, 4), (2, 8), (1, 8, 4)), + replayssm_dtypes=(torch.float16,) * 3, + mamba_type=MambaAttentionBackendEnum.MAMBA2, + mamba_cache_mode="align", + ) + + validate_mamba_state_copy_funcs({replayssm_spec: [0]}, _COPY_FUNCS) + + for invalid_funcs in ( + (get_conv_copy_spec,), + (*_DEFAULT_COPY_FUNCS, get_temporal_copy_spec), + ): + invalid_copy_funcs = { + **_COPY_FUNCS, + MambaAttentionBackendEnum.MAMBA2: invalid_funcs, + } + with pytest.raises(AssertionError, match="expects 2 state copy funcs"): + validate_mamba_state_copy_funcs({replayssm_spec: [0]}, invalid_copy_funcs) + + def test_mamba_groups_support_mixed_specs_in_uniform_group(): gdn_spec = MambaSpec( block_size=16, @@ -644,13 +803,15 @@ def test_mamba_groups_support_mixed_specs_in_uniform_group(): def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: - """Build a MambaSpecDecodeGPUContext stand-in exposing only the four + """Build a MambaSpecDecodeGPUContext stand-in exposing only the five per-request staging buffers touched by stage_postprocess_inputs_to_gpu.""" ctx = MagicMock() + ctx.block_size = 16 ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_scheduled_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_computed_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_draft_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) + ctx.is_prefilling_buf = _MockCpuGpuBuffer(max_num_reqs, torch.bool, device) return ctx @@ -690,6 +851,7 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): req_ids=req_ids, num_computed_tokens=[10, 20, 30], block_ids_per_req=[[0], [0], [0]], + num_prompt_tokens=[11, 20, 40], ) mamba_state_idx = {"req_a": 100, "req_b": 200, "req_c": 300} # A trailing entry past num_reqs must not be read. @@ -714,6 +876,9 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ctx.num_computed_tokens_buf.np[:num_reqs], [10, 20, 30] ) np.testing.assert_array_equal(ctx.num_draft_tokens_buf.np[:num_reqs], [2, 0, 4]) + np.testing.assert_array_equal( + ctx.is_prefilling_buf.np[:num_reqs], [True, False, True] + ) for buf in bufs: assert (buf.np[num_reqs:] == sentinel).all() @@ -725,6 +890,39 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ctx.num_draft_tokens_buf.gpu[:num_reqs], torch.tensor([2, 0, 4], dtype=torch.int32), ) + assert torch.equal( + ctx.is_prefilling_buf.gpu[:num_reqs], + torch.tensor([True, False, True]), + ) + assert ctx.replayssm_materialize_possible is True + + +def test_stage_postprocess_inputs_skips_impossible_materialization(): + device = torch.device("cpu") + ctx = _make_staging_ctx(max_num_reqs=4, device=device) + req_ids = ["req_a"] + scheduler_output = _make_postprocess_scheduler_output( + req_ids=req_ids, + num_scheduled_tokens={"req_a": 1}, + scheduled_spec_decode_tokens={"req_a": [1, 2]}, + ) + requests = _make_requests( + req_ids=req_ids, + num_computed_tokens=[8], + block_ids_per_req=[[0]], + num_prompt_tokens=[8], + ) + + stage_postprocess_inputs_to_gpu( + ctx, + scheduler_output, + req_ids, + 1, + requests, + {"req_a": 0}, + ) + + assert ctx.replayssm_materialize_possible is False def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): @@ -1054,6 +1252,7 @@ def test_no_copy_when_not_needed(self, device, test_config): # State should be unchanged torch.testing.assert_close(conv_state, conv_state_orig) torch.testing.assert_close(temporal_state, temporal_state_orig) + assert gpu_ctx.materialize_src_cols[0].item() == -1 @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16]) def test_various_batch_sizes(self, device, test_config, num_reqs): @@ -1376,6 +1575,10 @@ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( device=device, ) + assert gpu_ctx.materialize_src_cols[0].item() == 1 + assert gpu_ctx.materialize_dst_cols[0].item() == 1 + assert gpu_ctx.materialize_token_counts[0].item() == 1 + # --- Verify Python behavior (ground truth) --- # State should be unchanged (no copy when src_addr == dst_addr) torch.testing.assert_close( diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index a5589c5048ba..424efbc2d298 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -11,25 +11,31 @@ class _TestReplaySSMMixer(MambaMixer2): + _state_shapes = ((2,), (3,)) + _state_dtypes = (torch.float32, torch.float32) + def __init__(self): torch.nn.Module.__init__(self) self.use_replayssm = True self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER) self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) - self._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, ...], ...]: - return ((2,), (3,), (4,), (5,), (6,)) + return self._state_shapes def get_state_dtype(self) -> tuple[torch.dtype, ...]: - return (torch.float32,) * 5 + return self._state_dtypes + + def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: + return ((4,), (5,), (6,)) + def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: + return (torch.float32,) * 3 -def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: - return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8) + +def _packed_replayssm_cache(num_blocks: int, fill_value: int = 0) -> torch.Tensor: + return torch.full((num_blocks, 1, 1, 20), fill_value, dtype=torch.int8) def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): @@ -46,23 +52,50 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): SimpleNamespace(layer_names=[layer_names[0], layer_names[2]]), SimpleNamespace(layer_names=[layer_names[1]]), ] + replayssm_caches = { + name: [ + torch.zeros((4, *shape), dtype=torch.float32) + for shape in mixer.get_replayssm_state_shape() + ] + for name, mixer in ctx.items() + } - bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) + bind_kv_cache( + kv_cache, + ctx, + [], + kv_cache_groups=kv_cache_groups, + replayssm_caches={ + name: tuple(cache) for name, cache in replayssm_caches.items() + }, + ) + + assert all(len(mixer.kv_cache) == 5 for mixer in mixers) - tracker_names = ( - "_replayssm_ring_start", - "_replayssm_prev_num_accepted", - "_replayssm_prev_query_len", + assert ( + mixers[0]._replayssm_ring_start.data_ptr() + == mixers[2]._replayssm_ring_start.data_ptr() + ) + assert ( + mixers[0]._replayssm_prev_num_accepted.data_ptr() + == mixers[2]._replayssm_prev_num_accepted.data_ptr() + ) + assert ( + mixers[1]._replayssm_ring_start.data_ptr() + != mixers[0]._replayssm_ring_start.data_ptr() + ) + assert ( + mixers[1]._replayssm_prev_num_accepted.data_ptr() + != mixers[0]._replayssm_prev_num_accepted.data_ptr() ) - 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] + assert mixers[0]._replayssm_ring_start.shape == (4,) + assert mixers[0]._replayssm_prev_num_accepted.shape == (4,) + assert mixers[0]._replayssm_ring_start.dtype == torch.int32 + assert mixers[0]._replayssm_ring_start.is_contiguous() + assert torch.count_nonzero(mixers[0]._replayssm_ring_start) == 0 + assert torch.count_nonzero(mixers[0]._replayssm_prev_num_accepted) == 0 + assert not any(hasattr(m, "_commits_replayssm_trackers") for m in mixers) + assert not any(hasattr(m, "_updates_replayssm_trackers") for m in mixers) def test_bind_kv_cache(default_vllm_config): diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 6c6e8f4fe636..cc396d7b83f5 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -202,11 +202,10 @@ class CacheConfig: 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. 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.""" + Supports mamba_cache_mode 'none', 'align', and 'all'; 'all' requires the + FlashInfer backend. Mamba2 speculative decode also requires FlashInfer. + Prefix-boundary 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 e48a9178fe02..c85b0ae1055f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2908,10 +2908,10 @@ def validate_mamba_cached_kernel(self) -> "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": + if self.cache_config.mamba_cache_mode not in ("none", "align", "all"): raise ValueError( "FlashInfer ReplaySSM speculative decoding requires " - "--mamba-cache-mode none" + "--mamba-cache-mode none, align, or all" ) query_len = 1 + self.num_speculative_tokens if self.cache_config.replayssm_buffer_len < query_len: @@ -2926,22 +2926,25 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": "Mamba2 ReplaySSM speculative decoding requires " "--mamba-backend flashinfer" ) - elif self.cache_config.mamba_cache_mode == "all": + elif ( + self.cache_config.mamba_cache_mode == "all" + and self.mamba_config.backend != MambaBackendEnum.FLASHINFER + ): raise ValueError( - "--use-replayssm supports prefix caching only in align mode; " - "pass --mamba-cache-mode align" + "ReplaySSM prefix caching in all mode requires " + "--mamba-backend flashinfer" ) - elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: - if self.cache_config.mamba_cache_mode == "align": - raise ValueError( - "FlashInfer ReplaySSM does not support " - "--mamba-cache-mode align yet; use none" - ) - elif self.mamba_config.backend != MambaBackendEnum.TRITON: + elif self.mamba_config.backend not in ( + MambaBackendEnum.TRITON, + MambaBackendEnum.FLASHINFER, + ): raise ValueError( "--use-replayssm requires --mamba-backend triton or flashinfer" ) - elif self.use_v2_model_runner: + elif ( + self.mamba_config.backend == MambaBackendEnum.TRITON + and self.use_v2_model_runner + ): raise ValueError( "Triton ReplaySSM requires Model Runner V1; use " "--mamba-backend flashinfer or Model Runner V1" diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index 1dd942221b87..d12eb1dd4a39 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -42,6 +42,16 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: offset += nbytes self.kv_cache = tuple(states) + def bind_replayssm_cache(self, cache: tuple[torch.Tensor, ...]) -> None: + expected_shapes = self.get_replayssm_state_shape() + expected_dtypes = self.get_replayssm_state_dtype() + assert len(cache) == len(expected_shapes) == len(expected_dtypes) + assert all( + tuple(state.shape[1:]) == shape and state.dtype == dtype + for state, shape, dtype in zip(cache, expected_shapes, expected_dtypes) + ) + self.kv_cache = (*self.kv_cache, *cache) + @abstractmethod def get_state_shape(self) -> Iterable[tuple[int, ...]]: """ @@ -64,6 +74,12 @@ def is_kv_cache_tp_replicated(self) -> bool: def get_state_dtype(self) -> tuple[torch.dtype, ...]: pass + def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: + return () + + def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: + return () + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: mamba_block_size = vllm_config.cache_config.mamba_block_size assert mamba_block_size is not None @@ -71,6 +87,8 @@ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: return MambaSpec( shapes=tuple(self.get_state_shape()), dtypes=self.get_state_dtype(), + replayssm_shapes=self.get_replayssm_state_shape(), + replayssm_dtypes=self.get_replayssm_state_dtype(), block_size=mamba_block_size, page_size_padded=page_size_padded, mamba_type=self.mamba_type, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 29edb4751cd1..6a6917fc3caa 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -43,8 +43,6 @@ 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, ) @@ -530,9 +528,6 @@ def __init__( 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 if self.num_spec > 0: @@ -720,7 +715,7 @@ def conv_ssm_forward( 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 = prev_query_len = None + ring_start = prev_num_accepted = None attn_metadata: AttentionMetadata | None = None if attn_metadata_raw is not None: @@ -742,7 +737,6 @@ def conv_ssm_forward( 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 @@ -1002,20 +996,26 @@ def conv_ssm_forward( # tensor assert state_indices_tensor_p is not None ssm_state[state_indices_tensor_p] = varlen_states - if ring_start is not None and self._updates_replayssm_trackers: - assert prev_num_accepted is not None - reset_replayssm_ring_trackers( - ring_start, - prev_num_accepted, - prev_query_len, - state_indices_tensor_p, - ) # Process decode requests if has_decode: assert state_indices_tensor_d is not None if is_mamba_cache_all: - if self.num_spec > 0: + if self.use_replayssm: + # The ownership pre-copy seeds the last-scheduled page before + # forward. Keep both convolution and ReplaySSM on that private + # live page instead of touching the cached prefix source. + assert block_idx_last_scheduled_token_d is not None + live_indices = state_indices_tensor_d.gather( + 1, + block_idx_last_scheduled_token_d.to(torch.int64).unsqueeze(1), + ).squeeze(1) + state_indices_tensor_d_input = live_indices + state_indices_tensor_d_output = live_indices + block_idx_last_computed_token_d = ( + block_idx_last_scheduled_token_d + ) + elif self.num_spec > 0: assert block_idx_last_scheduled_token_prev_step_d is not None input_indices = ( block_idx_last_scheduled_token_prev_step_d.unsqueeze(1) @@ -1043,31 +1043,6 @@ def conv_ssm_forward( 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 self.num_spec > 0 - 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 - 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, @@ -1122,7 +1097,6 @@ def conv_ssm_forward( 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_x = hidden_states_d fi_dt = dt_d @@ -1158,16 +1132,11 @@ def conv_ssm_forward( 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=replayssm_state_indices_d, scratch=attn_metadata.replayssm_scratch, - update_trackers=( - self._updates_replayssm_trackers and self.num_spec == 0 - ), enable_stochastic_rounding=( self.mamba_config.enable_stochastic_rounding ), @@ -1227,20 +1196,15 @@ def conv_ssm_forward( def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None - base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( + return MambaStateDtypeCalculator.mamba2_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) - if self.use_replayssm: - return MambaStateDtypeCalculator.append_replayssm_ring( - base_dtype, self.model_config.dtype - ) - return base_dtype def get_state_shape(self) -> tuple[tuple[int, ...], ...]: tp_world_size = get_tensor_model_parallel_world_size() - base_shape = MambaStateShapeCalculator.mamba2_state_shape( + return MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=self.intermediate_size, tp_world_size=tp_world_size, n_groups=self.n_groups, @@ -1250,17 +1214,29 @@ def get_state_shape(self) -> tuple[tuple[int, ...], ...]: conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, ) - if self.use_replayssm: - assert self.replayssm_buffer_len is not None - return MambaStateShapeCalculator.append_replayssm_ring( - base_shapes=base_shape, - n_groups=self.n_groups, - tp_world_size=tp_world_size, - logical_window=self.replayssm_buffer_len, - backend=self.mamba_config.backend, - num_speculative_tokens=self.num_spec, - ) - return base_shape + + def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: + if not self.use_replayssm: + return () + assert self.model_config is not None + return MambaStateDtypeCalculator.append_replayssm_ring( + (), self.model_config.dtype + ) + + def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: + if not self.use_replayssm: + return () + assert self.replayssm_buffer_len is not None + tp_world_size = get_tensor_model_parallel_world_size() + base_shape = self.get_state_shape() + return MambaStateShapeCalculator.append_replayssm_ring( + base_shapes=base_shape, + n_groups=self.n_groups, + tp_world_size=tp_world_size, + logical_window=self.replayssm_buffer_len, + backend=self.mamba_config.backend, + num_speculative_tokens=self.num_spec, + )[len(base_shape) :] @property def mamba_type(self) -> MambaAttentionBackendEnum: @@ -1277,9 +1253,8 @@ 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. 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. + Tracker mutation is model-owned and runs once after the step; layer forwards + only consume the shared values. """ replayssm_mixers: dict[str, MambaMixer2] = {} @@ -1306,9 +1281,6 @@ 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]] first_state = first_mixer.kv_cache[1] num_blocks, device = first_state.shape[0], first_state.device @@ -1321,17 +1293,10 @@ 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._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/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index e64653afcbaf..fa05bc02b6cd 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -10,8 +10,10 @@ """ from abc import ABC, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from functools import cache +from typing import Any import torch @@ -25,198 +27,642 @@ logger = init_logger(__name__) -@triton.jit( - do_not_specialize=["n_slots", "state_batch_indices_stride"], - do_not_specialize_on_alignment=["state_batch_indices"], -) -def _update_replayssm_ring_trackers_kernel( - ring_start, - prev_num_accepted, - prev_query_len, - state_batch_indices, - state_batch_indices_stride, - n_slots, - num_states, - logical_window: tl.constexpr, - ring_buffer_len: tl.constexpr, - pad_slot_id: tl.constexpr, - RESET: tl.constexpr, - BLOCK: tl.constexpr, +@triton.jit(do_not_specialize=["num_reqs"]) +def _postprocess_replayssm_modelwide_kernel( + # Per-request step metadata. + idx_mapping, + query_metadata, + num_computed_tokens, + num_accepted_tokens, + is_prefilling, + live_cols, + materialize_src_cols, + materialize_dst_cols, + materialize_token_counts, + # Per-group address tables. + block_table_ptrs, + tracker_ring_start_ptrs, + tracker_num_committed_ptrs, + tracker_capacities, + group_layer_offsets, + # FlashInfer plan outputs. + src_slots, + dst_slots, + plan_ring_start, + plan_flush_count, + # Runtime sizes. + block_table_stride_req: tl.int64, + slot_table_stride_layer: tl.int64, + num_reqs, + # Compile-time model constants. + MAX_LAYERS_PER_GROUP: tl.constexpr, + MAMBA_BLOCK_SIZE: tl.constexpr, + LOGICAL_WINDOW: tl.constexpr, + RING_BUFFER_LEN: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + QUERY_IS_CUMULATIVE: tl.constexpr, + NUM_COMPUTED_IS_AFTER: tl.constexpr, + HAS_IDX_MAPPING: 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, + """Plan materialization and commit all ReplaySSM trackers in one launch. + + One CTA owns one ``(batch row, cache group)`` pair, and is therefore the + only writer of that group's tracker for the request. Group zero writes the + request-level ``ring_start``/``flush_count`` snapshot shared by every layer; + every group fills the layer rows belonging to its physical slot namespace. + """ + batch_idx = tl.program_id(0) + group_idx = tl.program_id(1) + active = batch_idx < num_reqs + req_idx = batch_idx + if HAS_IDX_MAPPING: + req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) + valid_req = active & (req_idx >= 0) + + if group_idx == 0: + # Always overwrite the request decision, including padded rows, so a + # fixed-capacity FlashInfer call never observes stale work. + tl.store(plan_ring_start + batch_idx, 0) + tl.store(plan_flush_count + batch_idx, -1) + + block_table = tl.load(block_table_ptrs + group_idx).to(tl.pointer_type(tl.int32)) + tracker_start = tl.load(tracker_ring_start_ptrs + group_idx).to( + tl.pointer_type(tl.int32) + ) + tracker_committed = tl.load(tracker_num_committed_ptrs + group_idx).to( + tl.pointer_type(tl.int32) + ) + tracker_capacity = tl.load(tracker_capacities + group_idx) + + live_col = tl.load(live_cols + req_idx, mask=valid_req, other=-1) + valid_live_col = valid_req & (live_col >= 0) + live_slot = tl.load( + block_table + batch_idx * block_table_stride_req + live_col, + mask=valid_live_col, + other=PAD_SLOT_ID, + ) + valid_live = ( + valid_live_col + & (live_slot != PAD_SLOT_ID) + & (live_slot >= 0) + & (live_slot < tracker_capacity) ) - valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states) - if RESET: - tl.store(ring_start + slots, 0, mask=valid) - tl.store(prev_num_accepted + slots, 0, mask=valid) - tl.store(prev_query_len + slots, 0, mask=valid) + + src_col = tl.load(materialize_src_cols + batch_idx, mask=active, other=-1) + dst_col = tl.load(materialize_dst_cols + batch_idx, mask=active, other=-1) + wants_materialize = valid_req & (src_col >= 0) & (dst_col >= 0) + materialize_src_slot = tl.load( + block_table + batch_idx * block_table_stride_req + src_col, + mask=wants_materialize, + other=PAD_SLOT_ID, + ) + materialize_dst_slot = tl.load( + block_table + batch_idx * block_table_stride_req + dst_col, + mask=wants_materialize, + other=PAD_SLOT_ID, + ) + valid_materialize = ( + wants_materialize + & (materialize_src_slot != PAD_SLOT_ID) + & (materialize_dst_slot != PAD_SLOT_ID) + & (materialize_src_slot >= 0) + & (materialize_dst_slot >= 0) + & (materialize_src_slot < tracker_capacity) + & (materialize_dst_slot < tracker_capacity) + ) + + # Fill every flattened layer row for this group. Invalid rows still receive + # the pad sentinel; request-level flush_count=-1 suppresses native writes. + layer_begin = tl.load(group_layer_offsets + group_idx) + layer_end = tl.load(group_layer_offsets + group_idx + 1) + for layer_offset in tl.static_range(0, MAX_LAYERS_PER_GROUP): + layer_idx = layer_begin + layer_offset + layer_valid = layer_idx < layer_end + table_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store( + src_slots + table_offset, + tl.where(valid_materialize, materialize_src_slot, PAD_SLOT_ID), + mask=layer_valid, + ) + tl.store( + dst_slots + table_offset, + tl.where(valid_materialize, materialize_dst_slot, PAD_SLOT_ID), + mask=layer_valid, + ) + + prefilling = tl.load(is_prefilling + batch_idx, mask=active, other=1) + if QUERY_IS_CUMULATIVE: + query_len = tl.load( + query_metadata + batch_idx + 1, mask=active, other=0 + ) - tl.load( + query_metadata + batch_idx, + mask=active, + other=0, + ) else: - prev = tl.load(prev_num_accepted + slots, mask=valid, other=0) - start = tl.load(ring_start + slots, mask=valid, other=0) - must_checkpoint = prev + 1 > logical_window - next_start = tl.where( - must_checkpoint, - (start + prev) % ring_buffer_len, - start, + query_len = tl.load(query_metadata + batch_idx, mask=active, other=0) + + if valid_req & prefilling: + computed = tl.load(num_computed_tokens + req_idx) + computed_before = tl.where( + NUM_COMPUTED_IS_AFTER, computed - query_len, computed ) - next_prev = tl.where(must_checkpoint, 1, prev + 1) - tl.store(ring_start + slots, next_start, mask=valid) - tl.store(prev_num_accepted + slots, next_prev, mask=valid) - - -@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, + computed_after = computed_before + query_len + first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) + last_col = tl.maximum( + (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1, + 0, + ) + # All-mode prefill writes every boundary state in this interval. Reset + # every corresponding cursor so any later prefix hit copies an exact + # canonical state instead of replaying rows from the slot's old owner. + for col in tl.range(first_col, last_col + 1): + prefill_slot = tl.load( + block_table + batch_idx * block_table_stride_req + col + ) + valid_prefill_slot = ( + (prefill_slot != PAD_SLOT_ID) + & (prefill_slot >= 0) + & (prefill_slot < tracker_capacity) + ) + tl.store(tracker_start + prefill_slot, 0, mask=valid_prefill_slot) + tl.store(tracker_committed + prefill_slot, 0, mask=valid_prefill_slot) + + if valid_live: + if prefilling: + if valid_materialize & (group_idx == 0): + # The prefill kernel already produced an exact canonical state; + # count zero asks FlashInfer to copy it byte-for-byte. + tl.store(plan_ring_start + batch_idx, 0) + tl.store(plan_flush_count + batch_idx, 0) + else: + old_start = tl.load(tracker_start + live_slot) + old_committed = tl.load(tracker_committed + live_slot) + accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) + checkpointed = old_committed + query_len > LOGICAL_WINDOW + next_start = tl.where( + checkpointed, + (old_start + old_committed) % RING_BUFFER_LEN, + old_start, + ) + next_committed = tl.where(checkpointed, accepted, old_committed + accepted) + + if valid_materialize & (group_idx == 0): + boundary_count = tl.load(materialize_token_counts + batch_idx) + flush_count = next_committed - (accepted - boundary_count) + tl.store(plan_ring_start + batch_idx, next_start) + tl.store(plan_flush_count + batch_idx, flush_count) + + tl.store(tracker_start + live_slot, next_start) + tl.store(tracker_committed + live_slot, next_committed) + + if valid_materialize: + # The immutable plan above preserves any in-place transition for the + # materializer; subsequent forwards see a canonical empty replay. + tl.store(tracker_start + materialize_dst_slot, 0) + tl.store(tracker_committed + materialize_dst_slot, 0) + + +@triton.jit(do_not_specialize=["num_reqs"]) +def _copy_reassigned_replayssm_slots_kernel( + idx_mapping, + src_cols, + dst_cols, + block_table_ptrs, + tracker_ring_start_ptrs, + tracker_num_committed_ptrs, + tracker_capacities, + group_layer_offsets, + src_slots, + dst_slots, + plan_ring_start, + plan_flush_count, + block_table_stride_req: tl.int64, + slot_table_stride_layer: tl.int64, + num_reqs, + MAX_LAYERS_PER_GROUP: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, + HAS_IDX_MAPPING: 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, + """Plan an exact copy when align reassigns a request's writable slot.""" + batch_idx = tl.program_id(0) + group_idx = tl.program_id(1) + active = batch_idx < num_reqs + req_idx = batch_idx + if HAS_IDX_MAPPING: + req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) + valid_req = active & (req_idx >= 0) + + if group_idx == 0: + tl.store(plan_ring_start + batch_idx, 0) + tl.store(plan_flush_count + batch_idx, -1) + + block_table = tl.load(block_table_ptrs + group_idx).to(tl.pointer_type(tl.int32)) + tracker_start = tl.load(tracker_ring_start_ptrs + group_idx).to( + tl.pointer_type(tl.int32) ) - valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states) - prev = tl.load(prev_num_accepted + slots, mask=valid, other=0) - start = tl.load(ring_start + slots, mask=valid, other=0) - previous_query_len = tl.load(prev_query_len + slots, mask=valid, other=0) - accepted = tl.load(num_accepted_tokens + offsets, mask=mask, other=0) - must_checkpoint = (previous_query_len > 0) & ( - prev + previous_query_len > logical_window + tracker_committed = tl.load(tracker_num_committed_ptrs + group_idx).to( + tl.pointer_type(tl.int32) ) - next_start = tl.where( - must_checkpoint, - (start + prev) % ring_buffer_len, - start, + tracker_capacity = tl.load(tracker_capacities + group_idx) + src_col = tl.load(src_cols + req_idx, mask=valid_req, other=-1) + dst_col = tl.load(dst_cols + req_idx, mask=valid_req, other=-1) + wants_copy = valid_req & (src_col >= 0) & (dst_col >= 0) & (src_col != dst_col) + src_slot = tl.load( + block_table + batch_idx * block_table_stride_req + src_col, + mask=wants_copy, + other=PAD_SLOT_ID, ) - next_prev = tl.where( - previous_query_len == 0, - 0, - tl.where(must_checkpoint, accepted, prev + accepted), + dst_slot = tl.load( + block_table + batch_idx * block_table_stride_req + dst_col, + mask=wants_copy, + other=PAD_SLOT_ID, ) - current_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_query_len, mask=valid) + valid_mapping = ( + wants_copy + & (src_slot != PAD_SLOT_ID) + & (dst_slot != PAD_SLOT_ID) + & (src_slot >= 0) + & (dst_slot >= 0) + & (src_slot < tracker_capacity) + & (dst_slot < tracker_capacity) + ) + needs_copy = valid_mapping & (src_slot != dst_slot) + if valid_mapping & (group_idx == 0): + # Snapshot the source cursor before resetting the distinct destination. + # The materializer uses it to copy the exact live state, including any + # committed replay rows that have not reached a prefix boundary. The + # logical migration activates the shared plan even when group 0 aliases; + # every group independently suppresses unchanged physical slots below. + tl.store(plan_ring_start + batch_idx, tl.load(tracker_start + src_slot)) + tl.store( + plan_flush_count + batch_idx, + tl.load(tracker_committed + src_slot), + ) + layer_begin = tl.load(group_layer_offsets + group_idx) + layer_end = tl.load(group_layer_offsets + group_idx + 1) + for layer_offset in tl.static_range(0, MAX_LAYERS_PER_GROUP): + layer_idx = layer_begin + layer_offset + layer_valid = layer_idx < layer_end + table_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store( + src_slots + table_offset, + tl.where(needs_copy, src_slot, PAD_SLOT_ID), + mask=layer_valid, + ) + tl.store( + dst_slots + table_offset, + tl.where(needs_copy, dst_slot, PAD_SLOT_ID), + mask=layer_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, - pad_slot_id: int = NULL_BLOCK_ID, -) -> None: - """Reset selected trackers, or advance them when a window is provided.""" - if state_batch_indices.dim() > 1: - state_batch_indices = state_batch_indices[:, 0] - n_slots = state_batch_indices.numel() - if n_slots == 0: - return - reset = logical_window is None - if reset: - logical_window = 0 - ring_buffer_len = 1 - else: - assert ring_buffer_len is not None - block = 128 - _update_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)]( - ring_start, - prev_num_accepted, - prev_query_len, - state_batch_indices, - 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, - RESET=reset, - BLOCK=block, - ) + if needs_copy: + # The reassigned destination must not inherit its prior owner's cursor. + tl.store(tracker_start + dst_slot, 0) + tl.store(tracker_committed + dst_slot, 0) + + +@dataclass +class ReplaySSMModelContext: + """Persistent all-layer tables for ReplaySSM post-step maintenance.""" + + mixers: list[Any] + group_layer_offsets: torch.Tensor + block_table_ptrs: torch.Tensor + tracker_ring_start_ptrs: torch.Tensor + tracker_num_committed_ptrs: torch.Tensor + tracker_capacities: torch.Tensor + state_ptrs: torch.Tensor + state_slot_strides: torch.Tensor + x_cache_ptrs: torch.Tensor + x_cache_slot_strides: torch.Tensor + b_cache_ptrs: torch.Tensor + b_cache_slot_strides: torch.Tensor + dt_cache_ptrs: torch.Tensor + dt_cache_slot_strides: torch.Tensor + a_ptrs: torch.Tensor + scale_ptrs: torch.Tensor + scale_slot_strides: torch.Tensor + src_slots: torch.Tensor + dst_slots: torch.Tensor + plan_ring_start: torch.Tensor + plan_flush_count: torch.Tensor + precopy_src_slots: torch.Tensor + precopy_dst_slots: torch.Tensor + precopy_ring_start: torch.Tensor + precopy_flush_count: torch.Tensor + block_table_stride_req: int + max_num_reqs: int + num_groups: int + max_layers_per_group: int + logical_window: int + ring_buffer_len: int + + @classmethod + def create( + cls, + kv_cache_config: KVCacheConfig, + mamba_group_ids: Sequence[int], + forward_context: Mapping[str, Any], + block_tables: Sequence[torch.Tensor], + max_num_reqs: int, + ) -> "ReplaySSMModelContext | None": + grouped = _flashinfer_replayssm_mixers_by_group( + kv_cache_config, mamba_group_ids, forward_context + ) + if not grouped: + return None + if len(block_tables) != len(mamba_group_ids): + raise ValueError( + f"expected {len(mamba_group_ids)} Mamba block tables, " + f"got {len(block_tables)}" + ) + block_table_by_gid = dict(zip(mamba_group_ids, block_tables)) + replayssm_block_tables = [block_table_by_gid[gid] for gid, _ in grouped] + + mixers = [mixer for _, group_mixers in grouped for mixer in group_mixers] + if not _replayssm_materialize_ready(mixers): + return None + first = mixers[0] + first_ssm = first.kv_cache[1] + first_x = first.kv_cache[2] + first_b = first.kv_cache[4] + compatibility = ( + first_ssm.dtype, + first_x.dtype, + first.A.dtype, + first_ssm.size(1), + first_ssm.size(2), + first_ssm.size(3), + first_ssm.size(1) // first_b.size(1), + int(first.replayssm_buffer_len), + first_x.size(2), + bool(first.mamba_config.enable_stochastic_rounding), + int(first.mamba_config.stochastic_rounding_philox_rounds or 0), + ) + for mixer in mixers[1:]: + ssm = mixer.kv_cache[1] + x_cache = mixer.kv_cache[2] + b_cache = mixer.kv_cache[4] + current = ( + ssm.dtype, + x_cache.dtype, + mixer.A.dtype, + ssm.size(1), + ssm.size(2), + ssm.size(3), + ssm.size(1) // b_cache.size(1), + int(mixer.replayssm_buffer_len), + x_cache.size(2), + bool(mixer.mamba_config.enable_stochastic_rounding), + int(mixer.mamba_config.stochastic_rounding_philox_rounds or 0), + ) + if current != compatibility: + raise ValueError( + "A single model-wide FlashInfer ReplaySSM materialization " + "launch requires identical layer specialization; got " + f"{compatibility} and {current}" + ) + + device = first_ssm.device + group_offsets = [0] + for _, group_mixers in grouped: + group_offsets.append(group_offsets[-1] + len(group_mixers)) + max_layers_per_group = max( + group_offsets[i + 1] - group_offsets[i] + for i in range(len(group_offsets) - 1) + ) + tracker_owners = [group_mixers[0] for _, group_mixers in grouped] + strides = {int(block_table.stride(0)) for block_table in replayssm_block_tables} + if len(strides) != 1: + raise ValueError( + "model-wide ReplaySSM requires one block-table row stride; " + f"got {sorted(strides)}" + ) + zero_table = torch.zeros(len(mixers), dtype=torch.int64, device=device) + return cls( + mixers=mixers, + group_layer_offsets=torch.tensor( + group_offsets, dtype=torch.int32, device=device + ), + block_table_ptrs=_cuda_i64_ptrs(replayssm_block_tables), + tracker_ring_start_ptrs=_cuda_i64_ptrs( + [m._replayssm_ring_start for m in tracker_owners] + ), + tracker_num_committed_ptrs=_cuda_i64_ptrs( + [m._replayssm_prev_num_accepted for m in tracker_owners] + ), + tracker_capacities=torch.tensor( + [m._replayssm_ring_start.numel() for m in tracker_owners], + dtype=torch.int32, + device=device, + ), + state_ptrs=_cuda_i64_ptrs([m.kv_cache[1] for m in mixers]), + state_slot_strides=_cuda_i64_slot_strides([m.kv_cache[1] for m in mixers]), + x_cache_ptrs=_cuda_i64_ptrs([m.kv_cache[2] for m in mixers]), + x_cache_slot_strides=_cuda_i64_slot_strides( + [m.kv_cache[2] for m in mixers] + ), + b_cache_ptrs=_cuda_i64_ptrs([m.kv_cache[4] for m in mixers]), + b_cache_slot_strides=_cuda_i64_slot_strides( + [m.kv_cache[4] for m in mixers] + ), + dt_cache_ptrs=_cuda_i64_ptrs([m.kv_cache[3] for m in mixers]), + dt_cache_slot_strides=_cuda_i64_slot_strides( + [m.kv_cache[3] for m in mixers] + ), + a_ptrs=_cuda_i64_ptrs([m.A for m in mixers]), + scale_ptrs=zero_table, + scale_slot_strides=zero_table.clone(), + src_slots=torch.full( + (len(mixers), max_num_reqs), + NULL_BLOCK_ID, + dtype=torch.int32, + device=device, + ), + dst_slots=torch.full( + (len(mixers), max_num_reqs), + NULL_BLOCK_ID, + dtype=torch.int32, + device=device, + ), + plan_ring_start=torch.zeros(max_num_reqs, dtype=torch.int32, device=device), + plan_flush_count=torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ), + precopy_src_slots=torch.full( + (len(mixers), max_num_reqs), + NULL_BLOCK_ID, + dtype=torch.int32, + device=device, + ), + precopy_dst_slots=torch.full( + (len(mixers), max_num_reqs), + NULL_BLOCK_ID, + dtype=torch.int32, + device=device, + ), + precopy_ring_start=torch.zeros( + max_num_reqs, dtype=torch.int32, device=device + ), + precopy_flush_count=torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ), + block_table_stride_req=next(iter(strides)), + max_num_reqs=max_num_reqs, + num_groups=len(grouped), + max_layers_per_group=max_layers_per_group, + logical_window=int(first.replayssm_buffer_len), + ring_buffer_len=first_x.size(2), + ) -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: - """Reset selected 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 postprocess_and_materialize( + self, + *, + idx_mapping: torch.Tensor | None, + query_metadata: torch.Tensor, + query_is_cumulative: bool, + num_computed_tokens: torch.Tensor, + num_computed_is_after: bool, + num_accepted_tokens: torch.Tensor, + is_prefilling: torch.Tensor, + live_cols: torch.Tensor, + materialize_src_cols: torch.Tensor, + materialize_dst_cols: torch.Tensor, + materialize_token_counts: torch.Tensor, + mamba_block_size: int, + num_reqs: int, + materialize_possible: bool = True, + ) -> None: + """Commit lifecycle metadata, then materialize all layers once.""" + if num_reqs == 0: + return + _postprocess_replayssm_modelwide_kernel[(self.max_num_reqs, self.num_groups)]( + idx_mapping, + query_metadata, + num_computed_tokens, + num_accepted_tokens, + is_prefilling, + live_cols, + materialize_src_cols, + materialize_dst_cols, + materialize_token_counts, + self.block_table_ptrs, + self.tracker_ring_start_ptrs, + self.tracker_num_committed_ptrs, + self.tracker_capacities, + self.group_layer_offsets, + self.src_slots, + self.dst_slots, + self.plan_ring_start, + self.plan_flush_count, + self.block_table_stride_req, + self.src_slots.stride(0), + num_reqs, + MAX_LAYERS_PER_GROUP=self.max_layers_per_group, + MAMBA_BLOCK_SIZE=mamba_block_size, + LOGICAL_WINDOW=self.logical_window, + RING_BUFFER_LEN=self.ring_buffer_len, + PAD_SLOT_ID=NULL_BLOCK_ID, + QUERY_IS_CUMULATIVE=query_is_cumulative, + NUM_COMPUTED_IS_AFTER=num_computed_is_after, + HAS_IDX_MAPPING=idx_mapping is not None, + ) + if materialize_possible: + self._materialize_planned( + self.src_slots, + self.dst_slots, + self.plan_ring_start, + self.plan_flush_count, + ) -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.""" - 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, - ) + def copy_reassigned_slots( + self, + *, + idx_mapping: torch.Tensor | None, + src_cols: torch.Tensor, + dst_cols: torch.Tensor, + num_reqs: int, + ) -> None: + """Copy exact live state when align assigns a new writable slot.""" + if num_reqs == 0: + return + _copy_reassigned_replayssm_slots_kernel[ + (self.max_num_reqs, self.num_groups) + ]( + idx_mapping, + src_cols, + dst_cols, + self.block_table_ptrs, + self.tracker_ring_start_ptrs, + self.tracker_num_committed_ptrs, + self.tracker_capacities, + self.group_layer_offsets, + self.precopy_src_slots, + self.precopy_dst_slots, + self.precopy_ring_start, + self.precopy_flush_count, + self.block_table_stride_req, + self.precopy_src_slots.stride(0), + num_reqs, + MAX_LAYERS_PER_GROUP=self.max_layers_per_group, + PAD_SLOT_ID=NULL_BLOCK_ID, + HAS_IDX_MAPPING=idx_mapping is not None, + ) + self._materialize_planned( + self.precopy_src_slots, + self.precopy_dst_slots, + self.precopy_ring_start, + self.precopy_flush_count, + ) + + def _materialize_planned( + self, + src_slots: torch.Tensor, + dst_slots: torch.Tensor, + ring_start: torch.Tensor, + flush_count: torch.Tensor, + ) -> None: + first = self.mixers[0] + mamba_config = first.mamba_config + rand_seed = None + philox_rounds = 0 + if mamba_config.enable_stochastic_rounding: + rand_seed = torch.randint( + 0, 2**32, (1,), device=src_slots.device, dtype=torch.int64 + ) + philox_rounds = mamba_config.stochastic_rounding_philox_rounds or 10 + _load_replayssm_materialize()( + self.state_ptrs, + self.state_slot_strides, + self.x_cache_ptrs, + self.x_cache_slot_strides, + self.b_cache_ptrs, + self.b_cache_slot_strides, + self.dt_cache_ptrs, + self.dt_cache_slot_strides, + self.a_ptrs, + self.scale_ptrs, + self.scale_slot_strides, + src_slots, + dst_slots, + ring_start, + flush_count, + state_dtype=first.kv_cache[1].dtype, + input_dtype=first.kv_cache[2].dtype, + matrixA_dtype=first.A.dtype, + dim=first.kv_cache[1].size(2), + dstate=first.kv_cache[1].size(3), + num_heads=first.kv_cache[1].size(1), + heads_per_group=(first.kv_cache[1].size(1) // first.kv_cache[4].size(1)), + max_window=self.logical_window, + ring_buffer_len=self.ring_buffer_len, + rand_seed=rand_seed, + philox_rounds=philox_rounds, + ) class MambaSSUBackend(ABC): @@ -487,22 +933,19 @@ 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, dt_softplus: bool = False, state_batch_indices: torch.Tensor | None = None, null_block_id: int = NULL_BLOCK_ID, scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, - update_trackers: bool = True, enable_stochastic_rounding: bool = False, stochastic_rounding_philox_rounds: int = 0, 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.""" + """Run FlashInfer checkpointing SSU with model-owned tracker metadata.""" if _flashinfer_replayssm_kernel is None: raise RuntimeError( "FlashInfer ReplaySSM has not been initialized. " @@ -530,7 +973,7 @@ def selective_state_update_replayssm_flashinfer( if enable_stochastic_rounding else None ) - result = _flashinfer_replayssm_kernel( + return _flashinfer_replayssm_kernel( state, x_cache, B_cache, @@ -557,17 +1000,98 @@ def selective_state_update_replayssm_flashinfer( cumAdt_vec=cumAdt_vec, cb_old=cb_old, ) - if update_trackers and indices is not None: - update_replayssm_ring_trackers( - ring_start, - prev_num_accepted_tokens, - prev_query_len, - indices, - logical_window=logical_window, - ring_buffer_len=x_cache.size(2), - pad_slot_id=null_block_id, + + +def _reinterpret_u64_as_i64(value: int) -> int: + """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" + return value if value < (1 << 63) else value - (1 << 64) + + +def _cuda_i64_ptrs(tensors: list[torch.Tensor]) -> torch.Tensor: + return torch.tensor( + [_reinterpret_u64_as_i64(t.data_ptr()) for t in tensors], + dtype=torch.int64, + device=tensors[0].device, + ) + + +def _cuda_i64_slot_strides(tensors: list[torch.Tensor]) -> torch.Tensor: + return torch.tensor( + [t.stride(0) for t in tensors], + dtype=torch.int64, + device=tensors[0].device, + ) + + +def _flashinfer_replayssm_mixers_by_group( + kv_cache_config: KVCacheConfig, + mamba_group_ids: Sequence[int], + forward_context: Mapping[str, Any], +) -> list[tuple[int, list[Any]]]: + grouped: list[tuple[int, list[Any]]] = [] + for gid in mamba_group_ids: + mixers: list[Any] = [] + for layer_name in kv_cache_config.kv_cache_groups[gid].layer_names: + layer = forward_context.get(layer_name) + if layer is None: + continue + kv_cache = getattr(layer, "kv_cache", ()) + mamba_config = getattr(layer, "mamba_config", None) + backend = getattr(mamba_config, "backend", None) + if ( + getattr(layer, "use_replayssm", False) + and backend == MambaBackendEnum.FLASHINFER + and len(kv_cache) >= 5 + ): + mixers.append(layer) + if mixers: + grouped.append((gid, mixers)) + return grouped + + +@cache +def _load_replayssm_materialize() -> Callable[..., None]: + try: + from flashinfer.mamba.replayssm_materialize import ( + replayssm_materialize, + ) + except ImportError as e: + raise ImportError( + "FlashInfer ReplaySSM prefix caching requires " + "flashinfer.mamba.replayssm_materialize" + ) from e + return replayssm_materialize + + +def _replayssm_materialize_ready(mixers: list[Any]) -> bool: + """False only before the caches are allocated; raises on a bad cache. + + A skip here is not free: ``state_skip_postprocess`` has already told the + fused postprocess kernel not to copy this temporal state, so silently + doing nothing would leave the destination block holding stale SSM state. + The empty-cache case (profiling and other pre-allocation runs) is the one + legitimate no-op; anything else is a misconfiguration and must be loud. + """ + ssm = mixers[0].kv_cache[1] + x_cache = mixers[0].kv_cache[2] + if ssm.numel() == 0: + return False + if not ssm.is_cuda: + raise RuntimeError( + "FlashInfer ReplaySSM prefix materialization requires a CUDA SSM " + f"state cache; got device {ssm.device}" ) - return result + if x_cache.numel() == 0 or mixers[0]._replayssm_ring_start.numel() == 0: + raise RuntimeError( + "FlashInfer ReplaySSM prefix materialization requires allocated " + "replay ring buffers and ring trackers" + ) + if not mixers[0].replayssm_buffer_len: + raise RuntimeError( + "FlashInfer ReplaySSM prefix materialization requires " + "--replayssm-buffer-len >= 1" + ) + return True def initialize_mamba_ssu_backend( diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index a3401dfcf522..daaf8cc9f88a 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -745,11 +745,6 @@ def get_mamba_state_dtype_from_config( cache_config.mamba_cache_dtype, cache_config.mamba_ssm_cache_dtype, ) - if cache_config.use_replayssm: - return MambaStateDtypeCalculator.append_replayssm_ring( - base_dtype, - vllm_config.model_config.dtype, - ) return base_dtype @classmethod @@ -766,10 +761,8 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache - - x_cache/dt_cache/B_cache ring-buffer shapes (use_replayssm only) """ parallel_config = vllm_config.parallel_config - cache_config = vllm_config.cache_config hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim @@ -783,15 +776,6 @@ def get_mamba_state_shape_from_config( conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, ) - if cache_config.use_replayssm: - return MambaStateShapeCalculator.append_replayssm_ring( - base_shapes=base_shape, - n_groups=hf_config.n_groups, - tp_world_size=parallel_config.tensor_parallel_size, - logical_window=cache_config.replayssm_buffer_len, - backend=vllm_config.mamba_config.backend, - num_speculative_tokens=vllm_config.num_speculative_tokens, - ) return base_shape @classmethod diff --git a/vllm/model_executor/warmup/replayssm_warmup.py b/vllm/model_executor/warmup/replayssm_warmup.py index 28c45804ca7e..701cf2b8ec09 100644 --- a/vllm/model_executor/warmup/replayssm_warmup.py +++ b/vllm/model_executor/warmup/replayssm_warmup.py @@ -75,37 +75,18 @@ def _temporary_replayssm_autotune_state( runner: "GPUModelRunner", max_num_reqs: int ) -> Iterator[None]: from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 - from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( - reset_replayssm_ring_trackers, - update_replayssm_ring_trackers, - ) reset_tensors: dict[int, torch.Tensor] = {} - tracker_specs: dict[ - int, tuple[torch.Tensor, torch.Tensor, 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), - ), - ) tensors = ( *module.kv_cache, ring_start, prev_num_accepted, - prev_query_len, ) for tensor in tensors: if tensor.numel(): @@ -124,35 +105,6 @@ def _temporary_replayssm_autotune_state( block_table.block_table.np[:max_num_reqs, 0] = dummy_block_ids runner.input_batch.block_table.commit_block_table(max_num_reqs) - first_tracker = next(iter(tracker_specs.values()), None) - if first_tracker is not None and first_tracker[0].is_cuda: - state_slots = torch.arange( - 1, max_num_reqs + 1, dtype=torch.int32, device=first_tracker[0].device - ) - for ( - ring_start, - prev_num_accepted, - 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, 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, prev_query_len, state_slots - ) - try: yield finally: diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index f7591c557eef..297fb5516ae9 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -182,6 +182,8 @@ def __init__( ) = None self.decode_replayssm_state_indices_d: torch.Tensor | None = None # ReplaySSM CUDA-graph buffers for the selected backend. + if self.use_replayssm: + assert len(kv_cache_spec.replayssm_shapes) == 3 if self.use_replayssm and not self.use_flashinfer_replayssm: self.decode_write_pos_d: torch.Tensor = torch.empty( (self.decode_cudagraph_max_bs,), @@ -193,9 +195,8 @@ def __init__( dtype=torch.int8, device=device, ) - # B_cache shape = (ngroups, replayssm_buffer_len, dstate); the page - # layout is (conv_state, ssm_state, x_cache, dt_cache, B_cache). - bc_ngroups = kv_cache_spec.shapes[4][0] + # B_cache shape = (ngroups, replayssm_buffer_len, dstate). + bc_ngroups = kv_cache_spec.replayssm_shapes[2][0] bc_scratch_bs = max( self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs ) @@ -213,7 +214,7 @@ def __init__( allocate_checkpointing_ssu_scratch, ) - nheads = kv_cache_spec.shapes[2][0] + nheads = kv_cache_spec.replayssm_shapes[0][0] self.decode_replayssm_scratch = allocate_checkpointing_ssu_scratch( batch_size=scheduler_config.max_num_seqs, num_heads=nheads, @@ -548,14 +549,14 @@ def _compute_common_metadata( ) = self._compute_prefix_caching_block_indices( common_attn_metadata, mamba_block_size ) - if self.use_spec_decode and prev_last_scheduled_idx is not None: - fallback = (num_computed_tokens - 1) // mamba_block_size - fallback.clamp_(min=0) - block_idx_last_scheduled_token_prev_step = torch.where( - prev_last_scheduled_idx >= 0, - prev_last_scheduled_idx, - fallback, - ) + if self.use_spec_decode: + block_idx_last_scheduled_token_prev_step = block_idx_last_computed_token + if prev_last_scheduled_idx is not None: + block_idx_last_scheduled_token_prev_step = torch.where( + prev_last_scheduled_idx >= 0, + prev_last_scheduled_idx, + block_idx_last_computed_token, + ) else: state_indices_tensor = mamba_get_block_table_tensor( common_attn_metadata.block_table_tensor, @@ -845,8 +846,12 @@ def _update_metadata_for_cudagraph_capture( cb_old[:padded_bs], ) assert self.decode_replayssm_state_indices_d is not None + live_state_indices = self._select_replayssm_state_indices( + state_indices_tensor_d, + block_idx_last_scheduled_token, + ) self.decode_replayssm_state_indices_d[:padded_bs].copy_( - state_indices_tensor_d[:, 0], non_blocking=True + live_state_indices, non_blocking=True ) replayssm_state_indices_d = self.decode_replayssm_state_indices_d[ :padded_bs @@ -857,7 +862,10 @@ def _update_metadata_for_cudagraph_capture( 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() + replayssm_state_indices_d = self._select_replayssm_state_indices( + state_indices_tensor_d, + block_idx_last_scheduled_token, + ).contiguous() return replace( metadata, @@ -876,6 +884,20 @@ def _update_metadata_for_cudagraph_capture( ), ) + def _select_replayssm_state_indices( + self, + state_indices_tensor_d: torch.Tensor, + block_idx_last_scheduled_token: torch.Tensor | None, + ) -> torch.Tensor: + if self.vllm_config.cache_config.mamba_cache_mode != "all": + return state_indices_tensor_d[:, 0] + + assert block_idx_last_scheduled_token is not None + live_cols = block_idx_last_scheduled_token[ + : state_indices_tensor_d.size(0) + ].to(torch.int64) + return state_indices_tensor_d.gather(1, live_cols.unsqueeze(1)).squeeze(1) + def update_block_table( self, metadata: M, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 1bbe002c97ef..ec24e987413a 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -983,9 +983,9 @@ def check_enough_kv_cache_memory( ) _check_enough_kv_cache_memory( check_memory, - lambda: max_memory_usage_bytes(vllm_config, kv_cache_spec.values()), + partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), vllm_config.model_config.max_model_len, - lambda am: estimate_max_model_len(vllm_config, kv_cache_spec, am), + partial(_estimate_max_model_len_from_groups, vllm_config, groups), ) @@ -1358,16 +1358,35 @@ def _get_per_layer_spec( def _get_kv_cache_bytes_per_block( kv_cache_groups: list[KVCacheGroupSpec], ) -> int: - """Return the largest cache group's bytes per block.""" - bytes_per_block = max( + """Return canonical KV plus ReplaySSM bytes per physical block.""" + return _get_kv_cache_main_bytes_per_block( + kv_cache_groups + ) + _get_replayssm_bytes_per_block(kv_cache_groups) + + +def _get_kv_cache_main_bytes_per_block( + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + return max( sum( _get_per_layer_spec(group, layer_name).page_size_bytes for layer_name in group.layer_names ) for group in kv_cache_groups ) - assert bytes_per_block > 0 - return bytes_per_block + + +def _get_replayssm_bytes_per_block( + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + # ReplaySSM ring tensors use standalone allocations, so unlike canonical + # cache groups their storage cannot overlay by physical block ID. + return sum( + spec.replayssm_size_bytes + for group in kv_cache_groups + for layer_name in group.layer_names + if isinstance((spec := _get_per_layer_spec(group, layer_name)), MambaSpec) + ) def validate_kv_cache_layout( @@ -1434,12 +1453,18 @@ def get_kv_cache_config_from_groups( layout = vllm_config.cache_config.get_resolved_kv_cache_layout() validate_kv_cache_layout(layout, kv_cache_groups) - bytes_per_block = _get_kv_cache_bytes_per_block(kv_cache_groups) - interleaved_block_stride = bytes_per_block if layout.is_block_outermost else None + main_bytes_per_block = _get_kv_cache_main_bytes_per_block(kv_cache_groups) + bytes_per_block = main_bytes_per_block + _get_replayssm_bytes_per_block( + kv_cache_groups + ) + assert bytes_per_block > 0 + interleaved_block_stride = ( + main_bytes_per_block if layout.is_block_outermost else None + ) num_blocks = available_memory // bytes_per_block num_blocks = may_override_num_blocks(vllm_config, num_blocks) - size = bytes_per_block * num_blocks + size = main_bytes_per_block * num_blocks # Groups alias from byte 0. Spec regions are laid out differently: # diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index fbb3a22611c0..a7d86a4088db 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -857,6 +857,8 @@ def is_uniform_with_collection( class MambaSpec(KVCacheSpec): shapes: tuple[tuple[int, ...], ...] dtypes: tuple[torch.dtype, ...] + replayssm_shapes: tuple[tuple[int, ...], ...] = () + replayssm_dtypes: tuple[torch.dtype, ...] = () page_size_padded: int | None = None mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 mamba_cache_mode: str = "none" @@ -868,6 +870,10 @@ class MambaSpec(KVCacheSpec): # rank holds the full state (e.g. the replicated PLE conv state). tp_replicated: bool = False + def __post_init__(self) -> None: + if len(self.replayssm_shapes) != len(self.replayssm_dtypes): + raise ValueError("ReplaySSM shapes and dtypes must have equal length") + @property def state_content_size_bytes(self) -> int: return sum( @@ -875,6 +881,13 @@ def state_content_size_bytes(self) -> int: for (shape, dtype) in zip(self.shapes, self.dtypes) ) + @property + def replayssm_size_bytes(self) -> int: + return sum( + prod(shape) * get_dtype_size(dtype) + for (shape, dtype) in zip(self.replayssm_shapes, self.replayssm_dtypes) + ) + @property def page_size_bytes(self) -> int: page_size = sum( @@ -919,6 +932,8 @@ def is_uniform_with_collection( and spec.num_speculative_blocks == self.num_speculative_blocks and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks and spec.page_size_bytes == self.page_size_bytes + and spec.replayssm_shapes == self.replayssm_shapes + and spec.replayssm_dtypes == self.replayssm_dtypes and spec.tp_replicated == self.tp_replicated for spec in kv_cache_specs.values() ) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 8a7043e54a14..6be50bcdff44 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -26,6 +26,7 @@ AttentionGroup, add_kv_sharing_layers_to_kv_cache_groups, allocate_kv_cache, + allocate_replayssm_caches, bind_kv_cache, prepare_kernel_block_sizes, ) @@ -219,6 +220,7 @@ def init_kv_cache( vllm_config.cache_config.get_resolved_kv_cache_layout(), kernel_block_sizes, ) + replayssm_caches = allocate_replayssm_caches(kv_cache_config, device) for layer_name, target in get_shared_kv_cache_layers(vllm_config).items(): kv_caches[layer_name] = kv_caches[target] # Dual-attention models (e.g. LongCat-Flash) put two Attention modules per @@ -235,6 +237,7 @@ def init_kv_cache( runner_kv_caches, num_attn_module, kv_cache_groups=kv_cache_config.kv_cache_groups, + replayssm_caches=replayssm_caches, ) return kv_caches diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 24e38d47a74b..b6cd6ea04190 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1496,6 +1496,7 @@ def postprocess_sampled( num_sampled: torch.Tensor, num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, ) -> None: # Update the number of computed tokens. if self.is_last_pp_rank: @@ -1517,7 +1518,11 @@ def postprocess_sampled( ) self.model_state.postprocess_state( - idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu + idx_mapping, + num_sampled, + self.req_states.num_computed_tokens.gpu, + query_start_loc, + is_prefilling, ) def _merge_ec_connector_no_forward( @@ -1609,7 +1614,7 @@ def execute_model( scheduler_output, batch_req_state, batch_desc ) block_tables, slot_mappings = self.prepare_attn(input_batch) - # Mamba "align" pre-copy: migrate recurrent state across block + # Mamba prefix-cache pre-copy: migrate recurrent state across block # boundaries before the forward. Runs only on real batches, and # before model_state.prepare_attn gathers num_accepted_tokens so the # boundary reset is visible to the attention metadata. diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index cd1b5b64da7d..0c5002b65398 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -152,6 +152,8 @@ def postprocess_state( idx_mapping: torch.Tensor, num_sampled: torch.Tensor, num_computed_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 6d07c1588ca0..123d89dbb1c2 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -9,6 +9,7 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.config.mamba import MambaBackendEnum from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFuncsByType from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder @@ -89,15 +90,26 @@ def __init__( self.num_accepted_tokens_gpu = torch.ones( self.max_num_reqs, dtype=torch.int32, device=self.device ) - # Pre-copy "align" prefix-cache state (V2). The migration of each - # request's mamba state across block boundaries runs as a fused GPU - # kernel reusing the postprocess copy machinery, so the per-step src - # columns and the running state_idx are kept GPU-resident. + self._is_prefilling_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.bool, device=self.device + ) + # Pre-copy prefix-cache state (V2). The migration of each request's + # mamba state across block boundaries runs as a fused GPU 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._use_flashinfer_replayssm = ( + self.cache_config.use_replayssm is True + and vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) + self._mamba_lifecycle_mode = self._align_mode or ( + self.cache_config.mamba_cache_mode == "all" + and self._use_flashinfer_replayssm + ) self.recoverssm = ( RecoverSSMState() if self.cache_config.use_kda_recoverssm else None ) - if self._align_mode: + if self._mamba_lifecycle_mode: self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device ) @@ -111,15 +123,21 @@ def __init__( self._mamba_group_ids: list[int] = [] self._mamba_spec: MambaSpec | None = None self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None + self._mamba_kv_cache_config: KVCacheConfig | None = None + self._mamba_block_tables: tuple[torch.Tensor, ...] | None = None def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) # Must reset the speculative acceptance count in this idx which could be stale. self.num_accepted_tokens_gpu[req_index].fill_(1) - if self._align_mode: + if self._mamba_lifecycle_mode: # Seed the running state block from the resumed/prefilled position. + state_block_size = self.cache_config.block_size + if self.cache_config.mamba_cache_mode == "all": + state_block_size = self.cache_config.mamba_block_size + assert state_block_size is not None self._mamba_state_idx_gpu[req_index].fill_( - (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + (new_req_data.num_computed_tokens - 1) // state_block_size ) def _get_mamba_group_info( @@ -144,6 +162,8 @@ def _ensure_align_ctx( mamba_group_ids: list[int], block_tables: tuple[torch.Tensor, ...], ) -> MambaSpecDecodeGPUContext: + self._mamba_kv_cache_config = kv_cache_config + self._mamba_block_tables = block_tables if self._mamba_state_copy_funcs is None: mamba_groups = get_mamba_groups(kv_cache_config) mamba_types = {spec.mamba_type for spec in mamba_groups} @@ -187,13 +207,13 @@ def preprocess_state( kv_cache_config: KVCacheConfig, num_computed_tokens: torch.Tensor, ) -> None: - """Migrate each request's mamba state across block boundaries before the - forward (V1 align semantics, done on GPU). Runs on real batches only - (dummy DP/profiling runs skip preprocess_state), and before + """Migrate each request's mamba state across block boundaries before + the forward (V1 lifecycle semantics, done on GPU). Runs on real batches + only (dummy DP/profiling runs skip preprocess_state), and before ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset is visible to the forward kernels. """ - if not self._align_mode: + if not self._mamba_lifecycle_mode: return num_reqs = input_batch.num_reqs if num_reqs == 0: @@ -221,6 +241,13 @@ def preprocess_state( BLOCK_SIZE=block, MAMBA_BLOCK_SIZE=mamba_spec.block_size, ) + if ctx.replayssm is not None: + ctx.replayssm.copy_reassigned_slots( + idx_mapping=input_batch.idx_mapping, + src_cols=self._mamba_src_col_gpu, + dst_cols=self._mamba_state_idx_gpu, + num_reqs=num_reqs, + ) ctx.run_fused_precopy( num_reqs, self._mamba_state_idx_gpu, @@ -258,6 +285,7 @@ def prepare_attn( is_prefilling[: input_batch.num_reqs] = torch.from_numpy( input_batch.is_prefilling_np ) + self._is_prefilling_gpu[:num_reqs].copy_(is_prefilling, non_blocking=True) # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only # compute them during actual (non-capture) forward execution. @@ -339,6 +367,8 @@ def postprocess_state( idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int, num_computed_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, ) -> None: # 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. @@ -371,12 +401,11 @@ def postprocess_state( if not num_reqs: return - # 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 - # post-step advanced count. + # Save the running state at a crossed block boundary when spec-decode + # acceptance leaves the sequence non-block-aligned. num_computed_tokens + # already holds the post-step advanced count. if ( - self._align_mode + self._mamba_lifecycle_mode and num_computed_tokens is not None and self._mamba_ctx is not None ): @@ -387,6 +416,40 @@ def postprocess_state( num_computed_tokens, idx_mapping, ) + # Must match the condition that sets ``state_skip_postprocess``: + # the fused kernel skips the temporal copy for every FlashInfer + # ReplaySSM layer, so the materializer has to cover every one of + # them. Gating this on spec decode would drop the SSM state on any + # non-spec boundary where src_col != dst_col. Rows that need no + # work carry the -1 src_col sentinel and cost nothing. + if self._use_flashinfer_replayssm: + replayssm = self._mamba_ctx.replayssm + assert replayssm is not None + if query_start_loc is None: + raise RuntimeError( + "ReplaySSM postprocess requires the query_start_loc from " + "the forward that produced this acceptance" + ) + if is_prefilling is None: + is_prefilling = self._is_prefilling_gpu[:num_reqs] + replayssm.postprocess_and_materialize( + idx_mapping=idx_mapping, + query_metadata=query_start_loc, + query_is_cumulative=True, + num_computed_tokens=num_computed_tokens, + num_computed_is_after=True, + # run_fused_postprocess_align can reset the live buffer to + # one for the next step. Its persistent snapshot still + # contains the acceptance produced by this forward. + num_accepted_tokens=self._mamba_ctx.num_accepted_tokens_out, + is_prefilling=is_prefilling, + live_cols=self._mamba_state_idx_gpu, + materialize_src_cols=self._mamba_ctx.materialize_src_cols, + materialize_dst_cols=self._mamba_ctx.materialize_dst_cols, + materialize_token_counts=(self._mamba_ctx.materialize_token_counts), + mamba_block_size=self._mamba_ctx.block_size, + num_reqs=num_reqs, + ) @triton.jit diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index cf52a6d3821e..20d588b81c32 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -24,6 +24,8 @@ class PendingRecv: num_sampled: torch.Tensor # [num_reqs] num_rejected: torch.Tensor # [num_reqs] idx_mapping: torch.Tensor # [num_reqs] + query_start_loc: torch.Tensor # [num_reqs + 1] + is_prefilling: torch.Tensor # [num_reqs] idx_mapping_np: np.ndarray # [num_reqs] # Records which rows need a deferred postprocess (bool). need_sampled_mask: np.ndarray # [num_reqs] @@ -117,6 +119,8 @@ def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: num_sampled=slot.num_sampled, num_rejected=slot.num_rejected, idx_mapping=idx_mapping, + query_start_loc=slot.query_start_loc, + is_prefilling=slot.is_prefilling, ) def receive(self, input_batch: InputBatch) -> bool: @@ -139,6 +143,13 @@ def receive(self, input_batch: InputBatch) -> bool: num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device) + # These input-buffer views are reused every step. Snapshot them in + # the same deferred slot as the sampled output so model-owned state + # postprocess observes the query that produced this acceptance. + query_start_loc = input_batch.query_start_loc[: num_reqs + 1].clone() + is_prefilling = async_copy_to_gpu( + input_batch.is_prefilling_np.copy(), device=self.device + ) torch.distributed.broadcast( sampled_tokens, src=self.last_rank, group=self.broadcast_group ) @@ -151,12 +162,16 @@ def receive(self, input_batch: InputBatch) -> bool: # later used on the main stream. sampled_tokens.record_stream(self.main_stream) combined.record_stream(self.main_stream) + query_start_loc.record_stream(self.main_stream) + is_prefilling.record_stream(self.main_stream) self.queue[-1] = PendingRecv( event, sampled_tokens, num_sampled, num_rejected, input_batch.idx_mapping, + query_start_loc, + is_prefilling, input_batch.idx_mapping_np, need_sampled_mask, gen_at_receive_np, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 1f79e6ef4460..8e985dc04cba 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -38,6 +38,7 @@ ) from vllm.config.cache import CacheConfig from vllm.config.ec_manager_config import EncoderCacheManagerMetadata +from vllm.config.mamba import MambaBackendEnum from vllm.config.model import PROCESSED_LOGPROBS_MODES from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.eplb.eplb_state import EplbState @@ -245,6 +246,7 @@ KVBlockZeroer, add_kv_sharing_layers_to_kv_cache_groups, allocate_kv_cache, + allocate_replayssm_caches, bind_kv_cache, copy_kv_cache_blocks_inplace, prepare_kernel_block_sizes, @@ -1001,7 +1003,11 @@ def __init__( self._mamba_bufs: mamba_utils.MambaBuffers | None = None self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None - if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0: + if ( + self.cache_config.mamba_cache_mode == "all" + and self.num_spec_tokens > 0 + and not self.cache_config.use_replayssm + ): self.mamba_prev_last_scheduled_idx = self._make_buffer( self.max_num_reqs, dtype=torch.int32 ) @@ -1066,10 +1072,13 @@ def _get_mamba_state_copy_funcs(self) -> MambaStateCopyFuncsByType: return self._mamba_state_copy_funcs def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: - # Only reachable on the ``mamba_cache_mode == "align"`` path. - # The postprocess sub-object is additionally gated on spec - # decode + hybrid model. - assert self.cache_config.mamba_cache_mode == "align" + # The postprocess sub-object is also the model-level owner of + # FlashInfer ReplaySSM trackers, including STP. + assert self.cache_config.mamba_cache_mode == "align" or ( + self.cache_config.mamba_cache_mode == "all" + and self.cache_config.use_replayssm + and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) if self._mamba_bufs is None: self._mamba_bufs = mamba_utils.MambaBuffers.create( max_num_reqs=self.max_num_reqs, @@ -1079,6 +1088,11 @@ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: device=self.device, with_postprocess_align=( self.speculative_config is not None and self.model_config.is_hybrid + ) + or ( + self.cache_config.use_replayssm + and self.vllm_config.mamba_config.backend + == MambaBackendEnum.FLASHINFER ), ) return self._mamba_bufs @@ -1587,7 +1601,14 @@ def _update_states_after_model_execute( each sequence, and a shifting is done during the next iteration based on the number of accepted tokens. """ - if not self.speculative_config or not self.model_config.is_hybrid: + modelwide_replayssm = ( + self.cache_config.mamba_cache_mode in ("align", "all") + and self.cache_config.use_replayssm + and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) + if not modelwide_replayssm and ( + not self.speculative_config or not self.model_config.is_hybrid + ): return # Count the number of accepted tokens for each sequence. @@ -1596,7 +1617,7 @@ def _update_states_after_model_execute( num_reqs = output_token_ids.size(0) self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) - if self.cache_config.mamba_cache_mode == "align": + if self.cache_config.mamba_cache_mode == "align" or modelwide_replayssm: # Fused GPU postprocess: state copies + per-request accepted-token # update without CPU-GPU sync. The metadata # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is @@ -1614,8 +1635,19 @@ def _update_states_after_model_execute( mamba_state_copy_funcs=self._get_mamba_state_copy_funcs(), ) - assert self.num_accepted_tokens_event is not None - self.num_accepted_tokens_event.record() + if self.num_accepted_tokens_event is not None: + self.num_accepted_tokens_event.record() + + if self.cache_config.mamba_cache_mode == "all": + mamba_utils.postprocess_mamba_all( + scheduler_output, + self.kv_cache_config, + self.input_batch, + self.requests, + self.mamba_state_idx, + self.num_spec_tokens, + num_reqs, + ) else: self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True @@ -2125,8 +2157,15 @@ def _prepare_inputs( # _update_states_after_model_execute for hybrid models). # Skipped under async scheduling (non-align): the CPU copy races with # the in-flight D2H copy and with input-batch row moves. - needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and not ( - self.use_async_scheduling and self.cache_config.mamba_cache_mode != "align" + modelwide_replayssm = ( + self.cache_config.mamba_cache_mode in ("align", "all") + and self.cache_config.use_replayssm + and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) + needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and ( + not self.use_async_scheduling + or self.cache_config.mamba_cache_mode == "align" + or modelwide_replayssm ) if needs_cpu_accepted_counts: assert self.num_accepted_tokens_event is not None @@ -4400,7 +4439,13 @@ def execute_model( ) pad_attn = cudagraph_mode == CUDAGraphMode.FULL - if self.cache_config.mamba_cache_mode == "align": + modelwide_replayssm = ( + self.cache_config.mamba_cache_mode in ("align", "all") + and self.cache_config.use_replayssm + and self.vllm_config.mamba_config.backend + == MambaBackendEnum.FLASHINFER + ) + if self.cache_config.mamba_cache_mode == "align" or modelwide_replayssm: # preprocess_mamba reads req_state.num_computed_tokens (CPU) # to decide copy operations, so we must apply deferred # corrections before it runs. @@ -4429,11 +4474,9 @@ def execute_model( ) self.num_accepted_tokens.copy_to_gpu(num_reqs) - # Stage per-request inputs for the fused postprocess kernel - # only when that kernel will actually run. The kernel is - # gated on spec-decode + hybrid (see MambaBuffers.create); - # without it, ``mamba_bufs.postprocess_align`` is None and - # the staging buffers don't exist. + # Stage inputs only when the fused postprocess will run. This + # includes spec-decode hybrid models and model-owned + # FlashInfer ReplaySSM lifecycle maintenance under STP. if mamba_bufs.postprocess_align is not None: mamba_utils.stage_postprocess_inputs_to_gpu( mamba_bufs.postprocess_align, @@ -7420,6 +7463,7 @@ def initialize_kv_cache_tensors( self.cache_config.get_resolved_kv_cache_layout(), kernel_block_sizes, ) + replayssm_caches = allocate_replayssm_caches(kv_cache_config, self.device) # Set up cross-layer KV cache sharing for layer_name, target_layer_name in self.shared_kv_cache_layers.items(): @@ -7435,6 +7479,7 @@ def initialize_kv_cache_tensors( self.kv_caches, num_attn_module, kv_cache_groups=kv_cache_config.kv_cache_groups, + replayssm_caches=replayssm_caches, ) return kv_caches diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 724c714057f7..cde95d033ff0 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -8,6 +8,7 @@ import torch from vllm.config import CacheConfig +from vllm.config.mamba import MambaBackendEnum from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFuncsByType, @@ -15,6 +16,9 @@ get_temporal_copy_spec, is_conv_state_dim_first, ) +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + ReplaySSMModelContext, +) from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.math_utils import cdiv @@ -387,11 +391,23 @@ def postprocess_mamba_fused_kernel( state_inner_sizes_ptr, # number of elements in inner dimensions state_conv_widths_ptr, # conv width for conv states (0 for temporal) state_group_indices_ptr, # maps state_idx to group index in block table + # Nonzero for temporal states reconstructed by an external materializer. + # The kernel still emits the request-level decision below, but does not + # overwrite the checkpoint with the generic speculative-column copy. + state_skip_postprocess_ptr, # DS conv row metadata. Zero keeps the single-region copy path. state_dim_row_count_ptr, # int32: per-block dim row count for DS conv state_dim_row_stride_ptr, # int64: bytes between rows for DS conv # Output: num_accepted_tokens update (for src==dst case) num_accepted_tokens_out_ptr, + # Batch-ordered ReplaySSM materialization decision. The caller initializes + # src_col to -1, which remains the no-op sentinel when no boundary is hit. + materialize_src_col_ptr, + # Logical column holding the block-aligned prefix checkpoint. + materialize_dst_col_ptr, + # Current-query rows through that boundary: accept_token_bias + 1. The + # materializer adds older pending rows from the source slot's tracker. + materialize_token_count_ptr, # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The # per-request decision arrays are in req-state-slot order; the block table # is in batch order, so HAS_IDX_MAPPING splits the two indexings. @@ -468,6 +484,14 @@ def postprocess_mamba_fused_kernel( accept_token_bias = aligned_new_computed - num_tokens_running_state dest_block_idx = aligned_new_computed // block_size - 1 + if state_idx == 0 and tile_idx == 0: + tl.store(materialize_src_col_ptr + batch_idx, src_block_idx) + tl.store(materialize_dst_col_ptr + batch_idx, dest_block_idx) + tl.store( + materialize_token_count_ptr + batch_idx, + accept_token_bias + 1, + ) + # Update accepted-token count before early exits (per-request, so only # state_idx == 0 writes). Also guard on tile_idx == 0 so tiles > 0 # (when TEMPORAL_TILES > 1) do not duplicate the store. @@ -478,6 +502,9 @@ def postprocess_mamba_fused_kernel( if src_block_idx == dest_block_idx and accept_token_bias == 0: return + if tl.load(state_skip_postprocess_ptr + state_idx): + return + bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx _copy_mamba_state_block( state_idx, @@ -564,6 +591,7 @@ def precopy_mamba_align_fused_kernel( state_inner_sizes_ptr, state_conv_widths_ptr, state_group_indices_ptr, + state_skip_precopy_ptr, state_dim_row_count_ptr, state_dim_row_stride_ptr, idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) @@ -608,6 +636,10 @@ def precopy_mamba_align_fused_kernel( # so there is nothing to copy. if src_col < 0 or src_col == dst_col: return + if tl.load(state_skip_precopy_ptr + state_idx): + # FlashInfer ReplaySSM materializes this temporal destination before the + # generic pre-copy launch. Copy only conv/other model-owned states here. + return token_bias = tl.load(token_bias_ptr + req_idx) _copy_mamba_state_block( @@ -713,10 +745,10 @@ def validate_mamba_state_copy_funcs( f"missing state copy funcs for {mamba_spec.mamba_type}" ) state_copy_funcs = copy_funcs[mamba_spec.mamba_type] - assert 0 < len(state_copy_funcs) <= len(mamba_spec.shapes), ( - f"{mamba_spec.mamba_type} declares {len(mamba_spec.shapes)} states, " - f"but provides {len(state_copy_funcs)} state copy funcs; expected " - "a non-empty copyable prefix" + assert len(state_copy_funcs) == len(mamba_spec.shapes), ( + f"{mamba_spec.mamba_type} expects {len(mamba_spec.shapes)} state copy " + "funcs for its canonical state tensors, but provides " + f"{len(state_copy_funcs)}" ) @@ -774,8 +806,8 @@ class MambaSpecDecodeGPUContext: Context for GPU-side Mamba state copy operations during the fused postprocess path. - Only used when speculative decoding is enabled on a hybrid model - (and the mamba_cache_config is in align mode). + Used for hybrid speculative state maintenance and for model-owned + FlashInfer ReplaySSM prefix-cache lifecycle maintenance. Precomputes memory layout metadata (base addresses, strides, element sizes) so the GPU kernel can perform state copies without CPU-GPU sync. @@ -794,6 +826,7 @@ class MambaSpecDecodeGPUContext: state_inner_sizes: torch.Tensor # int64: elements in inner dimensions state_conv_widths: torch.Tensor # int32: conv width (0 for temporal states) state_group_indices: torch.Tensor # int32: maps state_idx to group index + state_skip_postprocess: torch.Tensor # int32: materializer owns this state # DS conv row metadata. Zero keeps the single-region copy path. state_dim_row_count: torch.Tensor # int32: per-block dim row count state_dim_row_stride: torch.Tensor # int64: bytes between rows @@ -806,6 +839,9 @@ class MambaSpecDecodeGPUContext: # Output buffer for num_accepted_tokens updates num_accepted_tokens_out: torch.Tensor + materialize_src_cols: torch.Tensor + materialize_dst_cols: torch.Tensor + materialize_token_counts: torch.Tensor # Per-group block-table base addresses: int64[num_groups]. Populated in # initialize_from_forward_context from the persistent per-group block @@ -820,16 +856,29 @@ class MambaSpecDecodeGPUContext: # Per-request staging buffers (CPU+GPU mirrors). The runner stages # values into the CPU view in ``_prepare_inputs`` and the fused kernel # reads the GPU side. These only exist when the postprocess kernel is - # enabled (spec decode + hybrid + align mode). + # enabled (spec decode + hybrid, or FlashInfer ReplaySSM prefix caching). mamba_state_idx_buf: CpuGpuBuffer | None = None num_scheduled_tokens_buf: CpuGpuBuffer | None = None num_computed_tokens_buf: CpuGpuBuffer | None = None num_draft_tokens_buf: CpuGpuBuffer | None = None + is_prefilling_buf: CpuGpuBuffer | None = None precopy_src_col_buf: CpuGpuBuffer | None = None precopy_token_bias_buf: CpuGpuBuffer | None = None # Flag to track if metadata has been populated is_initialized: bool = False + # True when any temporal state is owned by the FlashInfer ReplaySSM + # materializer (i.e. some state_skip_postprocess entry is set). Cached at + # populate time so the per-step postprocess can skip the layer scan for + # Triton / non-ReplaySSM configs. + has_flashinfer_replayssm: bool = False + # Persistent all-layer ReplaySSM descriptors, populated with the cache + # addresses on first real forward. None for non-FlashInfer configurations. + replayssm: ReplaySSMModelContext | None = None + # Host-side upper bound computed while staging the current batch. False + # means no acceptance outcome can cross an align boundary, so the native + # materializer launch can be skipped after tracker maintenance. + replayssm_materialize_possible: bool = False @classmethod def create( @@ -887,6 +936,9 @@ def create( state_group_indices=torch.zeros( total_states, dtype=torch.int32, device=device ), + state_skip_postprocess=torch.zeros( + total_states, dtype=torch.int32, device=device + ), state_dim_row_count=torch.zeros( total_states, dtype=torch.int32, device=device ), @@ -900,6 +952,15 @@ def create( num_accepted_tokens_out=torch.zeros( max_num_reqs, dtype=torch.int32, device=device ), + materialize_src_cols=torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ), + materialize_dst_cols=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), + materialize_token_counts=torch.empty( + max_num_reqs, dtype=torch.int32, device=device + ), block_table_ptrs=torch.zeros( len(mamba_group_ids), dtype=torch.int64, device=device ), @@ -916,9 +977,11 @@ def create( num_scheduled_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32), num_computed_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32), num_draft_tokens_buf=make_buffer(max_num_reqs, dtype=torch.int32), + is_prefilling_buf=make_buffer(max_num_reqs, dtype=torch.bool), precopy_src_col_buf=make_buffer(max_num_reqs, dtype=torch.int32), precopy_token_bias_buf=make_buffer(max_num_reqs, dtype=torch.int32), is_initialized=False, + replayssm=None, ) def initialize_from_forward_context( @@ -989,6 +1052,10 @@ def _populate_metadata( state_copy_funcs = mamba_state_copy_funcs[mamba_spec.mamba_type] attention = forward_context[layer_name] kv_caches: list[torch.Tensor] = attention.kv_cache + is_flashinfer_replayssm = ( + getattr(attention, "use_replayssm", False) + and attention.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) if len(kv_caches) < len(state_copy_funcs): raise ValueError( f"Expected at least {len(state_copy_funcs)} Mamba state " @@ -1037,6 +1104,8 @@ def _populate_metadata( self.state_conv_widths[idx] = state.size(1) self.state_inner_sizes[idx] = state.stride(1) else: + self.state_skip_postprocess[idx] = is_flashinfer_replayssm + self.has_flashinfer_replayssm |= bool(is_flashinfer_replayssm) # Temporal state: inner_size = natural elements per # block (prod of inner dims). The kernel uses this # to compute copy_size = inner_size * elem_size, @@ -1090,6 +1159,20 @@ def _populate_metadata( for i, bt in enumerate(block_tables): self.block_table_ptrs[i] = _reinterpret_u64_as_i64(bt.data_ptr()) + if self.has_flashinfer_replayssm: + self.replayssm = ReplaySSMModelContext.create( + kv_cache_config, + self.mamba_group_ids, + forward_context, + block_tables, + self.num_accepted_tokens_out.numel(), + ) + if self.replayssm is None: + raise RuntimeError( + "FlashInfer ReplaySSM state was discovered but its model-wide " + "materialization context could not be initialized" + ) + self.is_initialized = True def compute_aligned_state_indices( @@ -1159,6 +1242,7 @@ def run_fused_postprocess( self.num_accepted_tokens_out[:num_reqs].copy_( num_accepted_tokens_gpu[:num_reqs] ) + self.materialize_src_cols[:num_reqs].fill_(-1) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) @@ -1177,9 +1261,13 @@ def run_fused_postprocess( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, + self.state_skip_postprocess, self.state_dim_row_count, self.state_dim_row_stride, self.num_accepted_tokens_out, + self.materialize_src_cols, + self.materialize_dst_cols, + self.materialize_token_counts, None, # idx_mapping: V1 decision arrays are already in req order num_reqs, block_size=self.block_size, @@ -1223,6 +1311,7 @@ def run_fused_precopy( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, + self.state_skip_postprocess, self.state_dim_row_count, self.state_dim_row_stride, idx_mapping, @@ -1258,6 +1347,7 @@ def run_fused_postprocess_align( # decision buffer rather than only [:num_reqs]. num_accepted_tokens_snapshot = self.num_accepted_tokens_out num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) + self.materialize_src_cols[:num_reqs].fill_(-1) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) @@ -1275,9 +1365,13 @@ def run_fused_postprocess_align( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, + self.state_skip_postprocess, self.state_dim_row_count, self.state_dim_row_stride, num_accepted_tokens_gpu, + self.materialize_src_cols, + self.materialize_dst_cols, + self.materialize_token_counts, idx_mapping, num_reqs, block_size=self.block_size, @@ -1294,9 +1388,9 @@ class MambaBuffers: """Single owner for all mamba-specific runner buffers. The two sub-objects have different gates: - ``preprocess`` is needed whenever ``mamba_cache_mode == "align"``; - ``postprocess_align`` is needed only when align is combined with - speculative decoding on a hybrid model, and is ``None`` otherwise. + ``preprocess`` handles prefix-cache state migration; ``postprocess_align`` + owns the fused lifecycle state used by hybrid speculative decoding and + FlashInfer ReplaySSM, and is ``None`` otherwise. """ preprocess: MambaCopyBuffers @@ -1458,6 +1552,8 @@ def preprocess_mamba( copy_bufs.offset = 0 num_reqs = len(input_batch.req_ids) + src_cols = [-1] * num_reqs + dst_cols = [-1] * num_reqs if fused is not None: if num_reqs == 0: @@ -1504,6 +1600,8 @@ def preprocess_mamba( fused.state_idx.np[i] = curr_state_idx if prev_state_idx != -1 and prev_state_idx != curr_state_idx: + src_cols[i] = prev_state_idx + dst_cols[i] = curr_state_idx accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 if fused is not None: assert accept_token_bias >= 0 @@ -1527,6 +1625,13 @@ def preprocess_mamba( fused.state_idx.copy_to_gpu(num_reqs) fused.src_col.copy_to_gpu(num_reqs) fused.token_bias.copy_to_gpu(num_reqs) + if fused.ctx.replayssm is not None and any(col >= 0 for col in src_cols): + fused.ctx.replayssm.copy_reassigned_slots( + idx_mapping=None, + src_cols=fused.src_col.gpu, + dst_cols=fused.state_idx.gpu, + num_reqs=num_reqs, + ) fused.ctx.run_fused_precopy( num_reqs=num_reqs, state_idx_gpu=fused.state_idx.gpu, @@ -1596,7 +1701,7 @@ def postprocess_mamba_align_gpu( forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, ) -> None: - """GPU-side mamba postprocess for spec decode + hybrid + align mode. + """GPU-side Mamba postprocess for fused align state maintenance. Lazily binds the fused-kernel context to the persistent block tables and forward-context state pointers on the first call, runs the fused kernel, @@ -1604,13 +1709,14 @@ def postprocess_mamba_align_gpu( batch's CPU tensor for the next iteration's preprocess. """ ctx = bufs.postprocess_align - # Caller is responsible for gating on spec decode + hybrid; this assert is - # a tripwire if those gates ever drift apart. + # The caller enables this context for spec-decode hybrid state copies or + # for model-owned FlashInfer ReplaySSM lifecycle maintenance under STP. assert ctx is not None assert ctx.mamba_state_idx_buf is not None assert ctx.num_scheduled_tokens_buf is not None assert ctx.num_computed_tokens_buf is not None assert ctx.num_draft_tokens_buf is not None + assert ctx.is_prefilling_buf is not None if not ctx.is_initialized: ctx.initialize_from_forward_context( @@ -1631,6 +1737,23 @@ def postprocess_mamba_align_gpu( num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu, num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu, ) + if ctx.replayssm is not None: + ctx.replayssm.postprocess_and_materialize( + idx_mapping=None, + query_metadata=ctx.num_scheduled_tokens_buf.gpu, + query_is_cumulative=False, + num_computed_tokens=ctx.num_computed_tokens_buf.gpu, + num_computed_is_after=False, + num_accepted_tokens=num_accepted_tokens_gpu, + is_prefilling=ctx.is_prefilling_buf.gpu, + live_cols=ctx.mamba_state_idx_buf.gpu, + materialize_src_cols=ctx.materialize_src_cols, + materialize_dst_cols=ctx.materialize_dst_cols, + materialize_token_counts=ctx.materialize_token_counts, + mamba_block_size=ctx.block_size, + num_reqs=num_reqs, + materialize_possible=ctx.replayssm_materialize_possible, + ) # ``num_accepted_tokens_out`` is pre-initialized from # ``num_accepted_tokens_gpu``; the kernel only overwrites entries to 1 @@ -1653,7 +1776,7 @@ def stage_postprocess_inputs_to_gpu( Walks ``req_ids[:num_reqs]`` once, writing each request's mamba block index and scheduled/computed/draft token counts into the matching pinned - numpy views, then issues four non-blocking H→D copies. The fused kernel + numpy views, then issues five non-blocking H→D copies. The fused kernel indexes the resulting GPU tensors by ``req_idx``. Buffers live on ``ctx`` and only exist when the postprocess kernel is enabled. @@ -1664,6 +1787,7 @@ def stage_postprocess_inputs_to_gpu( assert ctx.num_scheduled_tokens_buf is not None assert ctx.num_computed_tokens_buf is not None assert ctx.num_draft_tokens_buf is not None + assert ctx.is_prefilling_buf is not None scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens num_scheduled = scheduler_output.num_scheduled_tokens @@ -1671,6 +1795,8 @@ def stage_postprocess_inputs_to_gpu( scheduled_np = ctx.num_scheduled_tokens_buf.np computed_np = ctx.num_computed_tokens_buf.np draft_np = ctx.num_draft_tokens_buf.np + prefill_np = ctx.is_prefilling_buf.np + materialize_possible = False for i in range(num_reqs): req_id = req_ids[i] @@ -1680,11 +1806,31 @@ def stage_postprocess_inputs_to_gpu( "preprocess_mamba must run before stage_postprocess_inputs_to_gpu" ) state_idx_np[i] = state_idx - scheduled_np[i] = num_scheduled[req_id] - computed_np[i] = requests[req_id].num_computed_tokens - draft_np[i] = len(scheduled_spec_tokens.get(req_id, [])) + scheduled = num_scheduled[req_id] + computed = requests[req_id].num_computed_tokens + num_draft = len(scheduled_spec_tokens.get(req_id, [])) + scheduled_np[i] = scheduled + computed_np[i] = computed + draft_np[i] = num_draft + prefill_np[i] = ( + requests[req_id].num_computed_tokens < requests[req_id].num_prompt_tokens + ) + + # The actual accepted length is only known on GPU after sampling, but + # it cannot exceed one target token plus every scheduled draft. Skip + # the native model-wide launch only when even that upper bound cannot + # reach the next aligned checkpoint. + running_state_tokens = computed + scheduled - num_draft + max_new_computed = running_state_tokens + num_draft + aligned_max_new_computed = ( + max_new_computed // ctx.block_size + ) * ctx.block_size + materialize_possible |= aligned_max_new_computed >= running_state_tokens + + ctx.replayssm_materialize_possible = materialize_possible ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs) + ctx.is_prefilling_buf.copy_to_gpu(num_reqs) diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index d767a8fffd05..344dcad6cbad 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -440,6 +440,37 @@ def allocate_kv_cache( return kv_caches +def allocate_replayssm_caches( + kv_cache_config: KVCacheConfig, + device: torch.device, +) -> dict[str, tuple[torch.Tensor, ...]]: + """Allocate ReplaySSM ring state separately from canonical Mamba pages.""" + caches: dict[str, tuple[torch.Tensor, ...]] = {} + for group in kv_cache_config.kv_cache_groups: + group_spec = group.kv_cache_spec + layer_specs: Iterable[tuple[str, KVCacheSpec]] + if isinstance(group_spec, UniformTypeKVCacheSpecs): + layer_specs = group_spec.kv_cache_specs.items() + else: + layer_specs = ((name, group_spec) for name in group.layer_names) + + for layer_name, spec in layer_specs: + if not isinstance(spec, MambaSpec) or not spec.replayssm_shapes: + continue + assert layer_name not in caches + caches[layer_name] = tuple( + torch.zeros( + (kv_cache_config.num_blocks, *shape), + dtype=dtype, + device=device, + ) + for shape, dtype in zip( + spec.replayssm_shapes, spec.replayssm_dtypes, strict=True + ) + ) + return caches + + def prepare_kernel_block_sizes( kv_cache_config: KVCacheConfig, attn_groups: list[list[AttentionGroup]] ) -> list[int]: @@ -577,6 +608,7 @@ def bind_kv_cache( runner_kv_caches: list[torch.Tensor], num_attn_module: int = 1, kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, + replayssm_caches: Mapping[str, tuple[torch.Tensor, ...]] | None = None, ) -> None: """ Bind the allocated KV cache to both ModelRunner and forward context so @@ -624,6 +656,10 @@ def bind_kv_cache( # layer for the KV connector to register. for layer_name, kv_cache in kv_caches.items(): forward_context[layer_name].bind_kv_cache(kv_cache) + if replayssm_caches is not None and layer_name in replayssm_caches: + forward_context[layer_name].bind_replayssm_cache( + replayssm_caches[layer_name] + ) share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) From 65fac9b5f775fb5c3f8ce818dfca98d86059d404 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 18:40:46 +0200 Subject: [PATCH 06/53] test(mamba): skip ReplaySSM graph test without autotuning Signed-off-by: Andrii Skliar --- tests/v1/attention/test_replayssm_metadata_builder.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index e1e2875e2539..9283fb5d44f1 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -298,6 +298,10 @@ def test_spec_decode_single_token_chunk_synthesizes_acceptance_metadata(): 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, From 25b2f1dceb972d577db459304eff30225d04048a Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 18:46:24 +0200 Subject: [PATCH 07/53] Simplify ReplaySSM lifecycle integration Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 45 +--- tests/test_config.py | 134 +++++----- .../v1/e2e/general/test_mamba_prefix_cache.py | 38 ++- tests/v1/e2e/test_replayssm_decode.py | 228 +++++++++--------- tests/v1/worker/test_mamba_utils.py | 99 ++++---- tests/v1/worker/test_utils.py | 33 +-- .../layers/mamba/mamba_mixer2.py | 2 +- .../layers/mamba/ops/ssu_dispatch.py | 70 +++--- vllm/v1/attention/backends/mamba_attn.py | 4 +- vllm/v1/worker/gpu_model_runner.py | 48 ++-- vllm/v1/worker/mamba_utils.py | 28 +-- 11 files changed, 320 insertions(+), 409 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index da169d89d13a..8587fb1d6e3d 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -363,12 +363,6 @@ def test_replayssm_materialize_ready_rejects_incomplete_cache(): with pytest.raises(RuntimeError, match="ring trackers"): ssu_dispatch._replayssm_materialize_ready([mixer]) - mixer = _materialize_mixer(device="cuda") - mixer.replayssm_buffer_len = 0 - with pytest.raises(RuntimeError, match="buffer-len >= 1"): - ssu_dispatch._replayssm_materialize_ready([mixer]) - - def test_replayssm_materialize_ready_requires_cuda_ssm_state(): mixer = _materialize_mixer(device="cpu") @@ -536,11 +530,10 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), mamba_block_size=4, num_reqs=1, - materialize_possible=False, ) torch.cuda.synchronize() - assert kernel.call_count == 0 + assert kernel.call_count == 1 assert ctx.plan_flush_count.tolist() == [-1, -1] for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): for source_slot in source_slots: @@ -585,42 +578,6 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_does_not_copy_unchanged_physical_slots(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for block_table in block_tables: - block_table[0, 1] = block_table[0, 0] - for mixers, source_slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[source_slot] = 2 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 - kernel = Mock() - monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.copy_reassigned_slots( - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), - src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - torch.cuda.synchronize() - - assert kernel.call_count == 1 - assert ctx.precopy_flush_count.tolist() == [4, -1] - assert ctx.precopy_src_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 4 - assert ctx.precopy_dst_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 4 - for mixers, source_slot in zip(groups, (1, 4)): - assert mixers[0]._replayssm_ring_start[source_slot].item() == 2 - assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 4 - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() diff --git a/tests/test_config.py b/tests/test_config.py index b2ae877ed9ef..8f0da3941cbc 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -190,81 +190,75 @@ def _replayssm_config( ) -def test_v2_replayssm_requires_flashinfer(): - config = _replayssm_config( - backend=MambaBackendEnum.TRITON, - use_v2_model_runner=True, - ) - - with pytest.raises(ValueError, match="requires Model Runner V1"): - VllmConfig.validate_mamba_cached_kernel(config) - - -def test_v2_flashinfer_replayssm_is_supported(): - config = _replayssm_config( - backend=MambaBackendEnum.FLASHINFER, - use_v2_model_runner=True, - ) - - assert VllmConfig.validate_mamba_cached_kernel(config) is config - - -def test_flashinfer_replayssm_allows_align_prefix_caching(): - config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) - config.cache_config.mamba_cache_mode = "align" - - assert VllmConfig.validate_mamba_cached_kernel(config) is config - - -def test_flashinfer_replayssm_allows_all_prefix_caching(): - config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) - config.cache_config.mamba_cache_mode = "all" - - assert VllmConfig.validate_mamba_cached_kernel(config) is config - - -@pytest.mark.parametrize("use_v2_model_runner", [False, True]) -def test_flashinfer_replayssm_spec_decode_allows_align_prefix_caching( - use_v2_model_runner, -): - config = _replayssm_config( - backend=MambaBackendEnum.FLASHINFER, - use_v2_model_runner=use_v2_model_runner, - ) - config.num_speculative_tokens = 3 - config.cache_config.mamba_cache_mode = "align" - - assert VllmConfig.validate_mamba_cached_kernel(config) is config - - -@pytest.mark.parametrize("use_v2_model_runner", [False, True]) -def test_flashinfer_replayssm_spec_decode_allows_all_prefix_caching( - use_v2_model_runner, +@pytest.mark.parametrize( + ( + "backend", + "use_v2_model_runner", + "mamba_cache_mode", + "num_speculative_tokens", + "replayssm_buffer_len", + "error_match", + ), + [ + (MambaBackendEnum.TRITON, True, "none", 0, 16, "requires Model Runner V1"), + (MambaBackendEnum.FLASHINFER, True, "none", 0, 16, None), + (MambaBackendEnum.FLASHINFER, False, "align", 0, 16, None), + (MambaBackendEnum.FLASHINFER, False, "all", 0, 16, None), + (MambaBackendEnum.FLASHINFER, False, "align", 3, 16, None), + (MambaBackendEnum.FLASHINFER, True, "align", 3, 16, None), + (MambaBackendEnum.FLASHINFER, False, "all", 3, 16, None), + (MambaBackendEnum.FLASHINFER, True, "all", 3, 16, None), + ( + MambaBackendEnum.TRITON, + False, + "all", + 0, + 16, + "all mode requires.*flashinfer", + ), + ( + MambaBackendEnum.FLASHINFER, + False, + "none", + 0, + 17, + "replayssm-buffer-len <= 16", + ), + ], + ids=[ + "triton-v2-rejected", + "flashinfer-v2", + "flashinfer-align", + "flashinfer-all", + "flashinfer-align-spec-v1", + "flashinfer-align-spec-v2", + "flashinfer-all-spec-v1", + "flashinfer-all-spec-v2", + "triton-all-rejected", + "flashinfer-buffer-too-long", + ], +) +def test_replayssm_config_matrix( + backend: MambaBackendEnum, + use_v2_model_runner: bool, + mamba_cache_mode: str, + num_speculative_tokens: int, + replayssm_buffer_len: int, + error_match: str | None, ): config = _replayssm_config( - backend=MambaBackendEnum.FLASHINFER, + backend=backend, use_v2_model_runner=use_v2_model_runner, ) - config.num_speculative_tokens = 3 - config.cache_config.mamba_cache_mode = "all" - - assert VllmConfig.validate_mamba_cached_kernel(config) is config - - -def test_triton_replayssm_rejects_all_prefix_caching(): - config = _replayssm_config(backend=MambaBackendEnum.TRITON) - config.cache_config.mamba_cache_mode = "all" - - with pytest.raises(ValueError, match="all mode requires.*flashinfer"): - VllmConfig.validate_mamba_cached_kernel(config) - + config.cache_config.mamba_cache_mode = mamba_cache_mode + config.cache_config.replayssm_buffer_len = replayssm_buffer_len + config.num_speculative_tokens = num_speculative_tokens -def test_flashinfer_replayssm_rejects_unsupported_buffer_length(): - config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) - config.cache_config.replayssm_buffer_len = 17 - - with pytest.raises(ValueError, match="replayssm-buffer-len <= 16"): - VllmConfig.validate_mamba_cached_kernel(config) + if error_match is None: + assert VllmConfig.validate_mamba_cached_kernel(config) is config + else: + with pytest.raises(ValueError, match=error_match): + VllmConfig.validate_mamba_cached_kernel(config) def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 830bfccc9183..ea77d29cbe92 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -873,13 +873,9 @@ def get_mamba_prefix_cache_step_configs( return tests -def _run_mamba_prefix_cache_mrv1( +def _run_mamba_prefix_cache_mrv1_configured( monkeypatch: pytest.MonkeyPatch, async_scheduling: bool ): - # This test patches the V1 model runner, so pin V1 explicitly: MoE/hybrid - # models like Qwen3-Next now default to the V2 runner. - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") - envs.disable_envs_cache() global async_scheduling_mode async_scheduling_mode = async_scheduling run_ref_mamba_state_in_subprocess() @@ -938,6 +934,20 @@ def _run_mamba_prefix_cache_mrv1( cleanup_dist_env_and_memory() +def _run_mamba_prefix_cache_mrv1( + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool +): + # This test patches the V1 model runner, so pin V1 explicitly: MoE/hybrid + # models like Qwen3-Next now default to the V2 runner. + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + envs.disable_envs_cache() + _run_mamba_prefix_cache_mrv1_configured(patch, async_scheduling) + finally: + envs.disable_envs_cache() + + @create_new_process_for_each_test("spawn") def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch): _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False) @@ -948,14 +958,11 @@ def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch): _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=True) -def _run_mamba_prefix_cache_mrv2( +def _run_mamba_prefix_cache_mrv2_configured( monkeypatch: pytest.MonkeyPatch, async_scheduling: bool ): global async_scheduling_mode async_scheduling_mode = async_scheduling - monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - envs.disable_envs_cache() from vllm.v1.worker.gpu.model_runner import GPUModelRunner as MRV2GPUModelRunner from vllm.v1.worker.gpu.model_states.mamba_hybrid import ( @@ -1229,6 +1236,19 @@ def fake_sample( cleanup_dist_env_and_memory() +def _run_mamba_prefix_cache_mrv2( + monkeypatch: pytest.MonkeyPatch, async_scheduling: bool +): + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + envs.disable_envs_cache() + _run_mamba_prefix_cache_mrv2_configured(patch, async_scheduling) + finally: + envs.disable_envs_cache() + + @create_new_process_for_each_test() def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch): _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=False) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 690b57d8eb0e..13d3cf69f68b 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -13,10 +13,12 @@ try: from flashinfer.mamba.checkpointing_ssu import ( CheckpointingSSURunner, - allocate_checkpointing_ssu_scratch, # noqa: F401 + allocate_checkpointing_ssu_scratch, ) - HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) + HAS_FLASHINFER_CHECKPOINTING_SSU = callable( + CheckpointingSSURunner + ) and callable(allocate_checkpointing_ssu_scratch) except ImportError: HAS_FLASHINFER_CHECKPOINTING_SSU = False @@ -40,12 +42,10 @@ "Once upon a time, in a small village,", ] -try: - from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner - - HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) -except ImportError: - HAS_FLASHINFER_CHECKPOINTING_SSU = False +requires_flashinfer_replayssm_materialization = pytest.mark.skipif( + not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), + reason="FlashInfer ReplaySSM materialization APIs not available", +) def _check_replayssm_parity( @@ -58,39 +58,52 @@ def _check_replayssm_parity( require_v2: bool = False, monkeypatch: pytest.MonkeyPatch | None = None, ): - # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a - # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are - # common-mode and only ReplaySSM varies. - if require_v2: - assert monkeypatch is not None - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - envs.disable_envs_cache() + def run() -> 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. + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + tensor_parallel_size=tensor_parallel_size, + mamba_backend=mamba_backend, + ) + with vllm_runner(model_name, **common) as llm: + if require_v2: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner + baseline = llm.generate_greedy_logprobs( + PROMPTS, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + if require_v2: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner + replay = llm.generate_greedy_logprobs( + PROMPTS, max_tokens=32, num_logprobs=5 + ) + + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline", + name_1=name_1, + ) - common = dict( - max_model_len=1024, - trust_remote_code=True, - enable_prefix_caching=False, - mamba_cache_mode="none", - tensor_parallel_size=tensor_parallel_size, - mamba_backend=mamba_backend, - ) - with vllm_runner(model_name, **common) as llm: - if require_v2: - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner - baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) - with vllm_runner( - model_name, use_replayssm=True, replayssm_buffer_len=16, **common - ) as llm: - if require_v2: - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner - replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) + if not require_v2: + run() + return - check_logprobs_close( - outputs_0_lst=baseline, - outputs_1_lst=replay, - name_0="baseline", - name_1=name_1, - ) + assert monkeypatch is not None + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + envs.disable_envs_cache() + run() + finally: + envs.disable_envs_cache() @pytest.mark.parametrize("model_name", MODELS) @@ -301,73 +314,78 @@ def _check_flashinfer_replayssm_prefix_caching( use_v2: bool, tensor_parallel_size: int, ): - # ReplaySSM materializes the exact SSM state at each cacheable block - # boundary, so cached prefixes must match the always-materialized baseline. - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2 else "0") - envs.disable_envs_cache() - - common = dict( - max_model_len=8192, - trust_remote_code=True, - enable_prefix_caching=True, - enable_chunked_prefill=True, - mamba_cache_mode=mamba_cache_mode, - mamba_backend="flashinfer", - disable_log_stats=False, # required for llm.get_metrics() - tensor_parallel_size=tensor_parallel_size, - ) - if moe_backend is not None: - common["moe_backend"] = moe_backend - if use_ngram: - common["speculative_config"] = { - "method": "ngram", - "num_speculative_tokens": 3, - "prompt_lookup_max": 3, - } - - with vllm_runner(model_name, **common) as llm: - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 - baseline_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size - llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + def run() -> None: + # ReplaySSM materializes the exact SSM state at each cacheable block + # boundary, so cached prefixes must match the always-materialized baseline. + common = dict( + max_model_len=8192, + trust_remote_code=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + mamba_cache_mode=mamba_cache_mode, + mamba_backend="flashinfer", + disable_log_stats=False, # required for llm.get_metrics() + tensor_parallel_size=tensor_parallel_size, ) - baseline = llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + if moe_backend is not None: + common["moe_backend"] = moe_backend + if use_ngram: + common["speculative_config"] = { + "method": "ngram", + "num_speculative_tokens": 3, + "prompt_lookup_max": 3, + } + + with vllm_runner(model_name, **common) as llm: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 + baseline_block_size = ( + llm.llm.llm_engine.vllm_config.cache_config.block_size + ) + llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + baseline = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + baseline_hits = _prefix_cache_hits(llm) + + with vllm_runner( + model_name, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 + replay_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size + assert replay_block_size == baseline_block_size + llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + replay_hits = _prefix_cache_hits(llm) + + assert baseline_hits > 0 + assert replay_hits > 0, ( + f"ReplaySSM {mamba_cache_mode}-mode run produced no prefix-cache hits; " + "the shared prefix may be shorter than one mamba block, so prefix " + "caching is inert" ) - baseline_hits = _prefix_cache_hits(llm) - - with vllm_runner( - model_name, use_replayssm=True, replayssm_buffer_len=16, **common - ) as llm: - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 - replay_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size - assert replay_block_size == baseline_block_size - llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 - ) - replay = llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0=f"flashinfer_baseline_{mamba_cache_mode}_pc", + name_1=f"flashinfer_replayssm_{mamba_cache_mode}_pc", ) - replay_hits = _prefix_cache_hits(llm) - assert baseline_hits > 0 - assert replay_hits > 0, ( - f"ReplaySSM {mamba_cache_mode}-mode run produced no prefix-cache hits; " - "the shared prefix may be shorter than one mamba block, so prefix " - "caching is inert" - ) - check_logprobs_close( - outputs_0_lst=baseline, - outputs_1_lst=replay, - name_0=f"flashinfer_baseline_{mamba_cache_mode}_pc", - name_1=f"flashinfer_replayssm_{mamba_cache_mode}_pc", - ) + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_v2 else "0") + envs.disable_envs_cache() + run() + finally: + envs.disable_envs_cache() -@pytest.mark.skipif( - not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), - reason="FlashInfer ReplaySSM materialization APIs not available", -) +@requires_flashinfer_replayssm_materialization @pytest.mark.parametrize("model_name", MODELS) @pytest.mark.parametrize( ("mamba_cache_mode", "use_v2", "use_ngram"), @@ -396,10 +414,7 @@ def test_flashinfer_replayssm_prefix_cache_tp1( ) -@pytest.mark.skipif( - not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), - reason="FlashInfer ReplaySSM materialization APIs not available", -) +@requires_flashinfer_replayssm_materialization @large_gpu_mark(min_gb=40) def test_flashinfer_replayssm_all_prefix_cache_v2(vllm_runner, monkeypatch): _check_flashinfer_replayssm_prefix_caching( @@ -414,10 +429,7 @@ def test_flashinfer_replayssm_all_prefix_cache_v2(vllm_runner, monkeypatch): ) -@pytest.mark.skipif( - not (HAS_FLASHINFER_CHECKPOINTING_SSU and HAS_FLASHINFER_REPLAYSSM_MATERIALIZE), - reason="FlashInfer ReplaySSM materialization APIs not available", -) +@requires_flashinfer_replayssm_materialization @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) def test_flashinfer_replayssm_prefix_cache_v2_tp2( diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 1608c0d1d762..a7609518a7df 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -173,7 +173,9 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): ("with_replayssm", "num_computed_tokens", "expected_order"), [ pytest.param(True, 4, ["materialize", "copy"], id="replayssm-boundary"), - pytest.param(True, 3, ["copy"], id="replayssm-no-boundary"), + pytest.param( + True, 3, ["materialize", "copy"], id="replayssm-no-boundary" + ), pytest.param(False, 4, ["copy"], id="generic"), ], ) @@ -241,8 +243,6 @@ def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): ctx.replayssm.postprocess_and_materialize.side_effect = lambda **kwargs: ( order.append("materialize") ) - ctx.replayssm_materialize_possible = False - block_table = MagicMock() block_table.get_device_tensor.return_value = torch.zeros((1, 4), dtype=torch.int32) input_batch = MagicMock() @@ -263,12 +263,6 @@ def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): ) assert order == ["copy", "materialize"] - assert ( - ctx.replayssm.postprocess_and_materialize.call_args.kwargs[ - "materialize_possible" - ] - is False - ) assert accepted_cpu.tolist() == [3] @@ -339,14 +333,7 @@ def __getitem__(self, item): def test_reinterpret_u64_as_i64_preserves_pointer_bits(): - ptrs = [ - 0, - 1, - (1 << 63) - 1, - 1 << 63, - (1 << 63) + 1234, - (1 << 64) - 1, - ] + ptrs = [1 << 63, (1 << 64) - 1] ptr_tensor = torch.zeros(len(ptrs), dtype=torch.int64) for idx, ptr in enumerate(ptrs): @@ -396,7 +383,7 @@ def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): ] -def test_gpu_context_marks_flashinfer_replayssm_temporal_state(): +def test_gpu_context_initializes_flashinfer_replayssm_lifecycle(): cfg = _TestConfig(num_layers=1) device = torch.device("cpu") kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) @@ -408,14 +395,48 @@ def test_gpu_context_marks_flashinfer_replayssm_temporal_state(): attention.use_replayssm = True attention.mamba_config.backend = MambaBackendEnum.FLASHINFER - gpu_ctx.initialize_from_forward_context( - kv_cache_config, - {"layer_0": attention}, - _COPY_FUNCS, - [torch.empty(1, 4, dtype=torch.int32)], - ) + model_ctx = object() + with patch( + "vllm.v1.worker.mamba_utils.ReplaySSMModelContext.create", + return_value=model_ctx, + ) as create: + gpu_ctx.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, + [torch.empty(1, 4, dtype=torch.int32)], + ) assert gpu_ctx.state_skip_postprocess.tolist() == [0, 1] + assert gpu_ctx.replayssm is model_ctx + create.assert_called_once() + + +def test_gpu_context_rejects_missing_replayssm_lifecycle(): + cfg = _TestConfig(num_layers=1) + device = torch.device("cpu") + kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) + gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) + attention = _make_mock_attention( + torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), + torch.empty(cfg.num_blocks, cfg.temporal_state_dim), + ) + attention.use_replayssm = True + attention.mamba_config.backend = MambaBackendEnum.FLASHINFER + + with ( + patch( + "vllm.v1.worker.mamba_utils.ReplaySSMModelContext.create", + return_value=None, + ), + pytest.raises(RuntimeError, match="could not be initialized"), + ): + gpu_ctx.initialize_from_forward_context( + kv_cache_config, + {"layer_0": attention}, + _COPY_FUNCS, + [torch.empty(1, 4, dtype=torch.int32)], + ) def _make_postprocess_scheduler_output( @@ -806,7 +827,6 @@ def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: """Build a MambaSpecDecodeGPUContext stand-in exposing only the five per-request staging buffers touched by stage_postprocess_inputs_to_gpu.""" ctx = MagicMock() - ctx.block_size = 16 ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_scheduled_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_computed_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) @@ -894,35 +914,6 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ctx.is_prefilling_buf.gpu[:num_reqs], torch.tensor([True, False, True]), ) - assert ctx.replayssm_materialize_possible is True - - -def test_stage_postprocess_inputs_skips_impossible_materialization(): - device = torch.device("cpu") - ctx = _make_staging_ctx(max_num_reqs=4, device=device) - req_ids = ["req_a"] - scheduler_output = _make_postprocess_scheduler_output( - req_ids=req_ids, - num_scheduled_tokens={"req_a": 1}, - scheduled_spec_decode_tokens={"req_a": [1, 2]}, - ) - requests = _make_requests( - req_ids=req_ids, - num_computed_tokens=[8], - block_ids_per_req=[[0]], - num_prompt_tokens=[8], - ) - - stage_postprocess_inputs_to_gpu( - ctx, - scheduler_output, - req_ids, - 1, - requests, - {"req_a": 0}, - ) - - assert ctx.replayssm_materialize_possible is False def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 424efbc2d298..3f9b9542a3a4 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -72,30 +72,17 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): assert all(len(mixer.kv_cache) == 5 for mixer in mixers) - assert ( - mixers[0]._replayssm_ring_start.data_ptr() - == mixers[2]._replayssm_ring_start.data_ptr() - ) - assert ( - mixers[0]._replayssm_prev_num_accepted.data_ptr() - == mixers[2]._replayssm_prev_num_accepted.data_ptr() - ) - assert ( - mixers[1]._replayssm_ring_start.data_ptr() - != mixers[0]._replayssm_ring_start.data_ptr() - ) - assert ( - mixers[1]._replayssm_prev_num_accepted.data_ptr() - != mixers[0]._replayssm_prev_num_accepted.data_ptr() + tracker_names = ( + "_replayssm_ring_start", + "_replayssm_prev_num_accepted", ) - assert mixers[0]._replayssm_ring_start.shape == (4,) - assert mixers[0]._replayssm_prev_num_accepted.shape == (4,) - assert mixers[0]._replayssm_ring_start.dtype == torch.int32 - assert mixers[0]._replayssm_ring_start.is_contiguous() - assert torch.count_nonzero(mixers[0]._replayssm_ring_start) == 0 - assert torch.count_nonzero(mixers[0]._replayssm_prev_num_accepted) == 0 - assert not any(hasattr(m, "_commits_replayssm_trackers") for m in mixers) - assert not any(hasattr(m, "_updates_replayssm_trackers") for m in mixers) + 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 group_tracker.dtype == torch.int32 + assert torch.count_nonzero(group_tracker) == 0 def test_bind_kv_cache(default_vllm_config): diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 6a6917fc3caa..be3008cbef6d 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1252,7 +1252,7 @@ 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. For + block indices to the same request and must keep separate cursor tensors. Tracker mutation is model-owned and runs once after the step; layer forwards only consume the shared values. """ diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index fa05bc02b6cd..b6b999b2c034 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -323,6 +323,25 @@ def _copy_reassigned_replayssm_slots_kernel( tl.store(tracker_committed + dst_slot, 0) +def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: + ssm = mixer.kv_cache[1] + x_cache = mixer.kv_cache[2] + b_cache = mixer.kv_cache[4] + return ( + ssm.dtype, + x_cache.dtype, + mixer.A.dtype, + ssm.size(1), + ssm.size(2), + ssm.size(3), + ssm.size(1) // b_cache.size(1), + int(mixer.replayssm_buffer_len), + x_cache.size(2), + bool(mixer.mamba_config.enable_stochastic_rounding), + int(mixer.mamba_config.stochastic_rounding_philox_rounds or 0), + ) + + @dataclass class ReplaySSMModelContext: """Persistent all-layer tables for ReplaySSM post-step maintenance.""" @@ -387,37 +406,9 @@ def create( first = mixers[0] first_ssm = first.kv_cache[1] first_x = first.kv_cache[2] - first_b = first.kv_cache[4] - compatibility = ( - first_ssm.dtype, - first_x.dtype, - first.A.dtype, - first_ssm.size(1), - first_ssm.size(2), - first_ssm.size(3), - first_ssm.size(1) // first_b.size(1), - int(first.replayssm_buffer_len), - first_x.size(2), - bool(first.mamba_config.enable_stochastic_rounding), - int(first.mamba_config.stochastic_rounding_philox_rounds or 0), - ) + compatibility = _replayssm_specialization_key(first) for mixer in mixers[1:]: - ssm = mixer.kv_cache[1] - x_cache = mixer.kv_cache[2] - b_cache = mixer.kv_cache[4] - current = ( - ssm.dtype, - x_cache.dtype, - mixer.A.dtype, - ssm.size(1), - ssm.size(2), - ssm.size(3), - ssm.size(1) // b_cache.size(1), - int(mixer.replayssm_buffer_len), - x_cache.size(2), - bool(mixer.mamba_config.enable_stochastic_rounding), - int(mixer.mamba_config.stochastic_rounding_philox_rounds or 0), - ) + current = _replayssm_specialization_key(mixer) if current != compatibility: raise ValueError( "A single model-wide FlashInfer ReplaySSM materialization " @@ -534,7 +525,6 @@ def postprocess_and_materialize( materialize_token_counts: torch.Tensor, mamba_block_size: int, num_reqs: int, - materialize_possible: bool = True, ) -> None: """Commit lifecycle metadata, then materialize all layers once.""" if num_reqs == 0: @@ -571,13 +561,12 @@ def postprocess_and_materialize( HAS_IDX_MAPPING=idx_mapping is not None, ) - if materialize_possible: - self._materialize_planned( - self.src_slots, - self.dst_slots, - self.plan_ring_start, - self.plan_flush_count, - ) + self._materialize_planned( + self.src_slots, + self.dst_slots, + self.plan_ring_start, + self.plan_flush_count, + ) def copy_reassigned_slots( self, @@ -1086,11 +1075,6 @@ def _replayssm_materialize_ready(mixers: list[Any]) -> bool: "FlashInfer ReplaySSM prefix materialization requires allocated " "replay ring buffers and ring trackers" ) - if not mixers[0].replayssm_buffer_len: - raise RuntimeError( - "FlashInfer ReplaySSM prefix materialization requires " - "--replayssm-buffer-len >= 1" - ) return True diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 297fb5516ae9..32c680b788ad 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -183,7 +183,9 @@ def __init__( self.decode_replayssm_state_indices_d: torch.Tensor | None = None # ReplaySSM CUDA-graph buffers for the selected backend. if self.use_replayssm: - assert len(kv_cache_spec.replayssm_shapes) == 3 + assert len(kv_cache_spec.replayssm_shapes) == 3, ( + "FlashInfer ReplaySSM requires x, dt, and B ring-state tensors" + ) if self.use_replayssm and not self.use_flashinfer_replayssm: self.decode_write_pos_d: torch.Tensor = torch.empty( (self.decode_cudagraph_max_bs,), diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 8e985dc04cba..0590cc5e50ea 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1002,6 +1002,11 @@ def __init__( self.mamba_state_idx: dict[str, int] = {} self._mamba_bufs: mamba_utils.MambaBuffers | None = None self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None + self._use_modelwide_replayssm = ( + self.cache_config.mamba_cache_mode in ("align", "all") + and self.cache_config.use_replayssm + and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None if ( self.cache_config.mamba_cache_mode == "all" @@ -1074,10 +1079,9 @@ def _get_mamba_state_copy_funcs(self) -> MambaStateCopyFuncsByType: def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: # The postprocess sub-object is also the model-level owner of # FlashInfer ReplaySSM trackers, including STP. - assert self.cache_config.mamba_cache_mode == "align" or ( - self.cache_config.mamba_cache_mode == "all" - and self.cache_config.use_replayssm - and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + assert ( + self.cache_config.mamba_cache_mode == "align" + or self._use_modelwide_replayssm ) if self._mamba_bufs is None: self._mamba_bufs = mamba_utils.MambaBuffers.create( @@ -1089,11 +1093,7 @@ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: with_postprocess_align=( self.speculative_config is not None and self.model_config.is_hybrid ) - or ( - self.cache_config.use_replayssm - and self.vllm_config.mamba_config.backend - == MambaBackendEnum.FLASHINFER - ), + or self._use_modelwide_replayssm, ) return self._mamba_bufs @@ -1601,12 +1601,7 @@ def _update_states_after_model_execute( each sequence, and a shifting is done during the next iteration based on the number of accepted tokens. """ - modelwide_replayssm = ( - self.cache_config.mamba_cache_mode in ("align", "all") - and self.cache_config.use_replayssm - and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER - ) - if not modelwide_replayssm and ( + if not self._use_modelwide_replayssm and ( not self.speculative_config or not self.model_config.is_hybrid ): return @@ -1617,7 +1612,10 @@ def _update_states_after_model_execute( num_reqs = output_token_ids.size(0) self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) - if self.cache_config.mamba_cache_mode == "align" or modelwide_replayssm: + if ( + self.cache_config.mamba_cache_mode == "align" + or self._use_modelwide_replayssm + ): # Fused GPU postprocess: state copies + per-request accepted-token # update without CPU-GPU sync. The metadata # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is @@ -2157,15 +2155,10 @@ def _prepare_inputs( # _update_states_after_model_execute for hybrid models). # Skipped under async scheduling (non-align): the CPU copy races with # the in-flight D2H copy and with input-batch row moves. - modelwide_replayssm = ( - self.cache_config.mamba_cache_mode in ("align", "all") - and self.cache_config.use_replayssm - and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER - ) needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and ( not self.use_async_scheduling or self.cache_config.mamba_cache_mode == "align" - or modelwide_replayssm + or self._use_modelwide_replayssm ) if needs_cpu_accepted_counts: assert self.num_accepted_tokens_event is not None @@ -4439,13 +4432,10 @@ def execute_model( ) pad_attn = cudagraph_mode == CUDAGraphMode.FULL - modelwide_replayssm = ( - self.cache_config.mamba_cache_mode in ("align", "all") - and self.cache_config.use_replayssm - and self.vllm_config.mamba_config.backend - == MambaBackendEnum.FLASHINFER - ) - if self.cache_config.mamba_cache_mode == "align" or modelwide_replayssm: + if ( + self.cache_config.mamba_cache_mode == "align" + or self._use_modelwide_replayssm + ): # preprocess_mamba reads req_state.num_computed_tokens (CPU) # to decide copy operations, so we must apply deferred # corrections before it runs. diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index cde95d033ff0..94de5d1d4591 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -875,10 +875,6 @@ class MambaSpecDecodeGPUContext: # Persistent all-layer ReplaySSM descriptors, populated with the cache # addresses on first real forward. None for non-FlashInfer configurations. replayssm: ReplaySSMModelContext | None = None - # Host-side upper bound computed while staging the current batch. False - # means no acceptance outcome can cross an align boundary, so the native - # materializer launch can be skipped after tracker maintenance. - replayssm_materialize_possible: bool = False @classmethod def create( @@ -1552,9 +1548,6 @@ def preprocess_mamba( copy_bufs.offset = 0 num_reqs = len(input_batch.req_ids) - src_cols = [-1] * num_reqs - dst_cols = [-1] * num_reqs - if fused is not None: if num_reqs == 0: return @@ -1600,8 +1593,6 @@ def preprocess_mamba( fused.state_idx.np[i] = curr_state_idx if prev_state_idx != -1 and prev_state_idx != curr_state_idx: - src_cols[i] = prev_state_idx - dst_cols[i] = curr_state_idx accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 if fused is not None: assert accept_token_bias >= 0 @@ -1625,7 +1616,7 @@ def preprocess_mamba( fused.state_idx.copy_to_gpu(num_reqs) fused.src_col.copy_to_gpu(num_reqs) fused.token_bias.copy_to_gpu(num_reqs) - if fused.ctx.replayssm is not None and any(col >= 0 for col in src_cols): + if fused.ctx.replayssm is not None: fused.ctx.replayssm.copy_reassigned_slots( idx_mapping=None, src_cols=fused.src_col.gpu, @@ -1752,7 +1743,6 @@ def postprocess_mamba_align_gpu( materialize_token_counts=ctx.materialize_token_counts, mamba_block_size=ctx.block_size, num_reqs=num_reqs, - materialize_possible=ctx.replayssm_materialize_possible, ) # ``num_accepted_tokens_out`` is pre-initialized from @@ -1796,8 +1786,6 @@ def stage_postprocess_inputs_to_gpu( computed_np = ctx.num_computed_tokens_buf.np draft_np = ctx.num_draft_tokens_buf.np prefill_np = ctx.is_prefilling_buf.np - materialize_possible = False - for i in range(num_reqs): req_id = req_ids[i] state_idx = mamba_state_idx.get(req_id) @@ -1815,20 +1803,6 @@ def stage_postprocess_inputs_to_gpu( prefill_np[i] = ( requests[req_id].num_computed_tokens < requests[req_id].num_prompt_tokens ) - - # The actual accepted length is only known on GPU after sampling, but - # it cannot exceed one target token plus every scheduled draft. Skip - # the native model-wide launch only when even that upper bound cannot - # reach the next aligned checkpoint. - running_state_tokens = computed + scheduled - num_draft - max_new_computed = running_state_tokens + num_draft - aligned_max_new_computed = ( - max_new_computed // ctx.block_size - ) * ctx.block_size - materialize_possible |= aligned_max_new_computed >= running_state_tokens - - ctx.replayssm_materialize_possible = materialize_possible - ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) From 4d8d019c2efc8f71fcc67636724134c5a7da260e Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 18:55:54 +0200 Subject: [PATCH 08/53] Fix ReplaySSM pre-commit checks Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 3 ++- tests/v1/e2e/test_replayssm_decode.py | 10 ++++------ tests/v1/worker/test_mamba_utils.py | 4 +--- vllm/model_executor/layers/mamba/mamba_mixer2.py | 4 +--- vllm/model_executor/layers/mamba/ops/ssu_dispatch.py | 4 +--- vllm/v1/attention/backends/mamba_attn.py | 6 +++--- 6 files changed, 12 insertions(+), 19 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 8587fb1d6e3d..1b3366f62f99 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -363,6 +363,7 @@ def test_replayssm_materialize_ready_rejects_incomplete_cache(): with pytest.raises(RuntimeError, match="ring trackers"): ssu_dispatch._replayssm_materialize_ready([mixer]) + def test_replayssm_materialize_ready_requires_cuda_ssm_state(): mixer = _materialize_mixer(device="cpu") @@ -429,7 +430,7 @@ def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch) mamba_block_size=4, num_reqs=1, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert kernel.call_count == 1 args = kernel.call_args.args diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 13d3cf69f68b..38fb5ad182e8 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -16,9 +16,9 @@ allocate_checkpointing_ssu_scratch, ) - HAS_FLASHINFER_CHECKPOINTING_SSU = callable( - CheckpointingSSURunner - ) and callable(allocate_checkpointing_ssu_scratch) + HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) and callable( + allocate_checkpointing_ssu_scratch + ) except ImportError: HAS_FLASHINFER_CHECKPOINTING_SSU = False @@ -338,9 +338,7 @@ def run() -> None: with vllm_runner(model_name, **common) as llm: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 - baseline_block_size = ( - llm.llm.llm_engine.vllm_config.cache_config.block_size - ) + baseline_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size llm.generate_greedy_logprobs( PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 ) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a7609518a7df..ea063ce94b6a 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -173,9 +173,7 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): ("with_replayssm", "num_computed_tokens", "expected_order"), [ pytest.param(True, 4, ["materialize", "copy"], id="replayssm-boundary"), - pytest.param( - True, 3, ["materialize", "copy"], id="replayssm-no-boundary" - ), + pytest.param(True, 3, ["materialize", "copy"], id="replayssm-no-boundary"), pytest.param(False, 4, ["copy"], id="generic"), ], ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index be3008cbef6d..5981f76570ba 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1012,9 +1012,7 @@ def conv_ssm_forward( ).squeeze(1) state_indices_tensor_d_input = live_indices state_indices_tensor_d_output = live_indices - block_idx_last_computed_token_d = ( - block_idx_last_scheduled_token_d - ) + block_idx_last_computed_token_d = block_idx_last_scheduled_token_d elif self.num_spec > 0: assert block_idx_last_scheduled_token_prev_step_d is not None input_indices = ( diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index b6b999b2c034..f74da7e8e30d 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -579,9 +579,7 @@ def copy_reassigned_slots( """Copy exact live state when align assigns a new writable slot.""" if num_reqs == 0: return - _copy_reassigned_replayssm_slots_kernel[ - (self.max_num_reqs, self.num_groups) - ]( + _copy_reassigned_replayssm_slots_kernel[(self.max_num_reqs, self.num_groups)]( idx_mapping, src_cols, dst_cols, diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 32c680b788ad..1b63acfb7224 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -895,9 +895,9 @@ def _select_replayssm_state_indices( return state_indices_tensor_d[:, 0] assert block_idx_last_scheduled_token is not None - live_cols = block_idx_last_scheduled_token[ - : state_indices_tensor_d.size(0) - ].to(torch.int64) + live_cols = block_idx_last_scheduled_token[: state_indices_tensor_d.size(0)].to( + torch.int64 + ) return state_indices_tensor_d.gather(1, live_cols.unsqueeze(1)).squeeze(1) def update_block_table( From efbfcee5dbef5823349114018dbb8a7099dea0dd Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 19:06:56 +0200 Subject: [PATCH 09/53] Fix ReplaySSM type checks Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 8 ++++---- vllm/model_executor/models/diffusion_gemma.py | 7 ++++++- vllm/v1/worker/gpu/model_states/mamba_hybrid.py | 5 +++-- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 1b3366f62f99..c5a7d566ffff 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -487,7 +487,7 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch mamba_block_size=4, num_reqs=1, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert kernel.call_count == 1 assert ctx.plan_ring_start.tolist() == [15, 0] @@ -532,7 +532,7 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): mamba_block_size=4, num_reqs=1, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert kernel.call_count == 1 assert ctx.plan_flush_count.tolist() == [-1, -1] @@ -565,7 +565,7 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert kernel.call_count == 1 assert ctx.precopy_ring_start.tolist() == [2, 0] @@ -603,7 +603,7 @@ def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) - torch.cuda.synchronize() + torch.accelerator.synchronize() assert kernel.call_count == 1 assert ctx.precopy_ring_start.tolist() == [2, 0] diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 3a51e5aad678..52872f310445 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -975,7 +975,12 @@ def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any] return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]} def postprocess_state( - self, idx_mapping, num_sampled, num_computed_tokens=None + self, + idx_mapping, + num_sampled, + num_computed_tokens=None, + query_start_loc=None, + is_prefilling=None, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 123d89dbb1c2..538aad921da0 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -134,8 +134,9 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: # Seed the running state block from the resumed/prefilled position. state_block_size = self.cache_config.block_size if self.cache_config.mamba_cache_mode == "all": - state_block_size = self.cache_config.mamba_block_size - assert state_block_size is not None + mamba_block_size = self.cache_config.mamba_block_size + assert mamba_block_size is not None + state_block_size = mamba_block_size self._mamba_state_idx_gpu[req_index].fill_( (new_req_data.num_computed_tokens - 1) // state_block_size ) From 47a0133551c2c520bac79f7277941715a18d5051 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 19:36:18 +0200 Subject: [PATCH 10/53] fix(mamba): maintain FlashInfer ReplaySSM trackers in all cache modes Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 101 +++++++++++++++- .../worker/test_mamba_hybrid_model_state.py | 50 +++++++- tests/v1/worker/test_mamba_utils.py | 76 ++++++++++++ .../layers/mamba/ops/ssu_dispatch.py | 34 +++++- .../worker/gpu/model_states/mamba_hybrid.py | 109 +++++++++++------- vllm/v1/worker/gpu_model_runner.py | 61 +++++----- vllm/v1/worker/mamba_utils.py | 43 ++++--- 7 files changed, 377 insertions(+), 97 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index c5a7d566ffff..d2b1bb9b8708 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -371,7 +371,7 @@ def test_replayssm_materialize_ready_requires_cuda_ssm_state(): ssu_dispatch._replayssm_materialize_ready([mixer]) -def _modelwide_replayssm_fixture(): +def _modelwide_replayssm_fixture(cache_mode: str = "align"): groups = [ [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], @@ -390,7 +390,19 @@ def _modelwide_replayssm_fixture(): forward_context[name] = mixer config = Mock() - config.kv_cache_groups = [Mock(layer_names=names) for names in layer_names] + specs = [ + MambaSpec( + block_size=1024 if cache_mode == "none" else 4, + shapes=((4, 3, 5),), + dtypes=(torch.float32,), + mamba_cache_mode=cache_mode, + ) + for _ in layer_names + ] + config.kv_cache_groups = [ + Mock(layer_names=names, kv_cache_spec=spec) + for names, spec in zip(layer_names, specs) + ] block_tables = [ torch.tensor([[1, 2, 3], [0, 0, 0]], dtype=torch.int32, device="cuda"), torch.tensor([[4, 5, 6], [0, 0, 0]], dtype=torch.int32, device="cuda"), @@ -398,6 +410,91 @@ def _modelwide_replayssm_fixture(): return groups, config, forward_context, block_tables +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_none_commits_trackers_without_materialization( + monkeypatch, +): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture( + cache_mode="none" + ) + materializer = Mock() + monkeypatch.setattr( + ssu_dispatch, "_load_replayssm_materialize", lambda: materializer + ) + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + assert not ctx.materialize_prefixes + + query_len = torch.zeros(2, dtype=torch.int32, device="cuda") + num_computed = torch.zeros(2, dtype=torch.int32, device="cuda") + accepted = torch.ones(2, dtype=torch.int32, device="cuda") + is_prefilling = torch.zeros(2, dtype=torch.bool, device="cuda") + live_cols = torch.zeros(2, dtype=torch.int32, device="cuda") + no_materialize = torch.full((2,), -1, dtype=torch.int32, device="cuda") + materialize_counts = torch.zeros(2, dtype=torch.int32, device="cuda") + + def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None: + query_len[0] = scheduled + accepted[0] = num_accepted + is_prefilling[0] = prefilling + ctx.postprocess_and_materialize( + idx_mapping=None, + query_metadata=query_len, + query_is_cumulative=False, + num_computed_tokens=num_computed, + num_computed_is_after=False, + num_accepted_tokens=accepted, + is_prefilling=is_prefilling, + live_cols=live_cols, + materialize_src_cols=no_materialize, + materialize_dst_cols=no_materialize, + materialize_token_counts=materialize_counts, + mamba_block_size=1024, + num_reqs=1, + ) + num_computed[0] += num_accepted + + # Prefill canonicalizes the one live slot in mode none. + for mixers, slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[slot] = 17 + mixers[0]._replayssm_prev_num_accepted[slot] = 9 + step(scheduled=8, num_accepted=1, prefilling=True) + + # Mix STP and MTP-shaped updates. Several scheduled lengths exceed the + # accepted count; repeated crossings eventually wrap the 20-row ring. + for scheduled, num_accepted in ( + (4, 2), + (4, 1), + (14, 3), + (18, 2), + (16, 16), + (1, 1), + ): + step(scheduled=scheduled, num_accepted=num_accepted) + torch.accelerator.synchronize() + + assert materializer.call_count == 0 + assert ctx.plan_flush_count.tolist() == [-1, -1] + for mixers, live_slot in zip(groups, (1, 4)): + # Both layers share this group tracker. The expected single transition + # per step would differ if either layer committed it independently. + assert mixers[0]._replayssm_ring_start[live_slot].item() == 4 + assert mixers[0]._replayssm_prev_num_accepted[live_slot].item() == 1 + + # A later prefill resets the exact same live physical slot again. + step(scheduled=8, num_accepted=1, prefilling=True) + torch.accelerator.synchronize() + for mixers, live_slot in zip(groups, (1, 4)): + assert mixers[0]._replayssm_ring_start[live_slot].item() == 0 + assert mixers[0]._replayssm_prev_num_accepted[live_slot].item() == 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index f02c7be84217..67c6962b47a2 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -26,7 +26,8 @@ def test_postprocess_state_scalar_with_int32_mapping( (4,), 9, dtype=torch.int32, device="cuda" ) state._align_mode = False - state._mamba_lifecycle_mode = False + state._needs_prefix_state_migration = False + state._use_flashinfer_replayssm = False state.recoverssm = None state._mamba_ctx = None idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda") @@ -39,6 +40,50 @@ def test_postprocess_state_scalar_with_int32_mapping( torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: + state = object.__new__(MambaHybridModelState) + state._align_mode = False + state._needs_prefix_state_migration = False + state._use_flashinfer_replayssm = True + state.recoverssm = None + state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") + state._replayssm_live_cols_gpu = torch.zeros( + 4, dtype=torch.int32, device="cuda" + ) + state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") + replayssm = Mock() + ctx = Mock( + is_initialized=True, + replayssm=replayssm, + materialize_src_cols=torch.full( + (4,), -1, dtype=torch.int32, device="cuda" + ), + materialize_dst_cols=torch.full( + (4,), -1, dtype=torch.int32, device="cuda" + ), + materialize_token_counts=torch.zeros(4, dtype=torch.int32, device="cuda"), + block_size=1024, + ) + state._mamba_ctx = ctx + idx_mapping = torch.tensor([2], dtype=torch.int32, device="cuda") + num_computed = torch.tensor([0, 0, 20, 0], dtype=torch.int32, device="cuda") + query_start_loc = torch.tensor([0, 4], dtype=torch.int32, device="cuda") + + state.postprocess_state( + idx_mapping, + 2, + num_computed_tokens=num_computed, + query_start_loc=query_start_loc, + ) + + ctx.run_fused_postprocess_align.assert_not_called() + assert replayssm.postprocess_and_materialize.call_count == 1 + kwargs = replayssm.postprocess_and_materialize.call_args.kwargs + assert kwargs["num_accepted_tokens"] is state.num_accepted_tokens_gpu + assert kwargs["live_cols"] is state._replayssm_live_cols_gpu + + def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: state = RecoverSSMState() metadata = Mock(spec=RecoverSSMMetadata) @@ -69,7 +114,8 @@ def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None: state = object.__new__(MambaHybridModelState) state._align_mode = True - state._mamba_lifecycle_mode = True + state._needs_prefix_state_migration = True + state._use_flashinfer_replayssm = False state._mamba_ctx = None state._mamba_state_idx_gpu = torch.full((5,), -1, dtype=torch.int32, device="cuda") state.recoverssm = RecoverSSMState() diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index ea063ce94b6a..29e265632262 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -258,12 +258,59 @@ def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): kv_cache_config=kv_cache_config, forward_context={}, mamba_state_copy_funcs={}, + run_prefix_state_migration=True, ) assert order == ["copy", "materialize"] assert accepted_cpu.tolist() == [3] +def test_postprocess_mamba_none_skips_prefix_copy(): + ctx = MagicMock() + ctx.is_initialized = True + ctx.mamba_group_ids = [0] + ctx.mamba_state_idx_buf = MagicMock(gpu=torch.zeros(1, dtype=torch.int32)) + ctx.num_scheduled_tokens_buf = MagicMock( + gpu=torch.tensor([4], dtype=torch.int32) + ) + ctx.num_computed_tokens_buf = MagicMock( + gpu=torch.tensor([20], dtype=torch.int32) + ) + ctx.num_draft_tokens_buf = MagicMock(gpu=torch.tensor([3], dtype=torch.int32)) + ctx.is_prefilling_buf = MagicMock(gpu=torch.tensor([False])) + ctx.materialize_src_cols = torch.full((1,), -1, dtype=torch.int32) + ctx.materialize_dst_cols = torch.full((1,), -1, dtype=torch.int32) + ctx.materialize_token_counts = torch.zeros(1, dtype=torch.int32) + ctx.block_size = 1024 + ctx.replayssm = MagicMock() + input_batch = MagicMock() + input_batch.block_table = [] + accepted = torch.tensor([2], dtype=torch.int32) + accepted_cpu = torch.zeros(1, dtype=torch.int32) + + postprocess_mamba_align_gpu( + bufs=MagicMock(postprocess_align=ctx), + num_reqs=1, + num_accepted_tokens_gpu=accepted, + num_accepted_tokens_cpu_tensor=accepted_cpu, + input_batch=input_batch, + kv_cache_config=MagicMock(), + forward_context={}, + mamba_state_copy_funcs={}, + run_prefix_state_migration=False, + ) + + ctx.run_fused_postprocess.assert_not_called() + assert ctx.replayssm.postprocess_and_materialize.call_count == 1 + assert ( + ctx.replayssm.postprocess_and_materialize.call_args.kwargs[ + "num_accepted_tokens" + ] + is accepted + ) + assert accepted_cpu.tolist() == [2] + + # ----------------------------------------------------------------------------- # Golden tests for postprocess_mamba_fused_kernel # ----------------------------------------------------------------------------- @@ -942,6 +989,35 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): ) +def test_stage_postprocess_inputs_to_gpu_uses_fixed_none_mode_live_col(): + device = torch.device("cpu") + ctx = _make_staging_ctx(max_num_reqs=4, device=device) + scheduler_output = _make_postprocess_scheduler_output( + req_ids=["req_a", "req_b"], + num_scheduled_tokens={"req_a": 4, "req_b": 1}, + ) + requests = _make_requests( + ["req_a", "req_b"], + [20, 7], + [[0], [0]], + num_prompt_tokens=[10, 8], + ) + + stage_postprocess_inputs_to_gpu( + ctx, + scheduler_output, + ["req_a", "req_b"], + 2, + requests, + {}, + fixed_live_col=0, + ) + + np.testing.assert_array_equal(ctx.mamba_state_idx_buf.np[:2], [0, 0]) + np.testing.assert_array_equal(ctx.num_scheduled_tokens_buf.np[:2], [4, 1]) + np.testing.assert_array_equal(ctx.num_computed_tokens_buf.np[:2], [20, 7]) + + def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: device = torch.device("cpu") config = _TestConfig(num_layers=1) diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index f74da7e8e30d..cb8ed6de4edb 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -377,6 +377,7 @@ class ReplaySSMModelContext: max_layers_per_group: int logical_window: int ring_buffer_len: int + materialize_prefixes: bool @classmethod def create( @@ -399,6 +400,21 @@ def create( ) block_table_by_gid = dict(zip(mamba_group_ids, block_tables)) replayssm_block_tables = [block_table_by_gid[gid] for gid, _ in grouped] + cache_modes = set() + for gid, _ in grouped: + spec = kv_cache_config.kv_cache_groups[gid].kv_cache_spec + if not isinstance(spec, MambaSpec): + raise TypeError( + "FlashInfer ReplaySSM layers require a Mamba cache spec; " + f"got {type(spec).__name__}" + ) + cache_modes.add(spec.mamba_cache_mode) + if len(cache_modes) != 1: + raise ValueError( + "model-wide ReplaySSM requires one Mamba cache mode; " + f"got {sorted(cache_modes)}" + ) + materialize_prefixes = next(iter(cache_modes)) in ("align", "all") mixers = [mixer for _, group_mixers in grouped for mixer in group_mixers] if not _replayssm_materialize_ready(mixers): @@ -507,6 +523,7 @@ def create( max_layers_per_group=max_layers_per_group, logical_window=int(first.replayssm_buffer_len), ring_buffer_len=first_x.size(2), + materialize_prefixes=materialize_prefixes, ) def postprocess_and_materialize( @@ -561,12 +578,13 @@ def postprocess_and_materialize( HAS_IDX_MAPPING=idx_mapping is not None, ) - self._materialize_planned( - self.src_slots, - self.dst_slots, - self.plan_ring_start, - self.plan_flush_count, - ) + if self.materialize_prefixes: + self._materialize_planned( + self.src_slots, + self.dst_slots, + self.plan_ring_start, + self.plan_flush_count, + ) def copy_reassigned_slots( self, @@ -579,6 +597,10 @@ def copy_reassigned_slots( """Copy exact live state when align assigns a new writable slot.""" if num_reqs == 0: return + if not self.materialize_prefixes: + raise RuntimeError( + "ReplaySSM writable-slot materialization requires align or all mode" + ) _copy_reassigned_replayssm_slots_kernel[(self.max_num_reqs, self.num_groups)]( idx_mapping, src_cols, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 538aad921da0..1920788fcde1 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -102,14 +102,24 @@ def __init__( self.cache_config.use_replayssm is True and vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER ) - self._mamba_lifecycle_mode = self._align_mode or ( + self._needs_prefix_state_migration = self._align_mode or ( self.cache_config.mamba_cache_mode == "all" and self._use_flashinfer_replayssm ) self.recoverssm = ( RecoverSSMState() if self.cache_config.use_kda_recoverssm else None ) - if self._mamba_lifecycle_mode: + if self._needs_prefix_state_migration or self._use_flashinfer_replayssm: + self._replayssm_live_cols_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.int32, device=self.device + ) + self._mamba_ctx: MambaSpecDecodeGPUContext | None = None + self._mamba_group_ids: list[int] = [] + self._mamba_spec: MambaSpec | None = None + self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None + self._mamba_kv_cache_config: KVCacheConfig | None = None + self._mamba_block_tables: tuple[torch.Tensor, ...] | None = None + if self._needs_prefix_state_migration: self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device ) @@ -119,18 +129,12 @@ def __init__( self._mamba_src_off_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device ) - self._mamba_ctx: MambaSpecDecodeGPUContext | None = None - self._mamba_group_ids: list[int] = [] - self._mamba_spec: MambaSpec | None = None - self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None - self._mamba_kv_cache_config: KVCacheConfig | None = None - self._mamba_block_tables: tuple[torch.Tensor, ...] | None = None def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) # Must reset the speculative acceptance count in this idx which could be stale. self.num_accepted_tokens_gpu[req_index].fill_(1) - if self._mamba_lifecycle_mode: + if self._needs_prefix_state_migration: # Seed the running state block from the resumed/prefilled position. state_block_size = self.cache_config.block_size if self.cache_config.mamba_cache_mode == "all": @@ -214,7 +218,7 @@ def preprocess_state( ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset is visible to the forward kernels. """ - if not self._mamba_lifecycle_mode: + if not self._needs_prefix_state_migration: return num_reqs = input_batch.num_reqs if num_reqs == 0: @@ -314,6 +318,10 @@ def prepare_attn( ) num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) + if self._use_flashinfer_replayssm: + mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) + self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + if self._align_mode: mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) aligned_index_builders = [] @@ -406,7 +414,7 @@ def postprocess_state( # acceptance leaves the sequence non-block-aligned. num_computed_tokens # already holds the post-step advanced count. if ( - self._mamba_lifecycle_mode + self._needs_prefix_state_migration and num_computed_tokens is not None and self._mamba_ctx is not None ): @@ -417,40 +425,53 @@ def postprocess_state( num_computed_tokens, idx_mapping, ) - # Must match the condition that sets ``state_skip_postprocess``: - # the fused kernel skips the temporal copy for every FlashInfer - # ReplaySSM layer, so the materializer has to cover every one of - # them. Gating this on spec decode would drop the SSM state on any - # non-spec boundary where src_col != dst_col. Rows that need no - # work carry the -1 src_col sentinel and cost nothing. - if self._use_flashinfer_replayssm: - replayssm = self._mamba_ctx.replayssm - assert replayssm is not None - if query_start_loc is None: - raise RuntimeError( - "ReplaySSM postprocess requires the query_start_loc from " - "the forward that produced this acceptance" - ) - if is_prefilling is None: - is_prefilling = self._is_prefilling_gpu[:num_reqs] - replayssm.postprocess_and_materialize( - idx_mapping=idx_mapping, - query_metadata=query_start_loc, - query_is_cumulative=True, - num_computed_tokens=num_computed_tokens, - num_computed_is_after=True, - # run_fused_postprocess_align can reset the live buffer to - # one for the next step. Its persistent snapshot still - # contains the acceptance produced by this forward. - num_accepted_tokens=self._mamba_ctx.num_accepted_tokens_out, - is_prefilling=is_prefilling, - live_cols=self._mamba_state_idx_gpu, - materialize_src_cols=self._mamba_ctx.materialize_src_cols, - materialize_dst_cols=self._mamba_ctx.materialize_dst_cols, - materialize_token_counts=(self._mamba_ctx.materialize_token_counts), - mamba_block_size=self._mamba_ctx.block_size, - num_reqs=num_reqs, + + if self._use_flashinfer_replayssm: + if num_computed_tokens is None: + raise RuntimeError( + "ReplaySSM postprocess requires the post-step computed-token " + "counts from the forward that produced this acceptance" + ) + if query_start_loc is None: + raise RuntimeError( + "ReplaySSM postprocess requires the query_start_loc from " + "the forward that produced this acceptance" ) + ctx = self._mamba_ctx + if ctx is None or not ctx.is_initialized: + raise RuntimeError( + "ReplaySSM postprocess context was not initialized before forward" + ) + replayssm = ctx.replayssm + assert replayssm is not None + if is_prefilling is None: + is_prefilling = self._is_prefilling_gpu[:num_reqs] + replayssm.postprocess_and_materialize( + idx_mapping=idx_mapping, + query_metadata=query_start_loc, + query_is_cumulative=True, + num_computed_tokens=num_computed_tokens, + num_computed_is_after=True, + # Prefix migration can reset the live buffer to one for the + # next step; use its snapshot in that case. Mode none never + # runs the migration kernel, so the acceptance buffer is exact. + num_accepted_tokens=( + ctx.num_accepted_tokens_out + if self._needs_prefix_state_migration + else self.num_accepted_tokens_gpu + ), + is_prefilling=is_prefilling, + live_cols=( + self._mamba_state_idx_gpu + if self._needs_prefix_state_migration + else self._replayssm_live_cols_gpu + ), + materialize_src_cols=ctx.materialize_src_cols, + materialize_dst_cols=ctx.materialize_dst_cols, + materialize_token_counts=ctx.materialize_token_counts, + mamba_block_size=ctx.block_size, + num_reqs=num_reqs, + ) @triton.jit diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0590cc5e50ea..501885bd0d14 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1002,11 +1002,17 @@ def __init__( self.mamba_state_idx: dict[str, int] = {} self._mamba_bufs: mamba_utils.MambaBuffers | None = None self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None - self._use_modelwide_replayssm = ( - self.cache_config.mamba_cache_mode in ("align", "all") - and self.cache_config.use_replayssm + self._use_flashinfer_replayssm = ( + self.cache_config.use_replayssm and self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER ) + self._needs_prefix_state_migration = ( + self.cache_config.mamba_cache_mode == "align" + or ( + self._use_flashinfer_replayssm + and self.cache_config.mamba_cache_mode == "all" + ) + ) self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None if ( self.cache_config.mamba_cache_mode == "all" @@ -1080,8 +1086,7 @@ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: # The postprocess sub-object is also the model-level owner of # FlashInfer ReplaySSM trackers, including STP. assert ( - self.cache_config.mamba_cache_mode == "align" - or self._use_modelwide_replayssm + self._needs_prefix_state_migration or self._use_flashinfer_replayssm ) if self._mamba_bufs is None: self._mamba_bufs = mamba_utils.MambaBuffers.create( @@ -1093,7 +1098,7 @@ def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: with_postprocess_align=( self.speculative_config is not None and self.model_config.is_hybrid ) - or self._use_modelwide_replayssm, + or self._use_flashinfer_replayssm, ) return self._mamba_bufs @@ -1601,7 +1606,7 @@ def _update_states_after_model_execute( each sequence, and a shifting is done during the next iteration based on the number of accepted tokens. """ - if not self._use_modelwide_replayssm and ( + if not self._use_flashinfer_replayssm and ( not self.speculative_config or not self.model_config.is_hybrid ): return @@ -1613,8 +1618,7 @@ def _update_states_after_model_execute( self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) if ( - self.cache_config.mamba_cache_mode == "align" - or self._use_modelwide_replayssm + self._needs_prefix_state_migration or self._use_flashinfer_replayssm ): # Fused GPU postprocess: state copies + per-request accepted-token # update without CPU-GPU sync. The metadata @@ -1631,6 +1635,7 @@ def _update_states_after_model_execute( kv_cache_config=self.kv_cache_config, forward_context=self.compilation_config.static_forward_context, mamba_state_copy_funcs=self._get_mamba_state_copy_funcs(), + run_prefix_state_migration=self._needs_prefix_state_migration, ) if self.num_accepted_tokens_event is not None: @@ -2158,7 +2163,7 @@ def _prepare_inputs( needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and ( not self.use_async_scheduling or self.cache_config.mamba_cache_mode == "align" - or self._use_modelwide_replayssm + or self._needs_prefix_state_migration ) if needs_cpu_accepted_counts: assert self.num_accepted_tokens_event is not None @@ -4432,10 +4437,8 @@ def execute_model( ) pad_attn = cudagraph_mode == CUDAGraphMode.FULL - if ( - self.cache_config.mamba_cache_mode == "align" - or self._use_modelwide_replayssm - ): + mamba_bufs = None + if self._needs_prefix_state_migration: # preprocess_mamba reads req_state.num_computed_tokens (CPU) # to decide copy operations, so we must apply deferred # corrections before it runs. @@ -4464,18 +4467,24 @@ def execute_model( ) self.num_accepted_tokens.copy_to_gpu(num_reqs) - # Stage inputs only when the fused postprocess will run. This - # includes spec-decode hybrid models and model-owned - # FlashInfer ReplaySSM lifecycle maintenance under STP. - if mamba_bufs.postprocess_align is not None: - mamba_utils.stage_postprocess_inputs_to_gpu( - mamba_bufs.postprocess_align, - scheduler_output, - self.input_batch.req_ids, - num_reqs, - self.requests, - self.mamba_state_idx, - ) + elif self._use_flashinfer_replayssm: + mamba_bufs = self._get_mamba_bufs() + + # Stage inputs whenever the fused state-copy or model-wide + # ReplaySSM tracker postprocess will run. Mode none does not run + # prefix preprocessing and always uses logical live column zero. + if mamba_bufs is not None and mamba_bufs.postprocess_align is not None: + mamba_utils.stage_postprocess_inputs_to_gpu( + mamba_bufs.postprocess_align, + scheduler_output, + self.input_batch.req_ids, + num_reqs, + self.requests, + self.mamba_state_idx, + fixed_live_col=( + None if self._needs_prefix_state_migration else 0 + ), + ) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 94de5d1d4591..04d89028c789 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -1691,6 +1691,7 @@ def postprocess_mamba_align_gpu( kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], mamba_state_copy_funcs: MambaStateCopyFuncsByType, + run_prefix_state_migration: bool, ) -> None: """GPU-side Mamba postprocess for fused align state maintenance. @@ -1720,14 +1721,17 @@ def postprocess_mamba_align_gpu( ], ) - ctx.run_fused_postprocess( - num_reqs=num_reqs, - num_accepted_tokens_gpu=num_accepted_tokens_gpu, - mamba_state_idx_gpu=ctx.mamba_state_idx_buf.gpu, - num_scheduled_tokens_gpu=ctx.num_scheduled_tokens_buf.gpu, - num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu, - num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu, - ) + accepted_tokens_for_postprocess = num_accepted_tokens_gpu + if run_prefix_state_migration: + ctx.run_fused_postprocess( + num_reqs=num_reqs, + num_accepted_tokens_gpu=num_accepted_tokens_gpu, + mamba_state_idx_gpu=ctx.mamba_state_idx_buf.gpu, + num_scheduled_tokens_gpu=ctx.num_scheduled_tokens_buf.gpu, + num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu, + num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu, + ) + accepted_tokens_for_postprocess = ctx.num_accepted_tokens_out if ctx.replayssm is not None: ctx.replayssm.postprocess_and_materialize( idx_mapping=None, @@ -1735,7 +1739,7 @@ def postprocess_mamba_align_gpu( query_is_cumulative=False, num_computed_tokens=ctx.num_computed_tokens_buf.gpu, num_computed_is_after=False, - num_accepted_tokens=num_accepted_tokens_gpu, + num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, live_cols=ctx.mamba_state_idx_buf.gpu, materialize_src_cols=ctx.materialize_src_cols, @@ -1750,7 +1754,7 @@ def postprocess_mamba_align_gpu( # when src_block_idx == dest_block_idx (copy within the same block), so # the original count is preserved for everyone else. num_accepted_tokens_cpu_tensor[:num_reqs].copy_( - ctx.num_accepted_tokens_out[:num_reqs], non_blocking=True + accepted_tokens_for_postprocess[:num_reqs], non_blocking=True ) @@ -1761,6 +1765,8 @@ def stage_postprocess_inputs_to_gpu( num_reqs: int, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], + *, + fixed_live_col: int | None = None, ) -> None: """Stage all per-request inputs the fused mamba postprocess kernel reads. @@ -1770,8 +1776,9 @@ def stage_postprocess_inputs_to_gpu( indexes the resulting GPU tensors by ``req_idx``. Buffers live on ``ctx`` and only exist when the postprocess kernel is enabled. - Invariant: ``preprocess_mamba`` must have run first for the same batch so - that every ``req_ids[i]`` has an entry in ``mamba_state_idx``. + Prefix-migration modes read the live column populated by + ``preprocess_mamba``. Mode ``none`` instead passes ``fixed_live_col=0`` + because its one backend-owned live state always occupies logical column 0. """ assert ctx.mamba_state_idx_buf is not None assert ctx.num_scheduled_tokens_buf is not None @@ -1788,11 +1795,13 @@ def stage_postprocess_inputs_to_gpu( prefill_np = ctx.is_prefilling_buf.np for i in range(num_reqs): req_id = req_ids[i] - state_idx = mamba_state_idx.get(req_id) - assert state_idx is not None, ( - f"mamba_state_idx missing entry for {req_id!r}; " - "preprocess_mamba must run before stage_postprocess_inputs_to_gpu" - ) + state_idx = fixed_live_col + if state_idx is None: + state_idx = mamba_state_idx.get(req_id) + assert state_idx is not None, ( + f"mamba_state_idx missing entry for {req_id!r}; " + "preprocess_mamba must run before stage_postprocess_inputs_to_gpu" + ) state_idx_np[i] = state_idx scheduled = num_scheduled[req_id] computed = requests[req_id].num_computed_tokens From a9339be1dd86c88f90f77aabb060f8c9d8317c0a Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 2 Sep 2026 19:42:55 +0200 Subject: [PATCH 11/53] refactor(mamba): clarify ReplaySSM lifecycle and prefix materialization Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 45 +++--- tests/v1/e2e/test_replayssm_decode.py | 2 +- .../worker/test_mamba_hybrid_model_state.py | 16 +-- tests/v1/worker/test_mamba_utils.py | 23 +-- vllm/config/cache.py | 8 +- .../layers/mamba/ops/ssu_dispatch.py | 135 ++++++++++++------ .../worker/gpu/model_states/mamba_hybrid.py | 20 +-- vllm/v1/worker/gpu_model_runner.py | 12 +- vllm/v1/worker/mamba_utils.py | 17 +-- 9 files changed, 152 insertions(+), 126 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index d2b1bb9b8708..50957fe8596a 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -347,28 +347,29 @@ def _materialize_mixer(device: str = "cpu") -> Mock: @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_replayssm_materialize_ready_rejects_incomplete_cache(): +def test_validate_replayssm_cache_rejects_incomplete_cache(): mixer = _materialize_mixer(device="cuda") mixer.kv_cache[1] = torch.empty(0, device="cuda") - assert not ssu_dispatch._replayssm_materialize_ready([mixer]) + with pytest.raises(RuntimeError, match="cache tensors"): + ssu_dispatch._validate_replayssm_cache([mixer]) mixer = _materialize_mixer(device="cuda") mixer.kv_cache[2] = torch.empty(0, device="cuda") - with pytest.raises(RuntimeError, match="replay ring buffers"): - ssu_dispatch._replayssm_materialize_ready([mixer]) + with pytest.raises(RuntimeError, match="cache tensors"): + ssu_dispatch._validate_replayssm_cache([mixer]) mixer = _materialize_mixer(device="cuda") mixer._replayssm_ring_start = torch.empty(0, dtype=torch.int32, device="cuda") with pytest.raises(RuntimeError, match="ring trackers"): - ssu_dispatch._replayssm_materialize_ready([mixer]) + ssu_dispatch._validate_replayssm_cache([mixer]) -def test_replayssm_materialize_ready_requires_cuda_ssm_state(): +def test_validate_replayssm_cache_requires_cuda_state(): mixer = _materialize_mixer(device="cpu") - with pytest.raises(RuntimeError, match="requires a CUDA SSM state cache"): - ssu_dispatch._replayssm_materialize_ready([mixer]) + with pytest.raises(RuntimeError, match="requires CUDA cache tensors"): + ssu_dispatch._validate_replayssm_cache([mixer]) def _modelwide_replayssm_fixture(cache_mode: str = "align"): @@ -443,12 +444,12 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None query_len[0] = scheduled accepted[0] = num_accepted is_prefilling[0] = prefilling - ctx.postprocess_and_materialize( + ctx.postprocess( idx_mapping=None, query_metadata=query_len, - query_is_cumulative=False, + query_metadata_is_cumulative=False, num_computed_tokens=num_computed, - num_computed_is_after=False, + num_computed_is_post_step=False, num_accepted_tokens=accepted, is_prefilling=is_prefilling, live_cols=live_cols, @@ -512,12 +513,12 @@ def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch) max_num_reqs=2, ) assert ctx is not None - ctx.postprocess_and_materialize( + ctx.postprocess( idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), - query_is_cumulative=True, + query_metadata_is_cumulative=True, num_computed_tokens=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), - num_computed_is_after=True, + num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), @@ -569,12 +570,12 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch max_num_reqs=2, ) assert ctx is not None - ctx.postprocess_and_materialize( + ctx.postprocess( idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), - query_is_cumulative=True, + query_metadata_is_cumulative=True, num_computed_tokens=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), - num_computed_is_after=True, + num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([3, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), @@ -614,12 +615,12 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.postprocess_and_materialize( + ctx.postprocess( idx_mapping=None, query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), - query_is_cumulative=False, + query_metadata_is_cumulative=False, num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_after=False, + num_computed_is_post_step=False, num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), @@ -656,7 +657,7 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.copy_reassigned_slots( + ctx.materialize_reassigned_slots( idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), @@ -694,7 +695,7 @@ def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.copy_reassigned_slots( + ctx.materialize_reassigned_slots( idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 38fb5ad182e8..5407aa34ac84 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -271,7 +271,7 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): finally: envs.disable_envs_cache() - assert any(len(token_ids) > 16 for token_ids, _ in replay) + assert any(len(output[0]) > 16 for output in replay) assert draft_count > 0 check_logprobs_close( outputs_0_lst=baseline, diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 67c6962b47a2..513a28be14b9 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -48,20 +48,14 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: state._use_flashinfer_replayssm = True state.recoverssm = None state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") - state._replayssm_live_cols_gpu = torch.zeros( - 4, dtype=torch.int32, device="cuda" - ) + state._replayssm_live_cols_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") replayssm = Mock() ctx = Mock( is_initialized=True, replayssm=replayssm, - materialize_src_cols=torch.full( - (4,), -1, dtype=torch.int32, device="cuda" - ), - materialize_dst_cols=torch.full( - (4,), -1, dtype=torch.int32, device="cuda" - ), + materialize_src_cols=torch.full((4,), -1, dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.full((4,), -1, dtype=torch.int32, device="cuda"), materialize_token_counts=torch.zeros(4, dtype=torch.int32, device="cuda"), block_size=1024, ) @@ -78,8 +72,8 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: ) ctx.run_fused_postprocess_align.assert_not_called() - assert replayssm.postprocess_and_materialize.call_count == 1 - kwargs = replayssm.postprocess_and_materialize.call_args.kwargs + assert replayssm.postprocess.call_count == 1 + kwargs = replayssm.postprocess.call_args.kwargs assert kwargs["num_accepted_tokens"] is state.num_accepted_tokens_gpu assert kwargs["live_cols"] is state._replayssm_live_cols_gpu diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 29e265632262..c725cd0389ec 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -201,8 +201,8 @@ def test_preprocess_mamba_uses_modelwide_materializer_when_present( align_ctx.precopy_token_bias_buf = _MockCpuGpuBuffer(1, torch.int32, device) align_ctx.replayssm = MagicMock() if with_replayssm else None if align_ctx.replayssm is not None: - align_ctx.replayssm.copy_reassigned_slots.side_effect = lambda **kwargs: ( - order.append("materialize") + align_ctx.replayssm.materialize_reassigned_slots.side_effect = ( + lambda **kwargs: order.append("materialize") ) align_ctx.run_fused_precopy.side_effect = lambda **kwargs: order.append("copy") @@ -238,7 +238,7 @@ def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): ctx.materialize_token_counts = torch.tensor([2], dtype=torch.int32) ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") ctx.replayssm = MagicMock() - ctx.replayssm.postprocess_and_materialize.side_effect = lambda **kwargs: ( + ctx.replayssm.postprocess.side_effect = lambda **kwargs: ( order.append("materialize") ) block_table = MagicMock() @@ -270,12 +270,8 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ctx.is_initialized = True ctx.mamba_group_ids = [0] ctx.mamba_state_idx_buf = MagicMock(gpu=torch.zeros(1, dtype=torch.int32)) - ctx.num_scheduled_tokens_buf = MagicMock( - gpu=torch.tensor([4], dtype=torch.int32) - ) - ctx.num_computed_tokens_buf = MagicMock( - gpu=torch.tensor([20], dtype=torch.int32) - ) + ctx.num_scheduled_tokens_buf = MagicMock(gpu=torch.tensor([4], dtype=torch.int32)) + ctx.num_computed_tokens_buf = MagicMock(gpu=torch.tensor([20], dtype=torch.int32)) ctx.num_draft_tokens_buf = MagicMock(gpu=torch.tensor([3], dtype=torch.int32)) ctx.is_prefilling_buf = MagicMock(gpu=torch.tensor([False])) ctx.materialize_src_cols = torch.full((1,), -1, dtype=torch.int32) @@ -301,13 +297,8 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ) ctx.run_fused_postprocess.assert_not_called() - assert ctx.replayssm.postprocess_and_materialize.call_count == 1 - assert ( - ctx.replayssm.postprocess_and_materialize.call_args.kwargs[ - "num_accepted_tokens" - ] - is accepted - ) + assert ctx.replayssm.postprocess.call_count == 1 + assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is accepted assert accepted_cpu.tolist() == [2] diff --git a/vllm/config/cache.py b/vllm/config/cache.py index cc396d7b83f5..fdb688ac4404 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -202,10 +202,10 @@ class CacheConfig: 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. - Supports mamba_cache_mode 'none', 'align', and 'all'; 'all' requires the - FlashInfer backend. Mamba2 speculative decode also requires FlashInfer. - Prefix-boundary flushes are most efficient when mamba_block_size is a - multiple of replayssm_buffer_len, but this is not required.""" + Triton supports 'none' and 'align' on Model Runner V1. FlashInfer supports + 'none', 'align', and 'all' on Model Runner V1 and V2. Mamba2 speculative + decoding requires FlashInfer. With prefix caching enabled, Triton supports + 'align'; FlashInfer supports 'align' and 'all'.""" use_kda_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index cb8ed6de4edb..af2b3804f274 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -60,16 +60,28 @@ def _postprocess_replayssm_modelwide_kernel( LOGICAL_WINDOW: tl.constexpr, RING_BUFFER_LEN: tl.constexpr, PAD_SLOT_ID: tl.constexpr, - QUERY_IS_CUMULATIVE: tl.constexpr, - NUM_COMPUTED_IS_AFTER: tl.constexpr, + QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, + NUM_COMPUTED_IS_POST_STEP: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, ) -> None: - """Plan materialization and commit all ReplaySSM trackers in one launch. - - One CTA owns one ``(batch row, cache group)`` pair, and is therefore the - only writer of that group's tracker for the request. Group zero writes the - request-level ``ring_start``/``flush_count`` snapshot shared by every layer; - every group fills the layer rows belonging to its physical slot namespace. + """Commit all ReplaySSM trackers and optionally plan materialization. + + Request vectors have shape ``[batch_capacity]``. Cumulative query metadata + has shape ``[batch_capacity + 1]``; non-cumulative metadata has shape + ``[batch_capacity]``. Group pointer/capacity tables have shape + ``[num_groups]``, group-layer offsets have shape ``[num_groups + 1]``, and + source/destination plans have shape ``[num_layers, batch_capacity]``. + + One CTA owns one ``(request row, cache group)`` tracker transition. Every + group updates its distinct physical-slot namespace. Group zero alone writes + the request-level ``ring_start``/``flush_count`` plan consumed by the one + all-layer FlashInfer materializer call. Invalid and padded rows still write + sentinels so fixed-capacity plans cannot retain stale work. + + The two metadata flags specialize runner input representation only; they do + not change recurrence semantics. Prefill has already produced canonical SSM + state, so this kernel only resets the affected ReplaySSM cursors and, when a + prefix snapshot is requested, emits ``flush_count=0`` for an exact copy. """ batch_idx = tl.program_id(0) group_idx = tl.program_id(1) @@ -151,7 +163,7 @@ def _postprocess_replayssm_modelwide_kernel( ) prefilling = tl.load(is_prefilling + batch_idx, mask=active, other=1) - if QUERY_IS_CUMULATIVE: + if QUERY_METADATA_IS_CUMULATIVE: query_len = tl.load( query_metadata + batch_idx + 1, mask=active, other=0 ) - tl.load( @@ -165,7 +177,7 @@ def _postprocess_replayssm_modelwide_kernel( if valid_req & prefilling: computed = tl.load(num_computed_tokens + req_idx) computed_before = tl.where( - NUM_COMPUTED_IS_AFTER, computed - query_len, computed + NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed ) computed_after = computed_before + query_len first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) @@ -344,7 +356,12 @@ def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: @dataclass class ReplaySSMModelContext: - """Persistent all-layer tables for ReplaySSM post-step maintenance.""" + """One model-wide owner for ReplaySSM trackers and materialization plans. + + Layers in a cache group share tracker tensors; the first layer is only the + representative used to capture those shared addresses. Materialization + plans remain layer-shaped because FlashInfer receives one row per mixer. + """ mixers: list[Any] group_layer_offsets: torch.Tensor @@ -400,6 +417,16 @@ def create( ) block_table_by_gid = dict(zip(mamba_group_ids, block_tables)) replayssm_block_tables = [block_table_by_gid[gid] for gid, _ in grouped] + for block_table in replayssm_block_tables: + if ( + block_table.ndim != 2 + or block_table.dtype != torch.int32 + or not block_table.is_cuda + or block_table.numel() == 0 + ): + raise ValueError( + "model-wide ReplaySSM requires non-empty 2D CUDA int32 block tables" + ) cache_modes = set() for gid, _ in grouped: spec = kv_cache_config.kv_cache_groups[gid].kv_cache_spec @@ -417,8 +444,7 @@ def create( materialize_prefixes = next(iter(cache_modes)) in ("align", "all") mixers = [mixer for _, group_mixers in grouped for mixer in group_mixers] - if not _replayssm_materialize_ready(mixers): - return None + _validate_replayssm_cache(mixers) first = mixers[0] first_ssm = first.kv_cache[1] first_x = first.kv_cache[2] @@ -440,7 +466,7 @@ def create( group_offsets[i + 1] - group_offsets[i] for i in range(len(group_offsets) - 1) ) - tracker_owners = [group_mixers[0] for _, group_mixers in grouped] + tracker_representatives = [group_mixers[0] for _, group_mixers in grouped] strides = {int(block_table.stride(0)) for block_table in replayssm_block_tables} if len(strides) != 1: raise ValueError( @@ -456,13 +482,13 @@ def create( ), block_table_ptrs=_cuda_i64_ptrs(replayssm_block_tables), tracker_ring_start_ptrs=_cuda_i64_ptrs( - [m._replayssm_ring_start for m in tracker_owners] + [m._replayssm_ring_start for m in tracker_representatives] ), tracker_num_committed_ptrs=_cuda_i64_ptrs( - [m._replayssm_prev_num_accepted for m in tracker_owners] + [m._replayssm_prev_num_accepted for m in tracker_representatives] ), tracker_capacities=torch.tensor( - [m._replayssm_ring_start.numel() for m in tracker_owners], + [m._replayssm_ring_start.numel() for m in tracker_representatives], dtype=torch.int32, device=device, ), @@ -526,14 +552,14 @@ def create( materialize_prefixes=materialize_prefixes, ) - def postprocess_and_materialize( + def postprocess( self, *, idx_mapping: torch.Tensor | None, query_metadata: torch.Tensor, - query_is_cumulative: bool, + query_metadata_is_cumulative: bool, num_computed_tokens: torch.Tensor, - num_computed_is_after: bool, + num_computed_is_post_step: bool, num_accepted_tokens: torch.Tensor, is_prefilling: torch.Tensor, live_cols: torch.Tensor, @@ -543,7 +569,7 @@ def postprocess_and_materialize( mamba_block_size: int, num_reqs: int, ) -> None: - """Commit lifecycle metadata, then materialize all layers once.""" + """Commit trackers, then publish a prefix snapshot when configured.""" if num_reqs == 0: return _postprocess_replayssm_modelwide_kernel[(self.max_num_reqs, self.num_groups)]( @@ -573,8 +599,8 @@ def postprocess_and_materialize( LOGICAL_WINDOW=self.logical_window, RING_BUFFER_LEN=self.ring_buffer_len, PAD_SLOT_ID=NULL_BLOCK_ID, - QUERY_IS_CUMULATIVE=query_is_cumulative, - NUM_COMPUTED_IS_AFTER=num_computed_is_after, + QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, + NUM_COMPUTED_IS_POST_STEP=num_computed_is_post_step, HAS_IDX_MAPPING=idx_mapping is not None, ) @@ -586,7 +612,7 @@ def postprocess_and_materialize( self.plan_flush_count, ) - def copy_reassigned_slots( + def materialize_reassigned_slots( self, *, idx_mapping: torch.Tensor | None, @@ -1072,30 +1098,45 @@ def _load_replayssm_materialize() -> Callable[..., None]: return replayssm_materialize -def _replayssm_materialize_ready(mixers: list[Any]) -> bool: - """False only before the caches are allocated; raises on a bad cache. +def _validate_replayssm_cache(mixers: list[Any]) -> None: + """Validate the cache tensors required by model-wide tracker ownership.""" + for layer_idx, mixer in enumerate(mixers): + cache_tensors = mixer.kv_cache[1:5] + if any(tensor.numel() == 0 for tensor in cache_tensors): + raise RuntimeError( + "FlashInfer ReplaySSM requires allocated SSM and replay-ring " + f"cache tensors for every layer; layer {layer_idx} is empty" + ) + if any(not tensor.is_cuda for tensor in cache_tensors): + devices = [str(tensor.device) for tensor in cache_tensors] + raise RuntimeError( + "FlashInfer ReplaySSM requires CUDA cache tensors for every " + f"layer; layer {layer_idx} uses {devices}" + ) - A skip here is not free: ``state_skip_postprocess`` has already told the - fused postprocess kernel not to copy this temporal state, so silently - doing nothing would leave the destination block holding stale SSM state. - The empty-cache case (profiling and other pre-allocation runs) is the one - legitimate no-op; anything else is a misconfiguration and must be loud. - """ - ssm = mixers[0].kv_cache[1] - x_cache = mixers[0].kv_cache[2] - if ssm.numel() == 0: - return False - if not ssm.is_cuda: - raise RuntimeError( - "FlashInfer ReplaySSM prefix materialization requires a CUDA SSM " - f"state cache; got device {ssm.device}" - ) - if x_cache.numel() == 0 or mixers[0]._replayssm_ring_start.numel() == 0: - raise RuntimeError( - "FlashInfer ReplaySSM prefix materialization requires allocated " - "replay ring buffers and ring trackers" - ) - return True + ring_start = mixer._replayssm_ring_start + num_committed = mixer._replayssm_prev_num_accepted + if ring_start.numel() == 0 or num_committed.numel() == 0: + raise RuntimeError( + "FlashInfer ReplaySSM requires allocated ring trackers for " + f"every layer; layer {layer_idx} is empty" + ) + if not ring_start.is_cuda or not num_committed.is_cuda: + raise RuntimeError( + "FlashInfer ReplaySSM requires CUDA ring trackers for every layer" + ) + if ( + ring_start.ndim != 1 + or num_committed.ndim != 1 + or ring_start.dtype != torch.int32 + or num_committed.dtype != torch.int32 + ): + raise ValueError("ReplaySSM ring trackers must be 1D int32 tensors") + if ring_start.numel() != num_committed.numel(): + raise ValueError( + "ReplaySSM ring trackers must have equal capacities; got " + f"{ring_start.numel()} and {num_committed.numel()}" + ) def initialize_mamba_ssu_backend( diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 1920788fcde1..40cb4768c4eb 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -161,7 +161,7 @@ def _get_mamba_group_info( self._mamba_spec = mamba_spec return self._mamba_group_ids, self._mamba_spec - def _ensure_align_ctx( + def _ensure_mamba_postprocess_ctx( self, kv_cache_config: KVCacheConfig, mamba_group_ids: list[int], @@ -224,7 +224,9 @@ def preprocess_state( if num_reqs == 0: return mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config) - ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + ctx = self._ensure_mamba_postprocess_ctx( + kv_cache_config, mamba_group_ids, block_tables + ) # The state-advance + pre-copy kernels run every step; they fast-exit per # request when src_col < 0 or src_col == dst_col, so no copy happens on @@ -247,7 +249,7 @@ def preprocess_state( MAMBA_BLOCK_SIZE=mamba_spec.block_size, ) if ctx.replayssm is not None: - ctx.replayssm.copy_reassigned_slots( + ctx.replayssm.materialize_reassigned_slots( idx_mapping=input_batch.idx_mapping, src_cols=self._mamba_src_col_gpu, dst_cols=self._mamba_state_idx_gpu, @@ -320,7 +322,9 @@ def prepare_attn( if self._use_flashinfer_replayssm: mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) - self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables) + self._ensure_mamba_postprocess_ctx( + kv_cache_config, mamba_group_ids, block_tables + ) if self._align_mode: mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) @@ -331,7 +335,7 @@ def prepare_attn( if hasattr(builder, "mamba_aligned_state_indices"): aligned_index_builders.append((group_idx, builder)) if aligned_index_builders: - ctx = self._ensure_align_ctx( + ctx = self._ensure_mamba_postprocess_ctx( kv_cache_config, mamba_group_ids, block_tables ) all_group_indices = ctx.compute_aligned_state_indices( @@ -446,12 +450,12 @@ def postprocess_state( assert replayssm is not None if is_prefilling is None: is_prefilling = self._is_prefilling_gpu[:num_reqs] - replayssm.postprocess_and_materialize( + replayssm.postprocess( idx_mapping=idx_mapping, query_metadata=query_start_loc, - query_is_cumulative=True, + query_metadata_is_cumulative=True, num_computed_tokens=num_computed_tokens, - num_computed_is_after=True, + num_computed_is_post_step=True, # Prefix migration can reset the live buffer to one for the # next step; use its snapshot in that case. Mode none never # runs the migration kernel, so the acceptance buffer is exact. diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 501885bd0d14..445574bb1d4e 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1085,9 +1085,7 @@ def _get_mamba_state_copy_funcs(self) -> MambaStateCopyFuncsByType: def _get_mamba_bufs(self) -> mamba_utils.MambaBuffers: # The postprocess sub-object is also the model-level owner of # FlashInfer ReplaySSM trackers, including STP. - assert ( - self._needs_prefix_state_migration or self._use_flashinfer_replayssm - ) + assert self._needs_prefix_state_migration or self._use_flashinfer_replayssm if self._mamba_bufs is None: self._mamba_bufs = mamba_utils.MambaBuffers.create( max_num_reqs=self.max_num_reqs, @@ -1617,9 +1615,7 @@ def _update_states_after_model_execute( num_reqs = output_token_ids.size(0) self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) - if ( - self._needs_prefix_state_migration or self._use_flashinfer_replayssm - ): + if self._needs_prefix_state_migration or self._use_flashinfer_replayssm: # Fused GPU postprocess: state copies + per-request accepted-token # update without CPU-GPU sync. The metadata # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is @@ -4481,9 +4477,7 @@ def execute_model( num_reqs, self.requests, self.mamba_state_idx, - fixed_live_col=( - None if self._needs_prefix_state_migration else 0 - ), + fixed_live_col=(None if self._needs_prefix_state_migration else 0), ) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 04d89028c789..6d2aff0a6552 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -1617,7 +1617,7 @@ def preprocess_mamba( fused.src_col.copy_to_gpu(num_reqs) fused.token_bias.copy_to_gpu(num_reqs) if fused.ctx.replayssm is not None: - fused.ctx.replayssm.copy_reassigned_slots( + fused.ctx.replayssm.materialize_reassigned_slots( idx_mapping=None, src_cols=fused.src_col.gpu, dst_cols=fused.state_idx.gpu, @@ -1693,12 +1693,13 @@ def postprocess_mamba_align_gpu( mamba_state_copy_funcs: MambaStateCopyFuncsByType, run_prefix_state_migration: bool, ) -> None: - """GPU-side Mamba postprocess for fused align state maintenance. + """Run model-wide Mamba state maintenance after token acceptance. Lazily binds the fused-kernel context to the persistent block tables and - forward-context state pointers on the first call, runs the fused kernel, - and async-copies the per-request accepted-token counts back to the input - batch's CPU tensor for the next iteration's preprocess. + forward-context state pointers on the first call. Prefix modes run the + generic state-copy planner before committing ReplaySSM trackers; mode none + commits only the trackers. The accepted counts are then copied back for any + CPU-side consumer on the next iteration. """ ctx = bufs.postprocess_align # The caller enables this context for spec-decode hybrid state copies or @@ -1733,12 +1734,12 @@ def postprocess_mamba_align_gpu( ) accepted_tokens_for_postprocess = ctx.num_accepted_tokens_out if ctx.replayssm is not None: - ctx.replayssm.postprocess_and_materialize( + ctx.replayssm.postprocess( idx_mapping=None, query_metadata=ctx.num_scheduled_tokens_buf.gpu, - query_is_cumulative=False, + query_metadata_is_cumulative=False, num_computed_tokens=ctx.num_computed_tokens_buf.gpu, - num_computed_is_after=False, + num_computed_is_post_step=False, num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, live_cols=ctx.mamba_state_idx_buf.gpu, From fe2707688b0f77866167b00e69e113f39a2b9a84 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 00:31:54 +0200 Subject: [PATCH 12/53] fix(mamba): compact ReplaySSM materialization requests Build compact physical-request maps inside the existing model-wide lifecycle kernels and pass them to the current FlashInfer materializer API. Avoid request scanning when prefix materialization is disabled and retain the mixed-batch PP publication fix. Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 63 ++++++- tests/v1/e2e/test_replayssm_decode.py | 11 +- .../worker/test_gpu_model_runner_v2_eplb.py | 63 +++++++ .../layers/mamba/ops/ssu_dispatch.py | 164 +++++++++++++++++- vllm/v1/worker/gpu/model_runner.py | 18 +- 5 files changed, 304 insertions(+), 15 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 50957fe8596a..dcb07a0dec0f 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -482,6 +482,7 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None assert materializer.call_count == 0 assert ctx.plan_flush_count.tolist() == [-1, -1] + assert ctx.active_request_indices.tolist() == [-1, -1] for mixers, live_slot in zip(groups, (1, 4)): # Both layers share this group tracker. The expected single transition # per step would differ if either layer committed it independently. @@ -537,6 +538,7 @@ def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch) assert args[12] is ctx.dst_slots assert args[13] is ctx.plan_ring_start assert args[14] is ctx.plan_flush_count + assert args[15] is ctx.active_request_indices assert kwargs["num_heads"] == 4 assert kwargs["heads_per_group"] == 2 assert kwargs["max_window"] == 16 @@ -547,12 +549,61 @@ def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch) assert ctx.dst_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 4 assert ctx.plan_ring_start.tolist() == [2, 0] assert ctx.plan_flush_count.tolist() == [6, -1] + assert ctx.active_request_indices.tolist() == [0, -1] assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 6 assert groups[0][0]._replayssm_prev_num_accepted[2].item() == 0 assert groups[1][0]._replayssm_prev_num_accepted[4].item() == 6 assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatch): + _, config, forward_context, block_tables = _modelwide_replayssm_fixture() + block_tables[0][1] = torch.tensor([3, 2, 1], device="cuda") + block_tables[1][1] = torch.tensor([7, 6, 5], device="cuda") + materializer = Mock() + monkeypatch.setattr( + ssu_dispatch, "_load_replayssm_materialize", lambda: materializer + ) + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor([1, 4], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, + num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([False, False], device="cuda"), + live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([-1, 1], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.tensor([0, 1], dtype=torch.int32, device="cuda"), + mamba_block_size=4, + num_reqs=2, + ) + ctx.materialize_reassigned_slots( + idx_mapping=None, + src_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), + dst_cols=torch.tensor([-1, 1], dtype=torch.int32, device="cuda"), + num_reqs=2, + ) + torch.accelerator.synchronize() + + assert materializer.call_count == 2 + assert ctx.plan_flush_count.tolist() == [-1, 1] + assert ctx.active_request_indices.tolist() == [1, -1] + assert ctx.precopy_flush_count.tolist() == [-1, 2] + assert ctx.precopy_active_request_indices.tolist() == [1, -1] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() @@ -571,14 +622,14 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch ) assert ctx is not None ctx.postprocess( - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + idx_mapping=torch.tensor([1], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, - num_computed_tokens=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), num_computed_is_post_step=True, - num_accepted_tokens=torch.tensor([3, 1], dtype=torch.int32, device="cuda"), + num_accepted_tokens=torch.tensor([1, 3], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + live_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), @@ -590,6 +641,7 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch assert kernel.call_count == 1 assert ctx.plan_ring_start.tolist() == [15, 0] assert ctx.plan_flush_count.tolist() == [1, -1] + assert ctx.active_request_indices.tolist() == [0, -1] for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): assert mixers[0]._replayssm_ring_start[source_slot].item() == 15 assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 3 @@ -634,6 +686,7 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): assert kernel.call_count == 1 assert ctx.plan_flush_count.tolist() == [-1, -1] + assert ctx.active_request_indices.tolist() == [-1, -1] for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): for source_slot in source_slots: assert mixers[0]._replayssm_ring_start[source_slot].item() == 0 @@ -668,6 +721,7 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): assert kernel.call_count == 1 assert ctx.precopy_ring_start.tolist() == [2, 0] assert ctx.precopy_flush_count.tolist() == [4, -1] + assert ctx.precopy_active_request_indices.tolist() == [0, -1] assert ctx.precopy_src_slots[:, 0].tolist() == [1, 1, 4, 4] assert ctx.precopy_dst_slots[:, 0].tolist() == [2, 2, 5, 5] for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): @@ -706,6 +760,7 @@ def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): assert kernel.call_count == 1 assert ctx.precopy_ring_start.tolist() == [2, 0] assert ctx.precopy_flush_count.tolist() == [4, -1] + assert ctx.precopy_active_request_indices.tolist() == [0, -1] assert ctx.precopy_src_slots[:, 0].tolist() == [ NULL_BLOCK_ID, NULL_BLOCK_ID, diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 5407aa34ac84..56b357da7b4f 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Engine-level parity: ReplaySSM standard decode vs the baseline SSM kernel.""" +from inspect import signature + import pytest import vllm.envs as envs @@ -25,7 +27,9 @@ try: from flashinfer.mamba.replayssm_materialize import replayssm_materialize - HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = callable(replayssm_materialize) + HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = callable(replayssm_materialize) and ( + "active_request_indices" in signature(replayssm_materialize).parameters + ) except ImportError: HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = False @@ -414,7 +418,8 @@ def test_flashinfer_replayssm_prefix_cache_tp1( @requires_flashinfer_replayssm_materialization @large_gpu_mark(min_gb=40) -def test_flashinfer_replayssm_all_prefix_cache_v2(vllm_runner, monkeypatch): +@pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"]) +def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: bool): _check_flashinfer_replayssm_prefix_caching( vllm_runner, MAMBA2_PREFIX_MODEL, @@ -422,7 +427,7 @@ def test_flashinfer_replayssm_all_prefix_cache_v2(vllm_runner, monkeypatch): mamba_cache_mode="all", moe_backend="triton", use_ngram=False, - use_v2=True, + use_v2=use_v2, tensor_parallel_size=1, ) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 559af5572de0..e683e47aeb91 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -5,6 +5,7 @@ from types import SimpleNamespace from typing import Any +import numpy as np import torch from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry @@ -202,3 +203,65 @@ def fake_receive(*args, **kwargs): output = mrv2.GPUModelRunner.sample_tokens(runner, None) assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] + + +def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( + monkeypatch, +): + events = [] + runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) + idx_mapping = torch.tensor([3, 7], dtype=torch.int64) + query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32) + input_batch = SimpleNamespace( + num_reqs=2, + idx_mapping=idx_mapping, + idx_mapping_np=np.array([3, 7], dtype=np.intp), + # Row 0 finishes prefill and will be processed from the deferred PP + # receive. Row 1 is a non-final chunk and must be published now. + num_computed_tokens_np=np.array([3, 2], dtype=np.int32), + prefill_len_np=np.array([4, 6], dtype=np.int32), + num_scheduled_tokens=np.array([1, 1], dtype=np.int32), + query_start_loc=query_start_loc, + ) + runner.execute_model_state = SimpleNamespace( + input_batch=input_batch, + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=None, + aux_hidden_states=None, + dp_sync=None, + finished_req_ids=set(), + ec_connector_output=None, + routed_experts=None, + ) + num_computed_tokens = torch.zeros(8, dtype=torch.int32) + runner.req_states = SimpleNamespace( + num_computed_tokens=SimpleNamespace(gpu=num_computed_tokens) + ) + postprocess_args = [] + runner.model_state = SimpleNamespace( + postprocess_state=lambda *args: postprocess_args.append(args) + ) + runner.pp_handler = SimpleNamespace( + receive=lambda *_: events.append("receive") or False + ) + runner.postprocess_num_computed_tokens = lambda *_: events.append( + "postprocess_num_computed_tokens" + ) + runner.eplb.step = lambda *args, **kwargs: events.append("eplb") + monkeypatch.setattr( + mrv2, + "async_copy_to_gpu", + lambda value, *, device: torch.as_tensor(value, device=device), + ) + + output = mrv2.GPUModelRunner.sample_tokens(runner, None) + + assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) + assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] + assert len(postprocess_args) == 1 + published_mapping, num_sampled, computed, query_metadata = postprocess_args[0] + assert published_mapping.tolist() == [-1, 7] + assert num_sampled == 0 + assert computed is num_computed_tokens + assert query_metadata is query_start_loc diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index af2b3804f274..067daedfd244 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -13,6 +13,7 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from functools import cache +from inspect import signature from typing import Any import torch @@ -27,7 +28,7 @@ logger = init_logger(__name__) -@triton.jit(do_not_specialize=["num_reqs"]) +@triton.jit(do_not_specialize=["num_reqs", "num_materialize_reqs"]) def _postprocess_replayssm_modelwide_kernel( # Per-request step metadata. idx_mapping, @@ -50,10 +51,12 @@ def _postprocess_replayssm_modelwide_kernel( dst_slots, plan_ring_start, plan_flush_count, + active_request_indices, # Runtime sizes. block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, num_reqs, + num_materialize_reqs, # Compile-time model constants. MAX_LAYERS_PER_GROUP: tl.constexpr, MAMBA_BLOCK_SIZE: tl.constexpr, @@ -63,6 +66,7 @@ def _postprocess_replayssm_modelwide_kernel( QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, NUM_COMPUTED_IS_POST_STEP: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, + MAX_NUM_REQS: tl.constexpr, ) -> None: """Commit all ReplaySSM trackers and optionally plan materialization. @@ -71,15 +75,18 @@ def _postprocess_replayssm_modelwide_kernel( ``[batch_capacity]``. Group pointer/capacity tables have shape ``[num_groups]``, group-layer offsets have shape ``[num_groups + 1]``, and source/destination plans have shape ``[num_layers, batch_capacity]``. + ``active_request_indices`` has shape ``[batch_capacity]`` and contains the + compact physical-request prefix consumed by FlashInfer, followed by ``-1``. One CTA owns one ``(request row, cache group)`` tracker transition. Every group updates its distinct physical-slot namespace. Group zero alone writes - the request-level ``ring_start``/``flush_count`` plan consumed by the one - all-layer FlashInfer materializer call. Invalid and padded rows still write - sentinels so fixed-capacity plans cannot retain stale work. + the request-level ``ring_start``/``flush_count`` plan and one of its CTAs + compacts active request indices for the all-layer FlashInfer materializer. + Invalid and padded rows still write sentinels so fixed-capacity plans cannot + retain stale work. - The two metadata flags specialize runner input representation only; they do - not change recurrence semantics. Prefill has already produced canonical SSM + Metadata flags specialize runner input representation only; they do not + change recurrence semantics. Prefill has already produced canonical SSM state, so this kernel only resets the affected ReplaySSM cursors and, when a prefix snapshot is requested, emits ``flush_count=0`` for an exact copy. """ @@ -234,6 +241,72 @@ def _postprocess_replayssm_modelwide_kernel( tl.store(tracker_start + materialize_dst_slot, 0) tl.store(tracker_committed + materialize_dst_slot, 0) + if (group_idx == 0) & (batch_idx == 0): + # FlashInfer consumes a compact active prefix and stops at the + # first -1. Build it from planner inputs here so sparse flushes do + # not require a second Triton launch or a host synchronization. + active_count = 0 + for candidate_idx in tl.range(0, num_materialize_reqs): + candidate_req_idx = candidate_idx + if HAS_IDX_MAPPING: + candidate_req_idx = tl.load(idx_mapping + candidate_idx) + candidate_valid_req = candidate_req_idx >= 0 + candidate_live_col = tl.load( + live_cols + candidate_req_idx, + mask=candidate_valid_req, + other=-1, + ) + candidate_valid_live_col = candidate_valid_req & (candidate_live_col >= 0) + candidate_live_slot = tl.load( + block_table + + candidate_idx * block_table_stride_req + + candidate_live_col, + mask=candidate_valid_live_col, + other=PAD_SLOT_ID, + ) + candidate_valid_live = ( + candidate_valid_live_col + & (candidate_live_slot != PAD_SLOT_ID) + & (candidate_live_slot >= 0) + & (candidate_live_slot < tracker_capacity) + ) + candidate_src_col = tl.load(materialize_src_cols + candidate_idx) + candidate_dst_col = tl.load(materialize_dst_cols + candidate_idx) + candidate_wants_materialize = ( + candidate_valid_req + & (candidate_src_col >= 0) + & (candidate_dst_col >= 0) + ) + candidate_src_slot = tl.load( + block_table + + candidate_idx * block_table_stride_req + + candidate_src_col, + mask=candidate_wants_materialize, + other=PAD_SLOT_ID, + ) + candidate_dst_slot = tl.load( + block_table + + candidate_idx * block_table_stride_req + + candidate_dst_col, + mask=candidate_wants_materialize, + other=PAD_SLOT_ID, + ) + candidate_valid_materialize = ( + candidate_valid_live + & candidate_wants_materialize + & (candidate_src_slot != PAD_SLOT_ID) + & (candidate_dst_slot != PAD_SLOT_ID) + & (candidate_src_slot >= 0) + & (candidate_dst_slot >= 0) + & (candidate_src_slot < tracker_capacity) + & (candidate_dst_slot < tracker_capacity) + ) + if candidate_valid_materialize: + tl.store(active_request_indices + active_count, candidate_idx) + active_count += 1 + if active_count < MAX_NUM_REQS: + tl.store(active_request_indices + active_count, -1) + @triton.jit(do_not_specialize=["num_reqs"]) def _copy_reassigned_replayssm_slots_kernel( @@ -249,12 +322,14 @@ def _copy_reassigned_replayssm_slots_kernel( dst_slots, plan_ring_start, plan_flush_count, + active_request_indices, block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, num_reqs, MAX_LAYERS_PER_GROUP: tl.constexpr, PAD_SLOT_ID: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, + MAX_NUM_REQS: tl.constexpr, ) -> None: """Plan an exact copy when align reassigns a request's writable slot.""" batch_idx = tl.program_id(0) @@ -334,6 +409,61 @@ def _copy_reassigned_replayssm_slots_kernel( tl.store(tracker_start + dst_slot, 0) tl.store(tracker_committed + dst_slot, 0) + if (group_idx == 0) & (batch_idx == 0): + # The logical reassignment can require other cache groups to copy + # even when group zero's physical slots alias. Compact every valid + # group-zero plan into the ordered map required by FlashInfer. + active_count = 0 + for candidate_idx in tl.range(0, num_reqs): + candidate_req_idx = candidate_idx + if HAS_IDX_MAPPING: + candidate_req_idx = tl.load(idx_mapping + candidate_idx) + candidate_valid_req = candidate_req_idx >= 0 + candidate_src_col = tl.load( + src_cols + candidate_req_idx, + mask=candidate_valid_req, + other=-1, + ) + candidate_dst_col = tl.load( + dst_cols + candidate_req_idx, + mask=candidate_valid_req, + other=-1, + ) + candidate_wants_copy = ( + candidate_valid_req + & (candidate_src_col >= 0) + & (candidate_dst_col >= 0) + & (candidate_src_col != candidate_dst_col) + ) + candidate_src_slot = tl.load( + block_table + + candidate_idx * block_table_stride_req + + candidate_src_col, + mask=candidate_wants_copy, + other=PAD_SLOT_ID, + ) + candidate_dst_slot = tl.load( + block_table + + candidate_idx * block_table_stride_req + + candidate_dst_col, + mask=candidate_wants_copy, + other=PAD_SLOT_ID, + ) + candidate_valid_mapping = ( + candidate_wants_copy + & (candidate_src_slot != PAD_SLOT_ID) + & (candidate_dst_slot != PAD_SLOT_ID) + & (candidate_src_slot >= 0) + & (candidate_dst_slot >= 0) + & (candidate_src_slot < tracker_capacity) + & (candidate_dst_slot < tracker_capacity) + ) + if candidate_valid_mapping: + tl.store(active_request_indices + active_count, candidate_idx) + active_count += 1 + if active_count < MAX_NUM_REQS: + tl.store(active_request_indices + active_count, -1) + def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: ssm = mixer.kv_cache[1] @@ -384,10 +514,12 @@ class ReplaySSMModelContext: dst_slots: torch.Tensor plan_ring_start: torch.Tensor plan_flush_count: torch.Tensor + active_request_indices: torch.Tensor precopy_src_slots: torch.Tensor precopy_dst_slots: torch.Tensor precopy_ring_start: torch.Tensor precopy_flush_count: torch.Tensor + precopy_active_request_indices: torch.Tensor block_table_stride_req: int max_num_reqs: int num_groups: int @@ -525,6 +657,9 @@ def create( plan_flush_count=torch.full( (max_num_reqs,), -1, dtype=torch.int32, device=device ), + active_request_indices=torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ), precopy_src_slots=torch.full( (len(mixers), max_num_reqs), NULL_BLOCK_ID, @@ -543,6 +678,9 @@ def create( precopy_flush_count=torch.full( (max_num_reqs,), -1, dtype=torch.int32, device=device ), + precopy_active_request_indices=torch.full( + (max_num_reqs,), -1, dtype=torch.int32, device=device + ), block_table_stride_req=next(iter(strides)), max_num_reqs=max_num_reqs, num_groups=len(grouped), @@ -591,9 +729,11 @@ def postprocess( self.dst_slots, self.plan_ring_start, self.plan_flush_count, + self.active_request_indices, self.block_table_stride_req, self.src_slots.stride(0), num_reqs, + num_reqs if self.materialize_prefixes else 0, MAX_LAYERS_PER_GROUP=self.max_layers_per_group, MAMBA_BLOCK_SIZE=mamba_block_size, LOGICAL_WINDOW=self.logical_window, @@ -602,6 +742,7 @@ def postprocess( QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, NUM_COMPUTED_IS_POST_STEP=num_computed_is_post_step, HAS_IDX_MAPPING=idx_mapping is not None, + MAX_NUM_REQS=self.max_num_reqs, ) if self.materialize_prefixes: @@ -610,6 +751,7 @@ def postprocess( self.dst_slots, self.plan_ring_start, self.plan_flush_count, + self.active_request_indices, ) def materialize_reassigned_slots( @@ -640,18 +782,21 @@ def materialize_reassigned_slots( self.precopy_dst_slots, self.precopy_ring_start, self.precopy_flush_count, + self.precopy_active_request_indices, self.block_table_stride_req, self.precopy_src_slots.stride(0), num_reqs, MAX_LAYERS_PER_GROUP=self.max_layers_per_group, PAD_SLOT_ID=NULL_BLOCK_ID, HAS_IDX_MAPPING=idx_mapping is not None, + MAX_NUM_REQS=self.max_num_reqs, ) self._materialize_planned( self.precopy_src_slots, self.precopy_dst_slots, self.precopy_ring_start, self.precopy_flush_count, + self.precopy_active_request_indices, ) def _materialize_planned( @@ -660,6 +805,7 @@ def _materialize_planned( dst_slots: torch.Tensor, ring_start: torch.Tensor, flush_count: torch.Tensor, + active_request_indices: torch.Tensor, ) -> None: first = self.mixers[0] mamba_config = first.mamba_config @@ -686,6 +832,7 @@ def _materialize_planned( dst_slots, ring_start, flush_count, + active_request_indices, state_dtype=first.kv_cache[1].dtype, input_dtype=first.kv_cache[2].dtype, matrixA_dtype=first.A.dtype, @@ -1095,6 +1242,11 @@ def _load_replayssm_materialize() -> Callable[..., None]: "FlashInfer ReplaySSM prefix caching requires " "flashinfer.mamba.replayssm_materialize" ) from e + if "active_request_indices" not in signature(replayssm_materialize).parameters: + raise ImportError( + "FlashInfer ReplaySSM prefix caching requires the ordered " + "active_request_indices materialization API" + ) return replayssm_materialize diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b6cd6ea04190..eb5cef26bd22 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -133,7 +133,7 @@ from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu.model_states import init_model_state from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner -from vllm.v1.worker.gpu.pp_utils import PPHandler +from vllm.v1.worker.gpu.pp_utils import PPHandler, compute_need_sampled_mask from vllm.v1.worker.gpu.sample.batch_shard import ( BatchSharder, all_to_all_logits, @@ -1869,7 +1869,21 @@ def sample_tokens( if not all_decode_next: # Might contain non-final prefill chunks, which will be scheduled # in the immediate next step (rather than in pp_size steps). - self.model_state.postprocess_state(input_batch.idx_mapping, 0) + idx_mapping = input_batch.idx_mapping + need_sampled_mask = compute_need_sampled_mask(input_batch) + if need_sampled_mask is not None: + # Sampled rows are published from the deferred PP receive. + # Mask them here so a mixed batch commits every row once. + idx_mapping_np = np.where( + need_sampled_mask, -1, input_batch.idx_mapping_np + ) + idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) + self.model_state.postprocess_state( + idx_mapping, + 0, + self.req_states.num_computed_tokens.gpu, + input_batch.query_start_loc, + ) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) From 85463689ac29c391433bacb6c43e9f9a191fc994 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 00:39:28 +0200 Subject: [PATCH 13/53] refactor(mamba): simplify ReplaySSM cache lifecycle Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 301 +++--- .../worker/test_mamba_hybrid_model_state.py | 38 + tests/v1/worker/test_mamba_utils.py | 55 +- .../layers/mamba/ops/ssu_dispatch.py | 885 ++++++------------ .../worker/gpu/model_states/mamba_hybrid.py | 214 +++-- vllm/v1/worker/gpu_model_runner.py | 15 +- vllm/v1/worker/mamba_utils.py | 39 +- 7 files changed, 734 insertions(+), 813 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index dcb07a0dec0f..84222ee4fc87 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -444,19 +444,27 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None query_len[0] = scheduled accepted[0] = num_accepted is_prefilling[0] = prefilling - ctx.postprocess( + ctx.preprocess( idx_mapping=None, query_metadata=query_len, query_metadata_is_cumulative=False, num_computed_tokens=num_computed, - num_computed_is_post_step=False, + is_prefilling=is_prefilling, + src_cols=live_cols, + dst_cols=live_cols, + mamba_block_size=1024, + num_reqs=1, + ) + ctx.postprocess( + idx_mapping=None, + query_metadata=query_len, + query_metadata_is_cumulative=False, num_accepted_tokens=accepted, is_prefilling=is_prefilling, live_cols=live_cols, materialize_src_cols=no_materialize, materialize_dst_cols=no_materialize, materialize_token_counts=materialize_counts, - mamba_block_size=1024, num_reqs=1, ) num_computed[0] += num_accepted @@ -481,8 +489,7 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None torch.accelerator.synchronize() assert materializer.call_count == 0 - assert ctx.plan_flush_count.tolist() == [-1, -1] - assert ctx.active_request_indices.tolist() == [-1, -1] + assert all(group.plan_flush_count.tolist() == [-1, -1] for group in ctx.groups) for mixers, live_slot in zip(groups, (1, 4)): # Both layers share this group tracker. The expected single transition # per step would differ if either layer committed it independently. @@ -498,7 +505,7 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch): +def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() for mixers, source_slot in zip(groups, (1, 4)): mixers[0]._replayssm_ring_start[source_slot] = 2 @@ -518,92 +525,45 @@ def test_modelwide_replayssm_postprocess_launches_materializer_once(monkeypatch) idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, - num_computed_tokens=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), - mamba_block_size=4, num_reqs=1, ) + ctx.materialize() torch.accelerator.synchronize() - assert kernel.call_count == 1 - args = kernel.call_args.args - kwargs = kernel.call_args.kwargs - assert args[11] is ctx.src_slots - assert args[12] is ctx.dst_slots - assert args[13] is ctx.plan_ring_start - assert args[14] is ctx.plan_flush_count - assert args[15] is ctx.active_request_indices - assert kwargs["num_heads"] == 4 - assert kwargs["heads_per_group"] == 2 - assert kwargs["max_window"] == 16 - assert kwargs["ring_buffer_len"] == 20 - assert ctx.src_slots[:, 0].tolist() == [1, 1, 4, 4] - assert ctx.dst_slots[:, 0].tolist() == [2, 2, 5, 5] - assert ctx.src_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 4 - assert ctx.dst_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 4 - assert ctx.plan_ring_start.tolist() == [2, 0] - assert ctx.plan_flush_count.tolist() == [6, -1] - assert ctx.active_request_indices.tolist() == [0, -1] + assert kernel.call_count == 2 + for call, group_ctx, src_slot, dst_slot in zip( + kernel.call_args_list, ctx.groups, (1, 4), (2, 5) + ): + args = call.args + kwargs = call.kwargs + assert args[11] is group_ctx.src_slots + assert args[12] is group_ctx.dst_slots + assert args[13] is group_ctx.plan_ring_start + assert args[14] is group_ctx.plan_flush_count + assert args[15] is group_ctx.active_request_indices + assert kwargs["num_heads"] == 4 + assert kwargs["heads_per_group"] == 2 + assert kwargs["max_window"] == 16 + assert kwargs["ring_buffer_len"] == 20 + assert group_ctx.src_slots[:, 0].tolist() == [src_slot] * 2 + assert group_ctx.dst_slots[:, 0].tolist() == [dst_slot] * 2 + assert group_ctx.src_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 2 + assert group_ctx.dst_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 2 + assert group_ctx.plan_ring_start.tolist() == [2, 0] + assert group_ctx.plan_flush_count.tolist() == [6, -1] + assert group_ctx.active_request_indices.tolist() == [0, -1] assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 6 assert groups[0][0]._replayssm_prev_num_accepted[2].item() == 0 assert groups[1][0]._replayssm_prev_num_accepted[4].item() == 6 assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatch): - _, config, forward_context, block_tables = _modelwide_replayssm_fixture() - block_tables[0][1] = torch.tensor([3, 2, 1], device="cuda") - block_tables[1][1] = torch.tensor([7, 6, 5], device="cuda") - materializer = Mock() - monkeypatch.setattr( - ssu_dispatch, "_load_replayssm_materialize", lambda: materializer - ) - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([1, 4], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([False, False], device="cuda"), - live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([-1, 1], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.tensor([0, 1], dtype=torch.int32, device="cuda"), - mamba_block_size=4, - num_reqs=2, - ) - ctx.materialize_reassigned_slots( - idx_mapping=None, - src_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), - dst_cols=torch.tensor([-1, 1], dtype=torch.int32, device="cuda"), - num_reqs=2, - ) - torch.accelerator.synchronize() - - assert materializer.call_count == 2 - assert ctx.plan_flush_count.tolist() == [-1, 1] - assert ctx.active_request_indices.tolist() == [1, -1] - assert ctx.precopy_flush_count.tolist() == [-1, 2] - assert ctx.precopy_active_request_indices.tolist() == [1, -1] - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() @@ -625,23 +585,21 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch idx_mapping=torch.tensor([1], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, - num_computed_tokens=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([1, 3], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - mamba_block_size=4, num_reqs=1, ) + ctx.materialize() torch.accelerator.synchronize() - assert kernel.call_count == 1 - assert ctx.plan_ring_start.tolist() == [15, 0] - assert ctx.plan_flush_count.tolist() == [1, -1] - assert ctx.active_request_indices.tolist() == [0, -1] + assert kernel.call_count == 2 + for group_ctx in ctx.groups: + assert group_ctx.plan_ring_start.tolist() == [15, 0] + assert group_ctx.plan_flush_count.tolist() == [1, -1] for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): assert mixers[0]._replayssm_ring_start[source_slot].item() == 15 assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 3 @@ -650,7 +608,54 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): +def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): + _, config, forward_context, block_tables = _modelwide_replayssm_fixture() + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + + before = [] + for group_ctx in ctx.groups: + group_ctx.ring_start.copy_( + torch.arange(group_ctx.ring_start.numel(), device="cuda") + ) + group_ctx.num_committed.fill_(7) + before.append((group_ctx.ring_start.clone(), group_ctx.num_committed.clone())) + group_ctx.src_slots.fill_(3) + group_ctx.dst_slots.fill_(4) + group_ctx.plan_ring_start.fill_(5) + group_ctx.plan_flush_count.fill_(6) + + ctx.postprocess( + idx_mapping=torch.tensor([-1], dtype=torch.int32, device="cuda"), + query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=True, + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), + is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), + live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.ones(2, dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.ones(2, dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + torch.accelerator.synchronize() + + for group_ctx, (ring_start, num_committed) in zip(ctx.groups, before): + assert torch.equal(group_ctx.ring_start, ring_start) + assert torch.equal(group_ctx.num_committed, num_committed) + assert torch.all(group_ctx.src_slots == NULL_BLOCK_ID) + assert torch.all(group_ctx.dst_slots == NULL_BLOCK_ID) + assert group_ctx.plan_ring_start.tolist() == [0, 0] + assert group_ctx.plan_flush_count.tolist() == [-1, -1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_preprocess_resets_prefill_slots(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): for source_slot in source_slots: @@ -667,32 +672,70 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slot(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.postprocess( + no_change = torch.full((2,), -1, dtype=torch.int32, device="cuda") + ctx.preprocess( idx_mapping=None, query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.tensor([-1, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), + src_cols=no_change, + dst_cols=no_change, mamba_block_size=4, num_reqs=1, ) torch.accelerator.synchronize() - assert kernel.call_count == 1 - assert ctx.plan_flush_count.tolist() == [-1, -1] - assert ctx.active_request_indices.tolist() == [-1, -1] + assert kernel.call_count == 0 for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): for source_slot in source_slots: assert mixers[0]._replayssm_ring_start[source_slot].item() == 0 assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for mixers, destination_slot in zip(groups, (2, 5)): + mixers[0]._replayssm_ring_start[destination_slot] = 7 + mixers[0]._replayssm_prev_num_accepted[destination_slot] = 9 + materializer = Mock() + monkeypatch.setattr( + ssu_dispatch, "_load_replayssm_materialize", lambda: materializer + ) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([True, False], device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + ctx.materialize() + torch.accelerator.synchronize() + + assert materializer.call_count == 2 + for group_ctx in ctx.groups: + assert group_ctx.plan_ring_start.tolist() == [0, 0] + assert group_ctx.plan_flush_count.tolist() == [0, -1] + for mixers, destination_slot in zip(groups, (2, 5)): + assert mixers[0]._replayssm_ring_start[destination_slot].item() == 0 + assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() @@ -710,20 +753,27 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.materialize_reassigned_slots( + ctx.preprocess( idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), + query_metadata=torch.zeros(2, dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + mamba_block_size=4, num_reqs=1, ) + ctx.materialize() torch.accelerator.synchronize() - assert kernel.call_count == 1 - assert ctx.precopy_ring_start.tolist() == [2, 0] - assert ctx.precopy_flush_count.tolist() == [4, -1] - assert ctx.precopy_active_request_indices.tolist() == [0, -1] - assert ctx.precopy_src_slots[:, 0].tolist() == [1, 1, 4, 4] - assert ctx.precopy_dst_slots[:, 0].tolist() == [2, 2, 5, 5] + assert kernel.call_count == 2 + for group_ctx, source_slot, destination_slot in zip(ctx.groups, (1, 4), (2, 5)): + assert group_ctx.plan_ring_start.tolist() == [2, 0] + assert group_ctx.plan_flush_count.tolist() == [4, -1] + assert group_ctx.active_request_indices.tolist() == [0, -1] + assert group_ctx.src_slots[:, 0].tolist() == [source_slot] * 2 + assert group_ctx.dst_slots[:, 0].tolist() == [destination_slot] * 2 for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): assert mixers[0]._replayssm_ring_start[source_slot].item() == 2 assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 4 @@ -732,9 +782,8 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): +def test_modelwide_replayssm_postprocess_materializes_in_place(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - block_tables[0][0, 1] = block_tables[0][0, 0] for mixers, source_slot in zip(groups, (1, 4)): mixers[0]._replayssm_ring_start[source_slot] = 2 mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 @@ -749,33 +798,27 @@ def test_modelwide_replayssm_copies_only_reassigned_cache_group(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.materialize_reassigned_slots( - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), - src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), + is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) + ctx.materialize() torch.accelerator.synchronize() - assert kernel.call_count == 1 - assert ctx.precopy_ring_start.tolist() == [2, 0] - assert ctx.precopy_flush_count.tolist() == [4, -1] - assert ctx.precopy_active_request_indices.tolist() == [0, -1] - assert ctx.precopy_src_slots[:, 0].tolist() == [ - NULL_BLOCK_ID, - NULL_BLOCK_ID, - 4, - 4, - ] - assert ctx.precopy_dst_slots[:, 0].tolist() == [ - NULL_BLOCK_ID, - NULL_BLOCK_ID, - 5, - 5, - ] - assert groups[0][0]._replayssm_ring_start[1].item() == 2 - assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 4 - assert groups[1][0]._replayssm_ring_start[4].item() == 2 - assert groups[1][0]._replayssm_prev_num_accepted[4].item() == 4 - assert groups[1][0]._replayssm_ring_start[5].item() == 0 - assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 + assert kernel.call_count == 2 + for group_ctx, slot in zip(ctx.groups, (1, 4)): + assert group_ctx.plan_ring_start.tolist() == [2, 0] + assert group_ctx.plan_flush_count.tolist() == [6, -1] + assert group_ctx.src_slots[:, 0].tolist() == [slot] * 2 + assert group_ctx.dst_slots[:, 0].tolist() == [slot] * 2 + for mixers, slot in zip(groups, (1, 4)): + assert mixers[0]._replayssm_ring_start[slot].item() == 0 + assert mixers[0]._replayssm_prev_num_accepted[slot].item() == 0 diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 513a28be14b9..3ee51fcf7416 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -4,6 +4,7 @@ from types import SimpleNamespace from unittest.mock import Mock +import numpy as np import pytest import torch @@ -51,6 +52,7 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: state._replayssm_live_cols_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") replayssm = Mock() + replayssm.materialize_prefixes = False ctx = Mock( is_initialized=True, replayssm=replayssm, @@ -78,6 +80,42 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: assert kwargs["live_cols"] is state._replayssm_live_cols_gpu +def test_flashinfer_replayssm_preprocess_runs_before_v2_forward() -> None: + state = object.__new__(MambaHybridModelState) + state._needs_prefix_state_migration = False + state._use_flashinfer_replayssm = True + state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool) + state._replayssm_live_cols_gpu = torch.zeros(4, dtype=torch.int32) + state._get_mamba_group_info = Mock(return_value=([0], Mock())) + replayssm = Mock() + replayssm.materialize_prefixes = False + ctx = Mock(replayssm=replayssm, block_size=1024) + state._ensure_mamba_postprocess_ctx = Mock(return_value=ctx) + input_batch = Mock( + num_reqs=2, + is_prefilling_np=np.array([True, False, False, False]), + idx_mapping=torch.tensor([1, 3], dtype=torch.int32), + query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32), + ) + num_computed = torch.tensor([0, 2, 0, 7], dtype=torch.int32) + + state.preprocess_state(input_batch, (), Mock(), num_computed) + + replayssm.preprocess.assert_called_once_with( + idx_mapping=input_batch.idx_mapping, + query_metadata=input_batch.query_start_loc, + query_metadata_is_cumulative=True, + num_computed_tokens=num_computed, + is_prefilling=state._is_prefilling_gpu, + src_cols=state._replayssm_live_cols_gpu, + dst_cols=state._replayssm_live_cols_gpu, + mamba_block_size=1024, + num_reqs=2, + ) + replayssm.materialize.assert_not_called() + ctx.run_fused_precopy.assert_not_called() + + def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: state = RecoverSSMState() metadata = Mock(spec=RecoverSSMMetadata) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index c725cd0389ec..adb86ccc10e2 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -170,18 +170,15 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): @pytest.mark.parametrize( - ("with_replayssm", "num_computed_tokens", "expected_order"), + "num_computed_tokens", [ - pytest.param(True, 4, ["materialize", "copy"], id="replayssm-boundary"), - pytest.param(True, 3, ["materialize", "copy"], id="replayssm-no-boundary"), - pytest.param(False, 4, ["copy"], id="generic"), + pytest.param(4, id="boundary"), + pytest.param(3, id="no-boundary"), ], ) -def test_preprocess_mamba_uses_modelwide_materializer_when_present( - with_replayssm: bool, +def test_preprocess_mamba_leaves_replayssm_for_staging( num_computed_tokens: int, - expected_order: list[str], -): +) -> None: spec = MagicMock(block_size=4, num_speculative_blocks=0) cache_config = MagicMock(enable_prefix_caching=True, use_replayssm=True) input_batch = MagicMock() @@ -199,11 +196,7 @@ def test_preprocess_mamba_uses_modelwide_materializer_when_present( align_ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(1, torch.int32, device) align_ctx.precopy_src_col_buf = _MockCpuGpuBuffer(1, torch.int32, device) align_ctx.precopy_token_bias_buf = _MockCpuGpuBuffer(1, torch.int32, device) - align_ctx.replayssm = MagicMock() if with_replayssm else None - if align_ctx.replayssm is not None: - align_ctx.replayssm.materialize_reassigned_slots.side_effect = ( - lambda **kwargs: order.append("materialize") - ) + align_ctx.replayssm = MagicMock() align_ctx.run_fused_precopy.side_effect = lambda **kwargs: order.append("copy") preprocess_mamba( @@ -219,10 +212,14 @@ def test_preprocess_mamba_uses_modelwide_materializer_when_present( align_ctx=align_ctx, ) - assert order == expected_order + assert order == ["copy"] + align_ctx.replayssm.preprocess.assert_not_called() + align_ctx.replayssm.materialize.assert_not_called() -def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): +def test_postprocess_mamba_align_commits_then_materializes_after_fused_copy( + monkeypatch, +): order: list[str] = [] ctx = MagicMock() ctx.is_initialized = True @@ -238,9 +235,9 @@ def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): ctx.materialize_token_counts = torch.tensor([2], dtype=torch.int32) ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") ctx.replayssm = MagicMock() - ctx.replayssm.postprocess.side_effect = lambda **kwargs: ( - order.append("materialize") - ) + ctx.replayssm.materialize_prefixes = True + ctx.replayssm.postprocess.side_effect = lambda **kwargs: order.append("postprocess") + ctx.replayssm.materialize.side_effect = lambda: order.append("materialize") block_table = MagicMock() block_table.get_device_tensor.return_value = torch.zeros((1, 4), dtype=torch.int32) input_batch = MagicMock() @@ -261,7 +258,7 @@ def test_postprocess_mamba_align_materializes_after_fused_copy(monkeypatch): run_prefix_state_migration=True, ) - assert order == ["copy", "materialize"] + assert order == ["copy", "postprocess", "materialize"] assert accepted_cpu.tolist() == [3] @@ -279,6 +276,7 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ctx.materialize_token_counts = torch.zeros(1, dtype=torch.int32) ctx.block_size = 1024 ctx.replayssm = MagicMock() + ctx.replayssm.materialize_prefixes = False input_batch = MagicMock() input_batch.block_table = [] accepted = torch.tensor([2], dtype=torch.int32) @@ -299,6 +297,7 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ctx.run_fused_postprocess.assert_not_called() assert ctx.replayssm.postprocess.call_count == 1 assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is accepted + ctx.replayssm.materialize.assert_not_called() assert accepted_cpu.tolist() == [2] @@ -868,6 +867,8 @@ def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock: ctx.num_computed_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.num_draft_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) ctx.is_prefilling_buf = _MockCpuGpuBuffer(max_num_reqs, torch.bool, device) + ctx.precopy_src_col_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device) + ctx.replayssm = None return ctx @@ -878,6 +879,10 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): device = torch.device("cpu") max_num_reqs = 8 ctx = _make_staging_ctx(max_num_reqs, device) + ctx.block_size = 4 + ctx.replayssm = MagicMock() + ctx.replayssm.materialize_prefixes = True + ctx.precopy_src_col_buf.gpu[:3] = torch.tensor([99, 199, 299], dtype=torch.int32) # Any negative int32 works as a sentinel: all staged values (state_idx, # scheduled/computed/draft token counts) are non-negative, so a negative @@ -950,6 +955,18 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ctx.is_prefilling_buf.gpu[:num_reqs], torch.tensor([True, False, True]), ) + ctx.replayssm.preprocess.assert_called_once_with( + idx_mapping=None, + query_metadata=ctx.num_scheduled_tokens_buf.gpu, + query_metadata_is_cumulative=False, + num_computed_tokens=ctx.num_computed_tokens_buf.gpu, + is_prefilling=ctx.is_prefilling_buf.gpu, + src_cols=ctx.precopy_src_col_buf.gpu, + dst_cols=ctx.mamba_state_idx_buf.gpu, + mamba_block_size=4, + num_reqs=num_reqs, + ) + ctx.replayssm.materialize.assert_called_once_with() def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 067daedfd244..d287f199e8b9 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -13,7 +13,6 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from functools import cache -from inspect import signature from typing import Any import torch @@ -28,280 +27,112 @@ logger = init_logger(__name__) -@triton.jit(do_not_specialize=["num_reqs", "num_materialize_reqs"]) -def _postprocess_replayssm_modelwide_kernel( - # Per-request step metadata. +@triton.jit(do_not_specialize=["num_reqs"]) +def _preprocess_replayssm_kernel( idx_mapping, query_metadata, num_computed_tokens, - num_accepted_tokens, is_prefilling, - live_cols, - materialize_src_cols, - materialize_dst_cols, - materialize_token_counts, - # Per-group address tables. - block_table_ptrs, - tracker_ring_start_ptrs, - tracker_num_committed_ptrs, - tracker_capacities, - group_layer_offsets, - # FlashInfer plan outputs. + src_cols, + dst_cols, + block_table, + tracker_start, + tracker_committed, src_slots, dst_slots, plan_ring_start, plan_flush_count, active_request_indices, - # Runtime sizes. block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, num_reqs, - num_materialize_reqs, - # Compile-time model constants. - MAX_LAYERS_PER_GROUP: tl.constexpr, MAMBA_BLOCK_SIZE: tl.constexpr, - LOGICAL_WINDOW: tl.constexpr, - RING_BUFFER_LEN: tl.constexpr, + NUM_LAYERS: tl.constexpr, PAD_SLOT_ID: tl.constexpr, QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, - NUM_COMPUTED_IS_POST_STEP: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, + MATERIALIZE_PREFIXES: tl.constexpr, MAX_NUM_REQS: tl.constexpr, ) -> None: - """Commit all ReplaySSM trackers and optionally plan materialization. - - Request vectors have shape ``[batch_capacity]``. Cumulative query metadata - has shape ``[batch_capacity + 1]``; non-cumulative metadata has shape - ``[batch_capacity]``. Group pointer/capacity tables have shape - ``[num_groups]``, group-layer offsets have shape ``[num_groups + 1]``, and - source/destination plans have shape ``[num_layers, batch_capacity]``. - ``active_request_indices`` has shape ``[batch_capacity]`` and contains the - compact physical-request prefix consumed by FlashInfer, followed by ``-1``. - - One CTA owns one ``(request row, cache group)`` tracker transition. Every - group updates its distinct physical-slot namespace. Group zero alone writes - the request-level ``ring_start``/``flush_count`` plan and one of its CTAs - compacts active request indices for the all-layer FlashInfer materializer. - Invalid and padded rows still write sentinels so fixed-capacity plans cannot - retain stale work. - - Metadata flags specialize runner input representation only; they do not - change recurrence semantics. Prefill has already produced canonical SSM - state, so this kernel only resets the affected ReplaySSM cursors and, when a - prefix snapshot is requested, emits ``flush_count=0`` for an exact copy. - """ + """Reset prefill state and prepare an optional writable-slot move.""" batch_idx = tl.program_id(0) - group_idx = tl.program_id(1) active = batch_idx < num_reqs req_idx = batch_idx if HAS_IDX_MAPPING: req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) valid_req = active & (req_idx >= 0) - if group_idx == 0: - # Always overwrite the request decision, including padded rows, so a - # fixed-capacity FlashInfer call never observes stale work. - tl.store(plan_ring_start + batch_idx, 0) - tl.store(plan_flush_count + batch_idx, -1) - - block_table = tl.load(block_table_ptrs + group_idx).to(tl.pointer_type(tl.int32)) - tracker_start = tl.load(tracker_ring_start_ptrs + group_idx).to( - tl.pointer_type(tl.int32) - ) - tracker_committed = tl.load(tracker_num_committed_ptrs + group_idx).to( - tl.pointer_type(tl.int32) - ) - tracker_capacity = tl.load(tracker_capacities + group_idx) - - live_col = tl.load(live_cols + req_idx, mask=valid_req, other=-1) - valid_live_col = valid_req & (live_col >= 0) - live_slot = tl.load( - block_table + batch_idx * block_table_stride_req + live_col, - mask=valid_live_col, - other=PAD_SLOT_ID, - ) - valid_live = ( - valid_live_col - & (live_slot != PAD_SLOT_ID) - & (live_slot >= 0) - & (live_slot < tracker_capacity) - ) - - src_col = tl.load(materialize_src_cols + batch_idx, mask=active, other=-1) - dst_col = tl.load(materialize_dst_cols + batch_idx, mask=active, other=-1) - wants_materialize = valid_req & (src_col >= 0) & (dst_col >= 0) - materialize_src_slot = tl.load( + src_col = tl.load(src_cols + req_idx, mask=valid_req, other=-1) + dst_col = tl.load(dst_cols + req_idx, mask=valid_req, other=-1) + changed = valid_req & (src_col >= 0) & (src_col != dst_col) + src_slot = tl.load( block_table + batch_idx * block_table_stride_req + src_col, - mask=wants_materialize, + mask=changed, other=PAD_SLOT_ID, ) - materialize_dst_slot = tl.load( + dst_slot = tl.load( block_table + batch_idx * block_table_stride_req + dst_col, - mask=wants_materialize, + mask=changed, other=PAD_SLOT_ID, ) - valid_materialize = ( - wants_materialize - & (materialize_src_slot != PAD_SLOT_ID) - & (materialize_dst_slot != PAD_SLOT_ID) - & (materialize_src_slot >= 0) - & (materialize_dst_slot >= 0) - & (materialize_src_slot < tracker_capacity) - & (materialize_dst_slot < tracker_capacity) - ) - - # Fill every flattened layer row for this group. Invalid rows still receive - # the pad sentinel; request-level flush_count=-1 suppresses native writes. - layer_begin = tl.load(group_layer_offsets + group_idx) - layer_end = tl.load(group_layer_offsets + group_idx + 1) - for layer_offset in tl.static_range(0, MAX_LAYERS_PER_GROUP): - layer_idx = layer_begin + layer_offset - layer_valid = layer_idx < layer_end - table_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store( - src_slots + table_offset, - tl.where(valid_materialize, materialize_src_slot, PAD_SLOT_ID), - mask=layer_valid, - ) + for layer_idx in tl.static_range(0, NUM_LAYERS): + slot_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store(src_slots + slot_offset, src_slot) + tl.store(dst_slots + slot_offset, dst_slot) + + # Always clear the fixed-capacity plan row before deciding whether this + # request has work. FlashInfer treats flush_count < 0 as a no-op. + tl.store(plan_ring_start + batch_idx, 0) + tl.store(plan_flush_count + batch_idx, -1) + if changed: + # BlockManager guarantees that distinct live columns map to allocated, + # distinct physical slots. Snapshot the old owner before clearing the + # destination tracker. + tl.store(plan_ring_start + batch_idx, tl.load(tracker_start + src_slot)) tl.store( - dst_slots + table_offset, - tl.where(valid_materialize, materialize_dst_slot, PAD_SLOT_ID), - mask=layer_valid, - ) - - prefilling = tl.load(is_prefilling + batch_idx, mask=active, other=1) - if QUERY_METADATA_IS_CUMULATIVE: - query_len = tl.load( - query_metadata + batch_idx + 1, mask=active, other=0 - ) - tl.load( - query_metadata + batch_idx, - mask=active, - other=0, + plan_flush_count + batch_idx, + tl.load(tracker_committed + src_slot), ) - else: - query_len = tl.load(query_metadata + batch_idx, mask=active, other=0) + tl.store(tracker_start + dst_slot, 0) + tl.store(tracker_committed + dst_slot, 0) + prefilling = tl.load(is_prefilling + batch_idx, mask=active, other=0) if valid_req & prefilling: + if QUERY_METADATA_IS_CUMULATIVE: + query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( + query_metadata + batch_idx + ) + else: + query_len = tl.load(query_metadata + batch_idx) computed = tl.load(num_computed_tokens + req_idx) - computed_before = tl.where( - NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed - ) - computed_after = computed_before + query_len - first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) + first_col = tl.maximum(computed // MAMBA_BLOCK_SIZE, 0) last_col = tl.maximum( - (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1, + (computed + query_len + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1, 0, ) - # All-mode prefill writes every boundary state in this interval. Reset - # every corresponding cursor so any later prefix hit copies an exact - # canonical state instead of replaying rows from the slot's old owner. for col in tl.range(first_col, last_col + 1): - prefill_slot = tl.load( - block_table + batch_idx * block_table_stride_req + col - ) - valid_prefill_slot = ( - (prefill_slot != PAD_SLOT_ID) - & (prefill_slot >= 0) - & (prefill_slot < tracker_capacity) - ) - tl.store(tracker_start + prefill_slot, 0, mask=valid_prefill_slot) - tl.store(tracker_committed + prefill_slot, 0, mask=valid_prefill_slot) + slot = tl.load(block_table + batch_idx * block_table_stride_req + col) + tl.store(tracker_start + slot, 0) + tl.store(tracker_committed + slot, 0) - if valid_live: - if prefilling: - if valid_materialize & (group_idx == 0): - # The prefill kernel already produced an exact canonical state; - # count zero asks FlashInfer to copy it byte-for-byte. - tl.store(plan_ring_start + batch_idx, 0) - tl.store(plan_flush_count + batch_idx, 0) - else: - old_start = tl.load(tracker_start + live_slot) - old_committed = tl.load(tracker_committed + live_slot) - accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) - checkpointed = old_committed + query_len > LOGICAL_WINDOW - next_start = tl.where( - checkpointed, - (old_start + old_committed) % RING_BUFFER_LEN, - old_start, - ) - next_committed = tl.where(checkpointed, accepted, old_committed + accepted) - - if valid_materialize & (group_idx == 0): - boundary_count = tl.load(materialize_token_counts + batch_idx) - flush_count = next_committed - (accepted - boundary_count) - tl.store(plan_ring_start + batch_idx, next_start) - tl.store(plan_flush_count + batch_idx, flush_count) - - tl.store(tracker_start + live_slot, next_start) - tl.store(tracker_committed + live_slot, next_committed) - - if valid_materialize: - # The immutable plan above preserves any in-place transition for the - # materializer; subsequent forwards see a canonical empty replay. - tl.store(tracker_start + materialize_dst_slot, 0) - tl.store(tracker_committed + materialize_dst_slot, 0) - - if (group_idx == 0) & (batch_idx == 0): - # FlashInfer consumes a compact active prefix and stops at the - # first -1. Build it from planner inputs here so sparse flushes do - # not require a second Triton launch or a host synchronization. + if MATERIALIZE_PREFIXES & (batch_idx == 0): active_count = 0 - for candidate_idx in tl.range(0, num_materialize_reqs): + for candidate_idx in tl.range(0, num_reqs): candidate_req_idx = candidate_idx if HAS_IDX_MAPPING: candidate_req_idx = tl.load(idx_mapping + candidate_idx) - candidate_valid_req = candidate_req_idx >= 0 - candidate_live_col = tl.load( - live_cols + candidate_req_idx, - mask=candidate_valid_req, + valid_candidate = candidate_req_idx >= 0 + src_col = tl.load( + src_cols + candidate_req_idx, + mask=valid_candidate, other=-1, ) - candidate_valid_live_col = candidate_valid_req & (candidate_live_col >= 0) - candidate_live_slot = tl.load( - block_table - + candidate_idx * block_table_stride_req - + candidate_live_col, - mask=candidate_valid_live_col, - other=PAD_SLOT_ID, - ) - candidate_valid_live = ( - candidate_valid_live_col - & (candidate_live_slot != PAD_SLOT_ID) - & (candidate_live_slot >= 0) - & (candidate_live_slot < tracker_capacity) - ) - candidate_src_col = tl.load(materialize_src_cols + candidate_idx) - candidate_dst_col = tl.load(materialize_dst_cols + candidate_idx) - candidate_wants_materialize = ( - candidate_valid_req - & (candidate_src_col >= 0) - & (candidate_dst_col >= 0) - ) - candidate_src_slot = tl.load( - block_table - + candidate_idx * block_table_stride_req - + candidate_src_col, - mask=candidate_wants_materialize, - other=PAD_SLOT_ID, - ) - candidate_dst_slot = tl.load( - block_table - + candidate_idx * block_table_stride_req - + candidate_dst_col, - mask=candidate_wants_materialize, - other=PAD_SLOT_ID, - ) - candidate_valid_materialize = ( - candidate_valid_live - & candidate_wants_materialize - & (candidate_src_slot != PAD_SLOT_ID) - & (candidate_dst_slot != PAD_SLOT_ID) - & (candidate_src_slot >= 0) - & (candidate_dst_slot >= 0) - & (candidate_src_slot < tracker_capacity) - & (candidate_dst_slot < tracker_capacity) + dst_col = tl.load( + dst_cols + candidate_req_idx, + mask=valid_candidate, + other=-1, ) - if candidate_valid_materialize: + if valid_candidate & (src_col >= 0) & (src_col != dst_col): tl.store(active_request_indices + active_count, candidate_idx) active_count += 1 if active_count < MAX_NUM_REQS: @@ -309,15 +140,18 @@ def _postprocess_replayssm_modelwide_kernel( @triton.jit(do_not_specialize=["num_reqs"]) -def _copy_reassigned_replayssm_slots_kernel( +def _postprocess_replayssm_kernel( idx_mapping, - src_cols, - dst_cols, - block_table_ptrs, - tracker_ring_start_ptrs, - tracker_num_committed_ptrs, - tracker_capacities, - group_layer_offsets, + query_metadata, + num_accepted_tokens, + is_prefilling, + live_cols, + materialize_src_cols, + materialize_dst_cols, + materialize_token_counts, + block_table, + tracker_start, + tracker_committed, src_slots, dst_slots, plan_ring_start, @@ -326,139 +160,100 @@ def _copy_reassigned_replayssm_slots_kernel( block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, num_reqs, - MAX_LAYERS_PER_GROUP: tl.constexpr, + LOGICAL_WINDOW: tl.constexpr, + RING_BUFFER_LEN: tl.constexpr, + NUM_LAYERS: tl.constexpr, PAD_SLOT_ID: tl.constexpr, + QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, + MATERIALIZE_PREFIXES: tl.constexpr, MAX_NUM_REQS: tl.constexpr, ) -> None: - """Plan an exact copy when align reassigns a request's writable slot.""" + """Commit a completed step and prepare an optional prefix snapshot.""" batch_idx = tl.program_id(0) - group_idx = tl.program_id(1) active = batch_idx < num_reqs req_idx = batch_idx if HAS_IDX_MAPPING: req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) valid_req = active & (req_idx >= 0) - if group_idx == 0: - tl.store(plan_ring_start + batch_idx, 0) - tl.store(plan_flush_count + batch_idx, -1) - - block_table = tl.load(block_table_ptrs + group_idx).to(tl.pointer_type(tl.int32)) - tracker_start = tl.load(tracker_ring_start_ptrs + group_idx).to( - tl.pointer_type(tl.int32) - ) - tracker_committed = tl.load(tracker_num_committed_ptrs + group_idx).to( - tl.pointer_type(tl.int32) - ) - tracker_capacity = tl.load(tracker_capacities + group_idx) - src_col = tl.load(src_cols + req_idx, mask=valid_req, other=-1) - dst_col = tl.load(dst_cols + req_idx, mask=valid_req, other=-1) - wants_copy = valid_req & (src_col >= 0) & (dst_col >= 0) & (src_col != dst_col) + src_col = tl.load(materialize_src_cols + batch_idx, mask=valid_req, other=-1) + materialize = valid_req & (src_col >= 0) + dst_col = tl.load(materialize_dst_cols + batch_idx, mask=materialize, other=-1) src_slot = tl.load( block_table + batch_idx * block_table_stride_req + src_col, - mask=wants_copy, + mask=materialize, other=PAD_SLOT_ID, ) dst_slot = tl.load( block_table + batch_idx * block_table_stride_req + dst_col, - mask=wants_copy, + mask=materialize, other=PAD_SLOT_ID, ) - valid_mapping = ( - wants_copy - & (src_slot != PAD_SLOT_ID) - & (dst_slot != PAD_SLOT_ID) - & (src_slot >= 0) - & (dst_slot >= 0) - & (src_slot < tracker_capacity) - & (dst_slot < tracker_capacity) - ) - needs_copy = valid_mapping & (src_slot != dst_slot) - if valid_mapping & (group_idx == 0): - # Snapshot the source cursor before resetting the distinct destination. - # The materializer uses it to copy the exact live state, including any - # committed replay rows that have not reached a prefix boundary. The - # logical migration activates the shared plan even when group 0 aliases; - # every group independently suppresses unchanged physical slots below. - tl.store(plan_ring_start + batch_idx, tl.load(tracker_start + src_slot)) - tl.store( - plan_flush_count + batch_idx, - tl.load(tracker_committed + src_slot), - ) + for layer_idx in tl.static_range(0, NUM_LAYERS): + slot_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store(src_slots + slot_offset, src_slot) + tl.store(dst_slots + slot_offset, dst_slot) + + tl.store(plan_ring_start + batch_idx, 0) + tl.store(plan_flush_count + batch_idx, -1) + + if valid_req: + # The live column and any materialization destination are allocated by + # BlockManager. A bad mapping is an upstream lifecycle bug, not a + # recoverable per-request condition for this kernel to hide. + live_col = tl.load(live_cols + req_idx) + live_slot = tl.load(block_table + batch_idx * block_table_stride_req + live_col) + prefilling = tl.load(is_prefilling + batch_idx) + if prefilling: + if materialize: + # Prefill produced canonical state, so publish an exact copy. + tl.store(plan_flush_count + batch_idx, 0) + else: + if QUERY_METADATA_IS_CUMULATIVE: + query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( + query_metadata + batch_idx + ) + else: + query_len = tl.load(query_metadata + batch_idx) + accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) + old_start = tl.load(tracker_start + live_slot) + old_committed = tl.load(tracker_committed + live_slot) + checkpointed = old_committed + query_len > LOGICAL_WINDOW + next_start = tl.where( + checkpointed, + (old_start + old_committed) % RING_BUFFER_LEN, + old_start, + ) + next_committed = tl.where(checkpointed, accepted, old_committed + accepted) + tl.store(tracker_start + live_slot, next_start) + tl.store(tracker_committed + live_slot, next_committed) - layer_begin = tl.load(group_layer_offsets + group_idx) - layer_end = tl.load(group_layer_offsets + group_idx + 1) - for layer_offset in tl.static_range(0, MAX_LAYERS_PER_GROUP): - layer_idx = layer_begin + layer_offset - layer_valid = layer_idx < layer_end - table_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store( - src_slots + table_offset, - tl.where(needs_copy, src_slot, PAD_SLOT_ID), - mask=layer_valid, - ) - tl.store( - dst_slots + table_offset, - tl.where(needs_copy, dst_slot, PAD_SLOT_ID), - mask=layer_valid, - ) + if materialize: + boundary_count = tl.load(materialize_token_counts + batch_idx) + tl.store(plan_ring_start + batch_idx, next_start) + tl.store( + plan_flush_count + batch_idx, + next_committed - (accepted - boundary_count), + ) - if needs_copy: - # The reassigned destination must not inherit its prior owner's cursor. - tl.store(tracker_start + dst_slot, 0) - tl.store(tracker_committed + dst_slot, 0) + if materialize: + # src_slot == dst_slot is a valid in-place checkpoint. + tl.store(tracker_start + dst_slot, 0) + tl.store(tracker_committed + dst_slot, 0) - if (group_idx == 0) & (batch_idx == 0): - # The logical reassignment can require other cache groups to copy - # even when group zero's physical slots alias. Compact every valid - # group-zero plan into the ordered map required by FlashInfer. + if MATERIALIZE_PREFIXES & (batch_idx == 0): active_count = 0 for candidate_idx in tl.range(0, num_reqs): candidate_req_idx = candidate_idx if HAS_IDX_MAPPING: candidate_req_idx = tl.load(idx_mapping + candidate_idx) - candidate_valid_req = candidate_req_idx >= 0 - candidate_src_col = tl.load( - src_cols + candidate_req_idx, - mask=candidate_valid_req, + src_col = tl.load( + materialize_src_cols + candidate_idx, + mask=candidate_req_idx >= 0, other=-1, ) - candidate_dst_col = tl.load( - dst_cols + candidate_req_idx, - mask=candidate_valid_req, - other=-1, - ) - candidate_wants_copy = ( - candidate_valid_req - & (candidate_src_col >= 0) - & (candidate_dst_col >= 0) - & (candidate_src_col != candidate_dst_col) - ) - candidate_src_slot = tl.load( - block_table - + candidate_idx * block_table_stride_req - + candidate_src_col, - mask=candidate_wants_copy, - other=PAD_SLOT_ID, - ) - candidate_dst_slot = tl.load( - block_table - + candidate_idx * block_table_stride_req - + candidate_dst_col, - mask=candidate_wants_copy, - other=PAD_SLOT_ID, - ) - candidate_valid_mapping = ( - candidate_wants_copy - & (candidate_src_slot != PAD_SLOT_ID) - & (candidate_dst_slot != PAD_SLOT_ID) - & (candidate_src_slot >= 0) - & (candidate_dst_slot >= 0) - & (candidate_src_slot < tracker_capacity) - & (candidate_dst_slot < tracker_capacity) - ) - if candidate_valid_mapping: + if (candidate_req_idx >= 0) & (src_col >= 0): tl.store(active_request_indices + active_count, candidate_idx) active_count += 1 if active_count < MAX_NUM_REQS: @@ -485,45 +280,20 @@ def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: @dataclass -class ReplaySSMModelContext: - """One model-wide owner for ReplaySSM trackers and materialization plans. - - Layers in a cache group share tracker tensors; the first layer is only the - representative used to capture those shared addresses. Materialization - plans remain layer-shaped because FlashInfer receives one row per mixer. - """ +class _ReplaySSMGroupContext: + """ReplaySSM state sharing one physical cache-slot namespace.""" mixers: list[Any] - group_layer_offsets: torch.Tensor - block_table_ptrs: torch.Tensor - tracker_ring_start_ptrs: torch.Tensor - tracker_num_committed_ptrs: torch.Tensor - tracker_capacities: torch.Tensor - state_ptrs: torch.Tensor - state_slot_strides: torch.Tensor - x_cache_ptrs: torch.Tensor - x_cache_slot_strides: torch.Tensor - b_cache_ptrs: torch.Tensor - b_cache_slot_strides: torch.Tensor - dt_cache_ptrs: torch.Tensor - dt_cache_slot_strides: torch.Tensor - a_ptrs: torch.Tensor - scale_ptrs: torch.Tensor - scale_slot_strides: torch.Tensor + block_table: torch.Tensor + ring_start: torch.Tensor + num_committed: torch.Tensor + materialize_tables: tuple[torch.Tensor, ...] src_slots: torch.Tensor dst_slots: torch.Tensor plan_ring_start: torch.Tensor plan_flush_count: torch.Tensor active_request_indices: torch.Tensor - precopy_src_slots: torch.Tensor - precopy_dst_slots: torch.Tensor - precopy_ring_start: torch.Tensor - precopy_flush_count: torch.Tensor - precopy_active_request_indices: torch.Tensor - block_table_stride_req: int max_num_reqs: int - num_groups: int - max_layers_per_group: int logical_window: int ring_buffer_len: int materialize_prefixes: bool @@ -531,51 +301,18 @@ class ReplaySSMModelContext: @classmethod def create( cls, - kv_cache_config: KVCacheConfig, - mamba_group_ids: Sequence[int], - forward_context: Mapping[str, Any], - block_tables: Sequence[torch.Tensor], + mixers: list[Any], + block_table: torch.Tensor, + cache_mode: str, max_num_reqs: int, - ) -> "ReplaySSMModelContext | None": - grouped = _flashinfer_replayssm_mixers_by_group( - kv_cache_config, mamba_group_ids, forward_context - ) - if not grouped: - return None - if len(block_tables) != len(mamba_group_ids): - raise ValueError( - f"expected {len(mamba_group_ids)} Mamba block tables, " - f"got {len(block_tables)}" - ) - block_table_by_gid = dict(zip(mamba_group_ids, block_tables)) - replayssm_block_tables = [block_table_by_gid[gid] for gid, _ in grouped] - for block_table in replayssm_block_tables: - if ( - block_table.ndim != 2 - or block_table.dtype != torch.int32 - or not block_table.is_cuda - or block_table.numel() == 0 - ): - raise ValueError( - "model-wide ReplaySSM requires non-empty 2D CUDA int32 block tables" - ) - cache_modes = set() - for gid, _ in grouped: - spec = kv_cache_config.kv_cache_groups[gid].kv_cache_spec - if not isinstance(spec, MambaSpec): - raise TypeError( - "FlashInfer ReplaySSM layers require a Mamba cache spec; " - f"got {type(spec).__name__}" - ) - cache_modes.add(spec.mamba_cache_mode) - if len(cache_modes) != 1: - raise ValueError( - "model-wide ReplaySSM requires one Mamba cache mode; " - f"got {sorted(cache_modes)}" - ) - materialize_prefixes = next(iter(cache_modes)) in ("align", "all") - - mixers = [mixer for _, group_mixers in grouped for mixer in group_mixers] + ) -> "_ReplaySSMGroupContext": + if ( + block_table.ndim != 2 + or block_table.dtype != torch.int32 + or not block_table.is_cuda + or block_table.numel() == 0 + ): + raise ValueError("ReplaySSM requires a non-empty 2D CUDA int32 block table") _validate_replayssm_cache(mixers) first = mixers[0] first_ssm = first.kv_cache[1] @@ -585,62 +322,40 @@ def create( current = _replayssm_specialization_key(mixer) if current != compatibility: raise ValueError( - "A single model-wide FlashInfer ReplaySSM materialization " - "launch requires identical layer specialization; got " + "Layers in one ReplaySSM cache group require identical " + "materialization specialization; got " f"{compatibility} and {current}" ) + if ( + mixer._replayssm_ring_start.data_ptr() + != first._replayssm_ring_start.data_ptr() + or mixer._replayssm_prev_num_accepted.data_ptr() + != first._replayssm_prev_num_accepted.data_ptr() + ): + raise ValueError( + "Layers in one ReplaySSM cache group must share ring trackers" + ) device = first_ssm.device - group_offsets = [0] - for _, group_mixers in grouped: - group_offsets.append(group_offsets[-1] + len(group_mixers)) - max_layers_per_group = max( - group_offsets[i + 1] - group_offsets[i] - for i in range(len(group_offsets) - 1) - ) - tracker_representatives = [group_mixers[0] for _, group_mixers in grouped] - strides = {int(block_table.stride(0)) for block_table in replayssm_block_tables} - if len(strides) != 1: - raise ValueError( - "model-wide ReplaySSM requires one block-table row stride; " - f"got {sorted(strides)}" - ) - zero_table = torch.zeros(len(mixers), dtype=torch.int64, device=device) return cls( mixers=mixers, - group_layer_offsets=torch.tensor( - group_offsets, dtype=torch.int32, device=device + block_table=block_table, + ring_start=first._replayssm_ring_start, + num_committed=first._replayssm_prev_num_accepted, + materialize_tables=( + _cuda_i64_ptrs([m.kv_cache[1] for m in mixers]), + _cuda_i64_slot_strides([m.kv_cache[1] for m in mixers]), + _cuda_i64_ptrs([m.kv_cache[2] for m in mixers]), + _cuda_i64_slot_strides([m.kv_cache[2] for m in mixers]), + _cuda_i64_ptrs([m.kv_cache[4] for m in mixers]), + _cuda_i64_slot_strides([m.kv_cache[4] for m in mixers]), + _cuda_i64_ptrs([m.kv_cache[3] for m in mixers]), + _cuda_i64_slot_strides([m.kv_cache[3] for m in mixers]), + _cuda_i64_ptrs([m.A for m in mixers]), + zero_table, + zero_table.clone(), ), - block_table_ptrs=_cuda_i64_ptrs(replayssm_block_tables), - tracker_ring_start_ptrs=_cuda_i64_ptrs( - [m._replayssm_ring_start for m in tracker_representatives] - ), - tracker_num_committed_ptrs=_cuda_i64_ptrs( - [m._replayssm_prev_num_accepted for m in tracker_representatives] - ), - tracker_capacities=torch.tensor( - [m._replayssm_ring_start.numel() for m in tracker_representatives], - dtype=torch.int32, - device=device, - ), - state_ptrs=_cuda_i64_ptrs([m.kv_cache[1] for m in mixers]), - state_slot_strides=_cuda_i64_slot_strides([m.kv_cache[1] for m in mixers]), - x_cache_ptrs=_cuda_i64_ptrs([m.kv_cache[2] for m in mixers]), - x_cache_slot_strides=_cuda_i64_slot_strides( - [m.kv_cache[2] for m in mixers] - ), - b_cache_ptrs=_cuda_i64_ptrs([m.kv_cache[4] for m in mixers]), - b_cache_slot_strides=_cuda_i64_slot_strides( - [m.kv_cache[4] for m in mixers] - ), - dt_cache_ptrs=_cuda_i64_ptrs([m.kv_cache[3] for m in mixers]), - dt_cache_slot_strides=_cuda_i64_slot_strides( - [m.kv_cache[3] for m in mixers] - ), - a_ptrs=_cuda_i64_ptrs([m.A for m in mixers]), - scale_ptrs=zero_table, - scale_slot_strides=zero_table.clone(), src_slots=torch.full( (len(mixers), max_num_reqs), NULL_BLOCK_ID, @@ -660,179 +375,122 @@ def create( active_request_indices=torch.full( (max_num_reqs,), -1, dtype=torch.int32, device=device ), - precopy_src_slots=torch.full( - (len(mixers), max_num_reqs), - NULL_BLOCK_ID, - dtype=torch.int32, - device=device, - ), - precopy_dst_slots=torch.full( - (len(mixers), max_num_reqs), - NULL_BLOCK_ID, - dtype=torch.int32, - device=device, - ), - precopy_ring_start=torch.zeros( - max_num_reqs, dtype=torch.int32, device=device - ), - precopy_flush_count=torch.full( - (max_num_reqs,), -1, dtype=torch.int32, device=device - ), - precopy_active_request_indices=torch.full( - (max_num_reqs,), -1, dtype=torch.int32, device=device - ), - block_table_stride_req=next(iter(strides)), max_num_reqs=max_num_reqs, - num_groups=len(grouped), - max_layers_per_group=max_layers_per_group, logical_window=int(first.replayssm_buffer_len), ring_buffer_len=first_x.size(2), - materialize_prefixes=materialize_prefixes, + materialize_prefixes=cache_mode in ("align", "all"), ) - def postprocess( + def preprocess( self, *, idx_mapping: torch.Tensor | None, query_metadata: torch.Tensor, query_metadata_is_cumulative: bool, num_computed_tokens: torch.Tensor, - num_computed_is_post_step: bool, - num_accepted_tokens: torch.Tensor, is_prefilling: torch.Tensor, - live_cols: torch.Tensor, - materialize_src_cols: torch.Tensor, - materialize_dst_cols: torch.Tensor, - materialize_token_counts: torch.Tensor, + src_cols: torch.Tensor, + dst_cols: torch.Tensor, mamba_block_size: int, num_reqs: int, ) -> None: - """Commit trackers, then publish a prefix snapshot when configured.""" + """Reset prefill state and prepare an optional writable-slot move.""" if num_reqs == 0: return - _postprocess_replayssm_modelwide_kernel[(self.max_num_reqs, self.num_groups)]( + _preprocess_replayssm_kernel[(self.max_num_reqs,)]( idx_mapping, query_metadata, num_computed_tokens, - num_accepted_tokens, is_prefilling, - live_cols, - materialize_src_cols, - materialize_dst_cols, - materialize_token_counts, - self.block_table_ptrs, - self.tracker_ring_start_ptrs, - self.tracker_num_committed_ptrs, - self.tracker_capacities, - self.group_layer_offsets, + src_cols, + dst_cols, + self.block_table, + self.ring_start, + self.num_committed, self.src_slots, self.dst_slots, self.plan_ring_start, self.plan_flush_count, self.active_request_indices, - self.block_table_stride_req, + self.block_table.stride(0), self.src_slots.stride(0), num_reqs, - num_reqs if self.materialize_prefixes else 0, - MAX_LAYERS_PER_GROUP=self.max_layers_per_group, MAMBA_BLOCK_SIZE=mamba_block_size, - LOGICAL_WINDOW=self.logical_window, - RING_BUFFER_LEN=self.ring_buffer_len, + NUM_LAYERS=len(self.mixers), PAD_SLOT_ID=NULL_BLOCK_ID, QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, - NUM_COMPUTED_IS_POST_STEP=num_computed_is_post_step, HAS_IDX_MAPPING=idx_mapping is not None, + MATERIALIZE_PREFIXES=self.materialize_prefixes, MAX_NUM_REQS=self.max_num_reqs, ) - if self.materialize_prefixes: - self._materialize_planned( - self.src_slots, - self.dst_slots, - self.plan_ring_start, - self.plan_flush_count, - self.active_request_indices, - ) - - def materialize_reassigned_slots( + def postprocess( self, *, idx_mapping: torch.Tensor | None, - src_cols: torch.Tensor, - dst_cols: torch.Tensor, + query_metadata: torch.Tensor, + query_metadata_is_cumulative: bool, + num_accepted_tokens: torch.Tensor, + is_prefilling: torch.Tensor, + live_cols: torch.Tensor, + materialize_src_cols: torch.Tensor, + materialize_dst_cols: torch.Tensor, + materialize_token_counts: torch.Tensor, num_reqs: int, ) -> None: - """Copy exact live state when align assigns a new writable slot.""" + """Commit a completed step and prepare an optional prefix snapshot.""" if num_reqs == 0: return - if not self.materialize_prefixes: - raise RuntimeError( - "ReplaySSM writable-slot materialization requires align or all mode" - ) - _copy_reassigned_replayssm_slots_kernel[(self.max_num_reqs, self.num_groups)]( + _postprocess_replayssm_kernel[(self.max_num_reqs,)]( idx_mapping, - src_cols, - dst_cols, - self.block_table_ptrs, - self.tracker_ring_start_ptrs, - self.tracker_num_committed_ptrs, - self.tracker_capacities, - self.group_layer_offsets, - self.precopy_src_slots, - self.precopy_dst_slots, - self.precopy_ring_start, - self.precopy_flush_count, - self.precopy_active_request_indices, - self.block_table_stride_req, - self.precopy_src_slots.stride(0), + query_metadata, + num_accepted_tokens, + is_prefilling, + live_cols, + materialize_src_cols, + materialize_dst_cols, + materialize_token_counts, + self.block_table, + self.ring_start, + self.num_committed, + self.src_slots, + self.dst_slots, + self.plan_ring_start, + self.plan_flush_count, + self.active_request_indices, + self.block_table.stride(0), + self.src_slots.stride(0), num_reqs, - MAX_LAYERS_PER_GROUP=self.max_layers_per_group, + LOGICAL_WINDOW=self.logical_window, + RING_BUFFER_LEN=self.ring_buffer_len, + NUM_LAYERS=len(self.mixers), PAD_SLOT_ID=NULL_BLOCK_ID, + QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, HAS_IDX_MAPPING=idx_mapping is not None, + MATERIALIZE_PREFIXES=self.materialize_prefixes, MAX_NUM_REQS=self.max_num_reqs, ) - self._materialize_planned( - self.precopy_src_slots, - self.precopy_dst_slots, - self.precopy_ring_start, - self.precopy_flush_count, - self.precopy_active_request_indices, - ) - def _materialize_planned( - self, - src_slots: torch.Tensor, - dst_slots: torch.Tensor, - ring_start: torch.Tensor, - flush_count: torch.Tensor, - active_request_indices: torch.Tensor, - ) -> None: + def materialize(self) -> None: + """Execute the plan prepared by ``preprocess`` or ``postprocess``.""" + if not self.materialize_prefixes: + raise RuntimeError("ReplaySSM materialization requires align or all mode") first = self.mixers[0] mamba_config = first.mamba_config rand_seed = None philox_rounds = 0 if mamba_config.enable_stochastic_rounding: rand_seed = torch.randint( - 0, 2**32, (1,), device=src_slots.device, dtype=torch.int64 + 0, 2**32, (1,), device=self.src_slots.device, dtype=torch.int64 ) philox_rounds = mamba_config.stochastic_rounding_philox_rounds or 10 _load_replayssm_materialize()( - self.state_ptrs, - self.state_slot_strides, - self.x_cache_ptrs, - self.x_cache_slot_strides, - self.b_cache_ptrs, - self.b_cache_slot_strides, - self.dt_cache_ptrs, - self.dt_cache_slot_strides, - self.a_ptrs, - self.scale_ptrs, - self.scale_slot_strides, - src_slots, - dst_slots, - ring_start, - flush_count, - active_request_indices, + *self.materialize_tables, + self.src_slots, + self.dst_slots, + self.plan_ring_start, + self.plan_flush_count, + self.active_request_indices, state_dtype=first.kv_cache[1].dtype, input_dtype=first.kv_cache[2].dtype, matrixA_dtype=first.A.dtype, @@ -847,6 +505,74 @@ def _materialize_planned( ) +@dataclass +class ReplaySSMModelContext: + """ReplaySSM lifecycle split by physical cache-slot namespace.""" + + groups: list[_ReplaySSMGroupContext] + materialize_prefixes: bool + + @classmethod + def create( + cls, + kv_cache_config: KVCacheConfig, + mamba_group_ids: Sequence[int], + forward_context: Mapping[str, Any], + block_tables: Sequence[torch.Tensor], + max_num_reqs: int, + ) -> "ReplaySSMModelContext | None": + grouped = _flashinfer_replayssm_mixers_by_group( + kv_cache_config, mamba_group_ids, forward_context + ) + if not grouped: + return None + if len(block_tables) != len(mamba_group_ids): + raise ValueError( + f"expected {len(mamba_group_ids)} Mamba block tables, " + f"got {len(block_tables)}" + ) + + block_table_by_gid = dict(zip(mamba_group_ids, block_tables)) + modes = set() + group_args = [] + for gid, mixers in grouped: + spec = kv_cache_config.kv_cache_groups[gid].kv_cache_spec + if not isinstance(spec, MambaSpec): + raise TypeError( + "FlashInfer ReplaySSM layers require a Mamba cache spec; " + f"got {type(spec).__name__}" + ) + modes.add(spec.mamba_cache_mode) + group_args.append((mixers, block_table_by_gid[gid], spec.mamba_cache_mode)) + if len(modes) != 1: + raise ValueError( + "model-wide ReplaySSM requires one Mamba cache mode; " + f"got {sorted(modes)}" + ) + + groups = [ + _ReplaySSMGroupContext.create(*args, max_num_reqs) for args in group_args + ] + return cls( + groups=groups, + materialize_prefixes=next(iter(modes)) in ("align", "all"), + ) + + def preprocess(self, **kwargs: Any) -> None: + for group in self.groups: + group.preprocess(**kwargs) + + def postprocess(self, **kwargs: Any) -> None: + for group in self.groups: + group.postprocess(**kwargs) + + def materialize(self) -> None: + if not self.materialize_prefixes: + raise RuntimeError("ReplaySSM materialization requires align or all mode") + for group in self.groups: + group.materialize() + + class MambaSSUBackend(ABC): """Abstract base class for Mamba SSU backends.""" @@ -1242,11 +968,6 @@ def _load_replayssm_materialize() -> Callable[..., None]: "FlashInfer ReplaySSM prefix caching requires " "flashinfer.mamba.replayssm_materialize" ) from e - if "active_request_indices" not in signature(replayssm_materialize).parameters: - raise ImportError( - "FlashInfer ReplaySSM prefix caching requires the ordered " - "active_request_indices materialization API" - ) return replayssm_materialize diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 40cb4768c4eb..37a8783eb5a3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -218,7 +218,7 @@ def preprocess_state( ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset is visible to the forward kernels. """ - if not self._needs_prefix_state_migration: + if not (self._needs_prefix_state_migration or self._use_flashinfer_replayssm): return num_reqs = input_batch.num_reqs if num_reqs == 0: @@ -228,33 +228,53 @@ def preprocess_state( kv_cache_config, mamba_group_ids, block_tables ) - # The state-advance + pre-copy kernels run every step; they fast-exit per - # request when src_col < 0 or src_col == dst_col, so no copy happens on - # steps that don't cross a block boundary. (Skipping the launch entirely - # would need a V1-style async-D2H of the actual num_computed, since - # num_computed_tokens_np is an optimistic mirror under async scheduling; - # the launch cost is ~0.3% of TPOT, so the GPU fast-exit suffices.) - block = 256 - grid = (triton.cdiv(num_reqs, block),) - preprocess_mamba_align_fused_kernel[grid]( - input_batch.idx_mapping, - self._mamba_state_idx_gpu, - num_computed_tokens, - input_batch.query_start_loc, - self.num_accepted_tokens_gpu, - self._mamba_src_col_gpu, - self._mamba_src_off_gpu, - num_reqs, - BLOCK_SIZE=block, - MAMBA_BLOCK_SIZE=mamba_spec.block_size, - ) - if ctx.replayssm is not None: - ctx.replayssm.materialize_reassigned_slots( + replayssm = ctx.replayssm + if replayssm is not None: + self._is_prefilling_gpu[:num_reqs].copy_( + torch.from_numpy(input_batch.is_prefilling_np[:num_reqs]) + ) + + if self._needs_prefix_state_migration: + # This decision kernel fast-exits per request when no block boundary + # is crossed. Avoiding the launch would require a CPU sync under + # async scheduling. + block = 256 + grid = (triton.cdiv(num_reqs, block),) + preprocess_mamba_align_fused_kernel[grid]( + input_batch.idx_mapping, + self._mamba_state_idx_gpu, + num_computed_tokens, + input_batch.query_start_loc, + self.num_accepted_tokens_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + num_reqs, + BLOCK_SIZE=block, + MAMBA_BLOCK_SIZE=mamba_spec.block_size, + ) + + if replayssm is not None: + if self._needs_prefix_state_migration: + src_cols = self._mamba_src_col_gpu + dst_cols = self._mamba_state_idx_gpu + else: + src_cols = dst_cols = self._replayssm_live_cols_gpu + replayssm.preprocess( idx_mapping=input_batch.idx_mapping, - src_cols=self._mamba_src_col_gpu, - dst_cols=self._mamba_state_idx_gpu, + query_metadata=input_batch.query_start_loc, + query_metadata_is_cumulative=True, + num_computed_tokens=num_computed_tokens, + is_prefilling=self._is_prefilling_gpu, + src_cols=src_cols, + dst_cols=dst_cols, + mamba_block_size=ctx.block_size, num_reqs=num_reqs, ) + if replayssm.materialize_prefixes: + replayssm.materialize() + + if not self._needs_prefix_state_migration: + return ctx.run_fused_precopy( num_reqs, self._mamba_state_idx_gpu, @@ -382,6 +402,7 @@ def postprocess_state( num_computed_tokens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, + defer_after_drafting: bool = False, ) -> None: # 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. @@ -430,53 +451,106 @@ def postprocess_state( idx_mapping, ) - if self._use_flashinfer_replayssm: - if num_computed_tokens is None: - raise RuntimeError( - "ReplaySSM postprocess requires the post-step computed-token " - "counts from the forward that produced this acceptance" - ) - if query_start_loc is None: - raise RuntimeError( - "ReplaySSM postprocess requires the query_start_loc from " - "the forward that produced this acceptance" - ) - ctx = self._mamba_ctx - if ctx is None or not ctx.is_initialized: - raise RuntimeError( - "ReplaySSM postprocess context was not initialized before forward" - ) - replayssm = ctx.replayssm - assert replayssm is not None - if is_prefilling is None: - is_prefilling = self._is_prefilling_gpu[:num_reqs] - replayssm.postprocess( - idx_mapping=idx_mapping, - query_metadata=query_start_loc, - query_metadata_is_cumulative=True, - num_computed_tokens=num_computed_tokens, - num_computed_is_post_step=True, - # Prefix migration can reset the live buffer to one for the - # next step; use its snapshot in that case. Mode none never - # runs the migration kernel, so the acceptance buffer is exact. - num_accepted_tokens=( - ctx.num_accepted_tokens_out - if self._needs_prefix_state_migration - else self.num_accepted_tokens_gpu - ), - is_prefilling=is_prefilling, - live_cols=( - self._mamba_state_idx_gpu - if self._needs_prefix_state_migration - else self._replayssm_live_cols_gpu - ), - materialize_src_cols=ctx.materialize_src_cols, - materialize_dst_cols=ctx.materialize_dst_cols, - materialize_token_counts=ctx.materialize_token_counts, - mamba_block_size=ctx.block_size, - num_reqs=num_reqs, + if not defer_after_drafting: + self._publish_flashinfer_replayssm( + idx_mapping, + num_computed_tokens, + query_start_loc, + is_prefilling, ) + def postprocess_state_after_drafting( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor | int, + num_computed_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, + ) -> None: + """Publish ReplaySSM trackers after every forward in this step. + + MTP drafting reuses the tracker view consumed by the target forward. + The newly accepted transition belongs to the next target step, so it + must not become visible until the current draft pass has completed. + """ + num_reqs = idx_mapping.shape[0] + if ( + num_reqs + and self._use_flashinfer_replayssm + and not self._needs_prefix_state_migration + and not isinstance(num_sampled, int) + ): + # The MTP drafter reuses this request-state buffer after target + # acceptance was first scattered. Restore it before publication. + _scatter_num_accepted_kernel[(num_reqs,)]( + idx_mapping, + num_sampled, + self.num_accepted_tokens_gpu, + ) + self._publish_flashinfer_replayssm( + idx_mapping, + num_computed_tokens, + query_start_loc, + is_prefilling, + ) + + def _publish_flashinfer_replayssm( + self, + idx_mapping: torch.Tensor, + num_computed_tokens: torch.Tensor | None, + query_start_loc: torch.Tensor | None, + is_prefilling: torch.Tensor | None, + ) -> None: + """Commit the accepted target transition to ReplaySSM trackers.""" + num_reqs = idx_mapping.shape[0] + if not num_reqs or not self._use_flashinfer_replayssm: + return + + if num_computed_tokens is None: + raise RuntimeError( + "ReplaySSM postprocess requires the post-step computed-token " + "counts from the forward that produced this acceptance" + ) + if query_start_loc is None: + raise RuntimeError( + "ReplaySSM postprocess requires the query_start_loc from " + "the forward that produced this acceptance" + ) + ctx = self._mamba_ctx + if ctx is None or not ctx.is_initialized: + raise RuntimeError( + "ReplaySSM postprocess context was not initialized before forward" + ) + replayssm = ctx.replayssm + assert replayssm is not None + if is_prefilling is None: + is_prefilling = self._is_prefilling_gpu[:num_reqs] + replayssm.postprocess( + idx_mapping=idx_mapping, + query_metadata=query_start_loc, + query_metadata_is_cumulative=True, + # Prefix migration can reset the live buffer to one for the + # next step; use its snapshot in that case. Mode none never + # runs the migration kernel, so the acceptance buffer is exact. + num_accepted_tokens=( + ctx.num_accepted_tokens_out + if self._needs_prefix_state_migration + else self.num_accepted_tokens_gpu + ), + is_prefilling=is_prefilling, + live_cols=( + self._mamba_state_idx_gpu + if self._needs_prefix_state_migration + else self._replayssm_live_cols_gpu + ), + materialize_src_cols=ctx.materialize_src_cols, + materialize_dst_cols=ctx.materialize_dst_cols, + materialize_token_counts=ctx.materialize_token_counts, + num_reqs=num_reqs, + ) + if replayssm.materialize_prefixes: + replayssm.materialize() + @triton.jit def _scatter_num_accepted_kernel( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 445574bb1d4e..ae6990d8722c 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4470,8 +4470,21 @@ def execute_model( # ReplaySSM tracker postprocess will run. Mode none does not run # prefix preprocessing and always uses logical live column zero. if mamba_bufs is not None and mamba_bufs.postprocess_align is not None: + mamba_ctx = mamba_bufs.postprocess_align + if not mamba_ctx.is_initialized: + mamba_ctx.initialize_from_forward_context( + self.kv_cache_config, + self.compilation_config.static_forward_context, + self._get_mamba_state_copy_funcs(), + [ + self.input_batch.block_table[gid].get_device_tensor( + num_reqs + ) + for gid in mamba_ctx.mamba_group_ids + ], + ) mamba_utils.stage_postprocess_inputs_to_gpu( - mamba_bufs.postprocess_align, + mamba_ctx, scheduler_output, self.input_batch.req_ids, num_reqs, diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 6d2aff0a6552..aaf87cf45b68 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -637,8 +637,8 @@ def precopy_mamba_align_fused_kernel( if src_col < 0 or src_col == dst_col: return if tl.load(state_skip_precopy_ptr + state_idx): - # FlashInfer ReplaySSM materializes this temporal destination before the - # generic pre-copy launch. Copy only conv/other model-owned states here. + # FlashInfer ReplaySSM owns this temporal state. Copy only conv/other + # model-owned states here. return token_bias = tl.load(token_bias_ptr + req_idx) @@ -1616,13 +1616,6 @@ def preprocess_mamba( fused.state_idx.copy_to_gpu(num_reqs) fused.src_col.copy_to_gpu(num_reqs) fused.token_bias.copy_to_gpu(num_reqs) - if fused.ctx.replayssm is not None: - fused.ctx.replayssm.materialize_reassigned_slots( - idx_mapping=None, - src_cols=fused.src_col.gpu, - dst_cols=fused.state_idx.gpu, - num_reqs=num_reqs, - ) fused.ctx.run_fused_precopy( num_reqs=num_reqs, state_idx_gpu=fused.state_idx.gpu, @@ -1738,17 +1731,16 @@ def postprocess_mamba_align_gpu( idx_mapping=None, query_metadata=ctx.num_scheduled_tokens_buf.gpu, query_metadata_is_cumulative=False, - num_computed_tokens=ctx.num_computed_tokens_buf.gpu, - num_computed_is_post_step=False, num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, live_cols=ctx.mamba_state_idx_buf.gpu, materialize_src_cols=ctx.materialize_src_cols, materialize_dst_cols=ctx.materialize_dst_cols, materialize_token_counts=ctx.materialize_token_counts, - mamba_block_size=ctx.block_size, num_reqs=num_reqs, ) + if ctx.replayssm.materialize_prefixes: + ctx.replayssm.materialize() # ``num_accepted_tokens_out`` is pre-initialized from # ``num_accepted_tokens_gpu``; the kernel only overwrites entries to 1 @@ -1786,6 +1778,7 @@ def stage_postprocess_inputs_to_gpu( assert ctx.num_computed_tokens_buf is not None assert ctx.num_draft_tokens_buf is not None assert ctx.is_prefilling_buf is not None + assert ctx.precopy_src_col_buf is not None scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens num_scheduled = scheduler_output.num_scheduled_tokens @@ -1818,3 +1811,25 @@ def stage_postprocess_inputs_to_gpu( ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs) ctx.is_prefilling_buf.copy_to_gpu(num_reqs) + if ctx.replayssm is not None: + # Prefix modes use the source/destination columns prepared by + # preprocess_mamba. Mode none has no migration, so passing the live + # column as both endpoints makes that part of the kernel a no-op. + src_cols = ( + ctx.precopy_src_col_buf.gpu + if ctx.replayssm.materialize_prefixes + else ctx.mamba_state_idx_buf.gpu + ) + ctx.replayssm.preprocess( + idx_mapping=None, + query_metadata=ctx.num_scheduled_tokens_buf.gpu, + query_metadata_is_cumulative=False, + num_computed_tokens=ctx.num_computed_tokens_buf.gpu, + is_prefilling=ctx.is_prefilling_buf.gpu, + src_cols=src_cols, + dst_cols=ctx.mamba_state_idx_buf.gpu, + mamba_block_size=ctx.block_size, + num_reqs=num_reqs, + ) + if ctx.replayssm.materialize_prefixes: + ctx.replayssm.materialize() From 791a66597156331ba175000eeed022a752f6a4b9 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 09:36:24 +0200 Subject: [PATCH 14/53] refactor(mamba): minimize ReplaySSM postprocess planning Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 50 ++++-- .../worker/test_mamba_hybrid_model_state.py | 1 - tests/v1/worker/test_mamba_utils.py | 5 +- .../layers/mamba/ops/ssu_dispatch.py | 151 ++++++++---------- .../worker/gpu/model_states/mamba_hybrid.py | 1 - vllm/v1/worker/mamba_utils.py | 19 +-- 6 files changed, 112 insertions(+), 115 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 84222ee4fc87..a18c11e1d02f 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -462,7 +462,6 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None num_accepted_tokens=accepted, is_prefilling=is_prefilling, live_cols=live_cols, - materialize_src_cols=no_materialize, materialize_dst_cols=no_materialize, materialize_token_counts=materialize_counts, num_reqs=1, @@ -528,7 +527,6 @@ def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_reqs=1, @@ -564,6 +562,44 @@ def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatch): + _, config, forward_context, block_tables = _modelwide_replayssm_fixture() + block_tables[0][1] = torch.tensor([3, 2, 1], device="cuda") + block_tables[1][1] = torch.tensor([7, 6, 5], device="cuda") + materializer = Mock() + monkeypatch.setattr( + ssu_dispatch, "_load_replayssm_materialize", lambda: materializer + ) + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor([1, 4], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([False, False], device="cuda"), + live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([-1, 1], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.tensor([0, 1], dtype=torch.int32, device="cuda"), + num_reqs=2, + ) + ctx.materialize() + torch.accelerator.synchronize() + + assert materializer.call_count == 2 + for group_ctx in ctx.groups: + assert group_ctx.plan_flush_count.tolist() == [-1, 1] + assert group_ctx.active_request_indices.tolist() == [1, -1] + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() @@ -588,7 +624,6 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch num_accepted_tokens=torch.tensor([1, 3], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), num_reqs=1, @@ -626,9 +661,6 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): ) group_ctx.num_committed.fill_(7) before.append((group_ctx.ring_start.clone(), group_ctx.num_committed.clone())) - group_ctx.src_slots.fill_(3) - group_ctx.dst_slots.fill_(4) - group_ctx.plan_ring_start.fill_(5) group_ctx.plan_flush_count.fill_(6) ctx.postprocess( @@ -638,7 +670,6 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.ones(2, dtype=torch.int32, device="cuda"), materialize_token_counts=torch.ones(2, dtype=torch.int32, device="cuda"), num_reqs=1, @@ -648,9 +679,6 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): for group_ctx, (ring_start, num_committed) in zip(ctx.groups, before): assert torch.equal(group_ctx.ring_start, ring_start) assert torch.equal(group_ctx.num_committed, num_committed) - assert torch.all(group_ctx.src_slots == NULL_BLOCK_ID) - assert torch.all(group_ctx.dst_slots == NULL_BLOCK_ID) - assert group_ctx.plan_ring_start.tolist() == [0, 0] assert group_ctx.plan_flush_count.tolist() == [-1, -1] @@ -719,7 +747,6 @@ def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), num_reqs=1, @@ -805,7 +832,6 @@ def test_modelwide_replayssm_postprocess_materializes_in_place(monkeypatch): num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_reqs=1, diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 3ee51fcf7416..12c1c8d9a29f 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -56,7 +56,6 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: ctx = Mock( is_initialized=True, replayssm=replayssm, - materialize_src_cols=torch.full((4,), -1, dtype=torch.int32, device="cuda"), materialize_dst_cols=torch.full((4,), -1, dtype=torch.int32, device="cuda"), materialize_token_counts=torch.zeros(4, dtype=torch.int32, device="cuda"), block_size=1024, diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index adb86ccc10e2..2c2d4b5b1f56 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -230,7 +230,6 @@ def test_postprocess_mamba_align_commits_then_materializes_after_fused_copy( ctx.num_draft_tokens_buf = MagicMock() ctx.is_prefilling_buf = MagicMock() ctx.num_accepted_tokens_out = torch.tensor([3], dtype=torch.int32) - ctx.materialize_src_cols = torch.tensor([2], dtype=torch.int32) ctx.materialize_dst_cols = torch.tensor([1], dtype=torch.int32) ctx.materialize_token_counts = torch.tensor([2], dtype=torch.int32) ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") @@ -271,7 +270,6 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ctx.num_computed_tokens_buf = MagicMock(gpu=torch.tensor([20], dtype=torch.int32)) ctx.num_draft_tokens_buf = MagicMock(gpu=torch.tensor([3], dtype=torch.int32)) ctx.is_prefilling_buf = MagicMock(gpu=torch.tensor([False])) - ctx.materialize_src_cols = torch.full((1,), -1, dtype=torch.int32) ctx.materialize_dst_cols = torch.full((1,), -1, dtype=torch.int32) ctx.materialize_token_counts = torch.zeros(1, dtype=torch.int32) ctx.block_size = 1024 @@ -1325,7 +1323,7 @@ def test_no_copy_when_not_needed(self, device, test_config): # State should be unchanged torch.testing.assert_close(conv_state, conv_state_orig) torch.testing.assert_close(temporal_state, temporal_state_orig) - assert gpu_ctx.materialize_src_cols[0].item() == -1 + assert gpu_ctx.materialize_dst_cols[0].item() == -1 @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16]) def test_various_batch_sizes(self, device, test_config, num_reqs): @@ -1648,7 +1646,6 @@ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( device=device, ) - assert gpu_ctx.materialize_src_cols[0].item() == 1 assert gpu_ctx.materialize_dst_cols[0].item() == 1 assert gpu_ctx.materialize_token_counts[0].item() == 1 diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index d287f199e8b9..f90884770f3b 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -80,8 +80,8 @@ def _preprocess_replayssm_kernel( tl.store(src_slots + slot_offset, src_slot) tl.store(dst_slots + slot_offset, dst_slot) - # Always clear the fixed-capacity plan row before deciding whether this - # request has work. FlashInfer treats flush_count < 0 as a no-op. + # Clear the fixed-capacity plan row before deciding whether this request + # belongs in the compact materialization list. tl.store(plan_ring_start + batch_idx, 0) tl.store(plan_flush_count + batch_idx, -1) if changed: @@ -146,7 +146,6 @@ def _postprocess_replayssm_kernel( num_accepted_tokens, is_prefilling, live_cols, - materialize_src_cols, materialize_dst_cols, materialize_token_counts, block_table, @@ -163,7 +162,6 @@ def _postprocess_replayssm_kernel( LOGICAL_WINDOW: tl.constexpr, RING_BUFFER_LEN: tl.constexpr, NUM_LAYERS: tl.constexpr, - PAD_SLOT_ID: tl.constexpr, QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, MATERIALIZE_PREFIXES: tl.constexpr, @@ -171,94 +169,84 @@ def _postprocess_replayssm_kernel( ) -> None: """Commit a completed step and prepare an optional prefix snapshot.""" batch_idx = tl.program_id(0) - active = batch_idx < num_reqs - req_idx = batch_idx - if HAS_IDX_MAPPING: - req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) - valid_req = active & (req_idx >= 0) - - src_col = tl.load(materialize_src_cols + batch_idx, mask=valid_req, other=-1) - materialize = valid_req & (src_col >= 0) - dst_col = tl.load(materialize_dst_cols + batch_idx, mask=materialize, other=-1) - src_slot = tl.load( - block_table + batch_idx * block_table_stride_req + src_col, - mask=materialize, - other=PAD_SLOT_ID, - ) - dst_slot = tl.load( - block_table + batch_idx * block_table_stride_req + dst_col, - mask=materialize, - other=PAD_SLOT_ID, - ) - for layer_idx in tl.static_range(0, NUM_LAYERS): - slot_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store(src_slots + slot_offset, src_slot) - tl.store(dst_slots + slot_offset, dst_slot) - - tl.store(plan_ring_start + batch_idx, 0) - tl.store(plan_flush_count + batch_idx, -1) - - if valid_req: - # The live column and any materialization destination are allocated by - # BlockManager. A bad mapping is an upstream lifecycle bug, not a - # recoverable per-request condition for this kernel to hide. - live_col = tl.load(live_cols + req_idx) - live_slot = tl.load(block_table + batch_idx * block_table_stride_req + live_col) - prefilling = tl.load(is_prefilling + batch_idx) - if prefilling: - if materialize: - # Prefill produced canonical state, so publish an exact copy. - tl.store(plan_flush_count + batch_idx, 0) - else: - if QUERY_METADATA_IS_CUMULATIVE: - query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( - query_metadata + batch_idx - ) - else: - query_len = tl.load(query_metadata + batch_idx) - accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) - old_start = tl.load(tracker_start + live_slot) - old_committed = tl.load(tracker_committed + live_slot) - checkpointed = old_committed + query_len > LOGICAL_WINDOW - next_start = tl.where( - checkpointed, - (old_start + old_committed) % RING_BUFFER_LEN, - old_start, - ) - next_committed = tl.where(checkpointed, accepted, old_committed + accepted) - tl.store(tracker_start + live_slot, next_start) - tl.store(tracker_committed + live_slot, next_committed) - - if materialize: - boundary_count = tl.load(materialize_token_counts + batch_idx) - tl.store(plan_ring_start + batch_idx, next_start) - tl.store( - plan_flush_count + batch_idx, - next_committed - (accepted - boundary_count), - ) - - if materialize: - # src_slot == dst_slot is a valid in-place checkpoint. - tl.store(tracker_start + dst_slot, 0) - tl.store(tracker_committed + dst_slot, 0) + # FlashInfer requires active rows to be a compact prefix. The generic + # planner's destination is the sole materialization work sentinel. if MATERIALIZE_PREFIXES & (batch_idx == 0): active_count = 0 for candidate_idx in tl.range(0, num_reqs): candidate_req_idx = candidate_idx if HAS_IDX_MAPPING: candidate_req_idx = tl.load(idx_mapping + candidate_idx) - src_col = tl.load( - materialize_src_cols + candidate_idx, - mask=candidate_req_idx >= 0, - other=-1, - ) - if (candidate_req_idx >= 0) & (src_col >= 0): + dst_col = tl.load(materialize_dst_cols + candidate_idx) + if (candidate_req_idx >= 0) & (dst_col >= 0): tl.store(active_request_indices + active_count, candidate_idx) active_count += 1 if active_count < MAX_NUM_REQS: tl.store(active_request_indices + active_count, -1) + # Clear the fixed-capacity plan row before any per-request early exit. + tl.store(plan_flush_count + batch_idx, -1) + if batch_idx >= num_reqs: + return + + req_idx = batch_idx + if HAS_IDX_MAPPING: + req_idx = tl.load(idx_mapping + batch_idx) + if req_idx < 0: + return + + # BlockManager supplies allocated live/destination columns. The generic + # Mamba planner always materializes from this same live column. + live_col = tl.load(live_cols + req_idx) + live_slot = tl.load(block_table + batch_idx * block_table_stride_req + live_col) + dst_col = tl.load(materialize_dst_cols + batch_idx) + materialize = dst_col >= 0 + prefilling = tl.load(is_prefilling + batch_idx) + if prefilling: + if materialize: + # Prefill produced canonical state, so publish an exact copy. + tl.store(plan_flush_count + batch_idx, 0) + else: + if QUERY_METADATA_IS_CUMULATIVE: + query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( + query_metadata + batch_idx + ) + else: + query_len = tl.load(query_metadata + batch_idx) + accepted = tl.load(num_accepted_tokens + req_idx) + old_start = tl.load(tracker_start + live_slot) + old_committed = tl.load(tracker_committed + live_slot) + checkpointed = old_committed + query_len > LOGICAL_WINDOW + next_start = tl.where( + checkpointed, + (old_start + old_committed) % RING_BUFFER_LEN, + old_start, + ) + next_committed = tl.where(checkpointed, accepted, old_committed + accepted) + tl.store(tracker_start + live_slot, next_start) + tl.store(tracker_committed + live_slot, next_committed) + + if materialize: + boundary_count = tl.load(materialize_token_counts + batch_idx) + tl.store(plan_ring_start + batch_idx, next_start) + tl.store( + plan_flush_count + batch_idx, + boundary_count + tl.where(checkpointed, 0, old_committed), + ) + + if not materialize: + return + dst_slot = tl.load(block_table + batch_idx * block_table_stride_req + dst_col) + for layer_idx in tl.static_range(0, NUM_LAYERS): + slot_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store(src_slots + slot_offset, live_slot) + tl.store(dst_slots + slot_offset, dst_slot) + + # The published destination is canonical and therefore has no live replay. + tl.store(tracker_start + dst_slot, 0) + tl.store(tracker_committed + dst_slot, 0) + def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: ssm = mixer.kv_cache[1] @@ -433,7 +421,6 @@ def postprocess( num_accepted_tokens: torch.Tensor, is_prefilling: torch.Tensor, live_cols: torch.Tensor, - materialize_src_cols: torch.Tensor, materialize_dst_cols: torch.Tensor, materialize_token_counts: torch.Tensor, num_reqs: int, @@ -447,7 +434,6 @@ def postprocess( num_accepted_tokens, is_prefilling, live_cols, - materialize_src_cols, materialize_dst_cols, materialize_token_counts, self.block_table, @@ -464,7 +450,6 @@ def postprocess( LOGICAL_WINDOW=self.logical_window, RING_BUFFER_LEN=self.ring_buffer_len, NUM_LAYERS=len(self.mixers), - PAD_SLOT_ID=NULL_BLOCK_ID, QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, HAS_IDX_MAPPING=idx_mapping is not None, MATERIALIZE_PREFIXES=self.materialize_prefixes, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 37a8783eb5a3..6716a47cd4cd 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -543,7 +543,6 @@ def _publish_flashinfer_replayssm( if self._needs_prefix_state_migration else self._replayssm_live_cols_gpu ), - materialize_src_cols=ctx.materialize_src_cols, materialize_dst_cols=ctx.materialize_dst_cols, materialize_token_counts=ctx.materialize_token_counts, num_reqs=num_reqs, diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index aaf87cf45b68..0d627c79c1e5 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -401,9 +401,8 @@ def postprocess_mamba_fused_kernel( # Output: num_accepted_tokens update (for src==dst case) num_accepted_tokens_out_ptr, # Batch-ordered ReplaySSM materialization decision. The caller initializes - # src_col to -1, which remains the no-op sentinel when no boundary is hit. - materialize_src_col_ptr, - # Logical column holding the block-aligned prefix checkpoint. + # dst_col to -1, which remains the no-op sentinel when no boundary is hit. + # The source is always mamba_state_idx_ptr. materialize_dst_col_ptr, # Current-query rows through that boundary: accept_token_bias + 1. The # materializer adds older pending rows from the source slot's tracker. @@ -485,7 +484,6 @@ def postprocess_mamba_fused_kernel( dest_block_idx = aligned_new_computed // block_size - 1 if state_idx == 0 and tile_idx == 0: - tl.store(materialize_src_col_ptr + batch_idx, src_block_idx) tl.store(materialize_dst_col_ptr + batch_idx, dest_block_idx) tl.store( materialize_token_count_ptr + batch_idx, @@ -839,7 +837,6 @@ class MambaSpecDecodeGPUContext: # Output buffer for num_accepted_tokens updates num_accepted_tokens_out: torch.Tensor - materialize_src_cols: torch.Tensor materialize_dst_cols: torch.Tensor materialize_token_counts: torch.Tensor @@ -948,12 +945,9 @@ def create( num_accepted_tokens_out=torch.zeros( max_num_reqs, dtype=torch.int32, device=device ), - materialize_src_cols=torch.full( + materialize_dst_cols=torch.full( (max_num_reqs,), -1, dtype=torch.int32, device=device ), - materialize_dst_cols=torch.empty( - max_num_reqs, dtype=torch.int32, device=device - ), materialize_token_counts=torch.empty( max_num_reqs, dtype=torch.int32, device=device ), @@ -1238,7 +1232,7 @@ def run_fused_postprocess( self.num_accepted_tokens_out[:num_reqs].copy_( num_accepted_tokens_gpu[:num_reqs] ) - self.materialize_src_cols[:num_reqs].fill_(-1) + self.materialize_dst_cols[:num_reqs].fill_(-1) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) @@ -1261,7 +1255,6 @@ def run_fused_postprocess( self.state_dim_row_count, self.state_dim_row_stride, self.num_accepted_tokens_out, - self.materialize_src_cols, self.materialize_dst_cols, self.materialize_token_counts, None, # idx_mapping: V1 decision arrays are already in req order @@ -1343,7 +1336,7 @@ def run_fused_postprocess_align( # decision buffer rather than only [:num_reqs]. num_accepted_tokens_snapshot = self.num_accepted_tokens_out num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) - self.materialize_src_cols[:num_reqs].fill_(-1) + self.materialize_dst_cols[:num_reqs].fill_(-1) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) @@ -1365,7 +1358,6 @@ def run_fused_postprocess_align( self.state_dim_row_count, self.state_dim_row_stride, num_accepted_tokens_gpu, - self.materialize_src_cols, self.materialize_dst_cols, self.materialize_token_counts, idx_mapping, @@ -1734,7 +1726,6 @@ def postprocess_mamba_align_gpu( num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, live_cols=ctx.mamba_state_idx_buf.gpu, - materialize_src_cols=ctx.materialize_src_cols, materialize_dst_cols=ctx.materialize_dst_cols, materialize_token_counts=ctx.materialize_token_counts, num_reqs=num_reqs, From 30d904252a226dd6f48438f4d5a7a6e875e119a6 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 10:06:44 +0200 Subject: [PATCH 15/53] fix(mamba): publish ReplaySSM state after drafting Signed-off-by: Andrii Skliar --- .../worker/test_gpu_model_runner_v2_eplb.py | 87 ++++++++++++++++++- .../worker/test_mamba_hybrid_model_state.py | 22 ++++- vllm/model_executor/models/diffusion_gemma.py | 1 + vllm/v1/worker/gpu/model_runner.py | 10 +++ vllm/v1/worker/gpu/model_states/interface.py | 12 +++ 5 files changed, 127 insertions(+), 5 deletions(-) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index e683e47aeb91..e20e084c9dd0 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from types import SimpleNamespace from typing import Any @@ -205,6 +206,83 @@ def fake_receive(*args, **kwargs): assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] +def test_v2_sample_tokens_publishes_state_after_drafting(monkeypatch): + events: list[Any] = [] + runner = _make_runner() + input_batch = SimpleNamespace( + req_ids=["request"], + idx_mapping=torch.tensor([0], dtype=torch.int64), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + ) + hidden_states = torch.zeros(1, 1) + runner.execute_model_state = SimpleNamespace( + input_batch=input_batch, + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=hidden_states, + aux_hidden_states=None, + dp_sync=None, + finished_req_ids=set(), + ec_connector_output=None, + routed_experts=None, + ) + sampled_token_ids = torch.tensor([[1]]) + num_sampled = torch.tensor([1], dtype=torch.int32) + num_rejected = torch.tensor([0], dtype=torch.int32) + runner.sample = lambda *_: ( + SimpleNamespace(sampled_token_ids=sampled_token_ids), + num_sampled, + num_rejected, + ) + runner.pp_handler = None + runner.prompt_logprobs_worker = SimpleNamespace( + compute_prompt_logprobs=lambda *_: {} + ) + runner.model = SimpleNamespace(compute_logits=None) + runner.main_stream = None + runner.output_copy_stream = None + runner.check_ep_fault = None + runner.pcp_manager = None + runner._draft_workspace_lane = None + runner.adaptive_verification = None + runner.sampler = SimpleNamespace( + penalties_state=SimpleNamespace(output_bin_counts=None), + sampling_states=SimpleNamespace( + temperature=SimpleNamespace(gpu=None), + seeds=SimpleNamespace(gpu=None), + ), + ) + runner.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=None), + num_computed_tokens=SimpleNamespace(gpu=torch.zeros(1, dtype=torch.int32)), + prompt_len=SimpleNamespace(np=None), + last_sampled_tokens=None, + next_prefill_tokens=None, + total_len=SimpleNamespace(gpu=None), + draft_tokens=torch.zeros((1, 1), dtype=torch.int64), + ) + + def postprocess_state(*_, defer_after_drafting=False): + events.append(("postprocess", defer_after_drafting)) + + def propose(*_, **__): + events.append("draft") + return torch.tensor([[2]]) + + runner.speculator = SimpleNamespace(supports_mm_inputs=False, propose=propose) + runner.model_state = SimpleNamespace( + postprocess_state=postprocess_state, + postprocess_state_after_drafting=lambda *_: events.append("publish"), + ) + monkeypatch.setattr(mrv2, "AsyncOutput", lambda **_: object()) + monkeypatch.setattr(mrv2, "post_update", lambda *_: None) + monkeypatch.setattr(mrv2, "use_workspace_lane", lambda _: nullcontext()) + + mrv2.GPUModelRunner.sample_tokens(runner, None) + + assert events == [("postprocess", True), "draft", "publish"] + + def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( monkeypatch, ): @@ -242,9 +320,12 @@ def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( runner.model_state = SimpleNamespace( postprocess_state=lambda *args: postprocess_args.append(args) ) - runner.pp_handler = SimpleNamespace( - receive=lambda *_: events.append("receive") or False - ) + + def receive(*_: Any) -> bool: + events.append("receive") + return False + + runner.pp_handler = SimpleNamespace(receive=receive) runner.postprocess_num_computed_tokens = lambda *_: events.append( "postprocess_num_computed_tokens" ) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 12c1c8d9a29f..a52acd409565 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -42,7 +42,10 @@ def test_postprocess_state_scalar_with_int32_mapping( @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") -def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: +@pytest.mark.parametrize("defer_after_drafting", [False, True]) +def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration( + defer_after_drafting: bool, +) -> None: state = object.__new__(MambaHybridModelState) state._align_mode = False state._needs_prefix_state_migration = False @@ -62,19 +65,34 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: ) state._mamba_ctx = ctx idx_mapping = torch.tensor([2], dtype=torch.int32, device="cuda") + num_sampled = torch.tensor([2], dtype=torch.int32, device="cuda") num_computed = torch.tensor([0, 0, 20, 0], dtype=torch.int32, device="cuda") query_start_loc = torch.tensor([0, 4], dtype=torch.int32, device="cuda") state.postprocess_state( idx_mapping, - 2, + num_sampled, num_computed_tokens=num_computed, query_start_loc=query_start_loc, + defer_after_drafting=defer_after_drafting, ) ctx.run_fused_postprocess_align.assert_not_called() + if defer_after_drafting: + replayssm.postprocess.assert_not_called() + # The drafter reuses this request-state buffer. The post-draft commit + # must consume the sampler-owned batch tensor instead of the clobbered + # request-state value. + state.num_accepted_tokens_gpu[2] = 1 + state.postprocess_state_after_drafting( + idx_mapping, + num_sampled, + num_computed_tokens=num_computed, + query_start_loc=query_start_loc, + ) assert replayssm.postprocess.call_count == 1 kwargs = replayssm.postprocess.call_args.kwargs + assert state.num_accepted_tokens_gpu.tolist() == [1, 1, 2, 1] assert kwargs["num_accepted_tokens"] is state.num_accepted_tokens_gpu assert kwargs["live_cols"] is state._replayssm_live_cols_gpu diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 52872f310445..b12871b4356b 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -981,6 +981,7 @@ def postprocess_state( num_computed_tokens=None, query_start_loc=None, is_prefilling=None, + defer_after_drafting=False, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index eb5cef26bd22..802d4c775e7f 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1497,6 +1497,7 @@ def postprocess_sampled( num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, + defer_after_drafting: bool = False, ) -> None: # Update the number of computed tokens. if self.is_last_pp_rank: @@ -1523,6 +1524,7 @@ def postprocess_sampled( self.req_states.num_computed_tokens.gpu, query_start_loc, is_prefilling, + defer_after_drafting=defer_after_drafting, ) def _merge_ec_connector_no_forward( @@ -1962,6 +1964,7 @@ def sample_tokens( num_sampled, num_rejected, input_batch.query_start_loc, + defer_after_drafting=self.speculator is not None, ) if self.speculator is not None: @@ -1996,6 +1999,13 @@ def sample_tokens( self.speculator.draft_token_confidence_probs, input_batch ) + self.model_state.postprocess_state_after_drafting( + input_batch.idx_mapping, + num_sampled, + self.req_states.num_computed_tokens.gpu, + input_batch.query_start_loc, + ) + if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does # not have a speculator (i.e. self.speculator is None) diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 0c5002b65398..9638e05fb4ed 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -154,9 +154,21 @@ def postprocess_state( num_computed_tokens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, + defer_after_drafting: bool = False, ) -> None: return None + def postprocess_state_after_drafting( + self, + idx_mapping: torch.Tensor, + num_sampled: torch.Tensor, + num_computed_tokens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + is_prefilling: torch.Tensor | None = None, + ) -> None: + """Publish state that must remain hidden from the current draft pass.""" + return None + @abstractmethod def prepare_inputs_embeds( self, From 00a899821121b2656473034640da6d840174b34d Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 12:43:25 +0200 Subject: [PATCH 16/53] [Mamba] Reconcile ReplaySSM live-state prefix caching Port live ReplaySSM state migration onto the simplified prefix-cache lifecycle. Copy canonical state, replay rings, and shared cursors through scheduler block copies, reset only fresh slots, and publish canonical prefix snapshots after sampling. Co-authored-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 124 +++++--- .../core/test_single_type_kv_cache_manager.py | 34 ++ ...est_replayssm_mtp_compaction_diagnostic.py | 64 ++++ tests/v1/worker/test_mamba_utils.py | 153 +++++++-- tests/v1/worker/test_utils.py | 69 +++- .../layers/mamba/ops/ssu_dispatch.py | 295 ++++++++++-------- vllm/v1/core/single_type_kv_cache_manager.py | 55 +++- vllm/v1/worker/gpu/model_runner.py | 8 +- .../worker/gpu/model_states/mamba_hybrid.py | 76 ++--- vllm/v1/worker/gpu_model_runner.py | 15 +- vllm/v1/worker/mamba_utils.py | 85 +++-- vllm/v1/worker/utils.py | 68 ++++ 12 files changed, 758 insertions(+), 288 deletions(-) create mode 100644 tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index a18c11e1d02f..19f3af4d9376 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -444,21 +444,12 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None query_len[0] = scheduled accepted[0] = num_accepted is_prefilling[0] = prefilling - ctx.preprocess( - idx_mapping=None, - query_metadata=query_len, - query_metadata_is_cumulative=False, - num_computed_tokens=num_computed, - is_prefilling=is_prefilling, - src_cols=live_cols, - dst_cols=live_cols, - mamba_block_size=1024, - num_reqs=1, - ) ctx.postprocess( idx_mapping=None, query_metadata=query_len, query_metadata_is_cumulative=False, + num_computed_tokens=num_computed, + num_computed_is_post_step=False, num_accepted_tokens=accepted, is_prefilling=is_prefilling, live_cols=live_cols, @@ -524,6 +515,8 @@ def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), @@ -562,6 +555,54 @@ def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +def test_modelwide_replayssm_materialization_uses_independent_group_mappings( + monkeypatch, +): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + groups[0][0]._replayssm_prev_num_accepted[1] = 4 + groups[1][0]._replayssm_prev_num_accepted[5] = 9 + # The request has a source in group zero but a null source in group one. + block_tables[1][0, 0] = NULL_BLOCK_ID + materializer = Mock() + monkeypatch.setattr( + ssu_dispatch, "_load_replayssm_materialize", lambda: materializer + ) + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, + num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), + is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + ctx.materialize() + torch.accelerator.synchronize() + + assert materializer.call_count == 2 + assert ctx.groups[0].active_request_indices.tolist() == [0, -1] + assert ctx.groups[0].plan_flush_count.tolist() == [6, -1] + assert ctx.groups[1].active_request_indices.tolist() == [-1, -1] + assert ctx.groups[1].plan_flush_count.tolist() == [-1, -1] + assert ctx.groups[1].src_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 2 + assert ctx.groups[1].dst_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 2 + assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 9 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatch): _, config, forward_context, block_tables = _modelwide_replayssm_fixture() @@ -584,6 +625,8 @@ def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatc idx_mapping=None, query_metadata=torch.tensor([1, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), @@ -621,6 +664,8 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch idx_mapping=torch.tensor([1], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([1, 3], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), @@ -667,6 +712,8 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): idx_mapping=torch.tensor([-1], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=True, num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), @@ -683,7 +730,7 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_preprocess_resets_prefill_slots(monkeypatch): +def test_modelwide_replayssm_postprocess_resets_prefill_slots(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): for source_slot in source_slots: @@ -700,16 +747,17 @@ def test_modelwide_replayssm_preprocess_resets_prefill_slots(monkeypatch): max_num_reqs=2, ) assert ctx is not None - no_change = torch.full((2,), -1, dtype=torch.int32, device="cuda") - ctx.preprocess( + ctx.postprocess( idx_mapping=None, query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), - src_cols=no_change, - dst_cols=no_change, - mamba_block_size=4, + live_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.full((2,), -1, dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), num_reqs=1, ) torch.accelerator.synchronize() @@ -744,6 +792,8 @@ def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): idx_mapping=None, query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), @@ -764,11 +814,16 @@ def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): +def test_modelwide_replayssm_resets_only_group_specific_fresh_slot(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, source_slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[source_slot] = 2 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 + # Group zero aliases the two logical columns and therefore already has a + # valid source. Group one has no physical source for the same logical move. + block_tables[0][0, 1] = block_tables[0][0, 0] + block_tables[1][0, 0] = NULL_BLOCK_ID + groups[0][0]._replayssm_ring_start[1] = 2 + groups[0][0]._replayssm_prev_num_accepted[1] = 4 + groups[1][0]._replayssm_ring_start[5] = 7 + groups[1][0]._replayssm_prev_num_accepted[5] = 9 kernel = Mock() monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) @@ -780,32 +835,19 @@ def test_modelwide_replayssm_copies_reassigned_live_slot_once(monkeypatch): max_num_reqs=2, ) assert ctx is not None - ctx.preprocess( + ctx.reset_new_slots( idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), - query_metadata=torch.zeros(2, dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - mamba_block_size=4, num_reqs=1, ) - ctx.materialize() torch.accelerator.synchronize() - assert kernel.call_count == 2 - for group_ctx, source_slot, destination_slot in zip(ctx.groups, (1, 4), (2, 5)): - assert group_ctx.plan_ring_start.tolist() == [2, 0] - assert group_ctx.plan_flush_count.tolist() == [4, -1] - assert group_ctx.active_request_indices.tolist() == [0, -1] - assert group_ctx.src_slots[:, 0].tolist() == [source_slot] * 2 - assert group_ctx.dst_slots[:, 0].tolist() == [destination_slot] * 2 - for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): - assert mixers[0]._replayssm_ring_start[source_slot].item() == 2 - assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 4 - assert mixers[0]._replayssm_ring_start[destination_slot].item() == 0 - assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 + assert kernel.call_count == 0 + assert groups[0][0]._replayssm_ring_start[1].item() == 2 + assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 4 + assert groups[1][0]._replayssm_ring_start[5].item() == 0 + assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") @@ -829,6 +871,8 @@ def test_modelwide_replayssm_postprocess_materializes_in_place(monkeypatch): idx_mapping=None, query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 4be95aba52b1..4b985ccac98b 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -103,6 +103,40 @@ def test_mamba_speculative_block_relocation_requires_exclusive_ownership(): manager._relocate_speculative_block([pinned_block], 0) +@pytest.mark.parametrize("mamba_cache_mode", ["align", "all"]) +def test_replayssm_queues_live_copy_for_new_state_block(mamba_cache_mode: str): + spec = MambaSpec( + block_size=4, + shapes=((2,), (3,)), + dtypes=(torch.float32, torch.float32), + replayssm_shapes=((4,), (5,), (6,)), + replayssm_dtypes=(torch.float32,) * 3, + mamba_cache_mode=mamba_cache_mode, + ) + block_pool = BlockPool(num_gpu_blocks=6, enable_caching=True, hash_block_size=4) + manager = MambaManager( + spec, + block_pool=block_pool, + enable_caching=True, + kv_cache_group_id=0, + scheduler_block_size=4, + ) + + manager.allocate_new_blocks("request", num_tokens=3, num_tokens_main_model=3) + assert manager.take_pending_cow_copies() == [] + source = manager.req_to_blocks["request"][-1] + + manager.allocate_new_blocks("request", num_tokens=5, num_tokens_main_model=5) + destination = manager.req_to_blocks["request"][-1] + assert manager.take_pending_cow_copies() == [(source, destination)] + assert source.ref_cnt == 2 + assert destination.ref_cnt == 2 + + block_pool.free_blocks([source, destination]) + assert source.ref_cnt == 1 + assert destination.ref_cnt == 1 + + def get_sliding_window_manager( sliding_window_spec, block_pool, diff --git a/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py b/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py new file mode 100644 index 000000000000..1e3caf38e037 --- /dev/null +++ b/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Temporary V2 MTP diagnostic that keeps both request rows live.""" + +import os + +import vllm.envs as envs + +from ...models.utils import check_logprobs_close +from .test_replayssm_decode import PROMPTS + + +def test_replayssm_flashinfer_mtp_v2_without_batch_compaction(vllm_runner, monkeypatch): + model = os.environ["REPLAYSSM_MODEL"] + prompts = [PROMPTS[1], PROMPTS[1]] + common = dict( + max_model_len=1024, + trust_remote_code=True, + enable_prefix_caching=False, + mamba_cache_mode="none", + mamba_backend="flashinfer", + speculative_config={"method": "mtp", "num_speculative_tokens": 3}, + ) + + try: + with monkeypatch.context() as patch: + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + envs.disable_envs_cache() + with vllm_runner(model, **common) as llm: + baseline = llm.generate_greedy_logprobs( + prompts, max_tokens=32, num_logprobs=5 + ) + with vllm_runner( + model, use_replayssm=True, replayssm_buffer_len=16, **common + ) as llm: + replay = llm.generate_greedy_logprobs( + prompts, max_tokens=32, num_logprobs=5 + ) + finally: + envs.disable_envs_cache() + + for baseline_output, replay_output in zip(baseline, replay): + baseline_ids = baseline_output[0] + replay_ids = replay_output[0] + matching_prefix = 0 + for baseline_id, replay_id in zip(baseline_ids, replay_ids): + if baseline_id != replay_id: + break + matching_prefix += 1 + print( + "REPLAYSSM_COMPACTION_DIAGNOSTIC", + { + "matching_prefix": matching_prefix, + "baseline": baseline_ids, + "replay": replay_ids, + }, + ) + + check_logprobs_close( + outputs_0_lst=baseline, + outputs_1_lst=replay, + name_0="baseline_mtp_v2_no_compaction", + name_1="replayssm_mtp_v2_no_compaction", + ) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 2c2d4b5b1f56..f1dacb293a1b 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -34,6 +34,7 @@ get_mamba_groups, postprocess_mamba_align_gpu, preprocess_mamba, + preprocess_mamba_align_fused_kernel, stage_postprocess_inputs_to_gpu, validate_mamba_state_copy_funcs, ) @@ -50,6 +51,43 @@ } +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +@pytest.mark.parametrize( + ("preserve_accepted", "expected_accepted"), [(False, 1), (True, 3)] +) +def test_preprocess_mamba_preserves_replayssm_accepted_offset( + preserve_accepted: bool, expected_accepted: int +): + device = torch.device("cuda") + idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) + state_idx = torch.tensor([0], dtype=torch.int32, device=device) + num_computed = torch.tensor([4], dtype=torch.int32, device=device) + query_start = torch.tensor([0, 1], dtype=torch.int32, device=device) + num_accepted = torch.tensor([3], dtype=torch.int32, device=device) + src_col = torch.empty(1, dtype=torch.int32, device=device) + src_off = torch.empty(1, dtype=torch.int32, device=device) + + preprocess_mamba_align_fused_kernel[(1,)]( + idx_mapping, + state_idx, + num_computed, + query_start, + num_accepted, + src_col, + src_off, + 1, + BLOCK_SIZE=32, + MAMBA_BLOCK_SIZE=4, + PRESERVE_ACCEPTED=preserve_accepted, + ) + torch.accelerator.synchronize() + + assert state_idx.item() == 1 + assert src_col.item() == 0 + assert src_off.item() == 2 + assert num_accepted.item() == expected_accepted + + def postprocess_mamba( scheduler_output: "SchedulerOutput", kv_cache_config: "KVCacheConfig", @@ -170,20 +208,29 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): @pytest.mark.parametrize( - "num_computed_tokens", + ( + "with_replayssm", + "num_computed_tokens", + "expected_order", + "expected_accepted", + ), [ - pytest.param(4, id="boundary"), - pytest.param(3, id="no-boundary"), + pytest.param(True, 4, ["reset", "copy"], 3, id="replayssm-boundary"), + pytest.param(True, 3, ["reset", "copy"], 3, id="replayssm-no-boundary"), + pytest.param(False, 4, ["copy"], 1, id="generic"), ], ) -def test_preprocess_mamba_leaves_replayssm_for_staging( +def test_preprocess_mamba_preserves_live_replayssm_state( + with_replayssm: bool, num_computed_tokens: int, + expected_order: list[str], + expected_accepted: int, ) -> None: spec = MagicMock(block_size=4, num_speculative_blocks=0) cache_config = MagicMock(enable_prefix_caching=True, use_replayssm=True) input_batch = MagicMock() input_batch.req_ids = ["r0"] - input_batch.num_accepted_tokens_cpu = np.array([1], dtype=np.int32) + input_batch.num_accepted_tokens_cpu = np.array([3], dtype=np.int32) copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec) requests = {"r0": MagicMock(num_computed_tokens=num_computed_tokens)} mamba_state_idx: dict[str, int] = {"r0": 0} @@ -196,7 +243,11 @@ def test_preprocess_mamba_leaves_replayssm_for_staging( align_ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(1, torch.int32, device) align_ctx.precopy_src_col_buf = _MockCpuGpuBuffer(1, torch.int32, device) align_ctx.precopy_token_bias_buf = _MockCpuGpuBuffer(1, torch.int32, device) - align_ctx.replayssm = MagicMock() + align_ctx.replayssm = MagicMock() if with_replayssm else None + if align_ctx.replayssm is not None: + align_ctx.replayssm.reset_new_slots.side_effect = lambda **kwargs: order.append( + "reset" + ) align_ctx.run_fused_precopy.side_effect = lambda **kwargs: order.append("copy") preprocess_mamba( @@ -212,13 +263,21 @@ def test_preprocess_mamba_leaves_replayssm_for_staging( align_ctx=align_ctx, ) - assert order == ["copy"] - align_ctx.replayssm.preprocess.assert_not_called() - align_ctx.replayssm.materialize.assert_not_called() + assert order == expected_order + assert input_batch.num_accepted_tokens_cpu[0] == expected_accepted + assert align_ctx.precopy_src_col_buf.np[0] == 0 -def test_postprocess_mamba_align_commits_then_materializes_after_fused_copy( - monkeypatch, +@pytest.mark.parametrize( + ("materialize_possible", "expected_order"), + [ + (True, ["copy", "postprocess", "materialize"]), + (False, ["copy", "postprocess"]), + ], +) +def test_postprocess_mamba_align_gates_materialization( + materialize_possible: bool, + expected_order: list[str], ): order: list[str] = [] ctx = MagicMock() @@ -235,6 +294,7 @@ def test_postprocess_mamba_align_commits_then_materializes_after_fused_copy( ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") ctx.replayssm = MagicMock() ctx.replayssm.materialize_prefixes = True + ctx.replayssm_materialize_possible = materialize_possible ctx.replayssm.postprocess.side_effect = lambda **kwargs: order.append("postprocess") ctx.replayssm.materialize.side_effect = lambda: order.append("materialize") block_table = MagicMock() @@ -257,7 +317,7 @@ def test_postprocess_mamba_align_commits_then_materializes_after_fused_copy( run_prefix_state_migration=True, ) - assert order == ["copy", "postprocess", "materialize"] + assert order == expected_order assert accepted_cpu.tolist() == [3] @@ -441,6 +501,7 @@ def test_gpu_context_initializes_flashinfer_replayssm_lifecycle(): ) assert gpu_ctx.state_skip_postprocess.tolist() == [0, 1] + assert gpu_ctx.state_skip_precopy.tolist() == [1, 1] assert gpu_ctx.replayssm is model_ctx create.assert_called_once() @@ -472,6 +533,33 @@ def test_gpu_context_rejects_missing_replayssm_lifecycle(): ) +def test_gpu_context_rejects_mixed_replayssm_and_baseline_layers(): + cfg = _TestConfig(num_layers=2) + device = torch.device("cpu") + layer_names = ["layer_0", "layer_1"] + kv_cache_config = _make_kv_cache_config(cfg, layer_names) + gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) + forward_context = { + name: _make_mock_attention( + torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), + torch.empty(cfg.num_blocks, cfg.temporal_state_dim), + ) + for name in layer_names + } + forward_context["layer_0"].use_replayssm = True + forward_context["layer_0"].mamba_config.backend = MambaBackendEnum.FLASHINFER + forward_context["layer_1"].use_replayssm = False + forward_context["layer_1"].mamba_config.backend = MambaBackendEnum.TRITON + + with pytest.raises(ValueError, match="mixed FlashInfer ReplaySSM"): + gpu_ctx.initialize_from_forward_context( + kv_cache_config, + forward_context, + _COPY_FUNCS, + [torch.empty(1, 4, dtype=torch.int32)], + ) + + def _make_postprocess_scheduler_output( req_ids: list[str], num_scheduled_tokens: dict[str, int], @@ -953,18 +1041,37 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ctx.is_prefilling_buf.gpu[:num_reqs], torch.tensor([True, False, True]), ) - ctx.replayssm.preprocess.assert_called_once_with( - idx_mapping=None, - query_metadata=ctx.num_scheduled_tokens_buf.gpu, - query_metadata_is_cumulative=False, - num_computed_tokens=ctx.num_computed_tokens_buf.gpu, - is_prefilling=ctx.is_prefilling_buf.gpu, - src_cols=ctx.precopy_src_col_buf.gpu, - dst_cols=ctx.mamba_state_idx_buf.gpu, - mamba_block_size=4, - num_reqs=num_reqs, + ctx.replayssm.reset_new_slots.assert_not_called() + ctx.replayssm.materialize.assert_not_called() + assert ctx.replayssm_materialize_possible + + +def test_stage_postprocess_inputs_skips_impossible_materialization(): + device = torch.device("cpu") + ctx = _make_staging_ctx(max_num_reqs=2, device=device) + ctx.block_size = 4 + ctx.replayssm = MagicMock(materialize_prefixes=True) + scheduler_output = _make_postprocess_scheduler_output( + req_ids=["req_a"], + num_scheduled_tokens={"req_a": 1}, + ) + requests = _make_requests( + ["req_a"], + [10], + [[0]], + num_prompt_tokens=[10], + ) + + stage_postprocess_inputs_to_gpu( + ctx, + scheduler_output, + ["req_a"], + 1, + requests, + {"req_a": 2}, ) - ctx.replayssm.materialize.assert_called_once_with() + + assert not ctx.replayssm_materialize_possible def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 3f9b9542a3a4..80606ad4ce85 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -3,11 +3,17 @@ from types import SimpleNamespace +import pytest import torch from vllm.config.mamba import MambaBackendEnum, MambaConfig from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 -from vllm.v1.worker.utils import bind_kv_cache +from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy +from vllm.v1.worker.utils import ( + bind_kv_cache, + copy_kv_cache_blocks_inplace, + get_replayssm_block_copy_tensors, +) class _TestReplaySSMMixer(MambaMixer2): @@ -85,6 +91,67 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): assert torch.count_nonzero(group_tracker) == 0 +def test_replayssm_block_copy_includes_rings_and_group_trackers(monkeypatch): + monkeypatch.setattr( + "vllm.v1.worker.utils.async_tensor_h2d", + lambda array, *, device, **_: torch.from_numpy(array).to(device), + ) + mixers = [_TestReplaySSMMixer() for _ in range(3)] + layer_names = [f"layers.{i}.mixer" for i in range(3)] + ctx = dict(zip(layer_names, mixers)) + kv_cache = {name: _packed_replayssm_cache(4) for name in layer_names} + kv_cache_groups = [ + SimpleNamespace(layer_names=[layer_names[0], layer_names[2]]), + SimpleNamespace(layer_names=[layer_names[1]]), + ] + replayssm_caches = { + name: tuple( + torch.zeros((4, *shape), dtype=torch.float32) + for shape in mixer.get_replayssm_state_shape() + ) + for name, mixer in ctx.items() + } + runner_kv_caches: list[torch.Tensor] = [] + bind_kv_cache( + kv_cache, + ctx, + runner_kv_caches, + kv_cache_groups=kv_cache_groups, + replayssm_caches=replayssm_caches, + ) + + src, dst = 1, 2 + for layer_idx, mixer in enumerate(mixers): + for state_idx, state in enumerate(mixer.kv_cache): + state[src].fill_(10 * layer_idx + state_idx + 1) + state[dst].fill_(-1) + for group_idx, mixer in enumerate(mixers[:2]): + mixer._replayssm_ring_start[src] = 20 + group_idx + mixer._replayssm_prev_num_accepted[src] = 30 + group_idx + + copy_kv_cache_blocks_inplace( + [*runner_kv_caches, *get_replayssm_block_copy_tensors(ctx)], + 4, + [KVCacheBlockCopy(src, dst)], + ) + + for mixer in mixers: + for state in mixer.kv_cache: + torch.testing.assert_close(state[dst], state[src]) + assert mixers[0]._replayssm_ring_start[dst].item() == 20 + assert mixers[0]._replayssm_prev_num_accepted[dst].item() == 30 + assert mixers[1]._replayssm_ring_start[dst].item() == 21 + assert mixers[1]._replayssm_prev_num_accepted[dst].item() == 31 + + +def test_replayssm_block_copy_validates_exact_cache_roles(): + mixer = _TestReplaySSMMixer() + mixer.kv_cache = tuple(torch.zeros(4, 1) for _ in range(4)) + + with pytest.raises(ValueError, match="exactly 5 cache roles"): + get_replayssm_block_copy_tensors({"layers.0.mixer": mixer}) + + def test_bind_kv_cache(default_vllm_config): from vllm.model_executor.layers.attention import Attention diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index f90884770f3b..3d9d141a59fe 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -28,33 +28,20 @@ @triton.jit(do_not_specialize=["num_reqs"]) -def _preprocess_replayssm_kernel( +def _reset_new_replayssm_slots_kernel( idx_mapping, - query_metadata, - num_computed_tokens, - is_prefilling, src_cols, dst_cols, block_table, tracker_start, tracker_committed, - src_slots, - dst_slots, - plan_ring_start, - plan_flush_count, - active_request_indices, block_table_stride_req: tl.int64, - slot_table_stride_layer: tl.int64, + tracker_capacity, num_reqs, - MAMBA_BLOCK_SIZE: tl.constexpr, - NUM_LAYERS: tl.constexpr, PAD_SLOT_ID: tl.constexpr, - QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, - MATERIALIZE_PREFIXES: tl.constexpr, - MAX_NUM_REQS: tl.constexpr, ) -> None: - """Reset prefill state and prepare an optional writable-slot move.""" + """Reset cursors only when this cache group has no physical source.""" batch_idx = tl.program_id(0) active = batch_idx < num_reqs req_idx = batch_idx @@ -62,87 +49,44 @@ def _preprocess_replayssm_kernel( req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) valid_req = active & (req_idx >= 0) - src_col = tl.load(src_cols + req_idx, mask=valid_req, other=-1) dst_col = tl.load(dst_cols + req_idx, mask=valid_req, other=-1) - changed = valid_req & (src_col >= 0) & (src_col != dst_col) - src_slot = tl.load( - block_table + batch_idx * block_table_stride_req + src_col, - mask=changed, - other=PAD_SLOT_ID, - ) + valid_dst_col = valid_req & (dst_col >= 0) dst_slot = tl.load( block_table + batch_idx * block_table_stride_req + dst_col, - mask=changed, + mask=valid_dst_col, other=PAD_SLOT_ID, ) - for layer_idx in tl.static_range(0, NUM_LAYERS): - slot_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store(src_slots + slot_offset, src_slot) - tl.store(dst_slots + slot_offset, dst_slot) + valid_dst = ( + valid_dst_col + & (dst_slot != PAD_SLOT_ID) + & (dst_slot >= 0) + & (dst_slot < tracker_capacity) + ) - # Clear the fixed-capacity plan row before deciding whether this request - # belongs in the compact materialization list. - tl.store(plan_ring_start + batch_idx, 0) - tl.store(plan_flush_count + batch_idx, -1) - if changed: - # BlockManager guarantees that distinct live columns map to allocated, - # distinct physical slots. Snapshot the old owner before clearing the - # destination tracker. - tl.store(plan_ring_start + batch_idx, tl.load(tracker_start + src_slot)) - tl.store( - plan_flush_count + batch_idx, - tl.load(tracker_committed + src_slot), - ) - tl.store(tracker_start + dst_slot, 0) - tl.store(tracker_committed + dst_slot, 0) - - prefilling = tl.load(is_prefilling + batch_idx, mask=active, other=0) - if valid_req & prefilling: - if QUERY_METADATA_IS_CUMULATIVE: - query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( - query_metadata + batch_idx - ) - else: - query_len = tl.load(query_metadata + batch_idx) - computed = tl.load(num_computed_tokens + req_idx) - first_col = tl.maximum(computed // MAMBA_BLOCK_SIZE, 0) - last_col = tl.maximum( - (computed + query_len + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1, - 0, - ) - for col in tl.range(first_col, last_col + 1): - slot = tl.load(block_table + batch_idx * block_table_stride_req + col) - tl.store(tracker_start + slot, 0) - tl.store(tracker_committed + slot, 0) + src_col = tl.load(src_cols + req_idx, mask=valid_req, other=-1) + valid_src_col = valid_req & (src_col >= 0) + src_slot = tl.load( + block_table + batch_idx * block_table_stride_req + src_col, + mask=valid_src_col, + other=PAD_SLOT_ID, + ) + valid_src = ( + valid_src_col + & (src_slot != PAD_SLOT_ID) + & (src_slot >= 0) + & (src_slot < tracker_capacity) + ) - if MATERIALIZE_PREFIXES & (batch_idx == 0): - active_count = 0 - for candidate_idx in tl.range(0, num_reqs): - candidate_req_idx = candidate_idx - if HAS_IDX_MAPPING: - candidate_req_idx = tl.load(idx_mapping + candidate_idx) - valid_candidate = candidate_req_idx >= 0 - src_col = tl.load( - src_cols + candidate_req_idx, - mask=valid_candidate, - other=-1, - ) - dst_col = tl.load( - dst_cols + candidate_req_idx, - mask=valid_candidate, - other=-1, - ) - if valid_candidate & (src_col >= 0) & (src_col != dst_col): - tl.store(active_request_indices + active_count, candidate_idx) - active_count += 1 - if active_count < MAX_NUM_REQS: - tl.store(active_request_indices + active_count, -1) + fresh = valid_dst & ~valid_src + tl.store(tracker_start + dst_slot, 0, mask=fresh) + tl.store(tracker_committed + dst_slot, 0, mask=fresh) @triton.jit(do_not_specialize=["num_reqs"]) def _postprocess_replayssm_kernel( idx_mapping, query_metadata, + num_computed_tokens, num_accepted_tokens, is_prefilling, live_cols, @@ -158,11 +102,15 @@ def _postprocess_replayssm_kernel( active_request_indices, block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, + tracker_capacity, num_reqs, + MAMBA_BLOCK_SIZE: tl.constexpr, LOGICAL_WINDOW: tl.constexpr, RING_BUFFER_LEN: tl.constexpr, NUM_LAYERS: tl.constexpr, + PAD_SLOT_ID: tl.constexpr, QUERY_METADATA_IS_CUMULATIVE: tl.constexpr, + NUM_COMPUTED_IS_POST_STEP: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, MATERIALIZE_PREFIXES: tl.constexpr, MAX_NUM_REQS: tl.constexpr, @@ -170,22 +118,53 @@ def _postprocess_replayssm_kernel( """Commit a completed step and prepare an optional prefix snapshot.""" batch_idx = tl.program_id(0) - # FlashInfer requires active rows to be a compact prefix. The generic - # planner's destination is the sole materialization work sentinel. + # FlashInfer requires active rows to be a compact prefix. Resolve physical + # validity per group so divergent group mappings remain independent. if MATERIALIZE_PREFIXES & (batch_idx == 0): active_count = 0 for candidate_idx in tl.range(0, num_reqs): candidate_req_idx = candidate_idx if HAS_IDX_MAPPING: candidate_req_idx = tl.load(idx_mapping + candidate_idx) - dst_col = tl.load(materialize_dst_cols + candidate_idx) - if (candidate_req_idx >= 0) & (dst_col >= 0): + valid_candidate = candidate_req_idx >= 0 + live_col = tl.load( + live_cols + candidate_req_idx, + mask=valid_candidate, + other=-1, + ) + dst_col = tl.load( + materialize_dst_cols + candidate_idx, + mask=valid_candidate, + other=-1, + ) + wants_materialize = valid_candidate & (live_col >= 0) & (dst_col >= 0) + live_slot = tl.load( + block_table + candidate_idx * block_table_stride_req + live_col, + mask=wants_materialize, + other=PAD_SLOT_ID, + ) + dst_slot = tl.load( + block_table + candidate_idx * block_table_stride_req + dst_col, + mask=wants_materialize, + other=PAD_SLOT_ID, + ) + valid_materialize = ( + wants_materialize + & (live_slot != PAD_SLOT_ID) + & (dst_slot != PAD_SLOT_ID) + & (live_slot >= 0) + & (dst_slot >= 0) + & (live_slot < tracker_capacity) + & (dst_slot < tracker_capacity) + ) + if valid_materialize: tl.store(active_request_indices + active_count, candidate_idx) active_count += 1 if active_count < MAX_NUM_REQS: tl.store(active_request_indices + active_count, -1) # Clear the fixed-capacity plan row before any per-request early exit. + tl.store(plan_ring_start + batch_idx, 0) tl.store(plan_flush_count + batch_idx, -1) if batch_idx >= num_reqs: return @@ -196,25 +175,78 @@ def _postprocess_replayssm_kernel( if req_idx < 0: return - # BlockManager supplies allocated live/destination columns. The generic - # Mamba planner always materializes from this same live column. live_col = tl.load(live_cols + req_idx) - live_slot = tl.load(block_table + batch_idx * block_table_stride_req + live_col) + valid_live_col = live_col >= 0 + live_slot = tl.load( + block_table + batch_idx * block_table_stride_req + live_col, + mask=valid_live_col, + other=PAD_SLOT_ID, + ) + valid_live = ( + valid_live_col + & (live_slot != PAD_SLOT_ID) + & (live_slot >= 0) + & (live_slot < tracker_capacity) + ) dst_col = tl.load(materialize_dst_cols + batch_idx) - materialize = dst_col >= 0 + wants_materialize = valid_live & (dst_col >= 0) + dst_slot = tl.load( + block_table + batch_idx * block_table_stride_req + dst_col, + mask=wants_materialize, + other=PAD_SLOT_ID, + ) + materialize = ( + wants_materialize + & (dst_slot != PAD_SLOT_ID) + & (dst_slot >= 0) + & (dst_slot < tracker_capacity) + ) + for layer_idx in tl.static_range(0, NUM_LAYERS): + slot_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store( + src_slots + slot_offset, + tl.where(materialize, live_slot, PAD_SLOT_ID), + ) + tl.store( + dst_slots + slot_offset, + tl.where(materialize, dst_slot, PAD_SLOT_ID), + ) + prefilling = tl.load(is_prefilling + batch_idx) + if QUERY_METADATA_IS_CUMULATIVE: + query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( + query_metadata + batch_idx + ) + else: + query_len = tl.load(query_metadata + batch_idx) + if prefilling: + computed = tl.load(num_computed_tokens + req_idx) + computed_before = tl.where( + NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed + ) + computed_after = computed_before + query_len + first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) + last_col = tl.maximum( + (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1, + 0, + ) + for col in tl.range(first_col, last_col + 1): + prefill_slot = tl.load( + block_table + batch_idx * block_table_stride_req + col + ) + valid_prefill_slot = ( + (prefill_slot != PAD_SLOT_ID) + & (prefill_slot >= 0) + & (prefill_slot < tracker_capacity) + ) + tl.store(tracker_start + prefill_slot, 0, mask=valid_prefill_slot) + tl.store(tracker_committed + prefill_slot, 0, mask=valid_prefill_slot) if materialize: # Prefill produced canonical state, so publish an exact copy. tl.store(plan_flush_count + batch_idx, 0) - else: - if QUERY_METADATA_IS_CUMULATIVE: - query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( - query_metadata + batch_idx - ) - else: - query_len = tl.load(query_metadata + batch_idx) - accepted = tl.load(num_accepted_tokens + req_idx) + elif valid_live: + accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) old_start = tl.load(tracker_start + live_slot) old_committed = tl.load(tracker_committed + live_slot) checkpointed = old_committed + query_len > LOGICAL_WINDOW @@ -235,17 +267,9 @@ def _postprocess_replayssm_kernel( boundary_count + tl.where(checkpointed, 0, old_committed), ) - if not materialize: - return - dst_slot = tl.load(block_table + batch_idx * block_table_stride_req + dst_col) - for layer_idx in tl.static_range(0, NUM_LAYERS): - slot_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store(src_slots + slot_offset, live_slot) - tl.store(dst_slots + slot_offset, dst_slot) - # The published destination is canonical and therefore has no live replay. - tl.store(tracker_start + dst_slot, 0) - tl.store(tracker_committed + dst_slot, 0) + tl.store(tracker_start + dst_slot, 0, mask=materialize) + tl.store(tracker_committed + dst_slot, 0, mask=materialize) def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: @@ -282,6 +306,7 @@ class _ReplaySSMGroupContext: plan_flush_count: torch.Tensor active_request_indices: torch.Tensor max_num_reqs: int + mamba_block_size: int logical_window: int ring_buffer_len: int materialize_prefixes: bool @@ -292,6 +317,7 @@ def create( mixers: list[Any], block_table: torch.Tensor, cache_mode: str, + mamba_block_size: int, max_num_reqs: int, ) -> "_ReplaySSMGroupContext": if ( @@ -364,52 +390,35 @@ def create( (max_num_reqs,), -1, dtype=torch.int32, device=device ), max_num_reqs=max_num_reqs, + mamba_block_size=mamba_block_size, logical_window=int(first.replayssm_buffer_len), ring_buffer_len=first_x.size(2), materialize_prefixes=cache_mode in ("align", "all"), ) - def preprocess( + def reset_new_slots( self, *, idx_mapping: torch.Tensor | None, - query_metadata: torch.Tensor, - query_metadata_is_cumulative: bool, - num_computed_tokens: torch.Tensor, - is_prefilling: torch.Tensor, src_cols: torch.Tensor, dst_cols: torch.Tensor, - mamba_block_size: int, num_reqs: int, ) -> None: - """Reset prefill state and prepare an optional writable-slot move.""" + """Reset cursors for fresh physical slots in this cache group.""" if num_reqs == 0: return - _preprocess_replayssm_kernel[(self.max_num_reqs,)]( + _reset_new_replayssm_slots_kernel[(self.max_num_reqs,)]( idx_mapping, - query_metadata, - num_computed_tokens, - is_prefilling, src_cols, dst_cols, self.block_table, self.ring_start, self.num_committed, - self.src_slots, - self.dst_slots, - self.plan_ring_start, - self.plan_flush_count, - self.active_request_indices, self.block_table.stride(0), - self.src_slots.stride(0), + self.num_committed.numel(), num_reqs, - MAMBA_BLOCK_SIZE=mamba_block_size, - NUM_LAYERS=len(self.mixers), PAD_SLOT_ID=NULL_BLOCK_ID, - QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, HAS_IDX_MAPPING=idx_mapping is not None, - MATERIALIZE_PREFIXES=self.materialize_prefixes, - MAX_NUM_REQS=self.max_num_reqs, ) def postprocess( @@ -418,6 +427,8 @@ def postprocess( idx_mapping: torch.Tensor | None, query_metadata: torch.Tensor, query_metadata_is_cumulative: bool, + num_computed_tokens: torch.Tensor, + num_computed_is_post_step: bool, num_accepted_tokens: torch.Tensor, is_prefilling: torch.Tensor, live_cols: torch.Tensor, @@ -431,6 +442,7 @@ def postprocess( _postprocess_replayssm_kernel[(self.max_num_reqs,)]( idx_mapping, query_metadata, + num_computed_tokens, num_accepted_tokens, is_prefilling, live_cols, @@ -446,18 +458,22 @@ def postprocess( self.active_request_indices, self.block_table.stride(0), self.src_slots.stride(0), + self.num_committed.numel(), num_reqs, + MAMBA_BLOCK_SIZE=self.mamba_block_size, LOGICAL_WINDOW=self.logical_window, RING_BUFFER_LEN=self.ring_buffer_len, NUM_LAYERS=len(self.mixers), + PAD_SLOT_ID=NULL_BLOCK_ID, QUERY_METADATA_IS_CUMULATIVE=query_metadata_is_cumulative, + NUM_COMPUTED_IS_POST_STEP=num_computed_is_post_step, HAS_IDX_MAPPING=idx_mapping is not None, MATERIALIZE_PREFIXES=self.materialize_prefixes, MAX_NUM_REQS=self.max_num_reqs, ) def materialize(self) -> None: - """Execute the plan prepared by ``preprocess`` or ``postprocess``.""" + """Publish the canonical prefix snapshots prepared by ``postprocess``.""" if not self.materialize_prefixes: raise RuntimeError("ReplaySSM materialization requires align or all mode") first = self.mixers[0] @@ -528,7 +544,14 @@ def create( f"got {type(spec).__name__}" ) modes.add(spec.mamba_cache_mode) - group_args.append((mixers, block_table_by_gid[gid], spec.mamba_cache_mode)) + group_args.append( + ( + mixers, + block_table_by_gid[gid], + spec.mamba_cache_mode, + spec.block_size, + ) + ) if len(modes) != 1: raise ValueError( "model-wide ReplaySSM requires one Mamba cache mode; " @@ -543,9 +566,9 @@ def create( materialize_prefixes=next(iter(modes)) in ("align", "all"), ) - def preprocess(self, **kwargs: Any) -> None: + def reset_new_slots(self, **kwargs: Any) -> None: for group in self.groups: - group.preprocess(**kwargs) + group.reset_new_slots(**kwargs) def postprocess(self, **kwargs: Any) -> None: for group in self.groups: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2091be31c644..b72bc79e19e1 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1372,6 +1372,12 @@ def __init__( self.block_size = kv_cache_spec.block_size self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks + self._copy_replayssm_live_state = bool(kv_cache_spec.replayssm_shapes) + if self._copy_replayssm_live_state: + assert self.num_speculative_blocks == 0, ( + "ReplaySSM keeps speculative state in its replay rings, not " + "separate scheduler blocks" + ) self.cached_blocks_this_step: set[BlockHashWithGroupId] = set() if self.mamba_cache_mode == "align": # Mapping from request ID to the index of the block @@ -1662,13 +1668,24 @@ def allocate_new_blocks( ) -> list[KVCacheBlock]: assert isinstance(self.kv_cache_spec, MambaSpec) if self.mamba_cache_mode != "align": + req_blocks = self.req_to_blocks[request_id] + prev_block_len = len(req_blocks) + partial_hit = self._partial_hit_reqs.get(request_id) + live_source = ( + partial_hit[1] + if partial_hit is not None + else (req_blocks[-1] if req_blocks else None) + ) # Allocate extra `num_speculative_blocks` blocks for # speculative decoding (MTP/EAGLE) with linear attention. if self.num_speculative_blocks > 0: num_tokens += self.block_size * self.num_speculative_blocks - return super().allocate_new_blocks( + new_blocks = super().allocate_new_blocks( request_id, num_tokens, num_tokens_main_model ) + if len(req_blocks) > prev_block_len and live_source is not None: + self._queue_replayssm_live_copy(live_source, req_blocks[-1]) + return new_blocks else: # We don't allocate blocks for lookahead tokens in align mode, because if # x * block_size tokens are scheduled, num_tokens is @@ -1676,7 +1693,8 @@ def allocate_new_blocks( # We can ignore lookahead tokens because current draft models don't have # mamba layers. num_tokens = num_tokens_main_model - req_blocks: list[KVCacheBlock] = self.req_to_blocks[request_id] + req_blocks = self.req_to_blocks[request_id] + prev_block_len = len(req_blocks) # NOTE(tdouble): this is an over-estimate of how many blocks we need because # num_tokens can include draft tokens that will later be rejected. num_required_blocks = ( @@ -1685,13 +1703,22 @@ def allocate_new_blocks( checkpoint_block = self._num_checkpoint_blocks.get(request_id, 0) partial_hit = self._partial_hit_reqs.get(request_id) has_partial_hit = partial_hit is not None + live_source = None + if partial_hit is not None: + live_source = partial_hit[1] + elif prev_block_len > 0: + live_source_idx = ( + prev_block_len - 1 - self.num_speculative_blocks + if request_id in self._allocated_block_reqs + else prev_block_len - 1 + ) + live_source = req_blocks[live_source_idx] # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. if num_required_blocks <= len(req_blocks) and not has_partial_hit: self._allocated_block_reqs.add(request_id) return [] else: - prev_block_len = len(req_blocks) blocks_allocated = request_id in self._allocated_block_reqs # Record the last state block if blocks_allocated: @@ -1772,11 +1799,33 @@ def allocate_new_blocks( self._apply_cow(request_id, block_idx, source_block, cow_block) returned_blocks = [cow_block] + returned_blocks req_blocks.extend(new_blocks) + live_dest_idx = len(req_blocks) - 1 - self.num_speculative_blocks + live_dest = req_blocks[live_dest_idx] + if any(live_dest is block for block in new_blocks): + self._queue_replayssm_live_copy(live_source, live_dest) self._allocated_block_reqs.add(request_id) self._partial_hit_reqs.pop(request_id, None) returned_blocks.extend(new_blocks) return returned_blocks + def _queue_replayssm_live_copy( + self, + source_block: KVCacheBlock | None, + destination_block: KVCacheBlock, + ) -> None: + """Queue and retain one complete live ReplaySSM slot migration.""" + if ( + not self._copy_replayssm_live_state + or source_block is None + or source_block.is_null + or destination_block.is_null + or source_block is destination_block + ): + return + source_block.ref_cnt += 1 + destination_block.ref_cnt += 1 + self._pending_cow_copies.append((source_block, destination_block)) + def _relocate_speculative_block( self, req_blocks: list[KVCacheBlock], block_idx: int ) -> None: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 802d4c775e7f..b2cf6eeef75b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -164,6 +164,7 @@ KVBlockZeroer, clear_layer_kv_caches, copy_kv_cache_blocks_inplace, + get_replayssm_block_copy_tensors, get_uniform_decode_token_count, ) from vllm.v1.worker.workspace import use_workspace_lane @@ -1111,7 +1112,12 @@ def update_requests(self, scheduler_output: SchedulerOutput) -> None: # zeroing new blocks and before the forward pass reads them. if scheduler_output.kv_cache_block_copies: copy_kv_cache_blocks_inplace( - self.kv_caches, + [ + *self.kv_caches, + *get_replayssm_block_copy_tensors( + self.compilation_config.static_forward_context + ), + ], self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, ) diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 6716a47cd4cd..75982a1b2897 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -212,13 +212,13 @@ def preprocess_state( kv_cache_config: KVCacheConfig, num_computed_tokens: torch.Tensor, ) -> None: - """Migrate each request's mamba state across block boundaries before - the forward (V1 lifecycle semantics, done on GPU). Runs on real batches - only (dummy DP/profiling runs skip preprocess_state), and before - ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset - is visible to the forward kernels. + """Prepare each request's Mamba slot before the forward. + + ReplaySSM live state has already moved through the scheduler block-copy + path, so this phase only resets a fresh slot's cursors. Baseline Mamba + retains the fused state pre-copy. Dummy DP/profiling runs skip this. """ - if not (self._needs_prefix_state_migration or self._use_flashinfer_replayssm): + if not self._needs_prefix_state_migration: return num_reqs = input_batch.num_reqs if num_reqs == 0: @@ -229,52 +229,32 @@ def preprocess_state( ) replayssm = ctx.replayssm - if replayssm is not None: - self._is_prefilling_gpu[:num_reqs].copy_( - torch.from_numpy(input_batch.is_prefilling_np[:num_reqs]) - ) - - if self._needs_prefix_state_migration: - # This decision kernel fast-exits per request when no block boundary - # is crossed. Avoiding the launch would require a CPU sync under - # async scheduling. - block = 256 - grid = (triton.cdiv(num_reqs, block),) - preprocess_mamba_align_fused_kernel[grid]( - input_batch.idx_mapping, - self._mamba_state_idx_gpu, - num_computed_tokens, - input_batch.query_start_loc, - self.num_accepted_tokens_gpu, - self._mamba_src_col_gpu, - self._mamba_src_off_gpu, - num_reqs, - BLOCK_SIZE=block, - MAMBA_BLOCK_SIZE=mamba_spec.block_size, - ) + # This decision kernel fast-exits per request when no block boundary is + # crossed. Avoiding the launch would require a CPU sync under async + # scheduling. + block = 256 + grid = (triton.cdiv(num_reqs, block),) + preprocess_mamba_align_fused_kernel[grid]( + input_batch.idx_mapping, + self._mamba_state_idx_gpu, + num_computed_tokens, + input_batch.query_start_loc, + self.num_accepted_tokens_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + num_reqs, + BLOCK_SIZE=block, + MAMBA_BLOCK_SIZE=mamba_spec.block_size, + PRESERVE_ACCEPTED=replayssm is not None, + ) if replayssm is not None: - if self._needs_prefix_state_migration: - src_cols = self._mamba_src_col_gpu - dst_cols = self._mamba_state_idx_gpu - else: - src_cols = dst_cols = self._replayssm_live_cols_gpu - replayssm.preprocess( + replayssm.reset_new_slots( idx_mapping=input_batch.idx_mapping, - query_metadata=input_batch.query_start_loc, - query_metadata_is_cumulative=True, - num_computed_tokens=num_computed_tokens, - is_prefilling=self._is_prefilling_gpu, - src_cols=src_cols, - dst_cols=dst_cols, - mamba_block_size=ctx.block_size, + src_cols=self._mamba_src_col_gpu, + dst_cols=self._mamba_state_idx_gpu, num_reqs=num_reqs, ) - if replayssm.materialize_prefixes: - replayssm.materialize() - - if not self._needs_prefix_state_migration: - return ctx.run_fused_precopy( num_reqs, self._mamba_state_idx_gpu, @@ -529,6 +509,8 @@ def _publish_flashinfer_replayssm( idx_mapping=idx_mapping, query_metadata=query_start_loc, query_metadata_is_cumulative=True, + num_computed_tokens=num_computed_tokens, + num_computed_is_post_step=True, # Prefix migration can reset the live buffer to one for the # next step; use its snapshot in that case. Mode none never # runs the migration kernel, so the acceptance buffer is exact. diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index ae6990d8722c..8a0228616773 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -249,6 +249,7 @@ allocate_replayssm_caches, bind_kv_cache, copy_kv_cache_blocks_inplace, + get_replayssm_block_copy_tensors, prepare_kernel_block_sizes, sanity_check_mm_encoder_outputs, ) @@ -1250,7 +1251,12 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self._zero_block_ids(scheduler_output.new_block_ids_to_zero) if scheduler_output.kv_cache_block_copies: copy_kv_cache_blocks_inplace( - self.kv_caches, + [ + *self.kv_caches, + *get_replayssm_block_copy_tensors( + self.compilation_config.static_forward_context + ), + ], self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, ) @@ -4454,10 +4460,9 @@ def execute_model( mamba_bufs.preprocess, align_ctx=mamba_bufs.postprocess_align, ) - # preprocess_mamba resets num_accepted_tokens_cpu to 1 - # for requests whose state was copied to a new block. - # Re-sync to GPU so the mamba kernel reads from the - # correct initial state slot (init_token_idx = 0). + # Baseline Mamba may reset an accepted-token offset after + # shifting state. ReplaySSM preserves it with the exact live + # block copy. Re-sync either result to GPU. self.num_accepted_tokens.np[:num_reqs] = ( self.input_batch.num_accepted_tokens_cpu[:num_reqs] ) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 0d627c79c1e5..64e38b281d79 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -539,18 +539,19 @@ def preprocess_mamba_align_fused_kernel( num_reqs, BLOCK_SIZE: tl.constexpr, MAMBA_BLOCK_SIZE: tl.constexpr, + PRESERVE_ACCEPTED: tl.constexpr = False, ): """Fused align preprocess: emit the pre-copy src column/offset AND advance - state_idx (with accepted-token reset) in a single launch (V2 align). + state_idx in a single launch (V2 align). Per batch_idx (0..num_reqs-1), resolving req slot via idx_mapping: 1. Read pre-advance state_idx and num_accepted (last step's values). 2. Store the pre-copy src columns for ``precopy_mamba_align_fused_kernel``: - src_col = state_idx (the previous running block column) - src_off = max(num_accepted - 1, 0) (the accepted-token bias) - 3. Advance state_idx to the new running block, and reset num_accepted to 1 - when a block boundary is crossed (so the migrated state, now at the - start of the new block, is read with the neutral bias). + 3. Advance state_idx to the new running block. Baseline Mamba resets + num_accepted after shifting its state; exact ReplaySSM block copies + preserve it so the copied live slot keeps the same accepted position. """ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < num_reqs @@ -569,8 +570,9 @@ def preprocess_mamba_align_fused_kernel( computed_after = num_computed + query_end - query_start new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1 tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask) - should_reset = (state_idx >= 0) & (state_idx != new_state_idx) - tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) + if not PRESERVE_ACCEPTED: + should_reset = (state_idx >= 0) & (state_idx != new_state_idx) + tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset) @triton.jit(do_not_specialize=["num_reqs"]) @@ -825,6 +827,7 @@ class MambaSpecDecodeGPUContext: state_conv_widths: torch.Tensor # int32: conv width (0 for temporal states) state_group_indices: torch.Tensor # int32: maps state_idx to group index state_skip_postprocess: torch.Tensor # int32: materializer owns this state + state_skip_precopy: torch.Tensor # int32: scheduler block copy owns this state # DS conv row metadata. Zero keeps the single-region copy path. state_dim_row_count: torch.Tensor # int32: per-block dim row count state_dim_row_stride: torch.Tensor # int64: bytes between rows @@ -872,6 +875,9 @@ class MambaSpecDecodeGPUContext: # Persistent all-layer ReplaySSM descriptors, populated with the cache # addresses on first real forward. None for non-FlashInfer configurations. replayssm: ReplaySSMModelContext | None = None + # V1 can prove on the host that no accepted-token outcome reaches a + # boundary. V2 leaves this true because its sampled count is GPU-resident. + replayssm_materialize_possible: bool = True @classmethod def create( @@ -932,6 +938,9 @@ def create( state_skip_postprocess=torch.zeros( total_states, dtype=torch.int32, device=device ), + state_skip_precopy=torch.zeros( + total_states, dtype=torch.int32, device=device + ), state_dim_row_count=torch.zeros( total_states, dtype=torch.int32, device=device ), @@ -1034,6 +1043,8 @@ def _populate_metadata( block_tables: list[torch.Tensor], ) -> None: idx = 0 + has_replayssm_layer = False + has_baseline_layer = False for group_local_idx, mamba_group_id in enumerate(self.mamba_group_ids): kv_cache_group = kv_cache_config.kv_cache_groups[mamba_group_id] layer_names = kv_cache_group.layer_names @@ -1046,6 +1057,8 @@ def _populate_metadata( getattr(attention, "use_replayssm", False) and attention.mamba_config.backend == MambaBackendEnum.FLASHINFER ) + has_replayssm_layer |= bool(is_flashinfer_replayssm) + has_baseline_layer |= not is_flashinfer_replayssm if len(kv_caches) < len(state_copy_funcs): raise ValueError( f"Expected at least {len(state_copy_funcs)} Mamba state " @@ -1053,6 +1066,7 @@ def _populate_metadata( ) for state_type_idx, copy_func in enumerate(state_copy_funcs): state = kv_caches[state_type_idx] + self.state_skip_precopy[idx] = is_flashinfer_replayssm # Base address self.state_base_addrs[idx] = _reinterpret_u64_as_i64( state.data_ptr() @@ -1133,6 +1147,10 @@ def _populate_metadata( idx += 1 assert idx == self.num_states + if has_replayssm_layer and has_baseline_layer: + raise ValueError( + "mixed FlashInfer ReplaySSM and baseline Mamba layers are unsupported" + ) # Cache per-group block-table base addresses and per-request stride. # `block_tables[i]` is the persistent 2D int32 block-table tensor for @@ -1300,7 +1318,7 @@ def run_fused_precopy( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, - self.state_skip_postprocess, + self.state_skip_precopy, self.state_dim_row_count, self.state_dim_row_stride, idx_mapping, @@ -1583,6 +1601,11 @@ def preprocess_mamba( mamba_state_idx[req_id] = curr_state_idx if fused is not None: fused.state_idx.np[i] = curr_state_idx + if fused.ctx.replayssm is not None: + # A negative source means that this request has no live owner. + # Preserve same-column sources so ordinary decode steps are + # not mistaken for fresh slot ownership. + fused.src_col.np[i] = prev_state_idx if prev_state_idx != -1 and prev_state_idx != curr_state_idx: accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 @@ -1602,12 +1625,20 @@ def preprocess_mamba( req_state, forward_context, ) - input_batch.num_accepted_tokens_cpu[i] = 1 + if fused is None or fused.ctx.replayssm is None: + input_batch.num_accepted_tokens_cpu[i] = 1 if fused is not None: fused.state_idx.copy_to_gpu(num_reqs) fused.src_col.copy_to_gpu(num_reqs) fused.token_bias.copy_to_gpu(num_reqs) + if fused.ctx.replayssm is not None: + fused.ctx.replayssm.reset_new_slots( + idx_mapping=None, + src_cols=fused.src_col.gpu, + dst_cols=fused.state_idx.gpu, + num_reqs=num_reqs, + ) fused.ctx.run_fused_precopy( num_reqs=num_reqs, state_idx_gpu=fused.state_idx.gpu, @@ -1723,6 +1754,8 @@ def postprocess_mamba_align_gpu( idx_mapping=None, query_metadata=ctx.num_scheduled_tokens_buf.gpu, query_metadata_is_cumulative=False, + num_computed_tokens=ctx.num_computed_tokens_buf.gpu, + num_computed_is_post_step=False, num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, live_cols=ctx.mamba_state_idx_buf.gpu, @@ -1730,7 +1763,7 @@ def postprocess_mamba_align_gpu( materialize_token_counts=ctx.materialize_token_counts, num_reqs=num_reqs, ) - if ctx.replayssm.materialize_prefixes: + if ctx.replayssm.materialize_prefixes and ctx.replayssm_materialize_possible: ctx.replayssm.materialize() # ``num_accepted_tokens_out`` is pre-initialized from @@ -1769,7 +1802,6 @@ def stage_postprocess_inputs_to_gpu( assert ctx.num_computed_tokens_buf is not None assert ctx.num_draft_tokens_buf is not None assert ctx.is_prefilling_buf is not None - assert ctx.precopy_src_col_buf is not None scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens num_scheduled = scheduler_output.num_scheduled_tokens @@ -1778,6 +1810,9 @@ def stage_postprocess_inputs_to_gpu( computed_np = ctx.num_computed_tokens_buf.np draft_np = ctx.num_draft_tokens_buf.np prefill_np = ctx.is_prefilling_buf.np + materialize_possible = not ( + ctx.replayssm is not None and ctx.replayssm.materialize_prefixes + ) for i in range(num_reqs): req_id = req_ids[i] state_idx = fixed_live_col @@ -1797,30 +1832,16 @@ def stage_postprocess_inputs_to_gpu( prefill_np[i] = ( requests[req_id].num_computed_tokens < requests[req_id].num_prompt_tokens ) + if not materialize_possible: + running_state_tokens = computed + scheduled - num_draft + max_new_computed = running_state_tokens + num_draft + aligned_max_new_computed = ( + max_new_computed // ctx.block_size + ) * ctx.block_size + materialize_possible = aligned_max_new_computed >= running_state_tokens ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs) ctx.is_prefilling_buf.copy_to_gpu(num_reqs) - if ctx.replayssm is not None: - # Prefix modes use the source/destination columns prepared by - # preprocess_mamba. Mode none has no migration, so passing the live - # column as both endpoints makes that part of the kernel a no-op. - src_cols = ( - ctx.precopy_src_col_buf.gpu - if ctx.replayssm.materialize_prefixes - else ctx.mamba_state_idx_buf.gpu - ) - ctx.replayssm.preprocess( - idx_mapping=None, - query_metadata=ctx.num_scheduled_tokens_buf.gpu, - query_metadata_is_cumulative=False, - num_computed_tokens=ctx.num_computed_tokens_buf.gpu, - is_prefilling=ctx.is_prefilling_buf.gpu, - src_cols=src_cols, - dst_cols=ctx.mamba_state_idx_buf.gpu, - mamba_block_size=ctx.block_size, - num_reqs=num_reqs, - ) - if ctx.replayssm.materialize_prefixes: - ctx.replayssm.materialize() + ctx.replayssm_materialize_possible = materialize_possible diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 344dcad6cbad..f97e852d7a70 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -11,6 +11,7 @@ import torch from vllm.config import CacheConfig, VllmConfig +from vllm.config.mamba import MambaBackendEnum from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers @@ -734,6 +735,73 @@ def copy_kv_cache_blocks_inplace( blocks[dst] = blocks[src] +def get_replayssm_block_copy_tensors( + forward_context: Mapping[str, Any], +) -> list[torch.Tensor]: + """Return FlashInfer-owned ReplaySSM rings and shared cursors. + + The runner's ordinary cache list already covers the two canonical Mamba + roles. This validates the complete five-cache contract per ReplaySSM layer + and returns the three backend-owned rings plus the two group-shared cursor + tensors. ``copy_kv_cache_blocks_inplace`` deduplicates shared storage. + """ + extra_tensors: list[torch.Tensor] = [] + cache_roles = ( + "conv_state", + "ssm_state", + "x_cache", + "dt_cache", + "B_cache", + ) + cursor_roles = ( + ("ring_start", "_replayssm_ring_start"), + ("num_committed", "_replayssm_prev_num_accepted"), + ) + for layer_name, layer in forward_context.items(): + if not getattr(layer, "use_replayssm", False): + continue + mamba_config = getattr(layer, "mamba_config", None) + if getattr(mamba_config, "backend", None) != MambaBackendEnum.FLASHINFER: + continue + + kv_cache = getattr(layer, "kv_cache", ()) + if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) != len(cache_roles): + raise ValueError( + f"FlashInfer ReplaySSM layer {layer_name!r} must expose exactly " + f"{len(cache_roles)} cache roles {cache_roles}; got " + f"{len(kv_cache) if isinstance(kv_cache, (list, tuple)) else 0}" + ) + for role, tensor in zip(cache_roles, kv_cache, strict=True): + if not isinstance(tensor, torch.Tensor) or tensor.ndim == 0: + raise ValueError( + f"FlashInfer ReplaySSM layer {layer_name!r} has invalid " + f"{role} cache" + ) + capacity = kv_cache[0].shape[0] + for role, tensor in zip(cache_roles[1:], kv_cache[1:], strict=True): + if tensor.shape[0] != capacity: + raise ValueError( + f"FlashInfer ReplaySSM layer {layer_name!r} {role} capacity " + f"{tensor.shape[0]} does not match canonical capacity {capacity}" + ) + + extra_tensors.extend(kv_cache[2:5]) + for role, attr in cursor_roles: + cursor = getattr(layer, attr, None) + if ( + not isinstance(cursor, torch.Tensor) + or cursor.ndim != 1 + or cursor.dtype != torch.int32 + or cursor.numel() != capacity + ): + raise ValueError( + f"FlashInfer ReplaySSM layer {layer_name!r} has invalid " + f"{role} cursor for capacity {capacity}" + ) + extra_tensors.append(cursor) + return extra_tensors + + def is_uniform_query_len(num_reqs: int, num_tokens: int, max_query_len: int) -> bool: """Whether every request in the batch has the same query length. From 9d9c728624ab8bcbfbb7304d06883bdc9dda4403 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 13:58:45 +0200 Subject: [PATCH 17/53] test: trim ReplaySSM regression suite Remove temporary and duplicate coverage, reduce the config and engine matrices to distinct contracts, and share the pointer-bit conversion helper. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 13 ---- tests/test_config.py | 42 ++++++------ tests/v1/e2e/test_replayssm_decode.py | 54 +--------------- ...est_replayssm_mtp_compaction_diagnostic.py | 64 ------------------- .../worker/test_mamba_hybrid_model_state.py | 37 ----------- tests/v1/worker/test_mamba_utils.py | 11 +++- .../layers/mamba/mamba_utils.py | 5 ++ .../layers/mamba/ops/ssu_dispatch.py | 6 +- vllm/v1/worker/mamba_utils.py | 6 +- 9 files changed, 37 insertions(+), 201 deletions(-) delete mode 100644 tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 19f3af4d9376..c6040101f9df 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -310,19 +310,6 @@ def test_replayssm_physical_ring_shape( ) -@pytest.mark.parametrize( - ("value", "expected"), - [ - (0, 0), - ((1 << 63) - 1, (1 << 63) - 1), - (1 << 63, -(1 << 63)), - ((1 << 64) - 1, -1), - ], -) -def test_reinterpret_u64_as_i64(value: int, expected: int): - assert ssu_dispatch._reinterpret_u64_as_i64(value) == expected - - def _materialize_mixer(device: str = "cpu") -> Mock: mixer = Mock() mixer.kv_cache = [ diff --git a/tests/test_config.py b/tests/test_config.py index 8f0da3941cbc..accf68aaa4ab 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -86,16 +86,6 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) assert not config.cache_config.use_kda_recoverssm - config.mamba_config.backend = MambaBackendEnum.TRITON - with pytest.raises(ValueError, match="requires --mamba-backend flashinfer"): - VllmConfig.validate_mamba_cached_kernel(config) - - config.mamba_config.backend = MambaBackendEnum.FLASHINFER - config.cache_config.replayssm_buffer_len = 3 - with pytest.raises(ValueError, match="replayssm-buffer-len"): - VllmConfig.validate_mamba_cached_kernel(config) - config.cache_config.replayssm_buffer_len = 16 - config.model_config.architecture = "KimiLinearForCausalLM" config.mamba_config.backend = MambaBackendEnum.TRITON config.parallel_config.pipeline_parallel_size = 2 @@ -201,13 +191,24 @@ def _replayssm_config( ), [ (MambaBackendEnum.TRITON, True, "none", 0, 16, "requires Model Runner V1"), - (MambaBackendEnum.FLASHINFER, True, "none", 0, 16, None), - (MambaBackendEnum.FLASHINFER, False, "align", 0, 16, None), - (MambaBackendEnum.FLASHINFER, False, "all", 0, 16, None), - (MambaBackendEnum.FLASHINFER, False, "align", 3, 16, None), (MambaBackendEnum.FLASHINFER, True, "align", 3, 16, None), - (MambaBackendEnum.FLASHINFER, False, "all", 3, 16, None), - (MambaBackendEnum.FLASHINFER, True, "all", 3, 16, None), + ( + MambaBackendEnum.FLASHINFER, + True, + "align", + 3, + 3, + r"replayssm-buffer-len >= 1 \+ num_speculative_tokens", + ), + ( + MambaBackendEnum.TRITON, + False, + "align", + 3, + 16, + "requires --mamba-backend flashinfer", + ), + (MambaBackendEnum.FLASHINFER, False, "all", 0, 16, None), ( MambaBackendEnum.TRITON, False, @@ -227,13 +228,10 @@ def _replayssm_config( ], ids=[ "triton-v2-rejected", - "flashinfer-v2", - "flashinfer-align", - "flashinfer-all", - "flashinfer-align-spec-v1", "flashinfer-align-spec-v2", - "flashinfer-all-spec-v1", - "flashinfer-all-spec-v2", + "flashinfer-spec-buffer-too-short", + "triton-spec-rejected", + "flashinfer-all", "triton-all-rejected", "flashinfer-buffer-too-long", ], diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 56b357da7b4f..9a851f969d7b 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -36,7 +36,6 @@ # Mamba2 (Nemotron-3) hybrid. MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" MAMBA2_MTP_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4" -MAMBA2_PREFIX_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" MODELS = [ pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), ] @@ -201,37 +200,6 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na ) -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="flashinfer.mamba.checkpointing_ssu not available", -) -@pytest.mark.parametrize("model_name", MODELS) -def test_replayssm_flashinfer_matches_triton_replayssm(vllm_runner, model_name): - # Both backends implement ReplaySSM; compare them directly on V1 because - # Triton ReplaySSM is not supported on Model Runner V2. - common = dict( - max_model_len=1024, - trust_remote_code=True, - enable_prefix_caching=False, - mamba_cache_mode="none", - use_replayssm=True, - replayssm_buffer_len=16, - ) - with vllm_runner(model_name, mamba_backend="triton", **common) as llm: - triton = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5) - with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm: - flashinfer = llm.generate_greedy_logprobs( - PROMPTS, max_tokens=32, num_logprobs=5 - ) - - check_logprobs_close( - outputs_0_lst=triton, - outputs_1_lst=flashinfer, - name_0="replayssm_triton", - name_1="replayssm_flashinfer", - ) - - @pytest.mark.skipif( not HAS_FLASHINFER_CHECKPOINTING_SSU, reason="flashinfer.mamba.checkpointing_ssu not available", @@ -422,30 +390,10 @@ def test_flashinfer_replayssm_prefix_cache_tp1( def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: bool): _check_flashinfer_replayssm_prefix_caching( vllm_runner, - MAMBA2_PREFIX_MODEL, + MAMBA2_MODEL, monkeypatch, mamba_cache_mode="all", - moe_backend="triton", use_ngram=False, use_v2=use_v2, tensor_parallel_size=1, ) - - -@requires_flashinfer_replayssm_materialization -@multi_gpu_test(num_gpus=2) -@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) -def test_flashinfer_replayssm_prefix_cache_v2_tp2( - vllm_runner, - model_name, - monkeypatch: pytest.MonkeyPatch, -): - _check_flashinfer_replayssm_prefix_caching( - vllm_runner, - model_name, - monkeypatch, - mamba_cache_mode="align", - use_ngram=False, - use_v2=True, - tensor_parallel_size=2, - ) diff --git a/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py b/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py deleted file mode 100644 index 1e3caf38e037..000000000000 --- a/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py +++ /dev/null @@ -1,64 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Temporary V2 MTP diagnostic that keeps both request rows live.""" - -import os - -import vllm.envs as envs - -from ...models.utils import check_logprobs_close -from .test_replayssm_decode import PROMPTS - - -def test_replayssm_flashinfer_mtp_v2_without_batch_compaction(vllm_runner, monkeypatch): - model = os.environ["REPLAYSSM_MODEL"] - prompts = [PROMPTS[1], PROMPTS[1]] - common = dict( - max_model_len=1024, - trust_remote_code=True, - enable_prefix_caching=False, - mamba_cache_mode="none", - mamba_backend="flashinfer", - speculative_config={"method": "mtp", "num_speculative_tokens": 3}, - ) - - try: - with monkeypatch.context() as patch: - patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") - envs.disable_envs_cache() - with vllm_runner(model, **common) as llm: - baseline = llm.generate_greedy_logprobs( - prompts, max_tokens=32, num_logprobs=5 - ) - with vllm_runner( - model, use_replayssm=True, replayssm_buffer_len=16, **common - ) as llm: - replay = llm.generate_greedy_logprobs( - prompts, max_tokens=32, num_logprobs=5 - ) - finally: - envs.disable_envs_cache() - - for baseline_output, replay_output in zip(baseline, replay): - baseline_ids = baseline_output[0] - replay_ids = replay_output[0] - matching_prefix = 0 - for baseline_id, replay_id in zip(baseline_ids, replay_ids): - if baseline_id != replay_id: - break - matching_prefix += 1 - print( - "REPLAYSSM_COMPACTION_DIAGNOSTIC", - { - "matching_prefix": matching_prefix, - "baseline": baseline_ids, - "replay": replay_ids, - }, - ) - - check_logprobs_close( - outputs_0_lst=baseline, - outputs_1_lst=replay, - name_0="baseline_mtp_v2_no_compaction", - name_1="replayssm_mtp_v2_no_compaction", - ) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index a52acd409565..5789f0151425 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -4,7 +4,6 @@ from types import SimpleNamespace from unittest.mock import Mock -import numpy as np import pytest import torch @@ -97,42 +96,6 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration( assert kwargs["live_cols"] is state._replayssm_live_cols_gpu -def test_flashinfer_replayssm_preprocess_runs_before_v2_forward() -> None: - state = object.__new__(MambaHybridModelState) - state._needs_prefix_state_migration = False - state._use_flashinfer_replayssm = True - state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool) - state._replayssm_live_cols_gpu = torch.zeros(4, dtype=torch.int32) - state._get_mamba_group_info = Mock(return_value=([0], Mock())) - replayssm = Mock() - replayssm.materialize_prefixes = False - ctx = Mock(replayssm=replayssm, block_size=1024) - state._ensure_mamba_postprocess_ctx = Mock(return_value=ctx) - input_batch = Mock( - num_reqs=2, - is_prefilling_np=np.array([True, False, False, False]), - idx_mapping=torch.tensor([1, 3], dtype=torch.int32), - query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32), - ) - num_computed = torch.tensor([0, 2, 0, 7], dtype=torch.int32) - - state.preprocess_state(input_batch, (), Mock(), num_computed) - - replayssm.preprocess.assert_called_once_with( - idx_mapping=input_batch.idx_mapping, - query_metadata=input_batch.query_start_loc, - query_metadata_is_cumulative=True, - num_computed_tokens=num_computed, - is_prefilling=state._is_prefilling_gpu, - src_cols=state._replayssm_live_cols_gpu, - dst_cols=state._replayssm_live_cols_gpu, - mamba_block_size=1024, - num_reqs=2, - ) - replayssm.materialize.assert_not_called() - ctx.run_fused_precopy.assert_not_called() - - def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: state = RecoverSSMState() metadata = Mock(spec=RecoverSSMMetadata) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index f1dacb293a1b..bb453e50f5dc 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncsByType, + _reinterpret_u64_as_i64, get_conv_copy_spec, get_temporal_copy_spec, ) @@ -27,7 +28,6 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, - _reinterpret_u64_as_i64, batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, @@ -426,7 +426,14 @@ def __getitem__(self, item): def test_reinterpret_u64_as_i64_preserves_pointer_bits(): - ptrs = [1 << 63, (1 << 64) - 1] + ptrs = [ + 0, + 1, + (1 << 63) - 1, + 1 << 63, + (1 << 63) + 1234, + (1 << 64) - 1, + ] ptr_tensor = torch.zeros(len(ptrs), dtype=torch.int64) for idx, ptr in enumerate(ptrs): diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 9a91173abd84..78de733b36ed 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -50,6 +50,11 @@ def is_conv_state_dim_first() -> bool: return get_conv_state_layout() == "DS" +def _reinterpret_u64_as_i64(value: int) -> int: + """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" + return value if value < (1 << 63) else value - (1 << 64) + + class MambaStateDtypeCalculator: @classmethod def linear_attention_state_dtype( diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 3d9d141a59fe..bc90c7462f8c 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -19,6 +19,7 @@ from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.mamba_utils import _reinterpret_u64_as_i64 from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -918,11 +919,6 @@ def selective_state_update_replayssm_flashinfer( ) -def _reinterpret_u64_as_i64(value: int) -> int: - """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" - return value if value < (1 << 63) else value - (1 << 64) - - def _cuda_i64_ptrs(tensors: list[torch.Tensor]) -> torch.Tensor: return torch.tensor( [_reinterpret_u64_as_i64(t.data_ptr()) for t in tensors], diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 64e38b281d79..3f5c52687c2e 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -12,6 +12,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFuncsByType, + _reinterpret_u64_as_i64, get_conv_copy_spec, get_temporal_copy_spec, is_conv_state_dim_first, @@ -187,11 +188,6 @@ def _memcpy_u64_tiled( tl.store(dst_u8 + i + offsets, data, mask=mask) -def _reinterpret_u64_as_i64(value: int) -> int: - """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" - return value if value < (1 << 63) else value - (1 << 64) - - @triton.jit def _copy_mamba_state_block( state_idx, From c202d4b4c4c9fb3292064acdbac27f67f9fd7cc2 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 14:33:46 +0200 Subject: [PATCH 18/53] [Mamba] Strengthen ReplaySSM MTP cache tests Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 50 +++++++++++++++++++ ...est_replayssm_mtp_compaction_diagnostic.py | 48 ++++++++---------- 2 files changed, 72 insertions(+), 26 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 56b357da7b4f..bee2a882a307 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -432,6 +432,56 @@ def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: ) +@requires_flashinfer_replayssm_materialization +@large_gpu_mark(min_gb=40) +def test_flashinfer_replayssm_all_prefix_cache_mtp_v2(vllm_runner, monkeypatch): + common = dict( + max_model_len=8192, + trust_remote_code=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + mamba_cache_mode="all", + mamba_backend="flashinfer", + 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", "1") + envs.disable_envs_cache() + with vllm_runner( + MAMBA2_MTP_MODEL, + use_replayssm=True, + replayssm_buffer_len=16, + **common, + ) as llm: + first_pass = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + first_pass_hits = _prefix_cache_hits(llm) + cached = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + cached_hits = _prefix_cache_hits(llm) + 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() + + assert cached_hits > first_pass_hits + assert draft_count > 0 + check_logprobs_close( + outputs_0_lst=first_pass, + outputs_1_lst=cached, + name_0="replayssm_all_mtp_v2_first_pass", + name_1="replayssm_all_mtp_v2_cached", + ) + + @requires_flashinfer_replayssm_materialization @multi_gpu_test(num_gpus=2) @pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) diff --git a/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py b/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py index 1e3caf38e037..44c5af8ce383 100644 --- a/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py +++ b/tests/v1/e2e/test_replayssm_mtp_compaction_diagnostic.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Temporary V2 MTP diagnostic that keeps both request rows live.""" +"""V2 MTP diagnostic that keeps two identical ReplaySSM rows live.""" import os @@ -26,39 +26,35 @@ def test_replayssm_flashinfer_mtp_v2_without_batch_compaction(vllm_runner, monke with monkeypatch.context() as patch: patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") envs.disable_envs_cache() - with vllm_runner(model, **common) as llm: - baseline = llm.generate_greedy_logprobs( - prompts, max_tokens=32, num_logprobs=5 - ) with vllm_runner( model, use_replayssm=True, replayssm_buffer_len=16, **common ) as llm: - replay = llm.generate_greedy_logprobs( + outputs = llm.generate_greedy_logprobs( prompts, max_tokens=32, num_logprobs=5 ) finally: envs.disable_envs_cache() - for baseline_output, replay_output in zip(baseline, replay): - baseline_ids = baseline_output[0] - replay_ids = replay_output[0] - matching_prefix = 0 - for baseline_id, replay_id in zip(baseline_ids, replay_ids): - if baseline_id != replay_id: - break - matching_prefix += 1 - print( - "REPLAYSSM_COMPACTION_DIAGNOSTIC", - { - "matching_prefix": matching_prefix, - "baseline": baseline_ids, - "replay": replay_ids, - }, - ) + assert len(outputs) == 2 + row_0_ids = outputs[0][0] + row_1_ids = outputs[1][0] + matching_prefix = 0 + for row_0_id, row_1_id in zip(row_0_ids, row_1_ids): + if row_0_id != row_1_id: + break + matching_prefix += 1 + print( + "REPLAYSSM_COMPACTION_DIAGNOSTIC", + { + "matching_prefix": matching_prefix, + "row_0": row_0_ids, + "row_1": row_1_ids, + }, + ) check_logprobs_close( - outputs_0_lst=baseline, - outputs_1_lst=replay, - name_0="baseline_mtp_v2_no_compaction", - name_1="replayssm_mtp_v2_no_compaction", + outputs_0_lst=outputs[:1], + outputs_1_lst=outputs[1:], + name_0="replayssm_mtp_v2_row_0", + name_1="replayssm_mtp_v2_row_1", ) From 25890ec39d4e0722781b33b7a8e94a25012cfc1e Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 14:53:08 +0200 Subject: [PATCH 19/53] [Mamba] Respect MTP prefix-cache support boundary Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 51 --------------------------- 1 file changed, 51 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index ee16c6433f2f..9a851f969d7b 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -397,54 +397,3 @@ def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: use_v2=use_v2, tensor_parallel_size=1, ) - - -@requires_flashinfer_replayssm_materialization -@large_gpu_mark(min_gb=40) -def test_flashinfer_replayssm_all_prefix_cache_mtp_v2(vllm_runner, monkeypatch): - common = dict( - max_model_len=8192, - trust_remote_code=True, - enable_prefix_caching=True, - enable_chunked_prefill=True, - mamba_cache_mode="all", - mamba_backend="flashinfer", - 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", "1") - envs.disable_envs_cache() - with vllm_runner( - MAMBA2_MTP_MODEL, - use_replayssm=True, - replayssm_buffer_len=16, - **common, - ) as llm: - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner - first_pass = llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 - ) - first_pass_hits = _prefix_cache_hits(llm) - cached = llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 - ) - cached_hits = _prefix_cache_hits(llm) - 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() - - assert cached_hits > first_pass_hits - assert draft_count > 0 - check_logprobs_close( - outputs_0_lst=first_pass, - outputs_1_lst=cached, - name_0="replayssm_all_mtp_v2_first_pass", - name_1="replayssm_all_mtp_v2_cached", - ) From 5f6d99f95b920eb5a62bbeea5a8a1a953236dd47 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 15:07:45 +0200 Subject: [PATCH 20/53] [Mamba] Keep ReplaySSM prefix E2E on supported model Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 9a851f969d7b..a949f97ba56d 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -36,6 +36,7 @@ # Mamba2 (Nemotron-3) hybrid. MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16" MAMBA2_MTP_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4" +MAMBA2_PREFIX_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" MODELS = [ pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)), ] @@ -390,9 +391,10 @@ def test_flashinfer_replayssm_prefix_cache_tp1( def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: bool): _check_flashinfer_replayssm_prefix_caching( vllm_runner, - MAMBA2_MODEL, + MAMBA2_PREFIX_MODEL, monkeypatch, mamba_cache_mode="all", + moe_backend="triton", use_ngram=False, use_v2=use_v2, tensor_parallel_size=1, From ab96e0b21493ab89c4d414d6fb83e79233875085 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 15:48:06 +0200 Subject: [PATCH 21/53] [Mamba] Correct ReplaySSM cache lifecycle Restore non-empty prefix validation for canonical state copies and include Triton ReplaySSM rings in block migration. Align final-prefill tracker updates, reject unsupported pipeline parallelism, and remove unnecessary draft deferral and host materialization gating. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- .../kernels/mamba/test_precopy_mamba_align.py | 3 + tests/kernels/mamba/test_ssu_dispatch.py | 52 ++++++++++++++ tests/test_config.py | 9 +++ .../worker/test_gpu_model_runner_v2_eplb.py | 9 ++- .../worker/test_mamba_hybrid_model_state.py | 18 +---- tests/v1/worker/test_mamba_utils.py | 69 +++++++------------ tests/v1/worker/test_utils.py | 17 ++++- vllm/config/cache.py | 3 +- vllm/config/vllm.py | 6 +- .../layers/mamba/ops/ssu_dispatch.py | 12 ++-- vllm/model_executor/models/diffusion_gemma.py | 1 - vllm/v1/worker/gpu/model_runner.py | 10 --- vllm/v1/worker/gpu/model_states/interface.py | 12 ---- .../worker/gpu/model_states/mamba_hybrid.py | 37 ---------- vllm/v1/worker/mamba_utils.py | 20 +----- vllm/v1/worker/utils.py | 19 ++--- 16 files changed, 132 insertions(+), 165 deletions(-) diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py index 22b73c936ece..320fddaae81a 100644 --- a/tests/kernels/mamba/test_precopy_mamba_align.py +++ b/tests/kernels/mamba/test_precopy_mamba_align.py @@ -194,6 +194,7 @@ def test_precopy_matches_v1_copy_specs( base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( convs, ssms, device, conv_state_dim_first ) + state_skip_precopy = torch.zeros(NUM_LAYERS * 2, dtype=torch.int32, device=device) bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles) @@ -209,6 +210,7 @@ def test_precopy_matches_v1_copy_specs( inner, width, group, + state_skip_precopy, drc, drs, idx_mapping if has_idx_mapping else None, @@ -260,6 +262,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): diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index c6040101f9df..5333f0e74a65 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -756,6 +756,58 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slots(monkeypatch): assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 0 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") +@pytest.mark.parametrize( + ( + "query_metadata", + "query_metadata_is_cumulative", + "num_computed_is_post_step", + "num_computed", + ), + [([1, 0], False, False, 7), ([0, 1], True, True, 8)], + ids=["v1", "v2"], +) +def test_modelwide_replayssm_single_token_final_prefill_commits_as_decode( + query_metadata: list[int], + query_metadata_is_cumulative: bool, + num_computed_is_post_step: bool, + num_computed: int, +): + groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() + for mixers, source_slot in zip(groups, (1, 4)): + mixers[0]._replayssm_ring_start[source_slot] = 7 + mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 + + ctx = ReplaySSMModelContext.create( + config, + [0, 1], + forward_context, + block_tables, + max_num_reqs=2, + ) + assert ctx is not None + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor(query_metadata, dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=query_metadata_is_cumulative, + num_computed_tokens=torch.tensor( + [num_computed, 0], dtype=torch.int32, device="cuda" + ), + num_computed_is_post_step=num_computed_is_post_step, + num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), + is_prefilling=torch.tensor([True, False], device="cuda"), + live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), + materialize_dst_cols=torch.full((2,), -1, dtype=torch.int32, device="cuda"), + materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + torch.accelerator.synchronize() + + for mixers, source_slot in zip(groups, (1, 4)): + assert mixers[0]._replayssm_ring_start[source_slot].item() == 7 + assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 6 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() diff --git a/tests/test_config.py b/tests/test_config.py index accf68aaa4ab..4b6634bbd064 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -175,6 +175,7 @@ def _replayssm_config( model_config=None, num_speculative_tokens=0, mamba_config=SimpleNamespace(backend=backend), + parallel_config=SimpleNamespace(pipeline_parallel_size=1), use_v2_model_runner=use_v2_model_runner, kv_transfer_config=None, ) @@ -259,6 +260,14 @@ def test_replayssm_config_matrix( VllmConfig.validate_mamba_cached_kernel(config) +def test_replayssm_rejects_pipeline_parallelism(): + config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) + config.parallel_config.pipeline_parallel_size = 2 + + with pytest.raises(ValueError, match="pipeline_parallel_size=1"): + VllmConfig.validate_mamba_cached_kernel(config) + + def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index e20e084c9dd0..104dc7b1fca8 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -206,7 +206,7 @@ def fake_receive(*args, **kwargs): assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] -def test_v2_sample_tokens_publishes_state_after_drafting(monkeypatch): +def test_v2_sample_tokens_postprocesses_state_before_drafting(monkeypatch): events: list[Any] = [] runner = _make_runner() input_batch = SimpleNamespace( @@ -262,8 +262,8 @@ def test_v2_sample_tokens_publishes_state_after_drafting(monkeypatch): draft_tokens=torch.zeros((1, 1), dtype=torch.int64), ) - def postprocess_state(*_, defer_after_drafting=False): - events.append(("postprocess", defer_after_drafting)) + def postprocess_state(*_): + events.append("postprocess") def propose(*_, **__): events.append("draft") @@ -272,7 +272,6 @@ def propose(*_, **__): runner.speculator = SimpleNamespace(supports_mm_inputs=False, propose=propose) runner.model_state = SimpleNamespace( postprocess_state=postprocess_state, - postprocess_state_after_drafting=lambda *_: events.append("publish"), ) monkeypatch.setattr(mrv2, "AsyncOutput", lambda **_: object()) monkeypatch.setattr(mrv2, "post_update", lambda *_: None) @@ -280,7 +279,7 @@ def propose(*_, **__): mrv2.GPUModelRunner.sample_tokens(runner, None) - assert events == [("postprocess", True), "draft", "publish"] + assert events == ["postprocess", "draft"] def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 5789f0151425..ba84d6016a9b 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -41,10 +41,7 @@ def test_postprocess_state_scalar_with_int32_mapping( @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") -@pytest.mark.parametrize("defer_after_drafting", [False, True]) -def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration( - defer_after_drafting: bool, -) -> None: +def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: state = object.__new__(MambaHybridModelState) state._align_mode = False state._needs_prefix_state_migration = False @@ -73,22 +70,9 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration( num_sampled, num_computed_tokens=num_computed, query_start_loc=query_start_loc, - defer_after_drafting=defer_after_drafting, ) ctx.run_fused_postprocess_align.assert_not_called() - if defer_after_drafting: - replayssm.postprocess.assert_not_called() - # The drafter reuses this request-state buffer. The post-draft commit - # must consume the sampler-owned batch tensor instead of the clobbered - # request-state value. - state.num_accepted_tokens_gpu[2] = 1 - state.postprocess_state_after_drafting( - idx_mapping, - num_sampled, - num_computed_tokens=num_computed, - query_start_loc=query_start_loc, - ) assert replayssm.postprocess.call_count == 1 kwargs = replayssm.postprocess.call_args.kwargs assert state.num_accepted_tokens_gpu.tolist() == [1, 1, 2, 1] diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index bb453e50f5dc..794fd115dcf3 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -268,17 +268,7 @@ def test_preprocess_mamba_preserves_live_replayssm_state( assert align_ctx.precopy_src_col_buf.np[0] == 0 -@pytest.mark.parametrize( - ("materialize_possible", "expected_order"), - [ - (True, ["copy", "postprocess", "materialize"]), - (False, ["copy", "postprocess"]), - ], -) -def test_postprocess_mamba_align_gates_materialization( - materialize_possible: bool, - expected_order: list[str], -): +def test_postprocess_mamba_align_materializes_prefixes(): order: list[str] = [] ctx = MagicMock() ctx.is_initialized = True @@ -294,7 +284,6 @@ def test_postprocess_mamba_align_gates_materialization( ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") ctx.replayssm = MagicMock() ctx.replayssm.materialize_prefixes = True - ctx.replayssm_materialize_possible = materialize_possible ctx.replayssm.postprocess.side_effect = lambda **kwargs: order.append("postprocess") ctx.replayssm.materialize.side_effect = lambda: order.append("materialize") block_table = MagicMock() @@ -317,7 +306,7 @@ def test_postprocess_mamba_align_gates_materialization( run_prefix_state_migration=True, ) - assert order == expected_order + assert order == ["copy", "postprocess", "materialize"] assert accepted_cpu.tolist() == [3] @@ -837,7 +826,7 @@ def test_mamba_groups_support_different_state_specs(): assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] -def test_mamba_copy_funcs_ignore_replayssm_state_tensors(): +def test_mamba_copy_funcs_accept_nonempty_canonical_prefix(): replayssm_spec = MambaSpec( block_size=16, shapes=((4, 4), (2, 4, 4)), @@ -848,20 +837,37 @@ def test_mamba_copy_funcs_ignore_replayssm_state_tensors(): mamba_cache_mode="align", ) - validate_mamba_state_copy_funcs({replayssm_spec: [0]}, _COPY_FUNCS) + for valid_funcs in ((get_conv_copy_spec,), _DEFAULT_COPY_FUNCS): + copy_funcs = { + **_COPY_FUNCS, + MambaAttentionBackendEnum.MAMBA2: valid_funcs, + } + validate_mamba_state_copy_funcs({replayssm_spec: [0]}, copy_funcs) for invalid_funcs in ( - (get_conv_copy_spec,), + (), (*_DEFAULT_COPY_FUNCS, get_temporal_copy_spec), ): invalid_copy_funcs = { **_COPY_FUNCS, MambaAttentionBackendEnum.MAMBA2: invalid_funcs, } - with pytest.raises(AssertionError, match="expects 2 state copy funcs"): + with pytest.raises(AssertionError, match="non-empty copyable prefix"): validate_mamba_state_copy_funcs({replayssm_spec: [0]}, invalid_copy_funcs) +def test_gdn_copy_funcs_cover_copyable_prefix_only(): + gdn_spec = MambaSpec( + block_size=16, + shapes=((4, 4), (2, 4, 4), (2, 4), (2, 4)), + dtypes=(torch.float16,) * 4, + mamba_type=MambaAttentionBackendEnum.GDN_ATTN, + mamba_cache_mode="none", + ) + + validate_mamba_state_copy_funcs({gdn_spec: [0]}, _COPY_FUNCS) + + def test_mamba_groups_support_mixed_specs_in_uniform_group(): gdn_spec = MambaSpec( block_size=16, @@ -1050,35 +1056,6 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ) ctx.replayssm.reset_new_slots.assert_not_called() ctx.replayssm.materialize.assert_not_called() - assert ctx.replayssm_materialize_possible - - -def test_stage_postprocess_inputs_skips_impossible_materialization(): - device = torch.device("cpu") - ctx = _make_staging_ctx(max_num_reqs=2, device=device) - ctx.block_size = 4 - ctx.replayssm = MagicMock(materialize_prefixes=True) - scheduler_output = _make_postprocess_scheduler_output( - req_ids=["req_a"], - num_scheduled_tokens={"req_a": 1}, - ) - requests = _make_requests( - ["req_a"], - [10], - [[0]], - num_prompt_tokens=[10], - ) - - stage_postprocess_inputs_to_gpu( - ctx, - scheduler_output, - ["req_a"], - 1, - requests, - {"req_a": 2}, - ) - - assert not ctx.replayssm_materialize_possible def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 80606ad4ce85..9b7b22fbc574 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -20,10 +20,10 @@ class _TestReplaySSMMixer(MambaMixer2): _state_shapes = ((2,), (3,)) _state_dtypes = (torch.float32, torch.float32) - def __init__(self): + def __init__(self, backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER) -> None: torch.nn.Module.__init__(self) self.use_replayssm = True - self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER) + self.mamba_config = MambaConfig(backend=backend) self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) @@ -152,6 +152,19 @@ def test_replayssm_block_copy_validates_exact_cache_roles(): get_replayssm_block_copy_tensors({"layers.0.mixer": mixer}) +def test_replayssm_block_copy_includes_triton_rings_without_trackers(): + mixer = _TestReplaySSMMixer(MambaBackendEnum.TRITON) + mixer.kv_cache = tuple(torch.zeros(4, 1) for _ in range(5)) + + tensors = get_replayssm_block_copy_tensors({"layers.0.mixer": mixer}) + + assert len(tensors) == 3 + assert all( + actual is expected + for actual, expected in zip(tensors, mixer.kv_cache[2:5], strict=True) + ) + + def test_bind_kv_cache(default_vllm_config): from vllm.model_executor.layers.attention import Attention diff --git a/vllm/config/cache.py b/vllm/config/cache.py index fdb688ac4404..830025628082 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -205,7 +205,8 @@ class CacheConfig: Triton supports 'none' and 'align' on Model Runner V1. FlashInfer supports 'none', 'align', and 'all' on Model Runner V1 and V2. Mamba2 speculative decoding requires FlashInfer. With prefix caching enabled, Triton supports - 'align'; FlashInfer supports 'align' and 'all'.""" + 'align'; FlashInfer supports 'align' and 'all'. Pipeline parallelism is not + supported.""" 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 c85b0ae1055f..27e352e55ed0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2876,6 +2876,8 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": "--use-replayssm is not supported for architecture " f"{self.model_config.architecture!r}" ) + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError("ReplaySSM currently requires pipeline_parallel_size=1") if ( self.mamba_config.backend == MambaBackendEnum.FLASHINFER and self.cache_config.replayssm_buffer_len > 16 @@ -2901,10 +2903,6 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": 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" - ) if self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError("RecoverSSM requires --mamba-backend triton") elif use_mamba_replayssm_spec: diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index bc90c7462f8c..e620da3dd5fa 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -221,11 +221,15 @@ def _postprocess_replayssm_kernel( else: query_len = tl.load(query_metadata + batch_idx) + computed = tl.load(num_computed_tokens + req_idx) + computed_before = tl.where( + NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed + ) + # Mamba attention runs a one-token final prefill chunk with prior state as + # decode. Commit the same transition here instead of resetting its cursors. + prefilling = prefilling & ((query_len != 1) | (computed_before <= 0)) + if prefilling: - computed = tl.load(num_computed_tokens + req_idx) - computed_before = tl.where( - NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed - ) computed_after = computed_before + query_len first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) last_col = tl.maximum( diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index b12871b4356b..52872f310445 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -981,7 +981,6 @@ def postprocess_state( num_computed_tokens=None, query_start_loc=None, is_prefilling=None, - defer_after_drafting=False, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b2cf6eeef75b..5e3f90eb79ed 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1503,7 +1503,6 @@ def postprocess_sampled( num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, - defer_after_drafting: bool = False, ) -> None: # Update the number of computed tokens. if self.is_last_pp_rank: @@ -1530,7 +1529,6 @@ def postprocess_sampled( self.req_states.num_computed_tokens.gpu, query_start_loc, is_prefilling, - defer_after_drafting=defer_after_drafting, ) def _merge_ec_connector_no_forward( @@ -1970,7 +1968,6 @@ def sample_tokens( num_sampled, num_rejected, input_batch.query_start_loc, - defer_after_drafting=self.speculator is not None, ) if self.speculator is not None: @@ -2005,13 +2002,6 @@ def sample_tokens( self.speculator.draft_token_confidence_probs, input_batch ) - self.model_state.postprocess_state_after_drafting( - input_batch.idx_mapping, - num_sampled, - self.req_states.num_computed_tokens.gpu, - input_batch.query_start_loc, - ) - if self.num_speculative_steps > 0: # Spec-decode and diffusion LLMs both use draft tokens but the latter does # not have a speculator (i.e. self.speculator is None) diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 9638e05fb4ed..0c5002b65398 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -154,21 +154,9 @@ def postprocess_state( num_computed_tokens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, - defer_after_drafting: bool = False, ) -> None: return None - def postprocess_state_after_drafting( - self, - idx_mapping: torch.Tensor, - num_sampled: torch.Tensor, - num_computed_tokens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - is_prefilling: torch.Tensor | None = None, - ) -> None: - """Publish state that must remain hidden from the current draft pass.""" - return None - @abstractmethod def prepare_inputs_embeds( self, diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 75982a1b2897..89bde631c36c 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -382,7 +382,6 @@ def postprocess_state( num_computed_tokens: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, is_prefilling: torch.Tensor | None = None, - defer_after_drafting: bool = False, ) -> None: # 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. @@ -431,42 +430,6 @@ def postprocess_state( idx_mapping, ) - if not defer_after_drafting: - self._publish_flashinfer_replayssm( - idx_mapping, - num_computed_tokens, - query_start_loc, - is_prefilling, - ) - - def postprocess_state_after_drafting( - self, - idx_mapping: torch.Tensor, - num_sampled: torch.Tensor | int, - num_computed_tokens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - is_prefilling: torch.Tensor | None = None, - ) -> None: - """Publish ReplaySSM trackers after every forward in this step. - - MTP drafting reuses the tracker view consumed by the target forward. - The newly accepted transition belongs to the next target step, so it - must not become visible until the current draft pass has completed. - """ - num_reqs = idx_mapping.shape[0] - if ( - num_reqs - and self._use_flashinfer_replayssm - and not self._needs_prefix_state_migration - and not isinstance(num_sampled, int) - ): - # The MTP drafter reuses this request-state buffer after target - # acceptance was first scattered. Restore it before publication. - _scatter_num_accepted_kernel[(num_reqs,)]( - idx_mapping, - num_sampled, - self.num_accepted_tokens_gpu, - ) self._publish_flashinfer_replayssm( idx_mapping, num_computed_tokens, diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 3f5c52687c2e..cb50b6136373 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -741,10 +741,10 @@ def validate_mamba_state_copy_funcs( f"missing state copy funcs for {mamba_spec.mamba_type}" ) state_copy_funcs = copy_funcs[mamba_spec.mamba_type] - assert len(state_copy_funcs) == len(mamba_spec.shapes), ( + assert 0 < len(state_copy_funcs) <= len(mamba_spec.shapes), ( f"{mamba_spec.mamba_type} expects {len(mamba_spec.shapes)} state copy " "funcs for its canonical state tensors, but provides " - f"{len(state_copy_funcs)}" + f"{len(state_copy_funcs)}; expected a non-empty copyable prefix" ) @@ -871,9 +871,6 @@ class MambaSpecDecodeGPUContext: # Persistent all-layer ReplaySSM descriptors, populated with the cache # addresses on first real forward. None for non-FlashInfer configurations. replayssm: ReplaySSMModelContext | None = None - # V1 can prove on the host that no accepted-token outcome reaches a - # boundary. V2 leaves this true because its sampled count is GPU-resident. - replayssm_materialize_possible: bool = True @classmethod def create( @@ -1759,7 +1756,7 @@ def postprocess_mamba_align_gpu( materialize_token_counts=ctx.materialize_token_counts, num_reqs=num_reqs, ) - if ctx.replayssm.materialize_prefixes and ctx.replayssm_materialize_possible: + if ctx.replayssm.materialize_prefixes: ctx.replayssm.materialize() # ``num_accepted_tokens_out`` is pre-initialized from @@ -1806,9 +1803,6 @@ def stage_postprocess_inputs_to_gpu( computed_np = ctx.num_computed_tokens_buf.np draft_np = ctx.num_draft_tokens_buf.np prefill_np = ctx.is_prefilling_buf.np - materialize_possible = not ( - ctx.replayssm is not None and ctx.replayssm.materialize_prefixes - ) for i in range(num_reqs): req_id = req_ids[i] state_idx = fixed_live_col @@ -1828,16 +1822,8 @@ def stage_postprocess_inputs_to_gpu( prefill_np[i] = ( requests[req_id].num_computed_tokens < requests[req_id].num_prompt_tokens ) - if not materialize_possible: - running_state_tokens = computed + scheduled - num_draft - max_new_computed = running_state_tokens + num_draft - aligned_max_new_computed = ( - max_new_computed // ctx.block_size - ) * ctx.block_size - materialize_possible = aligned_max_new_computed >= running_state_tokens ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs) ctx.is_prefilling_buf.copy_to_gpu(num_reqs) - ctx.replayssm_materialize_possible = materialize_possible diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index f97e852d7a70..cbfc5c82828a 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -738,12 +738,13 @@ def copy_kv_cache_blocks_inplace( def get_replayssm_block_copy_tensors( forward_context: Mapping[str, Any], ) -> list[torch.Tensor]: - """Return FlashInfer-owned ReplaySSM rings and shared cursors. + """Return ReplaySSM rings and any backend-owned shared cursors. The runner's ordinary cache list already covers the two canonical Mamba roles. This validates the complete five-cache contract per ReplaySSM layer - and returns the three backend-owned rings plus the two group-shared cursor - tensors. ``copy_kv_cache_blocks_inplace`` deduplicates shared storage. + and returns the three backend-owned rings. FlashInfer additionally owns two + group-shared cursor tensors. ``copy_kv_cache_blocks_inplace`` deduplicates + shared storage. """ extra_tensors: list[torch.Tensor] = [] cache_roles = ( @@ -761,31 +762,31 @@ def get_replayssm_block_copy_tensors( if not getattr(layer, "use_replayssm", False): continue mamba_config = getattr(layer, "mamba_config", None) - if getattr(mamba_config, "backend", None) != MambaBackendEnum.FLASHINFER: - continue + backend = getattr(mamba_config, "backend", None) kv_cache = getattr(layer, "kv_cache", ()) if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) != len(cache_roles): raise ValueError( - f"FlashInfer ReplaySSM layer {layer_name!r} must expose exactly " + f"ReplaySSM layer {layer_name!r} must expose exactly " f"{len(cache_roles)} cache roles {cache_roles}; got " f"{len(kv_cache) if isinstance(kv_cache, (list, tuple)) else 0}" ) for role, tensor in zip(cache_roles, kv_cache, strict=True): if not isinstance(tensor, torch.Tensor) or tensor.ndim == 0: raise ValueError( - f"FlashInfer ReplaySSM layer {layer_name!r} has invalid " - f"{role} cache" + f"ReplaySSM layer {layer_name!r} has invalid {role} cache" ) capacity = kv_cache[0].shape[0] for role, tensor in zip(cache_roles[1:], kv_cache[1:], strict=True): if tensor.shape[0] != capacity: raise ValueError( - f"FlashInfer ReplaySSM layer {layer_name!r} {role} capacity " + f"ReplaySSM layer {layer_name!r} {role} capacity " f"{tensor.shape[0]} does not match canonical capacity {capacity}" ) extra_tensors.extend(kv_cache[2:5]) + if backend != MambaBackendEnum.FLASHINFER: + continue for role, attr in cursor_roles: cursor = getattr(layer, attr, None) if ( From 7a2c54093903033aad736464ce55de30386ee530 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 17:09:53 +0200 Subject: [PATCH 22/53] [Mamba] Simplify ReplaySSM prefix maintenance Remove orphaned PP state snapshots, per-state skip metadata, and per-step cache validation. Give ReplaySSM its own compact materialization plan while preserving canonical boundary arithmetic and accepted-count ownership across V1 and V2. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- .../kernels/mamba/test_precopy_mamba_align.py | 2 - tests/kernels/mamba/test_ssu_dispatch.py | 57 +++--- tests/v1/e2e/test_replayssm_decode.py | 51 ++++++ tests/v1/worker/test_gpu_model_runner.py | 111 ++++++++++++ .../worker/test_gpu_model_runner_v2_eplb.py | 22 +-- .../worker/test_mamba_hybrid_model_state.py | 40 ++++- tests/v1/worker/test_mamba_utils.py | 61 ++++--- .../layers/mamba/ops/ssu_dispatch.py | 165 +++++++++--------- vllm/v1/worker/gpu/model_runner.py | 32 ++-- .../worker/gpu/model_states/mamba_hybrid.py | 17 +- vllm/v1/worker/gpu/pp_utils.py | 15 -- vllm/v1/worker/gpu_model_runner.py | 28 +-- vllm/v1/worker/mamba_utils.py | 117 ++++--------- 13 files changed, 421 insertions(+), 297 deletions(-) diff --git a/tests/kernels/mamba/test_precopy_mamba_align.py b/tests/kernels/mamba/test_precopy_mamba_align.py index 320fddaae81a..160457d033d6 100644 --- a/tests/kernels/mamba/test_precopy_mamba_align.py +++ b/tests/kernels/mamba/test_precopy_mamba_align.py @@ -194,7 +194,6 @@ def test_precopy_matches_v1_copy_specs( base, blk_stride, elem, inner, width, group, drc, drs = _build_meta( convs, ssms, device, conv_state_dim_first ) - state_skip_precopy = torch.zeros(NUM_LAYERS * 2, dtype=torch.int32, device=device) bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device) idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) grid = (num_reqs, NUM_LAYERS * 2, temporal_tiles) @@ -210,7 +209,6 @@ def test_precopy_matches_v1_copy_specs( inner, width, group, - state_skip_precopy, drc, drs, idx_mapping if has_idx_mapping else None, diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 5333f0e74a65..1860c9747f1f 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -424,8 +424,6 @@ def test_modelwide_replayssm_none_commits_trackers_without_materialization( accepted = torch.ones(2, dtype=torch.int32, device="cuda") is_prefilling = torch.zeros(2, dtype=torch.bool, device="cuda") live_cols = torch.zeros(2, dtype=torch.int32, device="cuda") - no_materialize = torch.full((2,), -1, dtype=torch.int32, device="cuda") - materialize_counts = torch.zeros(2, dtype=torch.int32, device="cuda") def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None: query_len[0] = scheduled @@ -440,8 +438,6 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None num_accepted_tokens=accepted, is_prefilling=is_prefilling, live_cols=live_cols, - materialize_dst_cols=no_materialize, - materialize_token_counts=materialize_counts, num_reqs=1, ) num_computed[0] += num_accepted @@ -502,13 +498,11 @@ def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) ctx.materialize() @@ -568,13 +562,11 @@ def test_modelwide_replayssm_materialization_uses_independent_group_mappings( idx_mapping=None, query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([6, 0], dtype=torch.int32, device="cuda"), num_computed_is_post_step=False, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) ctx.materialize() @@ -612,13 +604,11 @@ def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatc idx_mapping=None, query_metadata=torch.tensor([1, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([0, 6], dtype=torch.int32, device="cuda"), num_computed_is_post_step=False, num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([-1, 1], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.tensor([0, 1], dtype=torch.int32, device="cuda"), num_reqs=2, ) ctx.materialize() @@ -629,6 +619,27 @@ def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatc assert group_ctx.plan_flush_count.tolist() == [-1, 1] assert group_ctx.active_request_indices.tolist() == [1, -1] + # A shorter batch must clear both the compacted active tail and every stale + # source/destination slot left by the prior materialization. + ctx.postprocess( + idx_mapping=None, + query_metadata=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), + query_metadata_is_cumulative=False, + num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_is_post_step=False, + num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), + is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), + live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_reqs=1, + ) + torch.accelerator.synchronize() + + for group_ctx in ctx.groups: + assert group_ctx.plan_flush_count.tolist() == [-1, -1] + assert group_ctx.active_request_indices.tolist() == [-1, -1] + assert torch.all(group_ctx.src_slots == NULL_BLOCK_ID) + assert torch.all(group_ctx.dst_slots == NULL_BLOCK_ID) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch): @@ -651,13 +662,11 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch idx_mapping=torch.tensor([1], dtype=torch.int32, device="cuda"), query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=True, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([0, 8], dtype=torch.int32, device="cuda"), num_computed_is_post_step=True, num_accepted_tokens=torch.tensor([1, 3], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([False, False], device="cuda"), live_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) ctx.materialize() @@ -666,7 +675,7 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch assert kernel.call_count == 2 for group_ctx in ctx.groups: assert group_ctx.plan_ring_start.tolist() == [15, 0] - assert group_ctx.plan_flush_count.tolist() == [1, -1] + assert group_ctx.plan_flush_count.tolist() == [3, -1] for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): assert mixers[0]._replayssm_ring_start[source_slot].item() == 15 assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 3 @@ -704,8 +713,6 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.ones(2, dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.ones(2, dtype=torch.int32, device="cuda"), num_reqs=1, ) torch.accelerator.synchronize() @@ -743,8 +750,6 @@ def test_modelwide_replayssm_postprocess_resets_prefill_slots(monkeypatch): num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), live_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.full((2,), -1, dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), num_reqs=1, ) torch.accelerator.synchronize() @@ -797,8 +802,6 @@ def test_modelwide_replayssm_single_token_final_prefill_commits_as_decode( num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.full((2,), -1, dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), num_reqs=1, ) torch.accelerator.synchronize() @@ -829,15 +832,13 @@ def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): assert ctx is not None ctx.postprocess( idx_mapping=None, - query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), + query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), num_computed_is_post_step=False, num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), is_prefilling=torch.tensor([True, False], device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.zeros(2, dtype=torch.int32, device="cuda"), num_reqs=1, ) ctx.materialize() @@ -910,13 +911,11 @@ def test_modelwide_replayssm_postprocess_materializes_in_place(monkeypatch): idx_mapping=None, query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_computed_is_post_step=False, num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_dst_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), num_reqs=1, ) ctx.materialize() diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index a949f97ba56d..561d45e4008a 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -399,3 +399,54 @@ def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: use_v2=use_v2, tensor_parallel_size=1, ) + + +@requires_flashinfer_replayssm_materialization +@large_gpu_mark(min_gb=40) +def test_flashinfer_replayssm_all_prefix_cache_mtp_v2(vllm_runner, monkeypatch): + common = dict( + max_model_len=8192, + trust_remote_code=True, + enable_prefix_caching=True, + enable_chunked_prefill=True, + mamba_cache_mode="all", + mamba_backend="flashinfer", + 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", "1") + envs.disable_envs_cache() + with vllm_runner( + MAMBA2_MTP_MODEL, + use_replayssm=True, + replayssm_buffer_len=16, + **common, + ) as llm: + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner + first_pass = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + first_pass_hits = _prefix_cache_hits(llm) + cached = llm.generate_greedy_logprobs( + PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + ) + cached_hits = _prefix_cache_hits(llm) + 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() + + assert cached_hits > first_pass_hits + assert draft_count > 0 + check_logprobs_close( + outputs_0_lst=first_pass, + outputs_1_lst=cached, + name_0="replayssm_all_mtp_v2_first_pass", + name_1="replayssm_all_mtp_v2_cached", + ) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 99331e7b354e..f610284a65f6 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1429,6 +1429,117 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch) assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16 +def test_replayssm_stp_keeps_accepted_one_without_cpu_copy(monkeypatch): + runner = object.__new__(GPUModelRunner) + runner._use_flashinfer_replayssm = True + runner._needs_prefix_state_migration = False + runner.num_spec_tokens = 0 + runner.speculative_config = None + runner.model_config = SimpleNamespace(is_hybrid=True) + runner.num_accepted_tokens = SimpleNamespace(gpu=torch.ones(2, dtype=torch.int32)) + accepted_cpu = torch.full((2,), 9, dtype=torch.int32) + runner.input_batch = SimpleNamespace(num_accepted_tokens_cpu_tensor=accepted_cpu) + runner.kv_cache_config = Mock() + runner.cache_config = SimpleNamespace(mamba_cache_mode="none") + runner.compilation_config = SimpleNamespace(static_forward_context={}) + runner._get_mamba_bufs = Mock(return_value=Mock()) + runner._get_mamba_state_copy_funcs = Mock(return_value={}) + runner.num_accepted_tokens_event = None + postprocess = Mock() + monkeypatch.setattr( + gpu_model_runner_module.mamba_utils, + "postprocess_mamba_align_gpu", + postprocess, + ) + + runner._update_states_after_model_execute( + torch.tensor([[42], [-1]], dtype=torch.int64), Mock() + ) + + assert runner.num_accepted_tokens.gpu.tolist() == [1, 1] + assert accepted_cpu.tolist() == [9, 9] + assert postprocess.call_args.kwargs["num_accepted_tokens_cpu_tensor"] is None + + +def test_v1_caches_replayssm_block_copy_tensors_after_binding(monkeypatch): + runner = object.__new__(GPUModelRunner) + runner.device = torch.device("cpu") + runner.cache_config = SimpleNamespace(get_resolved_kv_cache_layout=Mock()) + runner.shared_kv_cache_layers = {} + runner.model_config = SimpleNamespace(hf_config=SimpleNamespace(model_type="mamba")) + runner.compilation_config = SimpleNamespace(static_forward_context={}) + runner.kv_caches = [] + events = [] + extra = torch.empty(1) + monkeypatch.setattr( + gpu_model_runner_module, "allocate_kv_cache", lambda *_args, **_kwargs: {} + ) + monkeypatch.setattr( + gpu_model_runner_module, + "allocate_replayssm_caches", + lambda *_args, **_kwargs: {}, + ) + monkeypatch.setattr( + gpu_model_runner_module, + "bind_kv_cache", + lambda *_args, **_kwargs: events.append("bind"), + ) + + def get_extra(_context): + assert events == ["bind"] + events.append("validate") + return [extra] + + monkeypatch.setattr( + gpu_model_runner_module, "get_replayssm_block_copy_tensors", get_extra + ) + + runner.initialize_kv_cache_tensors( + SimpleNamespace(kv_cache_groups=[]), kernel_block_sizes=[] + ) + + assert events == ["bind", "validate"] + assert len(runner.replayssm_block_copy_tensors) == 1 + assert runner.replayssm_block_copy_tensors[0] is extra + + +def test_v2_block_copy_reuses_cached_replayssm_tensors(monkeypatch): + from vllm.v1.worker.gpu import model_runner as v2_model_runner_module + + runner = object.__new__(v2_model_runner_module.GPUModelRunner) + runner.req_states = SimpleNamespace( + num_computed_tokens_np=np.zeros(1, dtype=np.int32), + prefill_len=SimpleNamespace(np=np.zeros(1, dtype=np.int32)), + num_computed_prefill_tokens=np.zeros(1, dtype=np.int32), + ) + runner.block_tables = Mock() + runner.kv_block_zeroer = Mock() + canonical = torch.empty(1) + extra = torch.empty(1) + runner.kv_caches = [canonical] + runner.replayssm_block_copy_tensors = [extra] + runner.kv_cache_config = SimpleNamespace(num_blocks=4) + scheduler_output = SchedulerOutput.make_empty() + scheduler_output.kv_cache_block_copies = [Mock()] + copy_blocks = Mock() + monkeypatch.setattr( + v2_model_runner_module, "copy_kv_cache_blocks_inplace", copy_blocks + ) + monkeypatch.setattr( + v2_model_runner_module, + "get_replayssm_block_copy_tensors", + Mock(side_effect=AssertionError("hot path must use cached tensors")), + ) + + runner.update_requests(scheduler_output) + + assert copy_blocks.call_count == 1 + copied_tensors = copy_blocks.call_args.args[0] + assert len(copied_tensors) == 2 + assert copied_tensors[0] is canonical + assert copied_tensors[1] is extra + + def test_v2_runner_snapshots_late_interleave_adjustment(monkeypatch): from vllm.v1.worker.gpu import model_runner as v2_model_runner_module diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 104dc7b1fca8..ccc4740ca787 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -282,9 +282,7 @@ def propose(*_, **__): assert events == ["postprocess", "draft"] -def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( - monkeypatch, -): +def test_v2_sample_tokens_pp_mixed_batch_uses_ordinary_postprocess(monkeypatch): events = [] runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) idx_mapping = torch.tensor([3, 7], dtype=torch.int64) @@ -293,8 +291,6 @@ def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( num_reqs=2, idx_mapping=idx_mapping, idx_mapping_np=np.array([3, 7], dtype=np.intp), - # Row 0 finishes prefill and will be processed from the deferred PP - # receive. Row 1 is a non-final chunk and must be published now. num_computed_tokens_np=np.array([3, 2], dtype=np.int32), prefill_len_np=np.array([4, 6], dtype=np.int32), num_scheduled_tokens=np.array([1, 1], dtype=np.int32), @@ -311,10 +307,6 @@ def test_v2_sample_tokens_pp_mixed_batch_only_postprocesses_prefill_rows( ec_connector_output=None, routed_experts=None, ) - num_computed_tokens = torch.zeros(8, dtype=torch.int32) - runner.req_states = SimpleNamespace( - num_computed_tokens=SimpleNamespace(gpu=num_computed_tokens) - ) postprocess_args = [] runner.model_state = SimpleNamespace( postprocess_state=lambda *args: postprocess_args.append(args) @@ -329,19 +321,11 @@ def receive(*_: Any) -> bool: "postprocess_num_computed_tokens" ) runner.eplb.step = lambda *args, **kwargs: events.append("eplb") - monkeypatch.setattr( - mrv2, - "async_copy_to_gpu", - lambda value, *, device: torch.as_tensor(value, device=device), - ) - output = mrv2.GPUModelRunner.sample_tokens(runner, None) assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] assert len(postprocess_args) == 1 - published_mapping, num_sampled, computed, query_metadata = postprocess_args[0] - assert published_mapping.tolist() == [-1, 7] + published_mapping, num_sampled = postprocess_args[0] + assert published_mapping is idx_mapping assert num_sampled == 0 - assert computed is num_computed_tokens - assert query_metadata is query_start_loc diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index ba84d6016a9b..fb1a5e8d8588 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -55,8 +55,6 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: ctx = Mock( is_initialized=True, replayssm=replayssm, - materialize_dst_cols=torch.full((4,), -1, dtype=torch.int32, device="cuda"), - materialize_token_counts=torch.zeros(4, dtype=torch.int32, device="cuda"), block_size=1024, ) state._mamba_ctx = ctx @@ -80,6 +78,44 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: assert kwargs["live_cols"] is state._replayssm_live_cols_gpu +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_flashinfer_replayssm_prefix_uses_original_accepted_counts() -> None: + state = object.__new__(MambaHybridModelState) + state._align_mode = True + state._needs_prefix_state_migration = True + state._use_flashinfer_replayssm = True + state.recoverssm = None + state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") + state._mamba_state_idx_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") + state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") + replayssm = Mock(materialize_prefixes=True) + accepted_snapshot = torch.zeros(4, dtype=torch.int32, device="cuda") + ctx = Mock( + is_initialized=True, + replayssm=replayssm, + num_accepted_tokens_out=accepted_snapshot, + ) + + def normalize_live(*_args) -> None: + accepted_snapshot.copy_(state.num_accepted_tokens_gpu) + state.num_accepted_tokens_gpu[2] = 1 + + ctx.run_fused_postprocess_align.side_effect = normalize_live + state._mamba_ctx = ctx + + state.postprocess_state( + torch.tensor([2], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([0, 0, 8, 0], device="cuda"), + query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), + ) + + kwargs = replayssm.postprocess.call_args.kwargs + assert kwargs["num_accepted_tokens"] is accepted_snapshot + assert accepted_snapshot[2].item() == 3 + assert state.num_accepted_tokens_gpu[2].item() == 1 + + def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: state = RecoverSSMState() metadata = Mock(spec=RecoverSSMMetadata) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 794fd115dcf3..7b615a2209a5 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -215,8 +215,8 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): "expected_accepted", ), [ - pytest.param(True, 4, ["reset", "copy"], 3, id="replayssm-boundary"), - pytest.param(True, 3, ["reset", "copy"], 3, id="replayssm-no-boundary"), + pytest.param(True, 4, ["reset"], 3, id="replayssm-boundary"), + pytest.param(True, 3, ["reset"], 3, id="replayssm-no-boundary"), pytest.param(False, 4, ["copy"], 1, id="generic"), ], ) @@ -279,9 +279,13 @@ def test_postprocess_mamba_align_materializes_prefixes(): ctx.num_draft_tokens_buf = MagicMock() ctx.is_prefilling_buf = MagicMock() ctx.num_accepted_tokens_out = torch.tensor([3], dtype=torch.int32) - ctx.materialize_dst_cols = torch.tensor([1], dtype=torch.int32) - ctx.materialize_token_counts = torch.tensor([2], dtype=torch.int32) - ctx.run_fused_postprocess.side_effect = lambda **kwargs: order.append("copy") + accepted = torch.tensor([3], dtype=torch.int32) + + def run_fused_postprocess(**kwargs): + order.append("copy") + kwargs["num_accepted_tokens_gpu"].fill_(1) + + ctx.run_fused_postprocess.side_effect = run_fused_postprocess ctx.replayssm = MagicMock() ctx.replayssm.materialize_prefixes = True ctx.replayssm.postprocess.side_effect = lambda **kwargs: order.append("postprocess") @@ -297,7 +301,7 @@ def test_postprocess_mamba_align_materializes_prefixes(): postprocess_mamba_align_gpu( bufs=MagicMock(postprocess_align=ctx), num_reqs=1, - num_accepted_tokens_gpu=torch.tensor([3], dtype=torch.int32), + num_accepted_tokens_gpu=accepted, num_accepted_tokens_cpu_tensor=accepted_cpu, input_batch=input_batch, kv_cache_config=kv_cache_config, @@ -307,7 +311,10 @@ def test_postprocess_mamba_align_materializes_prefixes(): ) assert order == ["copy", "postprocess", "materialize"] - assert accepted_cpu.tolist() == [3] + assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is ( + ctx.num_accepted_tokens_out + ) + assert accepted_cpu.tolist() == [1] def test_postprocess_mamba_none_skips_prefix_copy(): @@ -319,8 +326,6 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ctx.num_computed_tokens_buf = MagicMock(gpu=torch.tensor([20], dtype=torch.int32)) ctx.num_draft_tokens_buf = MagicMock(gpu=torch.tensor([3], dtype=torch.int32)) ctx.is_prefilling_buf = MagicMock(gpu=torch.tensor([False])) - ctx.materialize_dst_cols = torch.full((1,), -1, dtype=torch.int32) - ctx.materialize_token_counts = torch.zeros(1, dtype=torch.int32) ctx.block_size = 1024 ctx.replayssm = MagicMock() ctx.replayssm.materialize_prefixes = False @@ -496,8 +501,7 @@ def test_gpu_context_initializes_flashinfer_replayssm_lifecycle(): [torch.empty(1, 4, dtype=torch.int32)], ) - assert gpu_ctx.state_skip_postprocess.tolist() == [0, 1] - assert gpu_ctx.state_skip_precopy.tolist() == [1, 1] + assert gpu_ctx.has_flashinfer_replayssm assert gpu_ctx.replayssm is model_ctx create.assert_called_once() @@ -1175,15 +1179,20 @@ def t(values): gpu_ctx.initialize_from_forward_context( kv_cache_config, forward_context, copy_funcs, [block_table] ) + accepted_gpu = t(num_accepted_tokens) gpu_ctx.run_fused_postprocess( num_reqs=len(req_ids), - num_accepted_tokens_gpu=t(num_accepted_tokens), + num_accepted_tokens_gpu=accepted_gpu, mamba_state_idx_gpu=t(mamba_state_idx), num_scheduled_tokens_gpu=t([num_scheduled_tokens[r] for r in req_ids]), num_computed_tokens_gpu=t(num_computed_tokens), num_draft_tokens_gpu=t([num_draft_tokens.get(r, 0) for r in req_ids]), ) torch.accelerator.synchronize() + # Keep the normalized live buffer available to the assertions below. The + # context-owned buffer now intentionally preserves the original counts for + # ReplaySSM. + gpu_ctx._test_num_accepted_tokens = accepted_gpu @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @@ -1349,7 +1358,7 @@ def test_matches_python_postprocess_mamba(self, device, test_config): device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch", ) @@ -1414,7 +1423,6 @@ def test_no_copy_when_not_needed(self, device, test_config): # State should be unchanged torch.testing.assert_close(conv_state, conv_state_orig) torch.testing.assert_close(temporal_state, temporal_state_orig) - assert gpu_ctx.materialize_dst_cols[0].item() == -1 @pytest.mark.parametrize("num_reqs", [1, 2, 8, 16]) def test_various_batch_sizes(self, device, test_config, num_reqs): @@ -1737,9 +1745,6 @@ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( device=device, ) - assert gpu_ctx.materialize_dst_cols[0].item() == 1 - assert gpu_ctx.materialize_token_counts[0].item() == 1 - # --- Verify Python behavior (ground truth) --- # State should be unchanged (no copy when src_addr == dst_addr) torch.testing.assert_close( @@ -1775,7 +1780,7 @@ def test_src_addr_equals_dst_addr_skips_copy_and_sets_accepted_to_1( device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python", ) @@ -1939,7 +1944,7 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python", ) @@ -2072,7 +2077,7 @@ def test_different_block_idx_copies_without_setting_accepted_to_1( device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python", ) @@ -2207,7 +2212,7 @@ def test_prefix_caching_shared_block_does_not_set_accepted_to_1( device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="GPU num_accepted_tokens should match Python (must NOT be 1)", ) @@ -2354,7 +2359,7 @@ def test_prefix_caching_nonsequential_block_ids_boundary(self, device, test_conf device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch with non-sequential block IDs", ) @@ -2512,7 +2517,7 @@ def test_prefix_caching_mixed_shared_and_distinct_blocks(self, device, test_conf device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch in mixed PC batch", ) @@ -2641,7 +2646,7 @@ def test_pc_aliased_blocks_skip_must_use_logical_idx_not_addr( # Old kernel (959ca0fd): `if src_addr == dst_addr` -> FAILS here (sets 1) # Fixed kernel (6466ce0d): `if src_block_idx == dest_block_idx and # accept_token_bias == 0` -> PASSES (preserves 3) - kernel_accepted = gpu_ctx.num_accepted_tokens_out[0].item() + kernel_accepted = gpu_ctx._test_num_accepted_tokens[0].item() assert kernel_accepted == 3, ( f"Kernel set num_accepted_tokens to {kernel_accepted} but expected 3. " f"The early-return guard likely compared physical addresses " @@ -2816,7 +2821,7 @@ def make_views(raw): device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch", ) @@ -2962,7 +2967,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): device=device, ) torch.testing.assert_close( - gpu_ctx.num_accepted_tokens_out[:num_reqs], + gpu_ctx._test_num_accepted_tokens[:num_reqs], expected_accepted, msg="num_accepted_tokens mismatch at accept_token_bias=2", ) @@ -3132,14 +3137,14 @@ def test_sd_and_ds_conv_layouts_match_snapshot( [expected_accepted], dtype=torch.int32, device=device ) torch.testing.assert_close( - gpu_ctx_sd.num_accepted_tokens_out[:num_reqs], + gpu_ctx_sd._test_num_accepted_tokens[:num_reqs], expected_accepted_tensor, rtol=0, atol=0, msg="SD num_accepted_tokens result is wrong", ) torch.testing.assert_close( - gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], + gpu_ctx_ds._test_num_accepted_tokens[:num_reqs], expected_accepted_tensor, rtol=0, atol=0, diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index e620da3dd5fa..8e87c8810594 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -28,6 +28,20 @@ logger = init_logger(__name__) +@triton.jit +def _mamba_state_copy_boundary( + num_tokens_running_state, + new_num_computed, + block_size: tl.constexpr, +): + """Return the canonical aligned Mamba state-copy decision.""" + aligned_new_computed = (new_num_computed // block_size) * block_size + needs_copy = aligned_new_computed >= num_tokens_running_state + accept_token_bias = aligned_new_computed - num_tokens_running_state + dest_col = aligned_new_computed // block_size - 1 + return needs_copy, accept_token_bias, dest_col + + @triton.jit(do_not_specialize=["num_reqs"]) def _reset_new_replayssm_slots_kernel( idx_mapping, @@ -91,8 +105,6 @@ def _postprocess_replayssm_kernel( num_accepted_tokens, is_prefilling, live_cols, - materialize_dst_cols, - materialize_token_counts, block_table, tracker_start, tracker_committed, @@ -100,7 +112,6 @@ def _postprocess_replayssm_kernel( dst_slots, plan_ring_start, plan_flush_count, - active_request_indices, block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, tracker_capacity, @@ -114,59 +125,18 @@ def _postprocess_replayssm_kernel( NUM_COMPUTED_IS_POST_STEP: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, MATERIALIZE_PREFIXES: tl.constexpr, - MAX_NUM_REQS: tl.constexpr, ) -> None: """Commit a completed step and prepare an optional prefix snapshot.""" batch_idx = tl.program_id(0) - # FlashInfer requires active rows to be a compact prefix. Resolve physical - # validity per group so divergent group mappings remain independent. - if MATERIALIZE_PREFIXES & (batch_idx == 0): - active_count = 0 - for candidate_idx in tl.range(0, num_reqs): - candidate_req_idx = candidate_idx - if HAS_IDX_MAPPING: - candidate_req_idx = tl.load(idx_mapping + candidate_idx) - valid_candidate = candidate_req_idx >= 0 - live_col = tl.load( - live_cols + candidate_req_idx, - mask=valid_candidate, - other=-1, - ) - dst_col = tl.load( - materialize_dst_cols + candidate_idx, - mask=valid_candidate, - other=-1, - ) - wants_materialize = valid_candidate & (live_col >= 0) & (dst_col >= 0) - live_slot = tl.load( - block_table + candidate_idx * block_table_stride_req + live_col, - mask=wants_materialize, - other=PAD_SLOT_ID, - ) - dst_slot = tl.load( - block_table + candidate_idx * block_table_stride_req + dst_col, - mask=wants_materialize, - other=PAD_SLOT_ID, - ) - valid_materialize = ( - wants_materialize - & (live_slot != PAD_SLOT_ID) - & (dst_slot != PAD_SLOT_ID) - & (live_slot >= 0) - & (dst_slot >= 0) - & (live_slot < tracker_capacity) - & (dst_slot < tracker_capacity) - ) - if valid_materialize: - tl.store(active_request_indices + active_count, candidate_idx) - active_count += 1 - if active_count < MAX_NUM_REQS: - tl.store(active_request_indices + active_count, -1) - - # Clear the fixed-capacity plan row before any per-request early exit. + # Clear all fixed-capacity outputs before any per-request early exit. This + # prevents a shorter batch from reusing stale materialization slots. tl.store(plan_ring_start + batch_idx, 0) tl.store(plan_flush_count + batch_idx, -1) + for layer_idx in tl.static_range(0, NUM_LAYERS): + slot_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store(src_slots + slot_offset, PAD_SLOT_ID) + tl.store(dst_slots + slot_offset, PAD_SLOT_ID) if batch_idx >= num_reqs: return @@ -176,6 +146,40 @@ def _postprocess_replayssm_kernel( if req_idx < 0: return + if QUERY_METADATA_IS_CUMULATIVE: + query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( + query_metadata + batch_idx + ) + else: + query_len = tl.load(query_metadata + batch_idx) + + computed = tl.load(num_computed_tokens + req_idx) + computed_before = tl.where( + NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed + ) + # Mamba attention runs a one-token final prefill chunk with prior state as + # decode. Commit the same transition here instead of resetting its cursors. + prefilling = tl.load(is_prefilling + batch_idx) + prefilling = prefilling & ((query_len != 1) | (computed_before <= 0)) + accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) + + # Derive this request's pre/post-step positions from ReplaySSM metadata, + # then share the canonical copy-boundary calculation with the generic + # Mamba state-copy kernel. + computed_after = tl.where( + prefilling, + computed_before + query_len, + computed if NUM_COMPUTED_IS_POST_STEP else computed_before + accepted, + ) + running_state_pos = tl.where( + prefilling, computed_after, computed_after - accepted + 1 + ) + boundary, accept_token_bias, dst_col = _mamba_state_copy_boundary( + running_state_pos, + computed_after, + MAMBA_BLOCK_SIZE, + ) + live_col = tl.load(live_cols + req_idx) valid_live_col = live_col >= 0 live_slot = tl.load( @@ -189,8 +193,7 @@ def _postprocess_replayssm_kernel( & (live_slot >= 0) & (live_slot < tracker_capacity) ) - dst_col = tl.load(materialize_dst_cols + batch_idx) - wants_materialize = valid_live & (dst_col >= 0) + wants_materialize = MATERIALIZE_PREFIXES & valid_live & boundary & (dst_col >= 0) dst_slot = tl.load( block_table + batch_idx * block_table_stride_req + dst_col, mask=wants_materialize, @@ -212,23 +215,6 @@ def _postprocess_replayssm_kernel( dst_slots + slot_offset, tl.where(materialize, dst_slot, PAD_SLOT_ID), ) - - prefilling = tl.load(is_prefilling + batch_idx) - if QUERY_METADATA_IS_CUMULATIVE: - query_len = tl.load(query_metadata + batch_idx + 1) - tl.load( - query_metadata + batch_idx - ) - else: - query_len = tl.load(query_metadata + batch_idx) - - computed = tl.load(num_computed_tokens + req_idx) - computed_before = tl.where( - NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed - ) - # Mamba attention runs a one-token final prefill chunk with prior state as - # decode. Commit the same transition here instead of resetting its cursors. - prefilling = prefilling & ((query_len != 1) | (computed_before <= 0)) - if prefilling: computed_after = computed_before + query_len first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) @@ -251,7 +237,6 @@ def _postprocess_replayssm_kernel( # Prefill produced canonical state, so publish an exact copy. tl.store(plan_flush_count + batch_idx, 0) elif valid_live: - accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) old_start = tl.load(tracker_start + live_slot) old_committed = tl.load(tracker_committed + live_slot) checkpointed = old_committed + query_len > LOGICAL_WINDOW @@ -265,11 +250,10 @@ def _postprocess_replayssm_kernel( tl.store(tracker_committed + live_slot, next_committed) if materialize: - boundary_count = tl.load(materialize_token_counts + batch_idx) tl.store(plan_ring_start + batch_idx, next_start) tl.store( plan_flush_count + batch_idx, - boundary_count + tl.where(checkpointed, 0, old_committed), + accept_token_bias + 1 + tl.where(checkpointed, 0, old_committed), ) # The published destination is canonical and therefore has no live replay. @@ -277,6 +261,29 @@ def _postprocess_replayssm_kernel( tl.store(tracker_committed + dst_slot, 0, mask=materialize) +@triton.jit +def _compact_replayssm_requests_kernel( + plan_flush_count, + active_request_indices, + num_reqs, + MAX_NUM_REQS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +) -> None: + """Build FlashInfer's active-prefix request list in one parallel program.""" + offsets = tl.arange(0, BLOCK_SIZE) + in_capacity = offsets < MAX_NUM_REQS + active = (offsets < num_reqs) & ( + tl.load(plan_flush_count + offsets, mask=in_capacity, other=-1) >= 0 + ) + candidates = tl.where(active, offsets, MAX_NUM_REQS) + compacted = tl.sort(candidates, dim=0) + tl.store( + active_request_indices + offsets, + tl.where(compacted < MAX_NUM_REQS, compacted, -1), + mask=in_capacity, + ) + + def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: ssm = mixer.kv_cache[1] x_cache = mixer.kv_cache[2] @@ -437,8 +444,6 @@ def postprocess( num_accepted_tokens: torch.Tensor, is_prefilling: torch.Tensor, live_cols: torch.Tensor, - materialize_dst_cols: torch.Tensor, - materialize_token_counts: torch.Tensor, num_reqs: int, ) -> None: """Commit a completed step and prepare an optional prefix snapshot.""" @@ -451,8 +456,6 @@ def postprocess( num_accepted_tokens, is_prefilling, live_cols, - materialize_dst_cols, - materialize_token_counts, self.block_table, self.ring_start, self.num_committed, @@ -460,7 +463,6 @@ def postprocess( self.dst_slots, self.plan_ring_start, self.plan_flush_count, - self.active_request_indices, self.block_table.stride(0), self.src_slots.stride(0), self.num_committed.numel(), @@ -474,8 +476,15 @@ def postprocess( NUM_COMPUTED_IS_POST_STEP=num_computed_is_post_step, HAS_IDX_MAPPING=idx_mapping is not None, MATERIALIZE_PREFIXES=self.materialize_prefixes, - MAX_NUM_REQS=self.max_num_reqs, ) + if self.materialize_prefixes: + _compact_replayssm_requests_kernel[(1,)]( + self.plan_flush_count, + self.active_request_indices, + num_reqs, + MAX_NUM_REQS=self.max_num_reqs, + BLOCK_SIZE=triton.next_power_of_2(self.max_num_reqs), + ) def materialize(self) -> None: """Publish the canonical prefix snapshots prepared by ``postprocess``.""" diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5e3f90eb79ed..bbd7982634f8 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -133,7 +133,7 @@ from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras from vllm.v1.worker.gpu.model_states import init_model_state from vllm.v1.worker.gpu.pool.pooling_runner import PoolingRunner -from vllm.v1.worker.gpu.pp_utils import PPHandler, compute_need_sampled_mask +from vllm.v1.worker.gpu.pp_utils import PPHandler from vllm.v1.worker.gpu.sample.batch_shard import ( BatchSharder, all_to_all_logits, @@ -685,6 +685,11 @@ def initialize_kv_cache( self.vllm_config, kv_cache_allocation_context=kv_cache_allocation_context, ) + # Validate and cache optional ReplaySSM-owned state once cache binding + # has populated every layer. Block copies are a per-step hot path. + self.replayssm_block_copy_tensors = get_replayssm_block_copy_tensors( + self.compilation_config.static_forward_context + ) if is_profiling: self.kv_connector = NO_OP_KV_CONNECTOR else: @@ -1112,12 +1117,7 @@ def update_requests(self, scheduler_output: SchedulerOutput) -> None: # zeroing new blocks and before the forward pass reads them. if scheduler_output.kv_cache_block_copies: copy_kv_cache_blocks_inplace( - [ - *self.kv_caches, - *get_replayssm_block_copy_tensors( - self.compilation_config.static_forward_context - ), - ], + [*self.kv_caches, *self.replayssm_block_copy_tensors], self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, ) @@ -1875,21 +1875,7 @@ def sample_tokens( if not all_decode_next: # Might contain non-final prefill chunks, which will be scheduled # in the immediate next step (rather than in pp_size steps). - idx_mapping = input_batch.idx_mapping - need_sampled_mask = compute_need_sampled_mask(input_batch) - if need_sampled_mask is not None: - # Sampled rows are published from the deferred PP receive. - # Mask them here so a mixed batch commits every row once. - idx_mapping_np = np.where( - need_sampled_mask, -1, input_batch.idx_mapping_np - ) - idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device) - self.model_state.postprocess_state( - idx_mapping, - 0, - self.req_states.num_computed_tokens.gpu, - input_batch.query_start_loc, - ) + self.model_state.postprocess_state(input_batch.idx_mapping, 0) # Post-step KV connector related operations. kv_connector_output = self.kv_connector.post_forward(finished_req_ids) @@ -2079,6 +2065,8 @@ def shutdown(self) -> None: self.cudagraph_manager = None if hasattr(self, "kv_caches"): self.kv_caches.clear() + if hasattr(self, "replayssm_block_copy_tensors"): + self.replayssm_block_copy_tensors.clear() if hasattr(self, "attn_groups"): self.attn_groups.clear() if hasattr(self, "kv_cache_config"): diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 89bde631c36c..3b7e7ee40f2f 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -255,13 +255,14 @@ def preprocess_state( dst_cols=self._mamba_state_idx_gpu, num_reqs=num_reqs, ) - ctx.run_fused_precopy( - num_reqs, - self._mamba_state_idx_gpu, - self._mamba_src_col_gpu, - self._mamba_src_off_gpu, - input_batch.idx_mapping, - ) + if replayssm is None: + ctx.run_fused_precopy( + num_reqs, + self._mamba_state_idx_gpu, + self._mamba_src_col_gpu, + self._mamba_src_off_gpu, + input_batch.idx_mapping, + ) def prepare_attn( self, @@ -488,8 +489,6 @@ def _publish_flashinfer_replayssm( if self._needs_prefix_state_migration else self._replayssm_live_cols_gpu ), - materialize_dst_cols=ctx.materialize_dst_cols, - materialize_token_counts=ctx.materialize_token_counts, num_reqs=num_reqs, ) if replayssm.materialize_prefixes: diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 20d588b81c32..cf52a6d3821e 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -24,8 +24,6 @@ class PendingRecv: num_sampled: torch.Tensor # [num_reqs] num_rejected: torch.Tensor # [num_reqs] idx_mapping: torch.Tensor # [num_reqs] - query_start_loc: torch.Tensor # [num_reqs + 1] - is_prefilling: torch.Tensor # [num_reqs] idx_mapping_np: np.ndarray # [num_reqs] # Records which rows need a deferred postprocess (bool). need_sampled_mask: np.ndarray # [num_reqs] @@ -119,8 +117,6 @@ def get_prev_sampled_outputs(self) -> dict[str, torch.Tensor] | None: num_sampled=slot.num_sampled, num_rejected=slot.num_rejected, idx_mapping=idx_mapping, - query_start_loc=slot.query_start_loc, - is_prefilling=slot.is_prefilling, ) def receive(self, input_batch: InputBatch) -> bool: @@ -143,13 +139,6 @@ def receive(self, input_batch: InputBatch) -> bool: num_reqs, self.max_sample_len, dtype=torch.int64, device=self.device ) combined = torch.empty(2, num_reqs, dtype=torch.int32, device=self.device) - # These input-buffer views are reused every step. Snapshot them in - # the same deferred slot as the sampled output so model-owned state - # postprocess observes the query that produced this acceptance. - query_start_loc = input_batch.query_start_loc[: num_reqs + 1].clone() - is_prefilling = async_copy_to_gpu( - input_batch.is_prefilling_np.copy(), device=self.device - ) torch.distributed.broadcast( sampled_tokens, src=self.last_rank, group=self.broadcast_group ) @@ -162,16 +151,12 @@ def receive(self, input_batch: InputBatch) -> bool: # later used on the main stream. sampled_tokens.record_stream(self.main_stream) combined.record_stream(self.main_stream) - query_start_loc.record_stream(self.main_stream) - is_prefilling.record_stream(self.main_stream) self.queue[-1] = PendingRecv( event, sampled_tokens, num_sampled, num_rejected, input_batch.idx_mapping, - query_start_loc, - is_prefilling, input_batch.idx_mapping_np, need_sampled_mask, gen_at_receive_np, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 8a0228616773..d0e1d452e2a9 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -617,6 +617,7 @@ def __init__( # self.model: nn.Module # Set after load_model # Initialize in initialize_kv_cache self.kv_caches: list[torch.Tensor] = [] + self.replayssm_block_copy_tensors: list[torch.Tensor] = [] # indexes: [kv_cache_group_id][attn_group] self.attn_groups: list[list[AttentionGroup]] = [] # self.kv_cache_config: KVCacheConfig @@ -1251,12 +1252,7 @@ def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None self._zero_block_ids(scheduler_output.new_block_ids_to_zero) if scheduler_output.kv_cache_block_copies: copy_kv_cache_blocks_inplace( - [ - *self.kv_caches, - *get_replayssm_block_copy_tensors( - self.compilation_config.static_forward_context - ), - ], + [*self.kv_caches, *self.replayssm_block_copy_tensors], self.kv_cache_config.num_blocks, scheduler_output.kv_cache_block_copies, ) @@ -1615,11 +1611,14 @@ def _update_states_after_model_execute( ): return - # Count the number of accepted tokens for each sequence. - # Valid tokens are contiguous from position 0, so counting non-(-1) - # tokens gives us the first -1 position (i.e., number of accepted). num_reqs = output_token_ids.size(0) - self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1) + if self.num_spec_tokens: + # Valid tokens are contiguous from position 0, so counting + # non-(-1) tokens gives the number accepted. STP is invariantly one + # and the input preparation already maintains that neutral value. + self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum( + dim=1 + ) if self._needs_prefix_state_migration or self._use_flashinfer_replayssm: # Fused GPU postprocess: state copies + per-request accepted-token @@ -1632,6 +1631,8 @@ def _update_states_after_model_execute( num_accepted_tokens_gpu=self.num_accepted_tokens.gpu, num_accepted_tokens_cpu_tensor=( self.input_batch.num_accepted_tokens_cpu_tensor + if self.num_spec_tokens + else None ), input_batch=self.input_batch, kv_cache_config=self.kv_cache_config, @@ -6703,6 +6704,8 @@ def _cleanup_profiling_kv_cache(self) -> None: for i in range(len(self.kv_caches)): self.kv_caches[i] = None # type: ignore self.kv_caches.clear() + if hasattr(self, "replayssm_block_copy_tensors"): + self.replayssm_block_copy_tensors.clear() if hasattr(self, "attn_groups"): self.attn_groups.clear() if hasattr(self, "kv_cache_config"): @@ -7492,6 +7495,11 @@ def initialize_kv_cache_tensors( kv_cache_groups=kv_cache_config.kv_cache_groups, replayssm_caches=replayssm_caches, ) + # Validate and cache optional ReplaySSM-owned state once cache binding + # has populated every layer. Block copies are a per-step hot path. + self.replayssm_block_copy_tensors = get_replayssm_block_copy_tensors( + self.compilation_config.static_forward_context + ) return kv_caches def maybe_add_kv_sharing_layers_to_kv_cache_groups( diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index cb50b6136373..9ab6a63e3228 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -19,6 +19,7 @@ ) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( ReplaySSMModelContext, + _mamba_state_copy_boundary, ) from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -387,22 +388,11 @@ def postprocess_mamba_fused_kernel( state_inner_sizes_ptr, # number of elements in inner dimensions state_conv_widths_ptr, # conv width for conv states (0 for temporal) state_group_indices_ptr, # maps state_idx to group index in block table - # Nonzero for temporal states reconstructed by an external materializer. - # The kernel still emits the request-level decision below, but does not - # overwrite the checkpoint with the generic speculative-column copy. - state_skip_postprocess_ptr, # DS conv row metadata. Zero keeps the single-region copy path. state_dim_row_count_ptr, # int32: per-block dim row count for DS conv state_dim_row_stride_ptr, # int64: bytes between rows for DS conv # Output: num_accepted_tokens update (for src==dst case) num_accepted_tokens_out_ptr, - # Batch-ordered ReplaySSM materialization decision. The caller initializes - # dst_col to -1, which remains the no-op sentinel when no boundary is hit. - # The source is always mamba_state_idx_ptr. - materialize_dst_col_ptr, - # Current-query rows through that boundary: accept_token_bias + 1. The - # materializer adds older pending rows from the source slot's tracker. - materialize_token_count_ptr, # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The # per-request decision arrays are in req-state-slot order; the block table # is in batch order, so HAS_IDX_MAPPING splits the two indexings. @@ -421,6 +411,9 @@ def postprocess_mamba_fused_kernel( # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds # the post-step new_num_computed value (V2 supplies the advanced count). PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, + # FlashInfer ReplaySSM owns temporal state, while this kernel continues to + # snapshot convolution state at the same block boundary. + SKIP_TEMPORAL_STATE_COPY: tl.constexpr = False, # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a # 3D grid (num_reqs, total_states, TEMPORAL_TILES). Default 1 preserves @@ -468,24 +461,15 @@ def postprocess_mamba_fused_kernel( num_tokens_running_state = num_computed + num_scheduled - num_draft new_num_computed = num_tokens_running_state + num_accepted - 1 - aligned_new_computed = (new_num_computed // block_size) * block_size - - needs_copy = aligned_new_computed >= num_tokens_running_state + needs_copy, accept_token_bias, dest_block_idx = _mamba_state_copy_boundary( + num_tokens_running_state, + new_num_computed, + block_size, + ) if not needs_copy: return - # Compute copy parameters - accept_token_bias = aligned_new_computed - num_tokens_running_state - dest_block_idx = aligned_new_computed // block_size - 1 - - if state_idx == 0 and tile_idx == 0: - tl.store(materialize_dst_col_ptr + batch_idx, dest_block_idx) - tl.store( - materialize_token_count_ptr + batch_idx, - accept_token_bias + 1, - ) - # Update accepted-token count before early exits (per-request, so only # state_idx == 0 writes). Also guard on tile_idx == 0 so tiles > 0 # (when TEMPORAL_TILES > 1) do not duplicate the store. @@ -496,7 +480,8 @@ def postprocess_mamba_fused_kernel( if src_block_idx == dest_block_idx and accept_token_bias == 0: return - if tl.load(state_skip_postprocess_ptr + state_idx): + conv_width = tl.load(state_conv_widths_ptr + state_idx) + if SKIP_TEMPORAL_STATE_COPY and conv_width == 0: return bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx @@ -587,7 +572,6 @@ def precopy_mamba_align_fused_kernel( state_inner_sizes_ptr, state_conv_widths_ptr, state_group_indices_ptr, - state_skip_precopy_ptr, state_dim_row_count_ptr, state_dim_row_stride_ptr, idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip) @@ -632,11 +616,6 @@ def precopy_mamba_align_fused_kernel( # so there is nothing to copy. if src_col < 0 or src_col == dst_col: return - if tl.load(state_skip_precopy_ptr + state_idx): - # FlashInfer ReplaySSM owns this temporal state. Copy only conv/other - # model-owned states here. - return - token_bias = tl.load(token_bias_ptr + req_idx) _copy_mamba_state_block( state_idx, @@ -822,8 +801,6 @@ class MambaSpecDecodeGPUContext: state_inner_sizes: torch.Tensor # int64: elements in inner dimensions state_conv_widths: torch.Tensor # int32: conv width (0 for temporal states) state_group_indices: torch.Tensor # int32: maps state_idx to group index - state_skip_postprocess: torch.Tensor # int32: materializer owns this state - state_skip_precopy: torch.Tensor # int32: scheduler block copy owns this state # DS conv row metadata. Zero keeps the single-region copy path. state_dim_row_count: torch.Tensor # int32: per-block dim row count state_dim_row_stride: torch.Tensor # int64: bytes between rows @@ -836,8 +813,6 @@ class MambaSpecDecodeGPUContext: # Output buffer for num_accepted_tokens updates num_accepted_tokens_out: torch.Tensor - materialize_dst_cols: torch.Tensor - materialize_token_counts: torch.Tensor # Per-group block-table base addresses: int64[num_groups]. Populated in # initialize_from_forward_context from the persistent per-group block @@ -863,10 +838,8 @@ class MambaSpecDecodeGPUContext: # Flag to track if metadata has been populated is_initialized: bool = False - # True when any temporal state is owned by the FlashInfer ReplaySSM - # materializer (i.e. some state_skip_postprocess entry is set). Cached at - # populate time so the per-step postprocess can skip the layer scan for - # Triton / non-ReplaySSM configs. + # True when the model-wide FlashInfer ReplaySSM lifecycle owns temporal + # state. Mixed ReplaySSM/baseline Mamba layers are rejected at populate time. has_flashinfer_replayssm: bool = False # Persistent all-layer ReplaySSM descriptors, populated with the cache # addresses on first real forward. None for non-FlashInfer configurations. @@ -928,12 +901,6 @@ def create( state_group_indices=torch.zeros( total_states, dtype=torch.int32, device=device ), - state_skip_postprocess=torch.zeros( - total_states, dtype=torch.int32, device=device - ), - state_skip_precopy=torch.zeros( - total_states, dtype=torch.int32, device=device - ), state_dim_row_count=torch.zeros( total_states, dtype=torch.int32, device=device ), @@ -947,12 +914,6 @@ def create( num_accepted_tokens_out=torch.zeros( max_num_reqs, dtype=torch.int32, device=device ), - materialize_dst_cols=torch.full( - (max_num_reqs,), -1, dtype=torch.int32, device=device - ), - materialize_token_counts=torch.empty( - max_num_reqs, dtype=torch.int32, device=device - ), block_table_ptrs=torch.zeros( len(mamba_group_ids), dtype=torch.int64, device=device ), @@ -1059,7 +1020,6 @@ def _populate_metadata( ) for state_type_idx, copy_func in enumerate(state_copy_funcs): state = kv_caches[state_type_idx] - self.state_skip_precopy[idx] = is_flashinfer_replayssm # Base address self.state_base_addrs[idx] = _reinterpret_u64_as_i64( state.data_ptr() @@ -1101,7 +1061,6 @@ def _populate_metadata( self.state_conv_widths[idx] = state.size(1) self.state_inner_sizes[idx] = state.stride(1) else: - self.state_skip_postprocess[idx] = is_flashinfer_replayssm self.has_flashinfer_replayssm |= bool(is_flashinfer_replayssm) # Temporal state: inner_size = natural elements per # block (prod of inner dims). The kernel uses this @@ -1239,17 +1198,17 @@ def run_fused_postprocess( if num_reqs == 0 or not self.is_initialized: return - # Initialize output to current values (unchanged unless src==dst) + # Preserve the original acceptance counts for ReplaySSM. The generic + # state-copy kernel normalizes the live buffer for the next iteration. self.num_accepted_tokens_out[:num_reqs].copy_( num_accepted_tokens_gpu[:num_reqs] ) - self.materialize_dst_cols[:num_reqs].fill_(-1) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) postprocess_mamba_fused_kernel[grid]( - num_accepted_tokens_gpu, + self.num_accepted_tokens_out, mamba_state_idx_gpu, num_scheduled_tokens_gpu, num_computed_tokens_gpu, @@ -1262,17 +1221,15 @@ def run_fused_postprocess( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, - self.state_skip_postprocess, self.state_dim_row_count, self.state_dim_row_stride, - self.num_accepted_tokens_out, - self.materialize_dst_cols, - self.materialize_token_counts, + num_accepted_tokens_gpu, None, # idx_mapping: V1 decision arrays are already in req order num_reqs, block_size=self.block_size, COPY_BLOCK_SIZE=1024, CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), + SKIP_TEMPORAL_STATE_COPY=self.has_flashinfer_replayssm, TEMPORAL_TILES=_TEMPORAL_TILES, ) @@ -1311,7 +1268,6 @@ def run_fused_precopy( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, - self.state_skip_precopy, self.state_dim_row_count, self.state_dim_row_stride, idx_mapping, @@ -1347,7 +1303,6 @@ def run_fused_postprocess_align( # decision buffer rather than only [:num_reqs]. num_accepted_tokens_snapshot = self.num_accepted_tokens_out num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) - self.materialize_dst_cols[:num_reqs].fill_(-1) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) @@ -1365,12 +1320,9 @@ def run_fused_postprocess_align( self.state_inner_sizes, self.state_conv_widths, self.state_group_indices, - self.state_skip_postprocess, self.state_dim_row_count, self.state_dim_row_stride, num_accepted_tokens_gpu, - self.materialize_dst_cols, - self.materialize_token_counts, idx_mapping, num_reqs, block_size=self.block_size, @@ -1378,6 +1330,7 @@ def run_fused_postprocess_align( CONV_STATE_DIM_FIRST=is_conv_state_dim_first(), HAS_IDX_MAPPING=True, PRECOMPUTED_NEW_COMPUTED=True, + SKIP_TEMPORAL_STATE_COPY=self.has_flashinfer_replayssm, TEMPORAL_TILES=_TEMPORAL_TILES, ) @@ -1632,13 +1585,14 @@ def preprocess_mamba( dst_cols=fused.state_idx.gpu, num_reqs=num_reqs, ) - fused.ctx.run_fused_precopy( - num_reqs=num_reqs, - state_idx_gpu=fused.state_idx.gpu, - src_col_gpu=fused.src_col.gpu, - token_bias_gpu=fused.token_bias.gpu, - idx_mapping=None, - ) + if fused.ctx.replayssm is None: + fused.ctx.run_fused_precopy( + num_reqs=num_reqs, + state_idx_gpu=fused.state_idx.gpu, + src_col_gpu=fused.src_col.gpu, + token_bias_gpu=fused.token_bias.gpu, + idx_mapping=None, + ) else: do_mamba_copy_block(copy_bufs) @@ -1695,7 +1649,7 @@ def postprocess_mamba_align_gpu( bufs: "MambaBuffers", num_reqs: int, num_accepted_tokens_gpu: torch.Tensor, - num_accepted_tokens_cpu_tensor: torch.Tensor, + num_accepted_tokens_cpu_tensor: torch.Tensor | None, input_batch: GPUInputBatch, kv_cache_config: KVCacheConfig, forward_context: dict[str, Any], @@ -1752,20 +1706,17 @@ def postprocess_mamba_align_gpu( num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, live_cols=ctx.mamba_state_idx_buf.gpu, - materialize_dst_cols=ctx.materialize_dst_cols, - materialize_token_counts=ctx.materialize_token_counts, num_reqs=num_reqs, ) if ctx.replayssm.materialize_prefixes: ctx.replayssm.materialize() - # ``num_accepted_tokens_out`` is pre-initialized from - # ``num_accepted_tokens_gpu``; the kernel only overwrites entries to 1 - # when src_block_idx == dest_block_idx (copy within the same block), so - # the original count is preserved for everyone else. - num_accepted_tokens_cpu_tensor[:num_reqs].copy_( - accepted_tokens_for_postprocess[:num_reqs], non_blocking=True - ) + if num_accepted_tokens_cpu_tensor is not None: + # CPU consumers need the normalized live counts for the next step, not + # the original snapshot retained for ReplaySSM. + num_accepted_tokens_cpu_tensor[:num_reqs].copy_( + num_accepted_tokens_gpu[:num_reqs], non_blocking=True + ) def stage_postprocess_inputs_to_gpu( From 2d75f40c324a1b032cd649c722d2dae8883326af Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 17:49:03 +0200 Subject: [PATCH 23/53] [Mamba] Bound ReplaySSM reset launch to active batch Launch one reset program per active request instead of the fixed maximum request capacity. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- vllm/model_executor/layers/mamba/ops/ssu_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 8e87c8810594..0079140969f9 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -419,7 +419,7 @@ def reset_new_slots( """Reset cursors for fresh physical slots in this cache group.""" if num_reqs == 0: return - _reset_new_replayssm_slots_kernel[(self.max_num_reqs,)]( + _reset_new_replayssm_slots_kernel[(num_reqs,)]( idx_mapping, src_cols, dst_cols, From ce2886852e7cef4b3261d49a6e6dfcb5220616fb Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 18:59:45 +0200 Subject: [PATCH 24/53] refactor(mamba): clarify ReplaySSM MTP bookkeeping Signed-off-by: Andrii Skliar --- .../layers/mamba/mamba_mixer2.py | 50 +++++++++++-------- .../layers/mamba/ops/ssu_dispatch.py | 47 ++++++++++++----- 2 files changed, 63 insertions(+), 34 deletions(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 29edb4751cd1..5bcdf0cd844e 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -720,6 +720,7 @@ def conv_ssm_forward( 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" + use_spec_decode = self.num_spec > 0 ring_start = prev_num_accepted = prev_query_len = None attn_metadata: AttentionMetadata | None = None @@ -1015,7 +1016,7 @@ def conv_ssm_forward( 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) @@ -1046,7 +1047,7 @@ def conv_ssm_forward( if ( self.use_replayssm and self.mamba_config.backend == MambaBackendEnum.FLASHINFER - and self.num_spec > 0 + and use_spec_decode and self._commits_replayssm_trackers ): assert ring_start is not None @@ -1057,6 +1058,9 @@ def conv_ssm_forward( 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, @@ -1084,7 +1088,7 @@ def conv_ssm_forward( # decode call still processes the full target + draft window. max_query_len=( 1 + self.num_spec - if self.use_replayssm and self.num_spec > 0 + if self.use_replayssm and use_spec_decode else state_indices_tensor_d.size(-1) ), ) @@ -1124,35 +1128,39 @@ def conv_ssm_forward( assert prev_num_accepted is not None assert prev_query_len is not None assert attn_metadata.replayssm_scratch is not None - fi_x = hidden_states_d - fi_dt = dt_d - fi_B = B_d - fi_C = C_d - fi_out = preallocated_ssm_out_d fi_cu_seqlens = query_start_loc_d fi_max_seqlen = None - if self.num_spec > 0: + 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: - fi_shape = (decode_batch, spec_query_len) - fi_x = fi_x.view(*fi_shape, *fi_x.shape[1:]) - fi_dt = fi_dt.view(*fi_shape, *fi_dt.shape[1:]) - fi_B = fi_B.view(*fi_shape, *fi_B.shape[1:]) - fi_C = fi_C.view(*fi_shape, *fi_C.shape[1:]) - fi_out = fi_out.view(*fi_shape, *fi_out.shape[1:]) + hidden_states_d = hidden_states_d.view( + decode_batch, + spec_query_len, + *hidden_states_d.shape[1:], + ) + dt_d = dt_d.view( + decode_batch, spec_query_len, *dt_d.shape[1:] + ) + B_d = B_d.view(decode_batch, spec_query_len, *B_d.shape[1:]) + C_d = C_d.view(decode_batch, spec_query_len, *C_d.shape[1:]) + preallocated_ssm_out_d = preallocated_ssm_out_d.view( + decode_batch, + spec_query_len, + *preallocated_ssm_out_d.shape[1:], + ) fi_cu_seqlens = None fi_max_seqlen = None selective_state_update_replayssm_flashinfer( ssm_state, - fi_x, - fi_dt, + hidden_states_d, + dt_d, A_d, - fi_B, - fi_C, - fi_out, + B_d, + C_d, + preallocated_ssm_out_d, x_cache, B_cache, dt_cache, @@ -1166,7 +1174,7 @@ def conv_ssm_forward( state_batch_indices=replayssm_state_indices_d, scratch=attn_metadata.replayssm_scratch, update_trackers=( - self._updates_replayssm_trackers and self.num_spec == 0 + self._updates_replayssm_trackers and not use_spec_decode ), enable_stochastic_rounding=( self.mamba_config.enable_stochastic_rounding diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index e64653afcbaf..90de3626b2c9 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -54,6 +54,8 @@ 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) @@ -96,29 +98,37 @@ def _commit_replayssm_ring_trackers_kernel( other=pad_slot_id, ) valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states) - prev = tl.load(prev_num_accepted + slots, mask=valid, other=0) - start = tl.load(ring_start + slots, mask=valid, other=0) - previous_query_len = tl.load(prev_query_len + slots, mask=valid, other=0) - accepted = tl.load(num_accepted_tokens + offsets, mask=mask, other=0) - must_checkpoint = (previous_query_len > 0) & ( - prev + previous_query_len > logical_window + 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, - (start + prev) % ring_buffer_len, - start, + (current_ring_start + accepted_since_checkpoint) % ring_buffer_len, + current_ring_start, ) next_prev = tl.where( - previous_query_len == 0, + previous_speculative_query_len == 0, 0, - tl.where(must_checkpoint, accepted, prev + accepted), + tl.where( + must_checkpoint, + accepted_from_previous_query, + accepted_since_checkpoint + accepted_from_previous_query, + ), ) - current_query_len = tl.load( + 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_query_len, mask=valid) + tl.store(prev_query_len + slots, current_speculative_query_len, mask=valid) def update_replayssm_ring_trackers( @@ -191,7 +201,18 @@ def commit_replayssm_ring_trackers( ring_buffer_len: int, pad_slot_id: int = NULL_BLOCK_ID, ) -> None: - """Commit the preceding speculative window and record the current one.""" + """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() From 1d828e076d923243550846c0a15f061799c52ff6 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 19:12:56 +0200 Subject: [PATCH 25/53] Refactor MTP decode tensor views Signed-off-by: Andrii Skliar --- .../layers/mamba/mamba_mixer2.py | 29 ++++++++++++------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 5bcdf0cd844e..af0b56882cb7 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -71,6 +71,13 @@ 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 @@ -1136,20 +1143,22 @@ def conv_ssm_forward( 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 = hidden_states_d.view( - decode_batch, - spec_query_len, - *hidden_states_d.shape[1:], + 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 ) - dt_d = dt_d.view( - decode_batch, spec_query_len, *dt_d.shape[1:] + C_d = _view_mtp_decode_tensor( + C_d, decode_batch, spec_query_len ) - B_d = B_d.view(decode_batch, spec_query_len, *B_d.shape[1:]) - C_d = C_d.view(decode_batch, spec_query_len, *C_d.shape[1:]) - preallocated_ssm_out_d = preallocated_ssm_out_d.view( + preallocated_ssm_out_d = _view_mtp_decode_tensor( + preallocated_ssm_out_d, decode_batch, spec_query_len, - *preallocated_ssm_out_d.shape[1:], ) fi_cu_seqlens = None fi_max_seqlen = None From 609e19e5d728d39705bfc95569c5891da45b0a5a Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 19:18:54 +0200 Subject: [PATCH 26/53] Run ReplaySSM end-to-end tests in CI Signed-off-by: Andrii Skliar --- .buildkite/test_areas/model_runner_v2.yaml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/.buildkite/test_areas/model_runner_v2.yaml b/.buildkite/test_areas/model_runner_v2.yaml index 60895f754add..ec4c0d16a921 100644 --- a/.buildkite/test_areas/model_runner_v2.yaml +++ b/.buildkite/test_areas/model_runner_v2.yaml @@ -24,6 +24,21 @@ steps: # Temporary hack filter to exclude ngram spec decoding based tests. - pytest -v -s entrypoints/llm/test_struct_output_generate.py -k "xgrammar and not speculative_config6 and not speculative_config7 and not speculative_config8 and not speculative_config0" +- 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/model_executor/layers/mamba/ + - vllm/v1/attention/backends/mamba_attn.py + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/ + commands: + - pytest -v -s v1/e2e/test_replayssm_decode.py + - label: ":nvidia: (H200 MIG 35GB) Model Runner V2 Examples" device: h200_35gb key: model-runner-v2-examples From 77dea00cf93e27be20f5d588cc7b7dfc23570b23 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 19:31:05 +0200 Subject: [PATCH 27/53] Expand ReplaySSM CI dependencies Signed-off-by: Andrii Skliar --- .buildkite/test_areas/engine.yaml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index fe1eae82d321..b802abfad745 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -103,10 +103,15 @@ steps: 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/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 From c1dd5a3272468422e15e9cf64714c959837746b7 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 19:38:11 +0200 Subject: [PATCH 28/53] test(mamba): run ReplaySSM MTP at TP2 Signed-off-by: Andrii Skliar --- .buildkite/test_areas/engine.yaml | 1 + tests/v1/e2e/test_replayssm_decode.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index b802abfad745..536b343e3c16 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -107,6 +107,7 @@ steps: - 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/ diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 1b61fd1f878e..269734bf3646 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -139,6 +139,7 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na ) +@multi_gpu_test(num_gpus=2) @pytest.mark.skipif( not HAS_FLASHINFER_CHECKPOINTING_SSU, reason="flashinfer.mamba.checkpointing_ssu not available", @@ -151,6 +152,7 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): 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}, ) From d4757fa1b023699be4620356f4f136481d82ab70 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Thu, 3 Sep 2026 19:39:39 +0200 Subject: [PATCH 29/53] [Mamba] Address ReplaySSM review feedback Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 10 +- tests/v1/core/test_contiguous_kv_packing.py | 11 ++ vllm/config/vllm.py | 112 ++++++++---------- .../layers/mamba/mamba_mixer2.py | 9 +- .../layers/mamba/mamba_utils.py | 15 +-- .../layers/mamba/ops/ssu_dispatch.py | 6 +- vllm/v1/attention/backends/mamba_attn.py | 9 +- vllm/v1/kv_cache_interface.py | 8 +- 8 files changed, 90 insertions(+), 90 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 1860c9747f1f..b2235649ea36 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -292,10 +292,10 @@ def test_replayssm_flashinfer_backend_init(): 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( - base_shapes, + shapes = MambaStateShapeCalculator.replayssm_ring_shapes( + num_heads=16, + head_dim=4, + state_size=16, n_groups=4, tp_world_size=2, logical_window=16, @@ -303,7 +303,7 @@ def test_replayssm_physical_ring_shape( 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), diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 069ce9641d7b..e1f286d6fa41 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -145,6 +145,17 @@ def test_replayssm_rings_do_not_expand_canonical_mamba_page(): assert torch.count_nonzero(replayssm_caches["mamba"][1]) == 0 +def test_mamba_spec_requires_exact_replayssm_ring_layout(): + with pytest.raises(ValueError, match="exactly three.*x, dt, and B"): + MambaSpec( + block_size=2, + shapes=((16,), (16,)), + dtypes=(torch.float32, torch.float32), + replayssm_shapes=((2,), (1,)), + replayssm_dtypes=(torch.float32,) * 2, + ) + + MAIN_KV_PAGE_BYTES = 2_048 COMPRESSED_PAGE_BYTES = 128 NUM_CACHE_TUPLES = 3 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 27e352e55ed0..d7d04a191f35 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2878,75 +2878,61 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": ) if self.parallel_config.pipeline_parallel_size > 1: raise ValueError("ReplaySSM currently requires pipeline_parallel_size=1") - 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.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.mamba_config.backend != MambaBackendEnum.TRITON: + if self.mamba_config.backend == MambaBackendEnum.TRITON: + if self.cache_config.use_kda_recoverssm: + 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" + ) + else: + if use_mamba_replayssm_spec: + raise ValueError( + "Mamba2 ReplaySSM speculative decoding requires " + "--mamba-backend flashinfer" + ) + if self.cache_config.mamba_cache_mode == "all": + raise ValueError( + "ReplaySSM prefix caching in all mode requires " + "--mamba-backend flashinfer" + ) + if self.use_v2_model_runner: + raise ValueError( + "Triton ReplaySSM requires Model Runner V1; use " + "--mamba-backend flashinfer or Model Runner V1" + ) + elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: + if self.cache_config.use_kda_recoverssm: raise ValueError("RecoverSSM requires --mamba-backend triton") - elif use_mamba_replayssm_spec: - if self.cache_config.mamba_cache_mode not in ("none", "align", "all"): + if self.cache_config.replayssm_buffer_len > 16: raise ValueError( - "FlashInfer ReplaySSM speculative decoding requires " - "--mamba-cache-mode none, align, or all" + "FlashInfer ReplaySSM requires --replayssm-buffer-len <= 16" ) - 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" - and self.mamba_config.backend != MambaBackendEnum.FLASHINFER - ): - raise ValueError( - "ReplaySSM prefix caching in all mode requires " - "--mamba-backend flashinfer" - ) - elif self.mamba_config.backend not in ( - MambaBackendEnum.TRITON, - MambaBackendEnum.FLASHINFER, - ): + if use_mamba_replayssm_spec: + 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}" + ) + else: raise ValueError( "--use-replayssm requires --mamba-backend triton or flashinfer" ) - elif ( - self.mamba_config.backend == MambaBackendEnum.TRITON - and self.use_v2_model_runner - ): - raise ValueError( - "Triton ReplaySSM requires Model Runner V1; use " - "--mamba-backend flashinfer or Model Runner V1" - ) if ( self.kv_transfer_config is not None and self.kv_transfer_config.is_kv_transfer_instance diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 5981f76570ba..95d134fdd9d9 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1226,15 +1226,16 @@ def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: return () assert self.replayssm_buffer_len is not None tp_world_size = get_tensor_model_parallel_world_size() - base_shape = self.get_state_shape() - return MambaStateShapeCalculator.append_replayssm_ring( - base_shapes=base_shape, + return MambaStateShapeCalculator.replayssm_ring_shapes( + num_heads=self.num_heads, + head_dim=self.head_dim, + state_size=self.ssm_state_size, n_groups=self.n_groups, tp_world_size=tp_world_size, logical_window=self.replayssm_buffer_len, backend=self.mamba_config.backend, num_speculative_tokens=self.num_spec, - )[len(base_shape) :] + ) @property def mamba_type(self) -> MambaAttentionBackendEnum: diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 78de733b36ed..06917e9767fc 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -215,28 +215,25 @@ def mamba2_state_shape( return conv_state_shape, temporal_state_shape @classmethod - def append_replayssm_ring( + def replayssm_ring_shapes( cls, - base_shapes: tuple[tuple[int, ...], ...], + num_heads: int, + head_dim: int, + state_size: int, n_groups: int, tp_world_size: int, logical_window: int, backend: MambaBackendEnum, num_speculative_tokens: int = 0, ) -> tuple[tuple[int, ...], ...]: - """Append the physical ReplaySSM ring shapes. - - ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; - B_cache uses the un-extended ``n_groups``. - """ + """Return the physical x, dt, and B ring shapes.""" ring_buffer_len = logical_window if backend == MambaBackendEnum.FLASHINFER: # 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_nheads = divide(num_heads, tp_world_size) local_ngroups = divide(n_groups, tp_world_size) return ( - *base_shapes, (local_nheads, ring_buffer_len, head_dim), (local_nheads, ring_buffer_len), (local_ngroups, ring_buffer_len, state_size), diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 0079140969f9..abb06086ec89 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -56,7 +56,11 @@ def _reset_new_replayssm_slots_kernel( PAD_SLOT_ID: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, ) -> None: - """Reset cursors only when this cache group has no physical source.""" + """Reset a fresh destination in this cache group's physical slot space. + + Source and destination are logical block-table columns. Their slot lookups + distinguish a missing group-local source from a continuation. + """ batch_idx = tl.program_id(0) active = batch_idx < num_reqs req_idx = batch_idx diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 1b63acfb7224..4401c0d94689 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -181,11 +181,9 @@ def __init__( 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: - assert len(kv_cache_spec.replayssm_shapes) == 3, ( - "FlashInfer ReplaySSM requires x, dt, and B ring-state tensors" - ) + # Canonical state is (conv, ssm). ReplaySSM auxiliaries are ordered as + # x=(nheads, ring, head_dim), dt=(nheads, ring), + # B=(ngroups, ring, dstate). if self.use_replayssm and not self.use_flashinfer_replayssm: self.decode_write_pos_d: torch.Tensor = torch.empty( (self.decode_cudagraph_max_bs,), @@ -197,7 +195,6 @@ def __init__( dtype=torch.int8, device=device, ) - # B_cache shape = (ngroups, replayssm_buffer_len, dstate). bc_ngroups = kv_cache_spec.replayssm_shapes[2][0] bc_scratch_bs = max( self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index a7d86a4088db..c6d45bc73e45 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -871,8 +871,12 @@ class MambaSpec(KVCacheSpec): tp_replicated: bool = False def __post_init__(self) -> None: - if len(self.replayssm_shapes) != len(self.replayssm_dtypes): - raise ValueError("ReplaySSM shapes and dtypes must have equal length") + if (self.replayssm_shapes or self.replayssm_dtypes) and ( + len(self.replayssm_shapes) != 3 or len(self.replayssm_dtypes) != 3 + ): + raise ValueError( + "ReplaySSM requires exactly three shape/dtype entries for x, dt, and B" + ) @property def state_content_size_bytes(self) -> int: From 57085cd88f2a0636eed4e516b8568adbc2673029 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 01:05:07 +0200 Subject: [PATCH 30/53] [Mamba] Simplify ReplaySSM cache ownership Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 44 +---- tests/model_executor/test_replayssm_warmup.py | 3 +- tests/v1/core/test_contiguous_kv_packing.py | 11 -- tests/v1/worker/test_gpu_model_runner.py | 4 +- .../worker/test_mamba_hybrid_model_state.py | 9 +- tests/v1/worker/test_utils.py | 21 +-- vllm/model_executor/layers/mamba/abstract.py | 11 +- .../layers/mamba/mamba_mixer2.py | 9 +- .../layers/mamba/ops/ssu_dispatch.py | 173 ++++++------------ vllm/model_executor/models/diffusion_gemma.py | 7 +- .../model_executor/warmup/replayssm_warmup.py | 1 + vllm/v1/core/single_type_kv_cache_manager.py | 43 +++-- vllm/v1/kv_cache_interface.py | 8 - vllm/v1/worker/gpu/model_runner.py | 3 - vllm/v1/worker/gpu/model_states/interface.py | 2 - .../worker/gpu/model_states/mamba_hybrid.py | 22 +-- vllm/v1/worker/gpu_model_runner.py | 15 +- vllm/v1/worker/utils.py | 64 +------ 18 files changed, 126 insertions(+), 324 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index b2235649ea36..9e8751506acb 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -315,6 +315,8 @@ def _materialize_mixer(device: str = "cpu") -> Mock: mixer.kv_cache = [ torch.empty(0, device=device), torch.empty(8, 4, 3, 5, device=device), + ] + mixer.replayssm_cache = [ torch.empty(8, 4, 20, 3, device=device), torch.empty(8, 4, 20, device=device), torch.empty(8, 2, 20, 5, device=device), @@ -333,32 +335,6 @@ def _materialize_mixer(device: str = "cpu") -> Mock: return mixer -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_validate_replayssm_cache_rejects_incomplete_cache(): - mixer = _materialize_mixer(device="cuda") - - mixer.kv_cache[1] = torch.empty(0, device="cuda") - with pytest.raises(RuntimeError, match="cache tensors"): - ssu_dispatch._validate_replayssm_cache([mixer]) - - mixer = _materialize_mixer(device="cuda") - mixer.kv_cache[2] = torch.empty(0, device="cuda") - with pytest.raises(RuntimeError, match="cache tensors"): - ssu_dispatch._validate_replayssm_cache([mixer]) - - mixer = _materialize_mixer(device="cuda") - mixer._replayssm_ring_start = torch.empty(0, dtype=torch.int32, device="cuda") - with pytest.raises(RuntimeError, match="ring trackers"): - ssu_dispatch._validate_replayssm_cache([mixer]) - - -def test_validate_replayssm_cache_requires_cuda_state(): - mixer = _materialize_mixer(device="cpu") - - with pytest.raises(RuntimeError, match="requires CUDA cache tensors"): - ssu_dispatch._validate_replayssm_cache([mixer]) - - def _modelwide_replayssm_fixture(cache_mode: str = "align"): groups = [ [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], @@ -423,7 +399,6 @@ def test_modelwide_replayssm_none_commits_trackers_without_materialization( num_computed = torch.zeros(2, dtype=torch.int32, device="cuda") accepted = torch.ones(2, dtype=torch.int32, device="cuda") is_prefilling = torch.zeros(2, dtype=torch.bool, device="cuda") - live_cols = torch.zeros(2, dtype=torch.int32, device="cuda") def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None: query_len[0] = scheduled @@ -437,7 +412,7 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None num_computed_is_post_step=False, num_accepted_tokens=accepted, is_prefilling=is_prefilling, - live_cols=live_cols, + live_cols=None, num_reqs=1, ) num_computed[0] += num_accepted @@ -619,8 +594,8 @@ def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatc assert group_ctx.plan_flush_count.tolist() == [-1, 1] assert group_ctx.active_request_indices.tolist() == [1, -1] - # A shorter batch must clear both the compacted active tail and every stale - # source/destination slot left by the prior materialization. + # A shorter batch clears the compacted active tail. Fixed-capacity plan and + # slot-table tails may stay stale because FlashInfer stops at the first -1. ctx.postprocess( idx_mapping=None, query_metadata=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), @@ -635,10 +610,10 @@ def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatc torch.accelerator.synchronize() for group_ctx in ctx.groups: - assert group_ctx.plan_flush_count.tolist() == [-1, -1] + assert group_ctx.plan_flush_count[0].item() == -1 assert group_ctx.active_request_indices.tolist() == [-1, -1] - assert torch.all(group_ctx.src_slots == NULL_BLOCK_ID) - assert torch.all(group_ctx.dst_slots == NULL_BLOCK_ID) + assert torch.all(group_ctx.src_slots[:, 0] == NULL_BLOCK_ID) + assert torch.all(group_ctx.dst_slots[:, 0] == NULL_BLOCK_ID) @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") @@ -720,7 +695,8 @@ def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): for group_ctx, (ring_start, num_committed) in zip(ctx.groups, before): assert torch.equal(group_ctx.ring_start, ring_start) assert torch.equal(group_ctx.num_committed, num_committed) - assert group_ctx.plan_flush_count.tolist() == [-1, -1] + assert group_ctx.plan_flush_count[0].item() == -1 + assert group_ctx.active_request_indices.tolist() == [-1, -1] @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index cb41a23f1336..b9dadec906d7 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -128,12 +128,13 @@ def test_replayssm_autotune_slots_restore_state_and_trackers(): mixer.kv_cache = ( torch.full((4, 2), 3.0), torch.full((4, 2), 3.0), - *(torch.full((4, 2, 17), 3.0) for _ in range(3)), ) + mixer.replayssm_cache = tuple(torch.full((4, 2, 17), 3.0) for _ in range(3)) mixer._replayssm_ring_start = torch.full((4,), 3, dtype=torch.int32) mixer._replayssm_prev_num_accepted = torch.full((4,), 3, dtype=torch.int32) tracked = ( *mixer.kv_cache, + *mixer.replayssm_cache, mixer._replayssm_ring_start, mixer._replayssm_prev_num_accepted, ) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index e1f286d6fa41..069ce9641d7b 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -145,17 +145,6 @@ def test_replayssm_rings_do_not_expand_canonical_mamba_page(): assert torch.count_nonzero(replayssm_caches["mamba"][1]) == 0 -def test_mamba_spec_requires_exact_replayssm_ring_layout(): - with pytest.raises(ValueError, match="exactly three.*x, dt, and B"): - MambaSpec( - block_size=2, - shapes=((16,), (16,)), - dtypes=(torch.float32, torch.float32), - replayssm_shapes=((2,), (1,)), - replayssm_dtypes=(torch.float32,) * 2, - ) - - MAIN_KV_PAGE_BYTES = 2_048 COMPRESSED_PAGE_BYTES = 128 NUM_CACHE_TUPLES = 3 diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index f610284a65f6..b9dc2cbf41f2 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1487,7 +1487,7 @@ def test_v1_caches_replayssm_block_copy_tensors_after_binding(monkeypatch): def get_extra(_context): assert events == ["bind"] - events.append("validate") + events.append("collect") return [extra] monkeypatch.setattr( @@ -1498,7 +1498,7 @@ def get_extra(_context): SimpleNamespace(kv_cache_groups=[]), kernel_block_sizes=[] ) - assert events == ["bind", "validate"] + assert events == ["bind", "collect"] assert len(runner.replayssm_block_copy_tensors) == 1 assert runner.replayssm_block_copy_tensors[0] is extra diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index fb1a5e8d8588..fdbdd0cbe345 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -48,7 +48,6 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: state._use_flashinfer_replayssm = True state.recoverssm = None state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") - state._replayssm_live_cols_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") replayssm = Mock() replayssm.materialize_prefixes = False @@ -62,12 +61,12 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: num_sampled = torch.tensor([2], dtype=torch.int32, device="cuda") num_computed = torch.tensor([0, 0, 20, 0], dtype=torch.int32, device="cuda") query_start_loc = torch.tensor([0, 4], dtype=torch.int32, device="cuda") + state._replayssm_query_start_loc = query_start_loc state.postprocess_state( idx_mapping, num_sampled, num_computed_tokens=num_computed, - query_start_loc=query_start_loc, ) ctx.run_fused_postprocess_align.assert_not_called() @@ -75,7 +74,7 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: kwargs = replayssm.postprocess.call_args.kwargs assert state.num_accepted_tokens_gpu.tolist() == [1, 1, 2, 1] assert kwargs["num_accepted_tokens"] is state.num_accepted_tokens_gpu - assert kwargs["live_cols"] is state._replayssm_live_cols_gpu + assert kwargs["live_cols"] is None @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @@ -88,6 +87,9 @@ def test_flashinfer_replayssm_prefix_uses_original_accepted_counts() -> None: state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") state._mamba_state_idx_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") + state._replayssm_query_start_loc = torch.tensor( + [0, 4], dtype=torch.int32, device="cuda" + ) replayssm = Mock(materialize_prefixes=True) accepted_snapshot = torch.zeros(4, dtype=torch.int32, device="cuda") ctx = Mock( @@ -107,7 +109,6 @@ def normalize_live(*_args) -> None: torch.tensor([2], dtype=torch.int32, device="cuda"), torch.tensor([3], dtype=torch.int32, device="cuda"), num_computed_tokens=torch.tensor([0, 0, 8, 0], device="cuda"), - query_start_loc=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), ) kwargs = replayssm.postprocess.call_args.kwargs diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 9b7b22fbc574..8c6acd825591 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -3,7 +3,6 @@ from types import SimpleNamespace -import pytest import torch from vllm.config.mamba import MambaBackendEnum, MambaConfig @@ -76,7 +75,8 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): }, ) - assert all(len(mixer.kv_cache) == 5 for mixer in mixers) + assert all(len(mixer.kv_cache) == 2 for mixer in mixers) + assert all(len(mixer.replayssm_cache) == 3 for mixer in mixers) tracker_names = ( "_replayssm_ring_start", @@ -122,7 +122,7 @@ def test_replayssm_block_copy_includes_rings_and_group_trackers(monkeypatch): src, dst = 1, 2 for layer_idx, mixer in enumerate(mixers): - for state_idx, state in enumerate(mixer.kv_cache): + for state_idx, state in enumerate((*mixer.kv_cache, *mixer.replayssm_cache)): state[src].fill_(10 * layer_idx + state_idx + 1) state[dst].fill_(-1) for group_idx, mixer in enumerate(mixers[:2]): @@ -136,7 +136,7 @@ def test_replayssm_block_copy_includes_rings_and_group_trackers(monkeypatch): ) for mixer in mixers: - for state in mixer.kv_cache: + for state in (*mixer.kv_cache, *mixer.replayssm_cache): torch.testing.assert_close(state[dst], state[src]) assert mixers[0]._replayssm_ring_start[dst].item() == 20 assert mixers[0]._replayssm_prev_num_accepted[dst].item() == 30 @@ -144,24 +144,17 @@ def test_replayssm_block_copy_includes_rings_and_group_trackers(monkeypatch): assert mixers[1]._replayssm_prev_num_accepted[dst].item() == 31 -def test_replayssm_block_copy_validates_exact_cache_roles(): - mixer = _TestReplaySSMMixer() - mixer.kv_cache = tuple(torch.zeros(4, 1) for _ in range(4)) - - with pytest.raises(ValueError, match="exactly 5 cache roles"): - get_replayssm_block_copy_tensors({"layers.0.mixer": mixer}) - - def test_replayssm_block_copy_includes_triton_rings_without_trackers(): mixer = _TestReplaySSMMixer(MambaBackendEnum.TRITON) - mixer.kv_cache = tuple(torch.zeros(4, 1) for _ in range(5)) + mixer.kv_cache = tuple(torch.zeros(4, 1) for _ in range(2)) + mixer.replayssm_cache = tuple(torch.zeros(4, 1) for _ in range(3)) tensors = get_replayssm_block_copy_tensors({"layers.0.mixer": mixer}) assert len(tensors) == 3 assert all( actual is expected - for actual, expected in zip(tensors, mixer.kv_cache[2:5], strict=True) + for actual, expected in zip(tensors, mixer.replayssm_cache, strict=True) ) diff --git a/vllm/model_executor/layers/mamba/abstract.py b/vllm/model_executor/layers/mamba/abstract.py index d12eb1dd4a39..037e8c2baa1f 100644 --- a/vllm/model_executor/layers/mamba/abstract.py +++ b/vllm/model_executor/layers/mamba/abstract.py @@ -24,6 +24,8 @@ class MambaBase(AttentionLayerBase): # Contains the KV cache (mamba state) for the layer # in the shape specified by `self.get_state_shape`. kv_cache: tuple[torch.Tensor, ...] + # ReplaySSM rings are auxiliary backend state, not canonical Mamba pages. + replayssm_cache: tuple[torch.Tensor, ...] = () supports_dcp: bool = False def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: @@ -43,14 +45,7 @@ def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: self.kv_cache = tuple(states) def bind_replayssm_cache(self, cache: tuple[torch.Tensor, ...]) -> None: - expected_shapes = self.get_replayssm_state_shape() - expected_dtypes = self.get_replayssm_state_dtype() - assert len(cache) == len(expected_shapes) == len(expected_dtypes) - assert all( - tuple(state.shape[1:]) == shape and state.dtype == dtype - for state, shape, dtype in zip(cache, expected_shapes, expected_dtypes) - ) - self.kv_cache = (*self.kv_cache, *cache) + self.replayssm_cache = cache @abstractmethod def get_state_shape(self) -> Iterable[tuple[int, ...]]: diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 95d134fdd9d9..d72f60d696be 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -523,9 +523,10 @@ def __init__( raise ValueError( "--use-replayssm requires tensor-parallel heads to divide evenly" ) - # ReplaySSM appends x/dt/B rings to (conv_state, ssm_state). - _n_state = 5 if self.use_replayssm else 2 - self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state)) + self.kv_cache = tuple(torch.tensor([]) for _ in range(2)) + self.replayssm_cache = ( + tuple(torch.tensor([]) for _ in range(3)) if self.use_replayssm else () + ) self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) @@ -733,7 +734,7 @@ def conv_ssm_forward( ) ssm_state = self.kv_cache[1] if self.use_replayssm: - x_cache, dt_cache, B_cache = self.kv_cache[2:5] + x_cache, dt_cache, B_cache = self.replayssm_cache if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: ring_start = self._replayssm_ring_start prev_num_accepted = self._replayssm_prev_num_accepted diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index abb06086ec89..33f7b8988fa6 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -42,7 +42,7 @@ def _mamba_state_copy_boundary( return needs_copy, accept_token_bias, dest_col -@triton.jit(do_not_specialize=["num_reqs"]) +@triton.jit def _reset_new_replayssm_slots_kernel( idx_mapping, src_cols, @@ -51,8 +51,6 @@ def _reset_new_replayssm_slots_kernel( tracker_start, tracker_committed, block_table_stride_req: tl.int64, - tracker_capacity, - num_reqs, PAD_SLOT_ID: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, ) -> None: @@ -62,11 +60,10 @@ def _reset_new_replayssm_slots_kernel( distinguish a missing group-local source from a continuation. """ batch_idx = tl.program_id(0) - active = batch_idx < num_reqs req_idx = batch_idx if HAS_IDX_MAPPING: - req_idx = tl.load(idx_mapping + batch_idx, mask=active, other=-1) - valid_req = active & (req_idx >= 0) + req_idx = tl.load(idx_mapping + batch_idx) + valid_req = req_idx >= 0 dst_col = tl.load(dst_cols + req_idx, mask=valid_req, other=-1) valid_dst_col = valid_req & (dst_col >= 0) @@ -75,12 +72,7 @@ def _reset_new_replayssm_slots_kernel( mask=valid_dst_col, other=PAD_SLOT_ID, ) - valid_dst = ( - valid_dst_col - & (dst_slot != PAD_SLOT_ID) - & (dst_slot >= 0) - & (dst_slot < tracker_capacity) - ) + valid_dst = valid_dst_col & (dst_slot != PAD_SLOT_ID) src_col = tl.load(src_cols + req_idx, mask=valid_req, other=-1) valid_src_col = valid_req & (src_col >= 0) @@ -89,19 +81,14 @@ def _reset_new_replayssm_slots_kernel( mask=valid_src_col, other=PAD_SLOT_ID, ) - valid_src = ( - valid_src_col - & (src_slot != PAD_SLOT_ID) - & (src_slot >= 0) - & (src_slot < tracker_capacity) - ) + valid_src = valid_src_col & (src_slot != PAD_SLOT_ID) fresh = valid_dst & ~valid_src tl.store(tracker_start + dst_slot, 0, mask=fresh) tl.store(tracker_committed + dst_slot, 0, mask=fresh) -@triton.jit(do_not_specialize=["num_reqs"]) +@triton.jit def _postprocess_replayssm_kernel( idx_mapping, query_metadata, @@ -118,8 +105,6 @@ def _postprocess_replayssm_kernel( plan_flush_count, block_table_stride_req: tl.int64, slot_table_stride_layer: tl.int64, - tracker_capacity, - num_reqs, MAMBA_BLOCK_SIZE: tl.constexpr, LOGICAL_WINDOW: tl.constexpr, RING_BUFFER_LEN: tl.constexpr, @@ -129,20 +114,13 @@ def _postprocess_replayssm_kernel( NUM_COMPUTED_IS_POST_STEP: tl.constexpr, HAS_IDX_MAPPING: tl.constexpr, MATERIALIZE_PREFIXES: tl.constexpr, + LIVE_COL_IS_ZERO: tl.constexpr, ) -> None: """Commit a completed step and prepare an optional prefix snapshot.""" batch_idx = tl.program_id(0) - # Clear all fixed-capacity outputs before any per-request early exit. This - # prevents a shorter batch from reusing stale materialization slots. tl.store(plan_ring_start + batch_idx, 0) tl.store(plan_flush_count + batch_idx, -1) - for layer_idx in tl.static_range(0, NUM_LAYERS): - slot_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store(src_slots + slot_offset, PAD_SLOT_ID) - tl.store(dst_slots + slot_offset, PAD_SLOT_ID) - if batch_idx >= num_reqs: - return req_idx = batch_idx if HAS_IDX_MAPPING: @@ -184,41 +162,35 @@ def _postprocess_replayssm_kernel( MAMBA_BLOCK_SIZE, ) - live_col = tl.load(live_cols + req_idx) + live_col = 0 if LIVE_COL_IS_ZERO else tl.load(live_cols + req_idx) valid_live_col = live_col >= 0 live_slot = tl.load( block_table + batch_idx * block_table_stride_req + live_col, mask=valid_live_col, other=PAD_SLOT_ID, ) - valid_live = ( - valid_live_col - & (live_slot != PAD_SLOT_ID) - & (live_slot >= 0) - & (live_slot < tracker_capacity) - ) + valid_live = valid_live_col & (live_slot != PAD_SLOT_ID) wants_materialize = MATERIALIZE_PREFIXES & valid_live & boundary & (dst_col >= 0) dst_slot = tl.load( block_table + batch_idx * block_table_stride_req + dst_col, mask=wants_materialize, other=PAD_SLOT_ID, ) - materialize = ( - wants_materialize - & (dst_slot != PAD_SLOT_ID) - & (dst_slot >= 0) - & (dst_slot < tracker_capacity) - ) - for layer_idx in tl.static_range(0, NUM_LAYERS): - slot_offset = layer_idx * slot_table_stride_layer + batch_idx - tl.store( - src_slots + slot_offset, - tl.where(materialize, live_slot, PAD_SLOT_ID), - ) - tl.store( - dst_slots + slot_offset, - tl.where(materialize, dst_slot, PAD_SLOT_ID), - ) + # Block-table writers emit either the null sentinel or an in-capacity ID. + materialize = wants_materialize & (dst_slot != PAD_SLOT_ID) + if MATERIALIZE_PREFIXES: + # FlashInfer's ABI requires packed [layer, batch] tables even though all + # layers in this cache group share the same physical slot namespace. + for layer_idx in tl.static_range(0, NUM_LAYERS): + slot_offset = layer_idx * slot_table_stride_layer + batch_idx + tl.store( + src_slots + slot_offset, + tl.where(materialize, live_slot, PAD_SLOT_ID), + ) + tl.store( + dst_slots + slot_offset, + tl.where(materialize, dst_slot, PAD_SLOT_ID), + ) if prefilling: computed_after = computed_before + query_len first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) @@ -230,11 +202,7 @@ def _postprocess_replayssm_kernel( prefill_slot = tl.load( block_table + batch_idx * block_table_stride_req + col ) - valid_prefill_slot = ( - (prefill_slot != PAD_SLOT_ID) - & (prefill_slot >= 0) - & (prefill_slot < tracker_capacity) - ) + valid_prefill_slot = prefill_slot != PAD_SLOT_ID tl.store(tracker_start + prefill_slot, 0, mask=valid_prefill_slot) tl.store(tracker_committed + prefill_slot, 0, mask=valid_prefill_slot) if materialize: @@ -279,19 +247,25 @@ def _compact_replayssm_requests_kernel( active = (offsets < num_reqs) & ( tl.load(plan_flush_count + offsets, mask=in_capacity, other=-1) >= 0 ) - candidates = tl.where(active, offsets, MAX_NUM_REQS) - compacted = tl.sort(candidates, dim=0) + active_i32 = active.to(tl.int32) + output_offsets = tl.cumsum(active_i32, axis=0) - 1 + num_active = tl.sum(active_i32, axis=0) + tl.store( + active_request_indices + output_offsets, + offsets, + mask=in_capacity & active, + ) tl.store( active_request_indices + offsets, - tl.where(compacted < MAX_NUM_REQS, compacted, -1), - mask=in_capacity, + -1, + mask=in_capacity & (offsets >= num_active), ) def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: ssm = mixer.kv_cache[1] - x_cache = mixer.kv_cache[2] - b_cache = mixer.kv_cache[4] + x_cache = mixer.replayssm_cache[0] + b_cache = mixer.replayssm_cache[2] return ( ssm.dtype, x_cache.dtype, @@ -343,10 +317,9 @@ def create( or block_table.numel() == 0 ): raise ValueError("ReplaySSM requires a non-empty 2D CUDA int32 block table") - _validate_replayssm_cache(mixers) first = mixers[0] first_ssm = first.kv_cache[1] - first_x = first.kv_cache[2] + first_x = first.replayssm_cache[0] compatibility = _replayssm_specialization_key(first) for mixer in mixers[1:]: current = _replayssm_specialization_key(mixer) @@ -376,12 +349,12 @@ def create( materialize_tables=( _cuda_i64_ptrs([m.kv_cache[1] for m in mixers]), _cuda_i64_slot_strides([m.kv_cache[1] for m in mixers]), - _cuda_i64_ptrs([m.kv_cache[2] for m in mixers]), - _cuda_i64_slot_strides([m.kv_cache[2] for m in mixers]), - _cuda_i64_ptrs([m.kv_cache[4] for m in mixers]), - _cuda_i64_slot_strides([m.kv_cache[4] for m in mixers]), - _cuda_i64_ptrs([m.kv_cache[3] for m in mixers]), - _cuda_i64_slot_strides([m.kv_cache[3] for m in mixers]), + _cuda_i64_ptrs([m.replayssm_cache[0] for m in mixers]), + _cuda_i64_slot_strides([m.replayssm_cache[0] for m in mixers]), + _cuda_i64_ptrs([m.replayssm_cache[2] for m in mixers]), + _cuda_i64_slot_strides([m.replayssm_cache[2] for m in mixers]), + _cuda_i64_ptrs([m.replayssm_cache[1] for m in mixers]), + _cuda_i64_slot_strides([m.replayssm_cache[1] for m in mixers]), _cuda_i64_ptrs([m.A for m in mixers]), zero_table, zero_table.clone(), @@ -431,8 +404,6 @@ def reset_new_slots( self.ring_start, self.num_committed, self.block_table.stride(0), - self.num_committed.numel(), - num_reqs, PAD_SLOT_ID=NULL_BLOCK_ID, HAS_IDX_MAPPING=idx_mapping is not None, ) @@ -447,13 +418,13 @@ def postprocess( num_computed_is_post_step: bool, num_accepted_tokens: torch.Tensor, is_prefilling: torch.Tensor, - live_cols: torch.Tensor, + live_cols: torch.Tensor | None, num_reqs: int, ) -> None: """Commit a completed step and prepare an optional prefix snapshot.""" if num_reqs == 0: return - _postprocess_replayssm_kernel[(self.max_num_reqs,)]( + _postprocess_replayssm_kernel[(num_reqs,)]( idx_mapping, query_metadata, num_computed_tokens, @@ -469,8 +440,6 @@ def postprocess( self.plan_flush_count, self.block_table.stride(0), self.src_slots.stride(0), - self.num_committed.numel(), - num_reqs, MAMBA_BLOCK_SIZE=self.mamba_block_size, LOGICAL_WINDOW=self.logical_window, RING_BUFFER_LEN=self.ring_buffer_len, @@ -480,6 +449,7 @@ def postprocess( NUM_COMPUTED_IS_POST_STEP=num_computed_is_post_step, HAS_IDX_MAPPING=idx_mapping is not None, MATERIALIZE_PREFIXES=self.materialize_prefixes, + LIVE_COL_IS_ZERO=live_cols is None, ) if self.materialize_prefixes: _compact_replayssm_requests_kernel[(1,)]( @@ -511,12 +481,14 @@ def materialize(self) -> None: self.plan_flush_count, self.active_request_indices, state_dtype=first.kv_cache[1].dtype, - input_dtype=first.kv_cache[2].dtype, + input_dtype=first.replayssm_cache[0].dtype, matrixA_dtype=first.A.dtype, dim=first.kv_cache[1].size(2), dstate=first.kv_cache[1].size(3), num_heads=first.kv_cache[1].size(1), - heads_per_group=(first.kv_cache[1].size(1) // first.kv_cache[4].size(1)), + heads_per_group=( + first.kv_cache[1].size(1) // first.replayssm_cache[2].size(1) + ), max_window=self.logical_window, ring_buffer_len=self.ring_buffer_len, rand_seed=rand_seed, @@ -964,13 +936,11 @@ def _flashinfer_replayssm_mixers_by_group( layer = forward_context.get(layer_name) if layer is None: continue - kv_cache = getattr(layer, "kv_cache", ()) mamba_config = getattr(layer, "mamba_config", None) backend = getattr(mamba_config, "backend", None) if ( getattr(layer, "use_replayssm", False) and backend == MambaBackendEnum.FLASHINFER - and len(kv_cache) >= 5 ): mixers.append(layer) if mixers: @@ -992,47 +962,6 @@ def _load_replayssm_materialize() -> Callable[..., None]: return replayssm_materialize -def _validate_replayssm_cache(mixers: list[Any]) -> None: - """Validate the cache tensors required by model-wide tracker ownership.""" - for layer_idx, mixer in enumerate(mixers): - cache_tensors = mixer.kv_cache[1:5] - if any(tensor.numel() == 0 for tensor in cache_tensors): - raise RuntimeError( - "FlashInfer ReplaySSM requires allocated SSM and replay-ring " - f"cache tensors for every layer; layer {layer_idx} is empty" - ) - if any(not tensor.is_cuda for tensor in cache_tensors): - devices = [str(tensor.device) for tensor in cache_tensors] - raise RuntimeError( - "FlashInfer ReplaySSM requires CUDA cache tensors for every " - f"layer; layer {layer_idx} uses {devices}" - ) - - ring_start = mixer._replayssm_ring_start - num_committed = mixer._replayssm_prev_num_accepted - if ring_start.numel() == 0 or num_committed.numel() == 0: - raise RuntimeError( - "FlashInfer ReplaySSM requires allocated ring trackers for " - f"every layer; layer {layer_idx} is empty" - ) - if not ring_start.is_cuda or not num_committed.is_cuda: - raise RuntimeError( - "FlashInfer ReplaySSM requires CUDA ring trackers for every layer" - ) - if ( - ring_start.ndim != 1 - or num_committed.ndim != 1 - or ring_start.dtype != torch.int32 - or num_committed.dtype != torch.int32 - ): - raise ValueError("ReplaySSM ring trackers must be 1D int32 tensors") - if ring_start.numel() != num_committed.numel(): - raise ValueError( - "ReplaySSM ring trackers must have equal capacities; got " - f"{ring_start.numel()} and {num_committed.numel()}" - ) - - def initialize_mamba_ssu_backend( mamba_config: MambaConfig, kv_cache_config: KVCacheConfig, diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 52872f310445..3a51e5aad678 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -975,12 +975,7 @@ def prepare_dummy_inputs(self, num_reqs: int, num_tokens: int) -> dict[str, Any] return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]} def postprocess_state( - self, - idx_mapping, - num_sampled, - num_computed_tokens=None, - query_start_loc=None, - is_prefilling=None, + self, idx_mapping, num_sampled, num_computed_tokens=None ) -> None: return None diff --git a/vllm/model_executor/warmup/replayssm_warmup.py b/vllm/model_executor/warmup/replayssm_warmup.py index 701cf2b8ec09..851dc85957d4 100644 --- a/vllm/model_executor/warmup/replayssm_warmup.py +++ b/vllm/model_executor/warmup/replayssm_warmup.py @@ -85,6 +85,7 @@ def _temporary_replayssm_autotune_state( prev_num_accepted = module._replayssm_prev_num_accepted tensors = ( *module.kv_cache, + *module.replayssm_cache, ring_start, prev_num_accepted, ) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index b72bc79e19e1..4e819d2c4578 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1668,6 +1668,14 @@ def allocate_new_blocks( ) -> list[KVCacheBlock]: assert isinstance(self.kv_cache_spec, MambaSpec) if self.mamba_cache_mode != "align": + # Allocate extra `num_speculative_blocks` blocks for + # speculative decoding (MTP/EAGLE) with linear attention. + if self.num_speculative_blocks > 0: + num_tokens += self.block_size * self.num_speculative_blocks + if not self._copy_replayssm_live_state: + return super().allocate_new_blocks( + request_id, num_tokens, num_tokens_main_model + ) req_blocks = self.req_to_blocks[request_id] prev_block_len = len(req_blocks) partial_hit = self._partial_hit_reqs.get(request_id) @@ -1676,10 +1684,6 @@ def allocate_new_blocks( if partial_hit is not None else (req_blocks[-1] if req_blocks else None) ) - # Allocate extra `num_speculative_blocks` blocks for - # speculative decoding (MTP/EAGLE) with linear attention. - if self.num_speculative_blocks > 0: - num_tokens += self.block_size * self.num_speculative_blocks new_blocks = super().allocate_new_blocks( request_id, num_tokens, num_tokens_main_model ) @@ -1704,15 +1708,16 @@ def allocate_new_blocks( partial_hit = self._partial_hit_reqs.get(request_id) has_partial_hit = partial_hit is not None live_source = None - if partial_hit is not None: - live_source = partial_hit[1] - elif prev_block_len > 0: - live_source_idx = ( - prev_block_len - 1 - self.num_speculative_blocks - if request_id in self._allocated_block_reqs - else prev_block_len - 1 - ) - live_source = req_blocks[live_source_idx] + if self._copy_replayssm_live_state: + if partial_hit is not None: + live_source = partial_hit[1] + elif prev_block_len > 0: + live_source_idx = ( + prev_block_len - 1 - self.num_speculative_blocks + if request_id in self._allocated_block_reqs + else prev_block_len - 1 + ) + live_source = req_blocks[live_source_idx] # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. if num_required_blocks <= len(req_blocks) and not has_partial_hit: @@ -1799,10 +1804,11 @@ def allocate_new_blocks( self._apply_cow(request_id, block_idx, source_block, cow_block) returned_blocks = [cow_block] + returned_blocks req_blocks.extend(new_blocks) - live_dest_idx = len(req_blocks) - 1 - self.num_speculative_blocks - live_dest = req_blocks[live_dest_idx] - if any(live_dest is block for block in new_blocks): - self._queue_replayssm_live_copy(live_source, live_dest) + if self._copy_replayssm_live_state: + live_dest_idx = len(req_blocks) - 1 - self.num_speculative_blocks + live_dest = req_blocks[live_dest_idx] + if any(live_dest is block for block in new_blocks): + self._queue_replayssm_live_copy(live_source, live_dest) self._allocated_block_reqs.add(request_id) self._partial_hit_reqs.pop(request_id, None) returned_blocks.extend(new_blocks) @@ -1815,8 +1821,7 @@ def _queue_replayssm_live_copy( ) -> None: """Queue and retain one complete live ReplaySSM slot migration.""" if ( - not self._copy_replayssm_live_state - or source_block is None + source_block is None or source_block.is_null or destination_block.is_null or source_block is destination_block diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index c6d45bc73e45..04ef876b9de4 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -870,14 +870,6 @@ class MambaSpec(KVCacheSpec): # rank holds the full state (e.g. the replicated PLE conv state). tp_replicated: bool = False - def __post_init__(self) -> None: - if (self.replayssm_shapes or self.replayssm_dtypes) and ( - len(self.replayssm_shapes) != 3 or len(self.replayssm_dtypes) != 3 - ): - raise ValueError( - "ReplaySSM requires exactly three shape/dtype entries for x, dt, and B" - ) - @property def state_content_size_bytes(self) -> int: return sum( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index bbd7982634f8..46e6cab3b58b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1502,7 +1502,6 @@ def postprocess_sampled( num_sampled: torch.Tensor, num_rejected: torch.Tensor, query_start_loc: torch.Tensor | None = None, - is_prefilling: torch.Tensor | None = None, ) -> None: # Update the number of computed tokens. if self.is_last_pp_rank: @@ -1527,8 +1526,6 @@ def postprocess_sampled( idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu, - query_start_loc, - is_prefilling, ) def _merge_ec_connector_no_forward( diff --git a/vllm/v1/worker/gpu/model_states/interface.py b/vllm/v1/worker/gpu/model_states/interface.py index 0c5002b65398..cd1b5b64da7d 100644 --- a/vllm/v1/worker/gpu/model_states/interface.py +++ b/vllm/v1/worker/gpu/model_states/interface.py @@ -152,8 +152,6 @@ def postprocess_state( idx_mapping: torch.Tensor, num_sampled: torch.Tensor, num_computed_tokens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - is_prefilling: torch.Tensor | None = None, ) -> None: return None diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 3b7e7ee40f2f..a8e7d6b45dae 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -93,6 +93,7 @@ def __init__( self._is_prefilling_gpu = torch.zeros( self.max_num_reqs, dtype=torch.bool, device=self.device ) + self._replayssm_query_start_loc: torch.Tensor | None = None # Pre-copy prefix-cache state (V2). The migration of each request's # mamba state across block boundaries runs as a fused GPU kernel reusing # the postprocess copy machinery, so the per-step src columns and the @@ -110,15 +111,10 @@ def __init__( RecoverSSMState() if self.cache_config.use_kda_recoverssm else None ) if self._needs_prefix_state_migration or self._use_flashinfer_replayssm: - self._replayssm_live_cols_gpu = torch.zeros( - self.max_num_reqs, dtype=torch.int32, device=self.device - ) self._mamba_ctx: MambaSpecDecodeGPUContext | None = None self._mamba_group_ids: list[int] = [] self._mamba_spec: MambaSpec | None = None self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None - self._mamba_kv_cache_config: KVCacheConfig | None = None - self._mamba_block_tables: tuple[torch.Tensor, ...] | None = None if self._needs_prefix_state_migration: self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device @@ -167,8 +163,6 @@ def _ensure_mamba_postprocess_ctx( mamba_group_ids: list[int], block_tables: tuple[torch.Tensor, ...], ) -> MambaSpecDecodeGPUContext: - self._mamba_kv_cache_config = kv_cache_config - self._mamba_block_tables = block_tables if self._mamba_state_copy_funcs is None: mamba_groups = get_mamba_groups(kv_cache_config) mamba_types = {spec.mamba_type for spec in mamba_groups} @@ -322,6 +316,7 @@ def prepare_attn( num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np) if self._use_flashinfer_replayssm: + self._replayssm_query_start_loc = input_batch.query_start_loc mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) self._ensure_mamba_postprocess_ctx( kv_cache_config, mamba_group_ids, block_tables @@ -381,8 +376,6 @@ def postprocess_state( idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int, num_computed_tokens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - is_prefilling: torch.Tensor | None = None, ) -> None: # 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. @@ -434,16 +427,12 @@ def postprocess_state( self._publish_flashinfer_replayssm( idx_mapping, num_computed_tokens, - query_start_loc, - is_prefilling, ) def _publish_flashinfer_replayssm( self, idx_mapping: torch.Tensor, num_computed_tokens: torch.Tensor | None, - query_start_loc: torch.Tensor | None, - is_prefilling: torch.Tensor | None, ) -> None: """Commit the accepted target transition to ReplaySSM trackers.""" num_reqs = idx_mapping.shape[0] @@ -455,6 +444,7 @@ def _publish_flashinfer_replayssm( "ReplaySSM postprocess requires the post-step computed-token " "counts from the forward that produced this acceptance" ) + query_start_loc = self._replayssm_query_start_loc if query_start_loc is None: raise RuntimeError( "ReplaySSM postprocess requires the query_start_loc from " @@ -467,8 +457,6 @@ def _publish_flashinfer_replayssm( ) replayssm = ctx.replayssm assert replayssm is not None - if is_prefilling is None: - is_prefilling = self._is_prefilling_gpu[:num_reqs] replayssm.postprocess( idx_mapping=idx_mapping, query_metadata=query_start_loc, @@ -483,11 +471,11 @@ def _publish_flashinfer_replayssm( if self._needs_prefix_state_migration else self.num_accepted_tokens_gpu ), - is_prefilling=is_prefilling, + is_prefilling=self._is_prefilling_gpu[:num_reqs], live_cols=( self._mamba_state_idx_gpu if self._needs_prefix_state_migration - else self._replayssm_live_cols_gpu + else None ), num_reqs=num_reqs, ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index d0e1d452e2a9..199e17f07cd4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -248,6 +248,7 @@ allocate_kv_cache, allocate_replayssm_caches, bind_kv_cache, + clear_layer_kv_caches, copy_kv_cache_blocks_inplace, get_replayssm_block_copy_tensors, prepare_kernel_block_sizes, @@ -6712,19 +6713,7 @@ def _cleanup_profiling_kv_cache(self) -> None: delattr(self, "kv_cache_config") self.cache_config.num_gpu_blocks = None - for layer in self.compilation_config.static_forward_context.values(): - if hasattr(layer, "kv_cache"): - kv_cache = layer.kv_cache - layer.kv_cache = ( - torch.tensor([]) if isinstance(kv_cache, torch.Tensor) else [] - ) - # Clean up quantized KV cache scale views - # (int8_per_token_head, fp8_per_token_head) - if hasattr(layer, "impl"): - if hasattr(layer.impl, "_k_scale_cache"): - layer.impl._k_scale_cache = None - if hasattr(layer.impl, "_v_scale_cache"): - layer.impl._v_scale_cache = None + clear_layer_kv_caches(self.compilation_config.static_forward_context.values()) gc.collect() torch.accelerator.empty_cache() diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index cbfc5c82828a..914c39b5dec4 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -684,6 +684,8 @@ def clear_layer_kv_caches(layers: Iterable[Any]) -> None: layer.impl._k_scale_cache = None if hasattr(layer.impl, "_v_scale_cache"): layer.impl._v_scale_cache = None + if hasattr(layer, "replayssm_cache"): + layer.replayssm_cache = () def copy_kv_cache_blocks_inplace( @@ -738,68 +740,18 @@ def copy_kv_cache_blocks_inplace( def get_replayssm_block_copy_tensors( forward_context: Mapping[str, Any], ) -> list[torch.Tensor]: - """Return ReplaySSM rings and any backend-owned shared cursors. - - The runner's ordinary cache list already covers the two canonical Mamba - roles. This validates the complete five-cache contract per ReplaySSM layer - and returns the three backend-owned rings. FlashInfer additionally owns two - group-shared cursor tensors. ``copy_kv_cache_blocks_inplace`` deduplicates - shared storage. - """ + """Return ReplaySSM rings and FlashInfer's group-shared cursors.""" extra_tensors: list[torch.Tensor] = [] - cache_roles = ( - "conv_state", - "ssm_state", - "x_cache", - "dt_cache", - "B_cache", - ) - cursor_roles = ( - ("ring_start", "_replayssm_ring_start"), - ("num_committed", "_replayssm_prev_num_accepted"), - ) - for layer_name, layer in forward_context.items(): + for layer in forward_context.values(): if not getattr(layer, "use_replayssm", False): continue + extra_tensors.extend(layer.replayssm_cache) mamba_config = getattr(layer, "mamba_config", None) backend = getattr(mamba_config, "backend", None) - - kv_cache = getattr(layer, "kv_cache", ()) - if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) != len(cache_roles): - raise ValueError( - f"ReplaySSM layer {layer_name!r} must expose exactly " - f"{len(cache_roles)} cache roles {cache_roles}; got " - f"{len(kv_cache) if isinstance(kv_cache, (list, tuple)) else 0}" + if backend == MambaBackendEnum.FLASHINFER: + extra_tensors.extend( + (layer._replayssm_ring_start, layer._replayssm_prev_num_accepted) ) - for role, tensor in zip(cache_roles, kv_cache, strict=True): - if not isinstance(tensor, torch.Tensor) or tensor.ndim == 0: - raise ValueError( - f"ReplaySSM layer {layer_name!r} has invalid {role} cache" - ) - capacity = kv_cache[0].shape[0] - for role, tensor in zip(cache_roles[1:], kv_cache[1:], strict=True): - if tensor.shape[0] != capacity: - raise ValueError( - f"ReplaySSM layer {layer_name!r} {role} capacity " - f"{tensor.shape[0]} does not match canonical capacity {capacity}" - ) - - extra_tensors.extend(kv_cache[2:5]) - if backend != MambaBackendEnum.FLASHINFER: - continue - for role, attr in cursor_roles: - cursor = getattr(layer, attr, None) - if ( - not isinstance(cursor, torch.Tensor) - or cursor.ndim != 1 - or cursor.dtype != torch.int32 - or cursor.numel() != capacity - ): - raise ValueError( - f"FlashInfer ReplaySSM layer {layer_name!r} has invalid " - f"{role} cursor for capacity {capacity}" - ) - extra_tensors.append(cursor) return extra_tensors From 9f93dcb87b5bb7f017f7657f23fe876d09596c48 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 03:00:10 +0200 Subject: [PATCH 31/53] [Mamba] Fix ReplaySSM materialization expectation The sparse compaction case advances from token 6 to the cache boundary at token 8, so materialization replays two transitions. Align the CUDA-only assertion with the planner and FlashInfer flush-count contract.\n\nAssisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 9e8751506acb..660dd4e2de72 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -591,7 +591,8 @@ def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatc assert materializer.call_count == 2 for group_ctx in ctx.groups: - assert group_ctx.plan_flush_count.tolist() == [-1, 1] + # Request 1 advances from token 6 to the block boundary at token 8. + assert group_ctx.plan_flush_count.tolist() == [-1, 2] assert group_ctx.active_request_indices.tolist() == [1, -1] # A shorter batch clears the compacted active tail. Fixed-capacity plan and From 0c4f84fca6bb53332f3550de507110602f12afd9 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 03:52:30 +0200 Subject: [PATCH 32/53] [Mamba] Preserve hybrid MTP prefix-cache hits Keep EAGLE block dropping on attention groups when a hybrid drafter cannot be identified, without applying the widened lookup window to Mamba state. Preserve that global classification through PP projection, use the Mamba block size for V2 state seeding, and avoid ReplaySSM-only prefill copies on the default path. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/v1/core/test_kv_cache_utils.py | 29 +++---- .../worker/test_mamba_hybrid_model_state.py | 19 +++++ vllm/v1/core/kv_cache_utils.py | 80 +++++++------------ .../worker/gpu/model_states/mamba_hybrid.py | 17 ++-- 4 files changed, 66 insertions(+), 79 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c5159e0d1627..9c98e04d4920 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1284,7 +1284,7 @@ def test_project_kv_cache_groups_to_worker(): spec_b = new_kv_cache_spec(num_kv_heads=4) global_groups = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), + KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a, is_eagle_group=True), ] worker_spec = {"layer1": spec_a, "layer2": spec_a} projected = kv_cache_utils._project_kv_cache_groups_to_worker( @@ -1300,6 +1300,7 @@ def test_project_kv_cache_groups_to_worker(): assert len(projected) == 1 assert projected[0].layer_names == [] assert projected[0].kv_cache_spec is spec_a + assert projected[0].is_eagle_group uniform_spec = UniformTypeKVCacheSpecs( block_size=16, @@ -3362,27 +3363,19 @@ def test_draft_group_not_annotated_without_spec_decode(): assert not any(g.is_eagle_group for g in groups) -def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): - # No group carries the draft marker, so every consumer falls back to - # flagging all groups -- including Mamba ones, which then can never report - # a hit. That is silent today; it must at least be visible. +def test_unidentifiable_draft_flags_only_non_mamba_groups(): + # When no group carries a draft marker, retain the conservative fallback + # for attention while excluding Mamba state from the widened lookup window. groups = get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) ) - assert not any(g.is_eagle_group for g in groups) - assert "no KV cache group could be identified as the draft model's" in ( - caplog_vllm.text - ) - assert "Mamba groups" in caplog_vllm.text - - -def test_no_warning_when_draft_group_is_identified(caplog_vllm): - get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) - ) - - assert "could be identified as the draft model's" not in caplog_vllm.text + for group in groups: + contains_mamba = any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + assert group.is_eagle_group is not contains_mamba def _deepseek_v4_specs(model_version="deepseek_v4"): diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index fdbdd0cbe345..9e94e497f6d3 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -16,6 +16,25 @@ from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState +def test_add_request_seeds_state_with_mamba_block_size() -> None: + state = object.__new__(MambaHybridModelState) + state.rope_state = None + state.prompt_embeds_state = None + state.cache_config = SimpleNamespace( + block_size=16, + mamba_block_size=8, + mamba_cache_mode="align", + ) + state._needs_prefix_state_migration = True + state.num_accepted_tokens_gpu = torch.full((2,), 9, dtype=torch.int32) + state._mamba_state_idx_gpu = torch.full((2,), -1, dtype=torch.int32) + + state.add_request(1, Mock(num_computed_tokens=17)) + + assert state.num_accepted_tokens_gpu.tolist() == [9, 1] + assert state._mamba_state_idx_gpu.tolist() == [-1, 2] + + @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @pytest.mark.parametrize(("num_sampled", "expected_value"), [(0, 1), (3, 3)]) def test_postprocess_state_scalar_with_int32_mapping( diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index ec24e987413a..ee35d95b3074 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1817,7 +1817,6 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: groups, use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config), ) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups @@ -1839,7 +1838,7 @@ def _annotate_eagle_groups( ) -> None: """Flag the KV cache groups that hold drafter attention layers. - Two detection rules, in order of preference: + Three detection rules, in order of preference: 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on MLAAttentionSpec and set by drafter attention layers that run a @@ -1857,6 +1856,10 @@ def _annotate_eagle_groups( ``use_deepseek_v4_fallback`` False. The caller gates this fallback on the configured model type. FIXME(yifan): avoid/generalize this hacky check. + 3. Hybrid fallback. If neither rule identifies a draft group and Mamba + groups are present, conservatively flag every non-Mamba group. This + preserves the existing all-groups fallback for attention caches without + applying its widened lookup window to Mamba state. Args: vllm_config: Config supplying the speculative method, if any. @@ -1876,55 +1879,25 @@ def _annotate_eagle_groups( ): group.is_eagle_group = True - if not use_deepseek_v4_fallback: - return - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - -def _warn_if_unannotated_eagle_mamba( - vllm_config: VllmConfig, - kv_cache_groups: list[KVCacheGroupSpec], -) -> None: - """Warn when the flag-all eagle fallback will silently disable reuse. - - With no group annotated, consumers flag every group as a draft group. That - widens a Mamba group's required lookup window to two consecutive chunks, - which align-mode checkpointing never produces, so reuse drops to zero with - no error and no metric to show it. + if use_deepseek_v4_fallback: + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: + if last_layer in group.layer_names: + group.is_eagle_group = True + break - Args: - vllm_config: Config supplying the speculative method, if any. - kv_cache_groups: Groups as they will be handed to consumers. - """ - spec_config = vllm_config.speculative_config - if spec_config is None or not spec_config.use_eagle(): - return - if any(group.is_eagle_group for group in kv_cache_groups): - return - mamba_groups = [ - idx - for idx, group in enumerate(kv_cache_groups) - if any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if not mamba_groups: - return - logger.warning( - "Speculative decoding (method=%s) is enabled but no KV cache group " - "could be identified as the draft model's, so every group -- " - "including Mamba groups %s -- will be treated as a draft group. A " - "Mamba group cannot satisfy the widened lookup window that implies, " - "so prefix-cache reuse across requests will be disabled and any " - "external KV offload tier will store without ever serving a hit.", - spec_config.method, - mamba_groups, - ) + if not any(group.is_eagle_group for group in kv_cache_groups): + non_mamba_groups = [ + group + for group in kv_cache_groups + if not any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if len(non_mamba_groups) < len(kv_cache_groups): + for group in non_mamba_groups: + group.is_eagle_group = True def _largest_divisor_at_most(value: int, limit: int) -> int: @@ -2017,7 +1990,6 @@ def get_kv_cache_groups( groups.append(KVCacheGroupSpec([name], aligned)) _annotate_eagle_groups(vllm_config, kv_cache_spec, groups) - _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups @@ -2240,7 +2212,11 @@ def _project_kv_cache_groups_to_worker( KVCacheGroupSpec( worker_layer_names, group_spec, - is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + # Empty projected groups preserve global group identity across + # PP ranks. Keep the annotation too, so a Mamba-only stage does + # not reinterpret "no local draft layers" as "all groups are + # draft groups" in local or external-cache coordinators. + is_eagle_group=group.is_eagle_group, ) ) return projected_groups diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index a8e7d6b45dae..fe72bdcbdf61 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -90,9 +90,6 @@ def __init__( self.num_accepted_tokens_gpu = torch.ones( self.max_num_reqs, dtype=torch.int32, device=self.device ) - self._is_prefilling_gpu = torch.zeros( - self.max_num_reqs, dtype=torch.bool, device=self.device - ) self._replayssm_query_start_loc: torch.Tensor | None = None # Pre-copy prefix-cache state (V2). The migration of each request's # mamba state across block boundaries runs as a fused GPU kernel reusing @@ -107,6 +104,10 @@ def __init__( self.cache_config.mamba_cache_mode == "all" and self._use_flashinfer_replayssm ) + if self._use_flashinfer_replayssm: + self._is_prefilling_gpu = torch.zeros( + self.max_num_reqs, dtype=torch.bool, device=self.device + ) self.recoverssm = ( RecoverSSMState() if self.cache_config.use_kda_recoverssm else None ) @@ -132,11 +133,8 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: self.num_accepted_tokens_gpu[req_index].fill_(1) if self._needs_prefix_state_migration: # Seed the running state block from the resumed/prefilled position. - state_block_size = self.cache_config.block_size - if self.cache_config.mamba_cache_mode == "all": - mamba_block_size = self.cache_config.mamba_block_size - assert mamba_block_size is not None - state_block_size = mamba_block_size + state_block_size = self.cache_config.mamba_block_size + assert state_block_size is not None self._mamba_state_idx_gpu[req_index].fill_( (new_req_data.num_computed_tokens - 1) // state_block_size ) @@ -287,7 +285,8 @@ def prepare_attn( is_prefilling[: input_batch.num_reqs] = torch.from_numpy( input_batch.is_prefilling_np ) - self._is_prefilling_gpu[:num_reqs].copy_(is_prefilling, non_blocking=True) + if self._use_flashinfer_replayssm: + self._is_prefilling_gpu[:num_reqs].copy_(is_prefilling, non_blocking=True) # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only # compute them during actual (non-capture) forward execution. From 27d0b5542d001d02c9f6e0c41f7cf0dabbe7d2fd Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 04:13:34 +0200 Subject: [PATCH 33/53] [Mamba] Fix all-mode MTP prefix-cache coverage Exercise two complete all-mode state blocks so MTP can drop the volatile trailing block and still reuse a valid prefix boundary. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 561d45e4008a..190b1e8f4170 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -266,6 +266,13 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): _PC_PREFIX + "Surprisingly, the experiments showed that", _PC_PREFIX + "The most important conclusion was that", ] +# All-mode MTP must cache at least two full state blocks: the drafter drops +# the volatile trailing block before resuming from the preceding boundary. +_PC_MTP_PREFIX = _PC_SENTENCE * 240 +MTP_PREFIX_CACHING_PROMPTS = [ + _PC_MTP_PREFIX + prompt.removeprefix(_PC_PREFIX) + for prompt in PREFIX_CACHING_PROMPTS +] def _prefix_cache_hits(llm) -> int: @@ -405,7 +412,7 @@ def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: @large_gpu_mark(min_gb=40) def test_flashinfer_replayssm_all_prefix_cache_mtp_v2(vllm_runner, monkeypatch): common = dict( - max_model_len=8192, + max_model_len=12288, trust_remote_code=True, enable_prefix_caching=True, enable_chunked_prefill=True, @@ -426,11 +433,11 @@ def test_flashinfer_replayssm_all_prefix_cache_mtp_v2(vllm_runner, monkeypatch): ) as llm: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner first_pass = llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + MTP_PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 ) first_pass_hits = _prefix_cache_hits(llm) cached = llm.generate_greedy_logprobs( - PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 + MTP_PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 ) cached_hits = _prefix_cache_hits(llm) draft_count = sum( From 38429ac56c36f6e79f714c901470b82ec70dc040 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 07:03:34 +0200 Subject: [PATCH 34/53] [Mamba] Finish ReplaySSM review cleanup Remove redundant validation and hot-path staging, clarify the shared state lifecycle, and reuse canonical metadata across both model runners. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 1 + tests/v1/worker/test_gpu_model_runner.py | 2 +- .../worker/test_mamba_hybrid_model_state.py | 4 +- tests/v1/worker/test_mamba_utils.py | 24 ++-- .../layers/mamba/mamba_mixer2.py | 19 ++- .../layers/mamba/mamba_utils.py | 9 +- .../layers/mamba/ops/ssu_dispatch.py | 112 +++++++----------- vllm/model_executor/models/nemotron_h.py | 6 +- vllm/v1/core/single_type_kv_cache_manager.py | 14 +-- .../worker/gpu/model_states/mamba_hybrid.py | 22 ++-- vllm/v1/worker/gpu_model_runner.py | 45 +++---- vllm/v1/worker/mamba_utils.py | 95 ++++++++------- vllm/v1/worker/utils.py | 2 + 13 files changed, 157 insertions(+), 198 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 660dd4e2de72..44bd98cf6f2a 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -489,6 +489,7 @@ def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): ): args = call.args kwargs = call.kwargs + assert args[9] is args[10] assert args[11] is group_ctx.src_slots assert args[12] is group_ctx.dst_slots assert args[13] is group_ctx.plan_ring_start diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index b9dc2cbf41f2..d4850c655ebf 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1448,7 +1448,7 @@ def test_replayssm_stp_keeps_accepted_one_without_cpu_copy(monkeypatch): postprocess = Mock() monkeypatch.setattr( gpu_model_runner_module.mamba_utils, - "postprocess_mamba_align_gpu", + "postprocess_mamba_gpu", postprocess, ) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 9e94e497f6d3..f07602c19eb5 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -94,6 +94,7 @@ def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: assert state.num_accepted_tokens_gpu.tolist() == [1, 1, 2, 1] assert kwargs["num_accepted_tokens"] is state.num_accepted_tokens_gpu assert kwargs["live_cols"] is None + assert state._replayssm_query_start_loc is None @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") @@ -114,7 +115,7 @@ def test_flashinfer_replayssm_prefix_uses_original_accepted_counts() -> None: ctx = Mock( is_initialized=True, replayssm=replayssm, - num_accepted_tokens_out=accepted_snapshot, + num_accepted_tokens_snapshot=accepted_snapshot, ) def normalize_live(*_args) -> None: @@ -134,6 +135,7 @@ def normalize_live(*_args) -> None: assert kwargs["num_accepted_tokens"] is accepted_snapshot assert accepted_snapshot[2].item() == 3 assert state.num_accepted_tokens_gpu[2].item() == 1 + assert state._replayssm_query_start_loc is None def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 7b615a2209a5..0423708d8fdf 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -32,7 +32,7 @@ collect_mamba_copy_meta, do_mamba_copy_block, get_mamba_groups, - postprocess_mamba_align_gpu, + postprocess_mamba_gpu, preprocess_mamba, preprocess_mamba_align_fused_kernel, stage_postprocess_inputs_to_gpu, @@ -99,7 +99,7 @@ def postprocess_mamba( ): """CPU reference for the align-mode postprocess. - Used as a golden against the GPU fused kernel (``postprocess_mamba_align_gpu``). + Used as a golden against the GPU fused kernel (``postprocess_mamba_gpu``). Mirrors what the production code did before the fused kernel replaced it; kept here because production no longer has a CPU implementation. """ @@ -268,7 +268,7 @@ def test_preprocess_mamba_preserves_live_replayssm_state( assert align_ctx.precopy_src_col_buf.np[0] == 0 -def test_postprocess_mamba_align_materializes_prefixes(): +def test_postprocess_mamba_materializes_prefixes(): order: list[str] = [] ctx = MagicMock() ctx.is_initialized = True @@ -278,7 +278,7 @@ def test_postprocess_mamba_align_materializes_prefixes(): ctx.num_computed_tokens_buf = MagicMock() ctx.num_draft_tokens_buf = MagicMock() ctx.is_prefilling_buf = MagicMock() - ctx.num_accepted_tokens_out = torch.tensor([3], dtype=torch.int32) + ctx.num_accepted_tokens_snapshot = torch.tensor([3], dtype=torch.int32) accepted = torch.tensor([3], dtype=torch.int32) def run_fused_postprocess(**kwargs): @@ -298,7 +298,7 @@ def run_fused_postprocess(**kwargs): kv_cache_config.kv_cache_groups = [MagicMock()] accepted_cpu = torch.zeros(1, dtype=torch.int32) - postprocess_mamba_align_gpu( + postprocess_mamba_gpu( bufs=MagicMock(postprocess_align=ctx), num_reqs=1, num_accepted_tokens_gpu=accepted, @@ -312,7 +312,7 @@ def run_fused_postprocess(**kwargs): assert order == ["copy", "postprocess", "materialize"] assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is ( - ctx.num_accepted_tokens_out + ctx.num_accepted_tokens_snapshot ) assert accepted_cpu.tolist() == [1] @@ -334,7 +334,7 @@ def test_postprocess_mamba_none_skips_prefix_copy(): accepted = torch.tensor([2], dtype=torch.int32) accepted_cpu = torch.zeros(1, dtype=torch.int32) - postprocess_mamba_align_gpu( + postprocess_mamba_gpu( bufs=MagicMock(postprocess_align=ctx), num_reqs=1, num_accepted_tokens_gpu=accepted, @@ -349,6 +349,7 @@ def test_postprocess_mamba_none_skips_prefix_copy(): ctx.run_fused_postprocess.assert_not_called() assert ctx.replayssm.postprocess.call_count == 1 assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is accepted + assert ctx.replayssm.postprocess.call_args.kwargs["live_cols"] is None ctx.replayssm.materialize.assert_not_called() assert accepted_cpu.tolist() == [2] @@ -1090,9 +1091,12 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): ) -def test_stage_postprocess_inputs_to_gpu_uses_fixed_none_mode_live_col(): +def test_stage_postprocess_inputs_to_gpu_skips_none_mode_live_col(): device = torch.device("cpu") ctx = _make_staging_ctx(max_num_reqs=4, device=device) + ctx.replayssm = MagicMock(materialize_prefixes=False) + ctx.mamba_state_idx_buf.cpu.fill_(17) + ctx.mamba_state_idx_buf.gpu.fill_(23) scheduler_output = _make_postprocess_scheduler_output( req_ids=["req_a", "req_b"], num_scheduled_tokens={"req_a": 4, "req_b": 1}, @@ -1111,10 +1115,10 @@ def test_stage_postprocess_inputs_to_gpu_uses_fixed_none_mode_live_col(): 2, requests, {}, - fixed_live_col=0, ) - np.testing.assert_array_equal(ctx.mamba_state_idx_buf.np[:2], [0, 0]) + np.testing.assert_array_equal(ctx.mamba_state_idx_buf.np[:2], [17, 17]) + assert ctx.mamba_state_idx_buf.gpu[:2].tolist() == [23, 23] np.testing.assert_array_equal(ctx.num_scheduled_tokens_buf.np[:2], [4, 1]) np.testing.assert_array_equal(ctx.num_computed_tokens_buf.np[:2], [20, 7]) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index d72f60d696be..ab59edf7f9ba 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -793,6 +793,7 @@ def conv_ssm_forward( dim=0, ) ) + conv_initial_state_idx_d = block_idx_last_computed_token_d block_idx_last_scheduled_token_d, block_idx_last_scheduled_token_p = ( torch.split( attn_metadata.block_idx_last_scheduled_token, @@ -819,6 +820,7 @@ def conv_ssm_forward( block_idx_first_scheduled_token_p = None block_idx_last_scheduled_token_d = None block_idx_last_computed_token_d = None + conv_initial_state_idx_d = None block_idx_last_scheduled_token_prev_step_d = None num_computed_tokens_p = None @@ -1007,13 +1009,10 @@ def conv_ssm_forward( # forward. Keep both convolution and ReplaySSM on that private # live page instead of touching the cached prefix source. assert block_idx_last_scheduled_token_d is not None - live_indices = state_indices_tensor_d.gather( - 1, - block_idx_last_scheduled_token_d.to(torch.int64).unsqueeze(1), - ).squeeze(1) - state_indices_tensor_d_input = live_indices - state_indices_tensor_d_output = live_indices - block_idx_last_computed_token_d = block_idx_last_scheduled_token_d + assert replayssm_state_indices_d is not None + state_indices_tensor_d_input = replayssm_state_indices_d + state_indices_tensor_d_output = replayssm_state_indices_d + conv_initial_state_idx_d = block_idx_last_scheduled_token_d elif self.num_spec > 0: assert block_idx_last_scheduled_token_prev_step_d is not None input_indices = ( @@ -1051,7 +1050,7 @@ def conv_ssm_forward( self.activation, conv_state_indices=state_indices_tensor_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token_d, - initial_state_idx=block_idx_last_computed_token_d, + initial_state_idx=conv_initial_state_idx_d, num_accepted_tokens=num_accepted_tokens, query_start_loc=query_start_loc_d, # ReplaySSM keeps one physical state block while a speculative @@ -1218,9 +1217,7 @@ def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: if not self.use_replayssm: return () assert self.model_config is not None - return MambaStateDtypeCalculator.append_replayssm_ring( - (), self.model_config.dtype - ) + return MambaStateDtypeCalculator.replayssm_ring_dtypes(self.model_config.dtype) def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: if not self.use_replayssm: diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 06917e9767fc..734823b9fb21 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -88,16 +88,13 @@ def mamba2_state_dtype( ) @classmethod - def append_replayssm_ring( + def replayssm_ring_dtypes( cls, - base_dtypes: tuple[torch.dtype, ...], model_dtype: ModelDType | torch.dtype, ) -> tuple[torch.dtype, ...]: - """Append the ReplaySSM ring dtypes to a base ``(conv, ssm)`` tuple: - ``(x_cache, dt_cache, B_cache)`` = ``(activation, fp32, activation)``. - """ + """Return ``(x_cache, dt_cache, B_cache)`` dtypes.""" activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype) - return (*base_dtypes, activation_dtype, torch.float32, activation_dtype) + return (activation_dtype, torch.float32, activation_dtype) @classmethod def _mamba_state_dtype( diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 33f7b8988fa6..4d912e99192d 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -13,7 +13,7 @@ from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass from functools import cache -from typing import Any +from typing import Any, NamedTuple import torch @@ -192,12 +192,14 @@ def _postprocess_replayssm_kernel( tl.where(materialize, dst_slot, PAD_SLOT_ID), ) if prefilling: - computed_after = computed_before + query_len - first_col = tl.maximum(computed_before // MAMBA_BLOCK_SIZE, 0) + first_col = computed_before // MAMBA_BLOCK_SIZE last_col = tl.maximum( (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1, 0, ) + # Any block touched by prefill can later become another request's live + # source through a prefix-cache hit. Clear every such block's cursors, + # including intermediate blocks that are neither live_slot nor dst_slot. for col in tl.range(first_col, last_col + 1): prefill_slot = tl.load( block_table + batch_idx * block_table_stride_req + col @@ -229,6 +231,8 @@ def _postprocess_replayssm_kernel( ) # The published destination is canonical and therefore has no live replay. + # When dst_slot aliases live_slot, these stores intentionally supersede the + # live tracker update above: the newly published snapshot is canonical. tl.store(tracker_start + dst_slot, 0, mask=materialize) tl.store(tracker_committed + dst_slot, 0, mask=materialize) @@ -262,23 +266,18 @@ def _compact_replayssm_requests_kernel( ) -def _replayssm_specialization_key(mixer: Any) -> tuple[Any, ...]: - ssm = mixer.kv_cache[1] - x_cache = mixer.replayssm_cache[0] - b_cache = mixer.replayssm_cache[2] - return ( - ssm.dtype, - x_cache.dtype, - mixer.A.dtype, - ssm.size(1), - ssm.size(2), - ssm.size(3), - ssm.size(1) // b_cache.size(1), - int(mixer.replayssm_buffer_len), - x_cache.size(2), - bool(mixer.mamba_config.enable_stochastic_rounding), - int(mixer.mamba_config.stochastic_rounding_philox_rounds or 0), - ) +class _ReplaySSMMaterializeTables(NamedTuple): + state_ptrs: torch.Tensor + state_slot_strides: torch.Tensor + x_cache_ptrs: torch.Tensor + x_cache_slot_strides: torch.Tensor + B_cache_ptrs: torch.Tensor + B_cache_slot_strides: torch.Tensor + dt_cache_ptrs: torch.Tensor + dt_cache_slot_strides: torch.Tensor + A_ptrs: torch.Tensor + state_scale_ptrs: torch.Tensor + state_scale_slot_strides: torch.Tensor @dataclass @@ -289,7 +288,7 @@ class _ReplaySSMGroupContext: block_table: torch.Tensor ring_start: torch.Tensor num_committed: torch.Tensor - materialize_tables: tuple[torch.Tensor, ...] + materialize_tables: _ReplaySSMMaterializeTables src_slots: torch.Tensor dst_slots: torch.Tensor plan_ring_start: torch.Tensor @@ -310,34 +309,9 @@ def create( mamba_block_size: int, max_num_reqs: int, ) -> "_ReplaySSMGroupContext": - if ( - block_table.ndim != 2 - or block_table.dtype != torch.int32 - or not block_table.is_cuda - or block_table.numel() == 0 - ): - raise ValueError("ReplaySSM requires a non-empty 2D CUDA int32 block table") first = mixers[0] first_ssm = first.kv_cache[1] first_x = first.replayssm_cache[0] - compatibility = _replayssm_specialization_key(first) - for mixer in mixers[1:]: - current = _replayssm_specialization_key(mixer) - if current != compatibility: - raise ValueError( - "Layers in one ReplaySSM cache group require identical " - "materialization specialization; got " - f"{compatibility} and {current}" - ) - if ( - mixer._replayssm_ring_start.data_ptr() - != first._replayssm_ring_start.data_ptr() - or mixer._replayssm_prev_num_accepted.data_ptr() - != first._replayssm_prev_num_accepted.data_ptr() - ): - raise ValueError( - "Layers in one ReplaySSM cache group must share ring trackers" - ) device = first_ssm.device zero_table = torch.zeros(len(mixers), dtype=torch.int64, device=device) @@ -346,18 +320,26 @@ def create( block_table=block_table, ring_start=first._replayssm_ring_start, num_committed=first._replayssm_prev_num_accepted, - materialize_tables=( - _cuda_i64_ptrs([m.kv_cache[1] for m in mixers]), - _cuda_i64_slot_strides([m.kv_cache[1] for m in mixers]), - _cuda_i64_ptrs([m.replayssm_cache[0] for m in mixers]), - _cuda_i64_slot_strides([m.replayssm_cache[0] for m in mixers]), - _cuda_i64_ptrs([m.replayssm_cache[2] for m in mixers]), - _cuda_i64_slot_strides([m.replayssm_cache[2] for m in mixers]), - _cuda_i64_ptrs([m.replayssm_cache[1] for m in mixers]), - _cuda_i64_slot_strides([m.replayssm_cache[1] for m in mixers]), - _cuda_i64_ptrs([m.A for m in mixers]), - zero_table, - zero_table.clone(), + materialize_tables=_ReplaySSMMaterializeTables( + state_ptrs=_cuda_i64_ptrs([m.kv_cache[1] for m in mixers]), + state_slot_strides=_cuda_i64_slot_strides( + [m.kv_cache[1] for m in mixers] + ), + x_cache_ptrs=_cuda_i64_ptrs([m.replayssm_cache[0] for m in mixers]), + x_cache_slot_strides=_cuda_i64_slot_strides( + [m.replayssm_cache[0] for m in mixers] + ), + B_cache_ptrs=_cuda_i64_ptrs([m.replayssm_cache[2] for m in mixers]), + B_cache_slot_strides=_cuda_i64_slot_strides( + [m.replayssm_cache[2] for m in mixers] + ), + dt_cache_ptrs=_cuda_i64_ptrs([m.replayssm_cache[1] for m in mixers]), + dt_cache_slot_strides=_cuda_i64_slot_strides( + [m.replayssm_cache[1] for m in mixers] + ), + A_ptrs=_cuda_i64_ptrs([m.A for m in mixers]), + state_scale_ptrs=zero_table, + state_scale_slot_strides=zero_table, ), src_slots=torch.full( (len(mixers), max_num_reqs), @@ -462,8 +444,6 @@ def postprocess( def materialize(self) -> None: """Publish the canonical prefix snapshots prepared by ``postprocess``.""" - if not self.materialize_prefixes: - raise RuntimeError("ReplaySSM materialization requires align or all mode") first = self.mixers[0] mamba_config = first.mamba_config rand_seed = None @@ -501,7 +481,10 @@ class ReplaySSMModelContext: """ReplaySSM lifecycle split by physical cache-slot namespace.""" groups: list[_ReplaySSMGroupContext] - materialize_prefixes: bool + + @property + def materialize_prefixes(self) -> bool: + return self.groups[0].materialize_prefixes @classmethod def create( @@ -551,10 +534,7 @@ def create( groups = [ _ReplaySSMGroupContext.create(*args, max_num_reqs) for args in group_args ] - return cls( - groups=groups, - materialize_prefixes=next(iter(modes)) in ("align", "all"), - ) + return cls(groups=groups) def reset_new_slots(self, **kwargs: Any) -> None: for group in self.groups: @@ -565,8 +545,6 @@ def postprocess(self, **kwargs: Any) -> None: group.postprocess(**kwargs) def materialize(self) -> None: - if not self.materialize_prefixes: - raise RuntimeError("ReplaySSM materialization requires align or all mode") for group in self.groups: group.materialize() diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index daaf8cc9f88a..a875c5c1da02 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -740,12 +740,11 @@ def get_mamba_state_dtype_from_config( vllm_config: "VllmConfig", ) -> tuple[torch.dtype, ...]: cache_config = vllm_config.cache_config - base_dtype = MambaStateDtypeCalculator.mamba2_state_dtype( + return MambaStateDtypeCalculator.mamba2_state_dtype( vllm_config.model_config.dtype, cache_config.mamba_cache_dtype, cache_config.mamba_ssm_cache_dtype, ) - return base_dtype @classmethod def get_mamba_state_shape_from_config( @@ -766,7 +765,7 @@ def get_mamba_state_shape_from_config( hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim - base_shape = MambaStateShapeCalculator.mamba2_state_shape( + return MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -776,7 +775,6 @@ def get_mamba_state_shape_from_config( conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, ) - return base_shape @classmethod def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 4e819d2c4578..717e3f8322ea 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1712,12 +1712,7 @@ def allocate_new_blocks( if partial_hit is not None: live_source = partial_hit[1] elif prev_block_len > 0: - live_source_idx = ( - prev_block_len - 1 - self.num_speculative_blocks - if request_id in self._allocated_block_reqs - else prev_block_len - 1 - ) - live_source = req_blocks[live_source_idx] + live_source = req_blocks[prev_block_len - 1] # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. if num_required_blocks <= len(req_blocks) and not has_partial_hit: @@ -1804,11 +1799,8 @@ def allocate_new_blocks( self._apply_cow(request_id, block_idx, source_block, cow_block) returned_blocks = [cow_block] + returned_blocks req_blocks.extend(new_blocks) - if self._copy_replayssm_live_state: - live_dest_idx = len(req_blocks) - 1 - self.num_speculative_blocks - live_dest = req_blocks[live_dest_idx] - if any(live_dest is block for block in new_blocks): - self._queue_replayssm_live_copy(live_source, live_dest) + if self._copy_replayssm_live_state and new_blocks: + self._queue_replayssm_live_copy(live_source, req_blocks[-1]) self._allocated_block_reqs.add(request_id) self._partial_hit_reqs.pop(request_id, None) returned_blocks.extend(new_blocks) diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index fe72bdcbdf61..8f6d64a1f48b 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -97,7 +97,7 @@ def __init__( # running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" self._use_flashinfer_replayssm = ( - self.cache_config.use_replayssm is True + self.cache_config.use_replayssm and vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER ) self._needs_prefix_state_migration = self._align_mode or ( @@ -247,7 +247,7 @@ def preprocess_state( dst_cols=self._mamba_state_idx_gpu, num_reqs=num_reqs, ) - if replayssm is None: + else: ctx.run_fused_precopy( num_reqs, self._mamba_state_idx_gpu, @@ -438,22 +438,18 @@ def _publish_flashinfer_replayssm( if not num_reqs or not self._use_flashinfer_replayssm: return - if num_computed_tokens is None: - raise RuntimeError( - "ReplaySSM postprocess requires the post-step computed-token " - "counts from the forward that produced this acceptance" - ) - query_start_loc = self._replayssm_query_start_loc + assert num_computed_tokens is not None + query_start_loc, self._replayssm_query_start_loc = ( + self._replayssm_query_start_loc, + None, + ) if query_start_loc is None: raise RuntimeError( "ReplaySSM postprocess requires the query_start_loc from " "the forward that produced this acceptance" ) ctx = self._mamba_ctx - if ctx is None or not ctx.is_initialized: - raise RuntimeError( - "ReplaySSM postprocess context was not initialized before forward" - ) + assert ctx is not None and ctx.is_initialized replayssm = ctx.replayssm assert replayssm is not None replayssm.postprocess( @@ -466,7 +462,7 @@ def _publish_flashinfer_replayssm( # next step; use its snapshot in that case. Mode none never # runs the migration kernel, so the acceptance buffer is exact. num_accepted_tokens=( - ctx.num_accepted_tokens_out + ctx.num_accepted_tokens_snapshot if self._needs_prefix_state_migration else self.num_accepted_tokens_gpu ), diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 199e17f07cd4..ffbd491e9f75 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1607,8 +1607,9 @@ def _update_states_after_model_execute( each sequence, and a shifting is done during the next iteration based on the number of accepted tokens. """ - if not self._use_flashinfer_replayssm and ( - not self.speculative_config or not self.model_config.is_hybrid + if not ( + self._use_flashinfer_replayssm + or (self.speculative_config and self.model_config.is_hybrid) ): return @@ -1626,7 +1627,7 @@ def _update_states_after_model_execute( # update without CPU-GPU sync. The metadata # (num_scheduled_tokens, num_draft_tokens, num_computed_tokens) is # pre-staged to GPU buffers in _prepare_inputs. - mamba_utils.postprocess_mamba_align_gpu( + mamba_utils.postprocess_mamba_gpu( bufs=self._get_mamba_bufs(), num_reqs=num_reqs, num_accepted_tokens_gpu=self.num_accepted_tokens.gpu, @@ -1643,18 +1644,9 @@ def _update_states_after_model_execute( ) if self.num_accepted_tokens_event is not None: + # STP ReplaySSM has no accepted-count D2H copy and therefore + # does not allocate an event. self.num_accepted_tokens_event.record() - - if self.cache_config.mamba_cache_mode == "all": - mamba_utils.postprocess_mamba_all( - scheduler_output, - self.kv_cache_config, - self.input_batch, - self.requests, - self.mamba_state_idx, - self.num_spec_tokens, - num_reqs, - ) else: self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True @@ -1662,16 +1654,16 @@ def _update_states_after_model_execute( assert self.num_accepted_tokens_event is not None self.num_accepted_tokens_event.record() - if self.cache_config.mamba_cache_mode == "all": - mamba_utils.postprocess_mamba_all( - scheduler_output, - self.kv_cache_config, - self.input_batch, - self.requests, - self.mamba_state_idx, - self.num_spec_tokens, - num_reqs, - ) + if self.cache_config.mamba_cache_mode == "all": + mamba_utils.postprocess_mamba_all( + scheduler_output, + self.kv_cache_config, + self.input_batch, + self.requests, + self.mamba_state_idx, + self.num_spec_tokens, + num_reqs, + ) def _update_streaming_request( self, req_id: str, new_req_data: NewRequestData @@ -2165,9 +2157,7 @@ def _prepare_inputs( # Skipped under async scheduling (non-align): the CPU copy races with # the in-flight D2H copy and with input-batch row moves. needs_cpu_accepted_counts = self.num_accepted_tokens_event is not None and ( - not self.use_async_scheduling - or self.cache_config.mamba_cache_mode == "align" - or self._needs_prefix_state_migration + not self.use_async_scheduling or self._needs_prefix_state_migration ) if needs_cpu_accepted_counts: assert self.num_accepted_tokens_event is not None @@ -4497,7 +4487,6 @@ def execute_model( num_reqs, self.requests, self.mamba_state_idx, - fixed_live_col=(None if self._needs_prefix_state_migration else 0), ) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 9ab6a63e3228..7e345ffcd24a 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -480,9 +480,12 @@ def postprocess_mamba_fused_kernel( if src_block_idx == dest_block_idx and accept_token_bias == 0: return - conv_width = tl.load(state_conv_widths_ptr + state_idx) - if SKIP_TEMPORAL_STATE_COPY and conv_width == 0: - return + if SKIP_TEMPORAL_STATE_COPY: + # state_conv_widths is also the state-kind metadata: convolution widths + # are positive, while zero explicitly denotes a temporal state. + conv_width = tl.load(state_conv_widths_ptr + state_idx) + if conv_width == 0: + return bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx _copy_mamba_state_block( @@ -811,8 +814,8 @@ class MambaSpecDecodeGPUContext: mamba_group_ids: list[int] num_groups: int - # Output buffer for num_accepted_tokens updates - num_accepted_tokens_out: torch.Tensor + # Snapshot retained while the live accepted-token buffer is normalized. + num_accepted_tokens_snapshot: torch.Tensor # Per-group block-table base addresses: int64[num_groups]. Populated in # initialize_from_forward_context from the persistent per-group block @@ -911,7 +914,7 @@ def create( num_states=total_states, mamba_group_ids=mamba_group_ids, num_groups=len(mamba_group_ids), - num_accepted_tokens_out=torch.zeros( + num_accepted_tokens_snapshot=torch.zeros( max_num_reqs, dtype=torch.int32, device=device ), block_table_ptrs=torch.zeros( @@ -1125,7 +1128,7 @@ def _populate_metadata( self.mamba_group_ids, forward_context, block_tables, - self.num_accepted_tokens_out.numel(), + self.num_accepted_tokens_snapshot.numel(), ) if self.replayssm is None: raise RuntimeError( @@ -1198,9 +1201,7 @@ def run_fused_postprocess( if num_reqs == 0 or not self.is_initialized: return - # Preserve the original acceptance counts for ReplaySSM. The generic - # state-copy kernel normalizes the live buffer for the next iteration. - self.num_accepted_tokens_out[:num_reqs].copy_( + self.num_accepted_tokens_snapshot[:num_reqs].copy_( num_accepted_tokens_gpu[:num_reqs] ) @@ -1208,7 +1209,7 @@ def run_fused_postprocess( grid = (num_reqs, total_states, _TEMPORAL_TILES) postprocess_mamba_fused_kernel[grid]( - self.num_accepted_tokens_out, + self.num_accepted_tokens_snapshot, mamba_state_idx_gpu, num_scheduled_tokens_gpu, num_computed_tokens_gpu, @@ -1301,13 +1302,13 @@ def run_fused_postprocess_align( # V2 reads non-contiguous idx_mapping positions, so snapshot the whole # decision buffer rather than only [:num_reqs]. - num_accepted_tokens_snapshot = self.num_accepted_tokens_out - num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu) + accepted_snapshot = self.num_accepted_tokens_snapshot + accepted_snapshot.copy_(num_accepted_tokens_gpu) total_states = self.num_states grid = (num_reqs, total_states, _TEMPORAL_TILES) postprocess_mamba_fused_kernel[grid]( - num_accepted_tokens_snapshot, + accepted_snapshot, state_idx_gpu, None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED new_num_computed_tokens_gpu, @@ -1519,7 +1520,8 @@ def preprocess_mamba( ) fused.src_col.np[:num_reqs] = -1 - fused.token_bias.np[:num_reqs] = 0 + if fused.ctx.replayssm is None: + fused.token_bias.np[:num_reqs] = 0 for i, req_id in enumerate(input_batch.req_ids): req_state = requests[req_id] @@ -1554,12 +1556,15 @@ def preprocess_mamba( fused.src_col.np[i] = prev_state_idx if prev_state_idx != -1 and prev_state_idx != curr_state_idx: - accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 if fused is not None: - assert accept_token_bias >= 0 - fused.src_col.np[i] = prev_state_idx - fused.token_bias.np[i] = accept_token_bias + if fused.ctx.replayssm is None: + accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 + assert accept_token_bias >= 0 + fused.src_col.np[i] = prev_state_idx + fused.token_bias.np[i] = accept_token_bias + input_batch.num_accepted_tokens_cpu[i] = 1 else: + accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 collect_mamba_copy_meta( copy_bufs, kv_cache_config, @@ -1571,13 +1576,11 @@ def preprocess_mamba( req_state, forward_context, ) - if fused is None or fused.ctx.replayssm is None: input_batch.num_accepted_tokens_cpu[i] = 1 if fused is not None: fused.state_idx.copy_to_gpu(num_reqs) fused.src_col.copy_to_gpu(num_reqs) - fused.token_bias.copy_to_gpu(num_reqs) if fused.ctx.replayssm is not None: fused.ctx.replayssm.reset_new_slots( idx_mapping=None, @@ -1585,7 +1588,8 @@ def preprocess_mamba( dst_cols=fused.state_idx.gpu, num_reqs=num_reqs, ) - if fused.ctx.replayssm is None: + else: + fused.token_bias.copy_to_gpu(num_reqs) fused.ctx.run_fused_precopy( num_reqs=num_reqs, state_idx_gpu=fused.state_idx.gpu, @@ -1644,7 +1648,7 @@ def preprocess_mamba_all_specdec( prev_last_scheduled_idx_buf.copy_to_gpu() -def postprocess_mamba_align_gpu( +def postprocess_mamba_gpu( *, bufs: "MambaBuffers", num_reqs: int, @@ -1668,7 +1672,6 @@ def postprocess_mamba_align_gpu( # The caller enables this context for spec-decode hybrid state copies or # for model-owned FlashInfer ReplaySSM lifecycle maintenance under STP. assert ctx is not None - assert ctx.mamba_state_idx_buf is not None assert ctx.num_scheduled_tokens_buf is not None assert ctx.num_computed_tokens_buf is not None assert ctx.num_draft_tokens_buf is not None @@ -1686,16 +1689,19 @@ def postprocess_mamba_align_gpu( ) accepted_tokens_for_postprocess = num_accepted_tokens_gpu + live_cols = None if run_prefix_state_migration: + assert ctx.mamba_state_idx_buf is not None + live_cols = ctx.mamba_state_idx_buf.gpu ctx.run_fused_postprocess( num_reqs=num_reqs, num_accepted_tokens_gpu=num_accepted_tokens_gpu, - mamba_state_idx_gpu=ctx.mamba_state_idx_buf.gpu, + mamba_state_idx_gpu=live_cols, num_scheduled_tokens_gpu=ctx.num_scheduled_tokens_buf.gpu, num_computed_tokens_gpu=ctx.num_computed_tokens_buf.gpu, num_draft_tokens_gpu=ctx.num_draft_tokens_buf.gpu, ) - accepted_tokens_for_postprocess = ctx.num_accepted_tokens_out + accepted_tokens_for_postprocess = ctx.num_accepted_tokens_snapshot if ctx.replayssm is not None: ctx.replayssm.postprocess( idx_mapping=None, @@ -1705,7 +1711,7 @@ def postprocess_mamba_align_gpu( num_computed_is_post_step=False, num_accepted_tokens=accepted_tokens_for_postprocess, is_prefilling=ctx.is_prefilling_buf.gpu, - live_cols=ctx.mamba_state_idx_buf.gpu, + live_cols=live_cols, num_reqs=num_reqs, ) if ctx.replayssm.materialize_prefixes: @@ -1726,22 +1732,14 @@ def stage_postprocess_inputs_to_gpu( num_reqs: int, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], - *, - fixed_live_col: int | None = None, ) -> None: """Stage all per-request inputs the fused mamba postprocess kernel reads. Walks ``req_ids[:num_reqs]`` once, writing each request's mamba block index and scheduled/computed/draft token counts into the matching pinned - numpy views, then issues five non-blocking H→D copies. The fused kernel - indexes the resulting GPU tensors by ``req_idx``. Buffers live on ``ctx`` - and only exist when the postprocess kernel is enabled. - - Prefix-migration modes read the live column populated by - ``preprocess_mamba``. Mode ``none`` instead passes ``fixed_live_col=0`` - because its one backend-owned live state always occupies logical column 0. + numpy views, then issues non-blocking H→D copies. Mode ``none`` does not + stage a live column because ReplaySSM's live state is always column zero. """ - assert ctx.mamba_state_idx_buf is not None assert ctx.num_scheduled_tokens_buf is not None assert ctx.num_computed_tokens_buf is not None assert ctx.num_draft_tokens_buf is not None @@ -1749,31 +1747,36 @@ def stage_postprocess_inputs_to_gpu( scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens num_scheduled = scheduler_output.num_scheduled_tokens - state_idx_np = ctx.mamba_state_idx_buf.np + stage_live_cols = ctx.replayssm is None or ctx.replayssm.materialize_prefixes + state_idx_np = None + if stage_live_cols: + assert ctx.mamba_state_idx_buf is not None + state_idx_np = ctx.mamba_state_idx_buf.np scheduled_np = ctx.num_scheduled_tokens_buf.np computed_np = ctx.num_computed_tokens_buf.np draft_np = ctx.num_draft_tokens_buf.np prefill_np = ctx.is_prefilling_buf.np for i in range(num_reqs): req_id = req_ids[i] - state_idx = fixed_live_col - if state_idx is None: + if stage_live_cols: state_idx = mamba_state_idx.get(req_id) assert state_idx is not None, ( f"mamba_state_idx missing entry for {req_id!r}; " "preprocess_mamba must run before stage_postprocess_inputs_to_gpu" ) - state_idx_np[i] = state_idx + assert state_idx_np is not None + state_idx_np[i] = state_idx scheduled = num_scheduled[req_id] - computed = requests[req_id].num_computed_tokens + req_state = requests[req_id] + computed = req_state.num_computed_tokens num_draft = len(scheduled_spec_tokens.get(req_id, [])) scheduled_np[i] = scheduled computed_np[i] = computed draft_np[i] = num_draft - prefill_np[i] = ( - requests[req_id].num_computed_tokens < requests[req_id].num_prompt_tokens - ) - ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) + prefill_np[i] = computed < req_state.num_prompt_tokens + if stage_live_cols: + assert ctx.mamba_state_idx_buf is not None + ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs) ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs) diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 914c39b5dec4..b11ad163c99c 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -749,6 +749,8 @@ def get_replayssm_block_copy_tensors( mamba_config = getattr(layer, "mamba_config", None) backend = getattr(mamba_config, "backend", None) if backend == MambaBackendEnum.FLASHINFER: + # Group-shared trackers appear once per layer; the block-copy helper + # deduplicates them by (device, data_ptr()). extra_tensors.extend( (layer._replayssm_ring_start, layer._replayssm_prev_num_accepted) ) From 4def0e57254f9bd3143020bcf487e72a94e210f8 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 09:23:18 +0200 Subject: [PATCH 35/53] [Mamba] Address ReplaySSM review follow-up Use one prefix-state migration predicate for V1 staging and postprocess, flatten the non-ReplaySSM state-copy branch, and pin external-offloader behavior for hybrid EAGLE group annotations. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- .../unit/offloading_connector/test_config.py | 30 +++++++++++++++++++ tests/v1/worker/test_mamba_utils.py | 3 ++ vllm/v1/worker/gpu_model_runner.py | 1 + vllm/v1/worker/mamba_utils.py | 27 +++++++++-------- 4 files changed, 48 insertions(+), 13 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 99d2ccfb82fc..548634e23a34 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -367,6 +367,36 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks(): ] == [1, 3] +def test_hybrid_eagle_annotation_does_not_change_partial_tail_support(): + config = _make_vllm_config() + config.cache_config.prefix_match_unit = 4 + config.speculative_config.use_eagle_block_drop.return_value = True + kv_cache_config = _make_mamba_hybrid_kv_cache_config() + offloading_config = build_offloading_config(config, kv_cache_config) + + unannotated = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), config, kv_cache_config + ) + # The offloader's existing fallback treats every group as an EAGLE group. + assert not unannotated.supports_partial_tail + assert [group.is_eagle_group for group in unannotated.kv_group_configs] == [ + True, + True, + ] + + kv_cache_config.kv_cache_groups[0].is_eagle_group = True + annotated = SchedulerOffloadConfig.from_spec( + MockOffloadingSpec(offloading_config), config, kv_cache_config + ) + # Explicit classification narrows the volatile set to attention. Partial + # tails remain disabled because at least one EAGLE group is still present. + assert not annotated.supports_partial_tail + assert [group.is_eagle_group for group in annotated.kv_group_configs] == [ + True, + False, + ] + + def test_preserves_data_parallel_config(): config = _make_vllm_config() config.parallel_config.data_parallel_index = 2 diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 0423708d8fdf..5ffb95698233 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -1029,6 +1029,7 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): num_reqs, requests, mamba_state_idx, + run_prefix_state_migration=True, ) np.testing.assert_array_equal( @@ -1088,6 +1089,7 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): 1, requests, mamba_state_idx, + run_prefix_state_migration=True, ) @@ -1115,6 +1117,7 @@ def test_stage_postprocess_inputs_to_gpu_skips_none_mode_live_col(): 2, requests, {}, + run_prefix_state_migration=False, ) np.testing.assert_array_equal(ctx.mamba_state_idx_buf.np[:2], [17, 17]) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index ffbd491e9f75..93dd09ae811b 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4487,6 +4487,7 @@ def execute_model( num_reqs, self.requests, self.mamba_state_idx, + run_prefix_state_migration=self._needs_prefix_state_migration, ) use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 7e345ffcd24a..8e89f58c48c1 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -1555,16 +1555,17 @@ def preprocess_mamba( # not mistaken for fresh slot ownership. fused.src_col.np[i] = prev_state_idx - if prev_state_idx != -1 and prev_state_idx != curr_state_idx: + if ( + prev_state_idx != -1 + and prev_state_idx != curr_state_idx + and (fused is None or fused.ctx.replayssm is None) + ): + accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 if fused is not None: - if fused.ctx.replayssm is None: - accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 - assert accept_token_bias >= 0 - fused.src_col.np[i] = prev_state_idx - fused.token_bias.np[i] = accept_token_bias - input_batch.num_accepted_tokens_cpu[i] = 1 + assert accept_token_bias >= 0 + fused.src_col.np[i] = prev_state_idx + fused.token_bias.np[i] = accept_token_bias else: - accept_token_bias = int(input_batch.num_accepted_tokens_cpu[i]) - 1 collect_mamba_copy_meta( copy_bufs, kv_cache_config, @@ -1576,7 +1577,7 @@ def preprocess_mamba( req_state, forward_context, ) - input_batch.num_accepted_tokens_cpu[i] = 1 + input_batch.num_accepted_tokens_cpu[i] = 1 if fused is not None: fused.state_idx.copy_to_gpu(num_reqs) @@ -1732,6 +1733,7 @@ def stage_postprocess_inputs_to_gpu( num_reqs: int, requests: dict[str, CachedRequestState], mamba_state_idx: dict[str, int], + run_prefix_state_migration: bool, ) -> None: """Stage all per-request inputs the fused mamba postprocess kernel reads. @@ -1747,9 +1749,8 @@ def stage_postprocess_inputs_to_gpu( scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens num_scheduled = scheduler_output.num_scheduled_tokens - stage_live_cols = ctx.replayssm is None or ctx.replayssm.materialize_prefixes state_idx_np = None - if stage_live_cols: + if run_prefix_state_migration: assert ctx.mamba_state_idx_buf is not None state_idx_np = ctx.mamba_state_idx_buf.np scheduled_np = ctx.num_scheduled_tokens_buf.np @@ -1758,7 +1759,7 @@ def stage_postprocess_inputs_to_gpu( prefill_np = ctx.is_prefilling_buf.np for i in range(num_reqs): req_id = req_ids[i] - if stage_live_cols: + if run_prefix_state_migration: state_idx = mamba_state_idx.get(req_id) assert state_idx is not None, ( f"mamba_state_idx missing entry for {req_id!r}; " @@ -1774,7 +1775,7 @@ def stage_postprocess_inputs_to_gpu( computed_np[i] = computed draft_np[i] = num_draft prefill_np[i] = computed < req_state.num_prompt_tokens - if stage_live_cols: + if run_prefix_state_migration: assert ctx.mamba_state_idx_buf is not None ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs) ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs) From 476dc96c3d713362718504aaf82b757e878ee66b Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 10:00:38 +0200 Subject: [PATCH 36/53] test: trim redundant ReplaySSM coverage Remove mock-level wiring and bookkeeping tests that duplicate retained lifecycle and end-to-end coverage. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 178 ------------- .../test_replayssm_metadata_builder.py | 20 -- .../unit/offloading_connector/test_config.py | 30 --- tests/v1/worker/test_gpu_model_runner.py | 111 -------- tests/v1/worker/test_mamba_utils.py | 249 ------------------ 5 files changed, 588 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 44bd98cf6f2a..409c987a0265 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -17,7 +17,6 @@ get_mamba_ssu_backend, initialize_mamba_ssu_backend, selective_state_update, - selective_state_update_replayssm_flashinfer, ) from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum @@ -35,16 +34,6 @@ except ImportError: HAS_FLASHINFER = False -try: - from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner - from flashinfer.mamba.checkpointing_ssu import ( - checkpointing_ssu as checkpointing_ssu_kernel, - ) - - HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) -except ImportError: - HAS_FLASHINFER_CHECKPOINTING_SSU = False - @pytest.fixture(autouse=True) def restore_backend_state(): @@ -215,72 +204,6 @@ def test_triton_basic_call(): assert not torch.isnan(out).any() -def test_replayssm_flashinfer_call_forwards_packed_mtp(monkeypatch): - import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod - - kernel = Mock(return_value=torch.empty(1, 6, 2, 4)) - monkeypatch.setattr(mod, "_flashinfer_replayssm_kernel", kernel) - - tokens, nheads, dim, dstate, ngroups = 6, 2, 4, 8, 1 - state = torch.empty(2, nheads, dim, dstate) - x = torch.empty(tokens, nheads, dim) - dt = torch.empty_like(x) - A = torch.empty(nheads, dim, dstate) - B = torch.empty(tokens, ngroups, dstate) - 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) - cu_seqlens = torch.tensor([0, 4, 6], 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, - state_batch_indices=torch.tensor([0, 1], dtype=torch.int32), - cu_seqlens=cu_seqlens, - max_seqlen=4, - ) - - args = kernel.call_args.args - kwargs = kernel.call_args.kwargs - assert args[6].shape == (1, tokens, nheads, dim) - assert args[7].shape == (1, tokens, nheads, dim) - assert args[9].shape == (1, tokens, ngroups, dstate) - assert args[10].shape == (1, tokens, ngroups, dstate) - assert args[11].shape == (1, tokens, nheads, dim) - assert kwargs["cu_seqlens"] is cu_seqlens - assert kwargs["max_seqlen"] == 4 - - -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="compatible flashinfer checkpointing_ssu not available", -) -def test_replayssm_flashinfer_backend_init(): - import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod - - initialize_mamba_ssu_backend( - MambaConfig(backend=MambaBackendEnum.FLASHINFER), - _kv_cache_config_with_ssu(), - use_replayssm=True, - ) - assert isinstance(get_mamba_ssu_backend(), FlashInferSSUBackend) - assert mod._flashinfer_replayssm_kernel is checkpointing_ssu_kernel - - @pytest.mark.parametrize( ("backend", "num_speculative_tokens", "expected_ring_len"), [ @@ -452,66 +375,6 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None assert mixers[0]._replayssm_prev_num_accepted[live_slot].item() == 0 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_materializes_each_cache_group(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, source_slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[source_slot] = 2 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 - kernel = Mock() - monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), - query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=True, - num_computed_tokens=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=True, - num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([False, False], device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - ctx.materialize() - torch.accelerator.synchronize() - - assert kernel.call_count == 2 - for call, group_ctx, src_slot, dst_slot in zip( - kernel.call_args_list, ctx.groups, (1, 4), (2, 5) - ): - args = call.args - kwargs = call.kwargs - assert args[9] is args[10] - assert args[11] is group_ctx.src_slots - assert args[12] is group_ctx.dst_slots - assert args[13] is group_ctx.plan_ring_start - assert args[14] is group_ctx.plan_flush_count - assert args[15] is group_ctx.active_request_indices - assert kwargs["num_heads"] == 4 - assert kwargs["heads_per_group"] == 2 - assert kwargs["max_window"] == 16 - assert kwargs["ring_buffer_len"] == 20 - assert group_ctx.src_slots[:, 0].tolist() == [src_slot] * 2 - assert group_ctx.dst_slots[:, 0].tolist() == [dst_slot] * 2 - assert group_ctx.src_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 2 - assert group_ctx.dst_slots[:, 1].tolist() == [NULL_BLOCK_ID] * 2 - assert group_ctx.plan_ring_start.tolist() == [2, 0] - assert group_ctx.plan_flush_count.tolist() == [6, -1] - assert group_ctx.active_request_indices.tolist() == [0, -1] - assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 6 - assert groups[0][0]._replayssm_prev_num_accepted[2].item() == 0 - assert groups[1][0]._replayssm_prev_num_accepted[4].item() == 6 - assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_materialization_uses_independent_group_mappings( monkeypatch, @@ -660,47 +523,6 @@ def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_postprocess_skips_filtered_pp_row(): - _, config, forward_context, block_tables = _modelwide_replayssm_fixture() - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - - before = [] - for group_ctx in ctx.groups: - group_ctx.ring_start.copy_( - torch.arange(group_ctx.ring_start.numel(), device="cuda") - ) - group_ctx.num_committed.fill_(7) - before.append((group_ctx.ring_start.clone(), group_ctx.num_committed.clone())) - group_ctx.plan_flush_count.fill_(6) - - ctx.postprocess( - idx_mapping=torch.tensor([-1], dtype=torch.int32, device="cuda"), - query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=True, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_post_step=True, - num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), - is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), - live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - torch.accelerator.synchronize() - - for group_ctx, (ring_start, num_committed) in zip(ctx.groups, before): - assert torch.equal(group_ctx.ring_start, ring_start) - assert torch.equal(group_ctx.num_committed, num_committed) - assert group_ctx.plan_flush_count[0].item() == -1 - assert group_ctx.active_request_indices.tolist() == [-1, -1] - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") def test_modelwide_replayssm_postprocess_resets_prefill_slots(monkeypatch): groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index 6d4704136685..252f64590c9b 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -374,23 +374,3 @@ def test_flashinfer_replayssm_all_uses_last_scheduled_state_page(): ).squeeze(1) assert torch.equal(metadata.replayssm_state_indices_d, expected) assert not torch.equal(expected, metadata.state_indices_tensor_d[:, 0]) - - -def test_flashinfer_replayssm_scratch_metadata_fresh_decode(): - checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu") - if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"): - pytest.skip("FlashInfer does not expose ReplaySSM scratch allocation") - - builder = _create_replayssm_builder(16, mamba_backend=MambaBackendEnum.FLASHINFER) - case = REPLAYSSM_BUILD_CASES["fresh_decode"] - meta = _build(builder, case) - - assert meta.write_pos_d is None - assert meta.is_flush_d is None - assert meta.bc_pre_scratch is None - assert meta.replayssm_scratch is not None - assert [tensor.shape for tensor in meta.replayssm_scratch] == [ - (1, 1, 32, 8), - (1, 1, 16), - (1, 1, 32, 8), - ] diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index 548634e23a34..99d2ccfb82fc 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -367,36 +367,6 @@ def test_dcp_scales_attention_but_not_mamba_group_blocks(): ] == [1, 3] -def test_hybrid_eagle_annotation_does_not_change_partial_tail_support(): - config = _make_vllm_config() - config.cache_config.prefix_match_unit = 4 - config.speculative_config.use_eagle_block_drop.return_value = True - kv_cache_config = _make_mamba_hybrid_kv_cache_config() - offloading_config = build_offloading_config(config, kv_cache_config) - - unannotated = SchedulerOffloadConfig.from_spec( - MockOffloadingSpec(offloading_config), config, kv_cache_config - ) - # The offloader's existing fallback treats every group as an EAGLE group. - assert not unannotated.supports_partial_tail - assert [group.is_eagle_group for group in unannotated.kv_group_configs] == [ - True, - True, - ] - - kv_cache_config.kv_cache_groups[0].is_eagle_group = True - annotated = SchedulerOffloadConfig.from_spec( - MockOffloadingSpec(offloading_config), config, kv_cache_config - ) - # Explicit classification narrows the volatile set to attention. Partial - # tails remain disabled because at least one EAGLE group is still present. - assert not annotated.supports_partial_tail - assert [group.is_eagle_group for group in annotated.kv_group_configs] == [ - True, - False, - ] - - def test_preserves_data_parallel_config(): config = _make_vllm_config() config.parallel_config.data_parallel_index = 2 diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index d4850c655ebf..99331e7b354e 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -1429,117 +1429,6 @@ def test_input_batch_reinitialized_after_late_interleave_adjustment(monkeypatch) assert input_batch_cls.call_args.kwargs["cp_kv_cache_interleave_size"] == 16 -def test_replayssm_stp_keeps_accepted_one_without_cpu_copy(monkeypatch): - runner = object.__new__(GPUModelRunner) - runner._use_flashinfer_replayssm = True - runner._needs_prefix_state_migration = False - runner.num_spec_tokens = 0 - runner.speculative_config = None - runner.model_config = SimpleNamespace(is_hybrid=True) - runner.num_accepted_tokens = SimpleNamespace(gpu=torch.ones(2, dtype=torch.int32)) - accepted_cpu = torch.full((2,), 9, dtype=torch.int32) - runner.input_batch = SimpleNamespace(num_accepted_tokens_cpu_tensor=accepted_cpu) - runner.kv_cache_config = Mock() - runner.cache_config = SimpleNamespace(mamba_cache_mode="none") - runner.compilation_config = SimpleNamespace(static_forward_context={}) - runner._get_mamba_bufs = Mock(return_value=Mock()) - runner._get_mamba_state_copy_funcs = Mock(return_value={}) - runner.num_accepted_tokens_event = None - postprocess = Mock() - monkeypatch.setattr( - gpu_model_runner_module.mamba_utils, - "postprocess_mamba_gpu", - postprocess, - ) - - runner._update_states_after_model_execute( - torch.tensor([[42], [-1]], dtype=torch.int64), Mock() - ) - - assert runner.num_accepted_tokens.gpu.tolist() == [1, 1] - assert accepted_cpu.tolist() == [9, 9] - assert postprocess.call_args.kwargs["num_accepted_tokens_cpu_tensor"] is None - - -def test_v1_caches_replayssm_block_copy_tensors_after_binding(monkeypatch): - runner = object.__new__(GPUModelRunner) - runner.device = torch.device("cpu") - runner.cache_config = SimpleNamespace(get_resolved_kv_cache_layout=Mock()) - runner.shared_kv_cache_layers = {} - runner.model_config = SimpleNamespace(hf_config=SimpleNamespace(model_type="mamba")) - runner.compilation_config = SimpleNamespace(static_forward_context={}) - runner.kv_caches = [] - events = [] - extra = torch.empty(1) - monkeypatch.setattr( - gpu_model_runner_module, "allocate_kv_cache", lambda *_args, **_kwargs: {} - ) - monkeypatch.setattr( - gpu_model_runner_module, - "allocate_replayssm_caches", - lambda *_args, **_kwargs: {}, - ) - monkeypatch.setattr( - gpu_model_runner_module, - "bind_kv_cache", - lambda *_args, **_kwargs: events.append("bind"), - ) - - def get_extra(_context): - assert events == ["bind"] - events.append("collect") - return [extra] - - monkeypatch.setattr( - gpu_model_runner_module, "get_replayssm_block_copy_tensors", get_extra - ) - - runner.initialize_kv_cache_tensors( - SimpleNamespace(kv_cache_groups=[]), kernel_block_sizes=[] - ) - - assert events == ["bind", "collect"] - assert len(runner.replayssm_block_copy_tensors) == 1 - assert runner.replayssm_block_copy_tensors[0] is extra - - -def test_v2_block_copy_reuses_cached_replayssm_tensors(monkeypatch): - from vllm.v1.worker.gpu import model_runner as v2_model_runner_module - - runner = object.__new__(v2_model_runner_module.GPUModelRunner) - runner.req_states = SimpleNamespace( - num_computed_tokens_np=np.zeros(1, dtype=np.int32), - prefill_len=SimpleNamespace(np=np.zeros(1, dtype=np.int32)), - num_computed_prefill_tokens=np.zeros(1, dtype=np.int32), - ) - runner.block_tables = Mock() - runner.kv_block_zeroer = Mock() - canonical = torch.empty(1) - extra = torch.empty(1) - runner.kv_caches = [canonical] - runner.replayssm_block_copy_tensors = [extra] - runner.kv_cache_config = SimpleNamespace(num_blocks=4) - scheduler_output = SchedulerOutput.make_empty() - scheduler_output.kv_cache_block_copies = [Mock()] - copy_blocks = Mock() - monkeypatch.setattr( - v2_model_runner_module, "copy_kv_cache_blocks_inplace", copy_blocks - ) - monkeypatch.setattr( - v2_model_runner_module, - "get_replayssm_block_copy_tensors", - Mock(side_effect=AssertionError("hot path must use cached tensors")), - ) - - runner.update_requests(scheduler_output) - - assert copy_blocks.call_count == 1 - copied_tensors = copy_blocks.call_args.args[0] - assert len(copied_tensors) == 2 - assert copied_tensors[0] is canonical - assert copied_tensors[1] is extra - - def test_v2_runner_snapshots_late_interleave_adjustment(monkeypatch): from vllm.v1.worker.gpu import model_runner as v2_model_runner_module diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 5ffb95698233..6067d5551e1f 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -32,7 +32,6 @@ collect_mamba_copy_meta, do_mamba_copy_block, get_mamba_groups, - postprocess_mamba_gpu, preprocess_mamba, preprocess_mamba_align_fused_kernel, stage_postprocess_inputs_to_gpu, @@ -207,153 +206,6 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): assert mamba_state_idx == {"keep": 99} -@pytest.mark.parametrize( - ( - "with_replayssm", - "num_computed_tokens", - "expected_order", - "expected_accepted", - ), - [ - pytest.param(True, 4, ["reset"], 3, id="replayssm-boundary"), - pytest.param(True, 3, ["reset"], 3, id="replayssm-no-boundary"), - pytest.param(False, 4, ["copy"], 1, id="generic"), - ], -) -def test_preprocess_mamba_preserves_live_replayssm_state( - with_replayssm: bool, - num_computed_tokens: int, - expected_order: list[str], - expected_accepted: int, -) -> None: - spec = MagicMock(block_size=4, num_speculative_blocks=0) - cache_config = MagicMock(enable_prefix_caching=True, use_replayssm=True) - input_batch = MagicMock() - input_batch.req_ids = ["r0"] - input_batch.num_accepted_tokens_cpu = np.array([3], dtype=np.int32) - copy_bufs = MagicMock(mamba_group_ids=[0], mamba_spec=spec) - requests = {"r0": MagicMock(num_computed_tokens=num_computed_tokens)} - mamba_state_idx: dict[str, int] = {"r0": 0} - sched = _make_scheduler_output(set(), None, set()) - sched.num_scheduled_tokens = {"r0": 1} - - order: list[str] = [] - device = torch.device("cpu") - align_ctx = MagicMock(is_initialized=True) - align_ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(1, torch.int32, device) - align_ctx.precopy_src_col_buf = _MockCpuGpuBuffer(1, torch.int32, device) - align_ctx.precopy_token_bias_buf = _MockCpuGpuBuffer(1, torch.int32, device) - align_ctx.replayssm = MagicMock() if with_replayssm else None - if align_ctx.replayssm is not None: - align_ctx.replayssm.reset_new_slots.side_effect = lambda **kwargs: order.append( - "reset" - ) - align_ctx.run_fused_precopy.side_effect = lambda **kwargs: order.append("copy") - - preprocess_mamba( - sched, - MagicMock(), - cache_config, - mamba_state_idx, - input_batch, - requests, - {}, - {}, - copy_bufs, - align_ctx=align_ctx, - ) - - assert order == expected_order - assert input_batch.num_accepted_tokens_cpu[0] == expected_accepted - assert align_ctx.precopy_src_col_buf.np[0] == 0 - - -def test_postprocess_mamba_materializes_prefixes(): - order: list[str] = [] - ctx = MagicMock() - ctx.is_initialized = True - ctx.mamba_group_ids = [0] - ctx.mamba_state_idx_buf = MagicMock() - ctx.num_scheduled_tokens_buf = MagicMock() - ctx.num_computed_tokens_buf = MagicMock() - ctx.num_draft_tokens_buf = MagicMock() - ctx.is_prefilling_buf = MagicMock() - ctx.num_accepted_tokens_snapshot = torch.tensor([3], dtype=torch.int32) - accepted = torch.tensor([3], dtype=torch.int32) - - def run_fused_postprocess(**kwargs): - order.append("copy") - kwargs["num_accepted_tokens_gpu"].fill_(1) - - ctx.run_fused_postprocess.side_effect = run_fused_postprocess - ctx.replayssm = MagicMock() - ctx.replayssm.materialize_prefixes = True - ctx.replayssm.postprocess.side_effect = lambda **kwargs: order.append("postprocess") - ctx.replayssm.materialize.side_effect = lambda: order.append("materialize") - block_table = MagicMock() - block_table.get_device_tensor.return_value = torch.zeros((1, 4), dtype=torch.int32) - input_batch = MagicMock() - input_batch.block_table = [block_table] - kv_cache_config = MagicMock() - kv_cache_config.kv_cache_groups = [MagicMock()] - accepted_cpu = torch.zeros(1, dtype=torch.int32) - - postprocess_mamba_gpu( - bufs=MagicMock(postprocess_align=ctx), - num_reqs=1, - num_accepted_tokens_gpu=accepted, - num_accepted_tokens_cpu_tensor=accepted_cpu, - input_batch=input_batch, - kv_cache_config=kv_cache_config, - forward_context={}, - mamba_state_copy_funcs={}, - run_prefix_state_migration=True, - ) - - assert order == ["copy", "postprocess", "materialize"] - assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is ( - ctx.num_accepted_tokens_snapshot - ) - assert accepted_cpu.tolist() == [1] - - -def test_postprocess_mamba_none_skips_prefix_copy(): - ctx = MagicMock() - ctx.is_initialized = True - ctx.mamba_group_ids = [0] - ctx.mamba_state_idx_buf = MagicMock(gpu=torch.zeros(1, dtype=torch.int32)) - ctx.num_scheduled_tokens_buf = MagicMock(gpu=torch.tensor([4], dtype=torch.int32)) - ctx.num_computed_tokens_buf = MagicMock(gpu=torch.tensor([20], dtype=torch.int32)) - ctx.num_draft_tokens_buf = MagicMock(gpu=torch.tensor([3], dtype=torch.int32)) - ctx.is_prefilling_buf = MagicMock(gpu=torch.tensor([False])) - ctx.block_size = 1024 - ctx.replayssm = MagicMock() - ctx.replayssm.materialize_prefixes = False - input_batch = MagicMock() - input_batch.block_table = [] - accepted = torch.tensor([2], dtype=torch.int32) - accepted_cpu = torch.zeros(1, dtype=torch.int32) - - postprocess_mamba_gpu( - bufs=MagicMock(postprocess_align=ctx), - num_reqs=1, - num_accepted_tokens_gpu=accepted, - num_accepted_tokens_cpu_tensor=accepted_cpu, - input_batch=input_batch, - kv_cache_config=MagicMock(), - forward_context={}, - mamba_state_copy_funcs={}, - run_prefix_state_migration=False, - ) - - ctx.run_fused_postprocess.assert_not_called() - assert ctx.replayssm.postprocess.call_count == 1 - assert ctx.replayssm.postprocess.call_args.kwargs["num_accepted_tokens"] is accepted - assert ctx.replayssm.postprocess.call_args.kwargs["live_cols"] is None - ctx.replayssm.materialize.assert_not_called() - assert accepted_cpu.tolist() == [2] - - # ----------------------------------------------------------------------------- # Golden tests for postprocess_mamba_fused_kernel # ----------------------------------------------------------------------------- @@ -478,62 +330,6 @@ def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): ] -def test_gpu_context_initializes_flashinfer_replayssm_lifecycle(): - cfg = _TestConfig(num_layers=1) - device = torch.device("cpu") - kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) - gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) - attention = _make_mock_attention( - torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), - torch.empty(cfg.num_blocks, cfg.temporal_state_dim), - ) - attention.use_replayssm = True - attention.mamba_config.backend = MambaBackendEnum.FLASHINFER - - model_ctx = object() - with patch( - "vllm.v1.worker.mamba_utils.ReplaySSMModelContext.create", - return_value=model_ctx, - ) as create: - gpu_ctx.initialize_from_forward_context( - kv_cache_config, - {"layer_0": attention}, - _COPY_FUNCS, - [torch.empty(1, 4, dtype=torch.int32)], - ) - - assert gpu_ctx.has_flashinfer_replayssm - assert gpu_ctx.replayssm is model_ctx - create.assert_called_once() - - -def test_gpu_context_rejects_missing_replayssm_lifecycle(): - cfg = _TestConfig(num_layers=1) - device = torch.device("cpu") - kv_cache_config = _make_kv_cache_config(cfg, ["layer_0"]) - gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) - attention = _make_mock_attention( - torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), - torch.empty(cfg.num_blocks, cfg.temporal_state_dim), - ) - attention.use_replayssm = True - attention.mamba_config.backend = MambaBackendEnum.FLASHINFER - - with ( - patch( - "vllm.v1.worker.mamba_utils.ReplaySSMModelContext.create", - return_value=None, - ), - pytest.raises(RuntimeError, match="could not be initialized"), - ): - gpu_ctx.initialize_from_forward_context( - kv_cache_config, - {"layer_0": attention}, - _COPY_FUNCS, - [torch.empty(1, 4, dtype=torch.int32)], - ) - - def test_gpu_context_rejects_mixed_replayssm_and_baseline_layers(): cfg = _TestConfig(num_layers=2) device = torch.device("cpu") @@ -861,18 +657,6 @@ def test_mamba_copy_funcs_accept_nonempty_canonical_prefix(): validate_mamba_state_copy_funcs({replayssm_spec: [0]}, invalid_copy_funcs) -def test_gdn_copy_funcs_cover_copyable_prefix_only(): - gdn_spec = MambaSpec( - block_size=16, - shapes=((4, 4), (2, 4, 4), (2, 4), (2, 4)), - dtypes=(torch.float16,) * 4, - mamba_type=MambaAttentionBackendEnum.GDN_ATTN, - mamba_cache_mode="none", - ) - - validate_mamba_state_copy_funcs({gdn_spec: [0]}, _COPY_FUNCS) - - def test_mamba_groups_support_mixed_specs_in_uniform_group(): gdn_spec = MambaSpec( block_size=16, @@ -1093,39 +877,6 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): ) -def test_stage_postprocess_inputs_to_gpu_skips_none_mode_live_col(): - device = torch.device("cpu") - ctx = _make_staging_ctx(max_num_reqs=4, device=device) - ctx.replayssm = MagicMock(materialize_prefixes=False) - ctx.mamba_state_idx_buf.cpu.fill_(17) - ctx.mamba_state_idx_buf.gpu.fill_(23) - scheduler_output = _make_postprocess_scheduler_output( - req_ids=["req_a", "req_b"], - num_scheduled_tokens={"req_a": 4, "req_b": 1}, - ) - requests = _make_requests( - ["req_a", "req_b"], - [20, 7], - [[0], [0]], - num_prompt_tokens=[10, 8], - ) - - stage_postprocess_inputs_to_gpu( - ctx, - scheduler_output, - ["req_a", "req_b"], - 2, - requests, - {}, - run_prefix_state_migration=False, - ) - - np.testing.assert_array_equal(ctx.mamba_state_idx_buf.np[:2], [17, 17]) - assert ctx.mamba_state_idx_buf.gpu[:2].tolist() == [23, 23] - np.testing.assert_array_equal(ctx.num_scheduled_tokens_buf.np[:2], [4, 1]) - np.testing.assert_array_equal(ctx.num_computed_tokens_buf.np[:2], [20, 7]) - - def test_gpu_context_ignores_auxiliary_cache_tensors() -> None: device = torch.device("cpu") config = _TestConfig(num_layers=1) From d2dae3d137d4052387fea73b8119482791cfc141 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 10:12:01 +0200 Subject: [PATCH 37/53] test: second trim pass on ReplaySSM coverage Remove unit tests whose failure mode is already loud in the retained end-to-end matrix, plus the mock scaffolding they required: - ssu_dispatch: 7 of 9 ReplaySSMModelContext tests. They drive a Mock materializer and assert on internal plan buffers; a wrong commit or plan shows up directly as a logprob mismatch in the none/align/all x STP/MTP e2e cells. Kept the none-mode ring-wrap lifecycle test and the v1/v2 metadata-convention parity test. - all-mode state-page and spec-decode anchor selection: covered by the all-mode and all+MTP prefix-cache e2es. - mixed-layer and copy-func validation asserts: the rejected configurations are not reachable from user-facing config. - PRESERVE_ACCEPTED kernel flag: subsumed by the align-mode accepted count test and the align+ngram e2e. - MRV2 PP mixed-batch postprocess: ordering is already asserted by test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank. - none-mode prefix-migration skip, Triton ring block-copy, and the autotune non-padding-slot skip: single-branch wiring checks. Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 308 ------------------ tests/model_executor/test_replayssm_warmup.py | 10 - .../test_mamba_update_block_table.py | 48 +-- .../test_replayssm_metadata_builder.py | 33 -- .../worker/test_gpu_model_runner_v2_eplb.py | 50 --- .../worker/test_mamba_hybrid_model_state.py | 38 --- tests/v1/worker/test_mamba_utils.py | 97 ------ tests/v1/worker/test_utils.py | 29 +- 8 files changed, 7 insertions(+), 606 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 409c987a0265..22112d5a7aa0 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -20,7 +20,6 @@ ) 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, @@ -375,192 +374,6 @@ def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None assert mixers[0]._replayssm_prev_num_accepted[live_slot].item() == 0 -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_materialization_uses_independent_group_mappings( - monkeypatch, -): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - groups[0][0]._replayssm_prev_num_accepted[1] = 4 - groups[1][0]._replayssm_prev_num_accepted[5] = 9 - # The request has a source in group zero but a null source in group one. - block_tables[1][0, 0] = NULL_BLOCK_ID - materializer = Mock() - monkeypatch.setattr( - ssu_dispatch, "_load_replayssm_materialize", lambda: materializer - ) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.tensor([6, 0], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), - is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - ctx.materialize() - torch.accelerator.synchronize() - - assert materializer.call_count == 2 - assert ctx.groups[0].active_request_indices.tolist() == [0, -1] - assert ctx.groups[0].plan_flush_count.tolist() == [6, -1] - assert ctx.groups[1].active_request_indices.tolist() == [-1, -1] - assert ctx.groups[1].plan_flush_count.tolist() == [-1, -1] - assert ctx.groups[1].src_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 2 - assert ctx.groups[1].dst_slots[:, 0].tolist() == [NULL_BLOCK_ID] * 2 - assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 9 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_compacts_sparse_materialization_requests(monkeypatch): - _, config, forward_context, block_tables = _modelwide_replayssm_fixture() - block_tables[0][1] = torch.tensor([3, 2, 1], device="cuda") - block_tables[1][1] = torch.tensor([7, 6, 5], device="cuda") - materializer = Mock() - monkeypatch.setattr( - ssu_dispatch, "_load_replayssm_materialize", lambda: materializer - ) - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([1, 4], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.tensor([0, 6], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.tensor([1, 2], dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([False, False], device="cuda"), - live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_reqs=2, - ) - ctx.materialize() - torch.accelerator.synchronize() - - assert materializer.call_count == 2 - for group_ctx in ctx.groups: - # Request 1 advances from token 6 to the block boundary at token 8. - assert group_ctx.plan_flush_count.tolist() == [-1, 2] - assert group_ctx.active_request_indices.tolist() == [1, -1] - - # A shorter batch clears the compacted active tail. Fixed-capacity plan and - # slot-table tails may stay stale because FlashInfer stops at the first -1. - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), - is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), - live_cols=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - torch.accelerator.synchronize() - - for group_ctx in ctx.groups: - assert group_ctx.plan_flush_count[0].item() == -1 - assert group_ctx.active_request_indices.tolist() == [-1, -1] - assert torch.all(group_ctx.src_slots[:, 0] == NULL_BLOCK_ID) - assert torch.all(group_ctx.dst_slots[:, 0] == NULL_BLOCK_ID) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_postprocess_commits_checkpoint_boundary(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, source_slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[source_slot] = 2 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 13 - kernel = Mock() - monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=torch.tensor([1], dtype=torch.int32, device="cuda"), - query_metadata=torch.tensor([0, 4], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=True, - num_computed_tokens=torch.tensor([0, 8], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=True, - num_accepted_tokens=torch.tensor([1, 3], dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([False, False], device="cuda"), - live_cols=torch.tensor([-1, 0], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - ctx.materialize() - torch.accelerator.synchronize() - - assert kernel.call_count == 2 - for group_ctx in ctx.groups: - assert group_ctx.plan_ring_start.tolist() == [15, 0] - assert group_ctx.plan_flush_count.tolist() == [3, -1] - for mixers, source_slot, destination_slot in zip(groups, (1, 4), (2, 5)): - assert mixers[0]._replayssm_ring_start[source_slot].item() == 15 - assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 3 - assert mixers[0]._replayssm_ring_start[destination_slot].item() == 0 - assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_postprocess_resets_prefill_slots(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): - for source_slot in source_slots: - mixers[0]._replayssm_ring_start[source_slot] = 7 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 9 - kernel = Mock() - monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([True, False], device="cuda"), - live_cols=torch.tensor([1, -1], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - torch.accelerator.synchronize() - - assert kernel.call_count == 0 - for mixers, source_slots in zip(groups, ((1, 2), (4, 5))): - for source_slot in source_slots: - assert mixers[0]._replayssm_ring_start[source_slot].item() == 0 - assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 0 - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") @pytest.mark.parametrize( ( @@ -609,124 +422,3 @@ def test_modelwide_replayssm_single_token_final_prefill_commits_as_decode( for mixers, source_slot in zip(groups, (1, 4)): assert mixers[0]._replayssm_ring_start[source_slot].item() == 7 assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 6 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_prefill_commit_publishes_exact_copy(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, destination_slot in zip(groups, (2, 5)): - mixers[0]._replayssm_ring_start[destination_slot] = 7 - mixers[0]._replayssm_prev_num_accepted[destination_slot] = 9 - materializer = Mock() - monkeypatch.setattr( - ssu_dispatch, "_load_replayssm_materialize", lambda: materializer - ) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([8, 0], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.zeros(2, dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.ones(2, dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([True, False], device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - ctx.materialize() - torch.accelerator.synchronize() - - assert materializer.call_count == 2 - for group_ctx in ctx.groups: - assert group_ctx.plan_ring_start.tolist() == [0, 0] - assert group_ctx.plan_flush_count.tolist() == [0, -1] - for mixers, destination_slot in zip(groups, (2, 5)): - assert mixers[0]._replayssm_ring_start[destination_slot].item() == 0 - assert mixers[0]._replayssm_prev_num_accepted[destination_slot].item() == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_resets_only_group_specific_fresh_slot(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - # Group zero aliases the two logical columns and therefore already has a - # valid source. Group one has no physical source for the same logical move. - block_tables[0][0, 1] = block_tables[0][0, 0] - block_tables[1][0, 0] = NULL_BLOCK_ID - groups[0][0]._replayssm_ring_start[1] = 2 - groups[0][0]._replayssm_prev_num_accepted[1] = 4 - groups[1][0]._replayssm_ring_start[5] = 7 - groups[1][0]._replayssm_prev_num_accepted[5] = 9 - kernel = Mock() - monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.reset_new_slots( - idx_mapping=torch.tensor([0], dtype=torch.int32, device="cuda"), - src_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - dst_cols=torch.tensor([1, 0], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - torch.accelerator.synchronize() - - assert kernel.call_count == 0 - assert groups[0][0]._replayssm_ring_start[1].item() == 2 - assert groups[0][0]._replayssm_prev_num_accepted[1].item() == 4 - assert groups[1][0]._replayssm_ring_start[5].item() == 0 - assert groups[1][0]._replayssm_prev_num_accepted[5].item() == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_postprocess_materializes_in_place(monkeypatch): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, source_slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[source_slot] = 2 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 - kernel = Mock() - monkeypatch.setattr(ssu_dispatch, "_load_replayssm_materialize", lambda: kernel) - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor([4, 0], dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=False, - num_computed_tokens=torch.tensor([2, 0], dtype=torch.int32, device="cuda"), - num_computed_is_post_step=False, - num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), - is_prefilling=torch.zeros(2, dtype=torch.bool, device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - ctx.materialize() - torch.accelerator.synchronize() - - assert kernel.call_count == 2 - for group_ctx, slot in zip(ctx.groups, (1, 4)): - assert group_ctx.plan_ring_start.tolist() == [2, 0] - assert group_ctx.plan_flush_count.tolist() == [6, -1] - assert group_ctx.src_slots[:, 0].tolist() == [slot] * 2 - assert group_ctx.dst_slots[:, 0].tolist() == [slot] * 2 - for mixers, slot in zip(groups, (1, 4)): - assert mixers[0]._replayssm_ring_start[slot].item() == 0 - assert mixers[0]._replayssm_prev_num_accepted[slot].item() == 0 diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index b9dadec906d7..594999ac1e30 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -110,16 +110,6 @@ def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): assert result is None -def test_replayssm_autotune_kwargs_skipped_without_non_padding_slot(): - with patch.object( - warmup, "flashinfer_replayssm_autotune_supported", return_value=True - ): - result = warmup._replayssm_autotune_kwargs( - _autotune_runner(num_blocks=1), PREFILL_KWARGS - ) - assert result is None - - def test_replayssm_autotune_slots_restore_state_and_trackers(): mixer = MambaMixer2.__new__(MambaMixer2) torch.nn.Module.__init__(mixer) diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 9a909952ce0f..4ec138270203 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -16,11 +16,7 @@ import torch -from tests.v1.attention.utils import ( - BatchSpec, - MockMambaBuilder, - create_common_attn_metadata, -) +from tests.v1.attention.utils import MockMambaBuilder from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadata from vllm.v1.kv_cache_interface import MambaSpec @@ -325,48 +321,6 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): assert builder.block_idx_last_scheduled_token_prev_step.dtype == torch.int32 -def test_all_spec_decode_without_drafts_uses_computed_state_anchor(): - """A step with no scheduled drafts still needs a valid input-state table. - - This occurs after a prefix hit when speculative decoding is configured but - the current step contains only one target token. The previous-step buffer - is not passed for that step, so the last computed block is the input anchor. - """ - block_size = 16 - seq_lens = [33, 49] - query_lens = [1, 1] - config = _make_vllm_config( - max_model_len=256, - max_num_seqs=len(seq_lens), - num_speculative_tokens=3, - ) - config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - spec = MambaSpec( - block_size=block_size, - shapes=((1,), (1,)), - dtypes=(torch.float32,), - mamba_cache_mode="all", - num_speculative_blocks=2, - ) - builder = MockMambaBuilder(spec, ["layer0"], config, torch.device("cpu")) - common = create_common_attn_metadata( - BatchSpec(seq_lens=seq_lens, query_lens=query_lens), - block_size, - torch.device("cpu"), - arange_block_indices=True, - ).replace(is_prefilling=torch.zeros(len(seq_lens), dtype=torch.bool)) - - metadata = builder.build(0, common) - - assert metadata.num_decodes == len(seq_lens) - assert metadata.block_idx_last_scheduled_token_prev_step is not None - expected = torch.tensor([1, 2], dtype=torch.int32) - torch.testing.assert_close( - metadata.block_idx_last_scheduled_token_prev_step, - expected, - ) - - def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode(): """Without spec decode, the prev-step buffer is unused and must not be allocated — the input anchor reduces to last_computed_token.""" diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index 252f64590c9b..740988c81352 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -341,36 +341,3 @@ def test_flashinfer_replayssm_state_indices_are_stable_for_full_cudagraph(): assert second_indices is not None assert second_indices.data_ptr() == first_ptr assert torch.equal(second_indices, second.state_indices_tensor_d[:, 0]) - - -def test_flashinfer_replayssm_all_uses_last_scheduled_state_page(): - builder = _create_replayssm_builder( - 16, - mamba_cache_mode="all", - mamba_backend=MambaBackendEnum.FLASHINFER, - num_speculative_tokens=3, - ) - builder.compilation_config.cudagraph_mode = CUDAGraphMode.NONE - - metadata = _build( - builder, - ReplaySSMBuildCase( - seq_lens=[34, 50], - query_lens=[1, 1], - is_prefilling=[False, False], - decode_base=[33, 49], - buffer_len=16, - expected_write_pos=[], - expected_is_flush=[], - mamba_cache_mode="all", - ), - ) - - assert metadata.replayssm_state_indices_d is not None - assert metadata.block_idx_last_scheduled_token is not None - live_cols = metadata.block_idx_last_scheduled_token[:2].to(torch.int64) - expected = metadata.state_indices_tensor_d.gather( - 1, live_cols.unsqueeze(1) - ).squeeze(1) - assert torch.equal(metadata.replayssm_state_indices_d, expected) - assert not torch.equal(expected, metadata.state_indices_tensor_d[:, 0]) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index ccc4740ca787..4d3457f6e781 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -6,7 +6,6 @@ from types import SimpleNamespace from typing import Any -import numpy as np import torch from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry @@ -280,52 +279,3 @@ def propose(*_, **__): mrv2.GPUModelRunner.sample_tokens(runner, None) assert events == ["postprocess", "draft"] - - -def test_v2_sample_tokens_pp_mixed_batch_uses_ordinary_postprocess(monkeypatch): - events = [] - runner = _make_runner(is_last_pp_rank=False, num_speculative_steps=0) - idx_mapping = torch.tensor([3, 7], dtype=torch.int64) - query_start_loc = torch.tensor([0, 1, 2], dtype=torch.int32) - input_batch = SimpleNamespace( - num_reqs=2, - idx_mapping=idx_mapping, - idx_mapping_np=np.array([3, 7], dtype=np.intp), - num_computed_tokens_np=np.array([3, 2], dtype=np.int32), - prefill_len_np=np.array([4, 6], dtype=np.int32), - num_scheduled_tokens=np.array([1, 1], dtype=np.int32), - query_start_loc=query_start_loc, - ) - runner.execute_model_state = SimpleNamespace( - input_batch=input_batch, - attn_metadata=None, - slot_mappings_by_layer=None, - hidden_states=None, - aux_hidden_states=None, - dp_sync=None, - finished_req_ids=set(), - ec_connector_output=None, - routed_experts=None, - ) - postprocess_args = [] - runner.model_state = SimpleNamespace( - postprocess_state=lambda *args: postprocess_args.append(args) - ) - - def receive(*_: Any) -> bool: - events.append("receive") - return False - - runner.pp_handler = SimpleNamespace(receive=receive) - runner.postprocess_num_computed_tokens = lambda *_: events.append( - "postprocess_num_computed_tokens" - ) - runner.eplb.step = lambda *args, **kwargs: events.append("eplb") - output = mrv2.GPUModelRunner.sample_tokens(runner, None) - - assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) - assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] - assert len(postprocess_args) == 1 - published_mapping, num_sampled = postprocess_args[0] - assert published_mapping is idx_mapping - assert num_sampled == 0 diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index f07602c19eb5..214f795da8aa 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -59,44 +59,6 @@ def test_postprocess_state_scalar_with_int32_mapping( torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") -def test_flashinfer_replayssm_none_postprocess_skips_prefix_migration() -> None: - state = object.__new__(MambaHybridModelState) - state._align_mode = False - state._needs_prefix_state_migration = False - state._use_flashinfer_replayssm = True - state.recoverssm = None - state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") - state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") - replayssm = Mock() - replayssm.materialize_prefixes = False - ctx = Mock( - is_initialized=True, - replayssm=replayssm, - block_size=1024, - ) - state._mamba_ctx = ctx - idx_mapping = torch.tensor([2], dtype=torch.int32, device="cuda") - num_sampled = torch.tensor([2], dtype=torch.int32, device="cuda") - num_computed = torch.tensor([0, 0, 20, 0], dtype=torch.int32, device="cuda") - query_start_loc = torch.tensor([0, 4], dtype=torch.int32, device="cuda") - state._replayssm_query_start_loc = query_start_loc - - state.postprocess_state( - idx_mapping, - num_sampled, - num_computed_tokens=num_computed, - ) - - ctx.run_fused_postprocess_align.assert_not_called() - assert replayssm.postprocess.call_count == 1 - kwargs = replayssm.postprocess.call_args.kwargs - assert state.num_accepted_tokens_gpu.tolist() == [1, 1, 2, 1] - assert kwargs["num_accepted_tokens"] is state.num_accepted_tokens_gpu - assert kwargs["live_cols"] is None - assert state._replayssm_query_start_loc is None - - @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") def test_flashinfer_replayssm_prefix_uses_original_accepted_counts() -> None: state = object.__new__(MambaHybridModelState) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 6067d5551e1f..a40df81db0ec 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -8,7 +8,6 @@ import pytest import torch -from vllm.config.mamba import MambaBackendEnum from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncsByType, @@ -33,9 +32,7 @@ do_mamba_copy_block, get_mamba_groups, preprocess_mamba, - preprocess_mamba_align_fused_kernel, stage_postprocess_inputs_to_gpu, - validate_mamba_state_copy_funcs, ) # Conv + temporal copy specs, in the order the tests' MambaSpec shapes expect. @@ -50,43 +47,6 @@ } -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -@pytest.mark.parametrize( - ("preserve_accepted", "expected_accepted"), [(False, 1), (True, 3)] -) -def test_preprocess_mamba_preserves_replayssm_accepted_offset( - preserve_accepted: bool, expected_accepted: int -): - device = torch.device("cuda") - idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) - state_idx = torch.tensor([0], dtype=torch.int32, device=device) - num_computed = torch.tensor([4], dtype=torch.int32, device=device) - query_start = torch.tensor([0, 1], dtype=torch.int32, device=device) - num_accepted = torch.tensor([3], dtype=torch.int32, device=device) - src_col = torch.empty(1, dtype=torch.int32, device=device) - src_off = torch.empty(1, dtype=torch.int32, device=device) - - preprocess_mamba_align_fused_kernel[(1,)]( - idx_mapping, - state_idx, - num_computed, - query_start, - num_accepted, - src_col, - src_off, - 1, - BLOCK_SIZE=32, - MAMBA_BLOCK_SIZE=4, - PRESERVE_ACCEPTED=preserve_accepted, - ) - torch.accelerator.synchronize() - - assert state_idx.item() == 1 - assert src_col.item() == 0 - assert src_off.item() == 2 - assert num_accepted.item() == expected_accepted - - def postprocess_mamba( scheduler_output: "SchedulerOutput", kv_cache_config: "KVCacheConfig", @@ -330,33 +290,6 @@ def test_gpu_context_reinterprets_high_data_ptrs_for_int64_metadata(): ] -def test_gpu_context_rejects_mixed_replayssm_and_baseline_layers(): - cfg = _TestConfig(num_layers=2) - device = torch.device("cpu") - layer_names = ["layer_0", "layer_1"] - kv_cache_config = _make_kv_cache_config(cfg, layer_names) - gpu_ctx = _make_gpu_ctx(cfg, kv_cache_config, device) - forward_context = { - name: _make_mock_attention( - torch.empty(cfg.num_blocks, cfg.conv_width, cfg.conv_inner_dim), - torch.empty(cfg.num_blocks, cfg.temporal_state_dim), - ) - for name in layer_names - } - forward_context["layer_0"].use_replayssm = True - forward_context["layer_0"].mamba_config.backend = MambaBackendEnum.FLASHINFER - forward_context["layer_1"].use_replayssm = False - forward_context["layer_1"].mamba_config.backend = MambaBackendEnum.TRITON - - with pytest.raises(ValueError, match="mixed FlashInfer ReplaySSM"): - gpu_ctx.initialize_from_forward_context( - kv_cache_config, - forward_context, - _COPY_FUNCS, - [torch.empty(1, 4, dtype=torch.int32)], - ) - - def _make_postprocess_scheduler_output( req_ids: list[str], num_scheduled_tokens: dict[str, int], @@ -627,36 +560,6 @@ def test_mamba_groups_support_different_state_specs(): assert ctx.state_conv_widths.tolist() == [4, 0, 4, 0, 12] -def test_mamba_copy_funcs_accept_nonempty_canonical_prefix(): - replayssm_spec = MambaSpec( - block_size=16, - shapes=((4, 4), (2, 4, 4)), - dtypes=(torch.float16,) * 2, - replayssm_shapes=((2, 8, 4), (2, 8), (1, 8, 4)), - replayssm_dtypes=(torch.float16,) * 3, - mamba_type=MambaAttentionBackendEnum.MAMBA2, - mamba_cache_mode="align", - ) - - for valid_funcs in ((get_conv_copy_spec,), _DEFAULT_COPY_FUNCS): - copy_funcs = { - **_COPY_FUNCS, - MambaAttentionBackendEnum.MAMBA2: valid_funcs, - } - validate_mamba_state_copy_funcs({replayssm_spec: [0]}, copy_funcs) - - for invalid_funcs in ( - (), - (*_DEFAULT_COPY_FUNCS, get_temporal_copy_spec), - ): - invalid_copy_funcs = { - **_COPY_FUNCS, - MambaAttentionBackendEnum.MAMBA2: invalid_funcs, - } - with pytest.raises(AssertionError, match="non-empty copyable prefix"): - validate_mamba_state_copy_funcs({replayssm_spec: [0]}, invalid_copy_funcs) - - def test_mamba_groups_support_mixed_specs_in_uniform_group(): gdn_spec = MambaSpec( block_size=16, diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 8c6acd825591..ffbcc7a43d44 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -16,21 +16,18 @@ class _TestReplaySSMMixer(MambaMixer2): - _state_shapes = ((2,), (3,)) - _state_dtypes = (torch.float32, torch.float32) - - def __init__(self, backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER) -> None: + def __init__(self) -> None: torch.nn.Module.__init__(self) self.use_replayssm = True - self.mamba_config = MambaConfig(backend=backend) + 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) def get_state_shape(self) -> tuple[tuple[int, ...], ...]: - return self._state_shapes + return ((2,), (3,)) def get_state_dtype(self) -> tuple[torch.dtype, ...]: - return self._state_dtypes + return (torch.float32, torch.float32) def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: return ((4,), (5,), (6,)) @@ -39,8 +36,8 @@ def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: return (torch.float32,) * 3 -def _packed_replayssm_cache(num_blocks: int, fill_value: int = 0) -> torch.Tensor: - return torch.full((num_blocks, 1, 1, 20), fill_value, dtype=torch.int8) +def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: + return torch.zeros((num_blocks, 1, 1, 20), dtype=torch.int8) def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): @@ -144,20 +141,6 @@ def test_replayssm_block_copy_includes_rings_and_group_trackers(monkeypatch): assert mixers[1]._replayssm_prev_num_accepted[dst].item() == 31 -def test_replayssm_block_copy_includes_triton_rings_without_trackers(): - mixer = _TestReplaySSMMixer(MambaBackendEnum.TRITON) - mixer.kv_cache = tuple(torch.zeros(4, 1) for _ in range(2)) - mixer.replayssm_cache = tuple(torch.zeros(4, 1) for _ in range(3)) - - tensors = get_replayssm_block_copy_tensors({"layers.0.mixer": mixer}) - - assert len(tensors) == 3 - assert all( - actual is expected - for actual, expected in zip(tensors, mixer.replayssm_cache, strict=True) - ) - - def test_bind_kv_cache(default_vllm_config): from vllm.model_executor.layers.attention import Attention From 2409c0c825229a762ae959b6e4d0ac6ee37375e7 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 10:25:44 +0200 Subject: [PATCH 38/53] test: drop superficial ReplaySSM assertions - test_utils: drop the tracker dtype assertion; the test is about which layers share a tracker, and the allocation dtype is explicit. - test_replayssm_warmup: drop the is_profile assertion. PREFILL_KWARGS already sets it, so it only checks a constant passes through. - test_replayssm_decode: callable() on a name that just imported is always true, so it added nothing over the import itself. Keep the replayssm_materialize signature probe, which does gate on a newer FlashInfer API. Signed-off-by: Andrii Skliar --- tests/model_executor/test_replayssm_warmup.py | 1 - tests/v1/e2e/test_replayssm_decode.py | 8 +++----- tests/v1/worker/test_utils.py | 1 - 3 files changed, 3 insertions(+), 7 deletions(-) diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py index 594999ac1e30..121571baaccb 100644 --- a/tests/model_executor/test_replayssm_warmup.py +++ b/tests/model_executor/test_replayssm_warmup.py @@ -76,7 +76,6 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): assert max_num_reqs == expected_num_reqs assert decode_kwargs["num_tokens"] == expected_num_reqs * query_len assert decode_kwargs["uniform_decode"] is True - assert decode_kwargs["is_profile"] is True if runner_kwargs.get("use_v2_model_runner"): assert decode_kwargs["valid_dummy_state_slots"] is True assert "profile_seq_lens" not in decode_kwargs diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 190b1e8f4170..dad52ba866af 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -13,21 +13,19 @@ from ...utils import large_gpu_mark, multi_gpu_test try: - from flashinfer.mamba.checkpointing_ssu import ( + from flashinfer.mamba.checkpointing_ssu import ( # noqa: F401 CheckpointingSSURunner, allocate_checkpointing_ssu_scratch, ) - HAS_FLASHINFER_CHECKPOINTING_SSU = callable(CheckpointingSSURunner) and callable( - allocate_checkpointing_ssu_scratch - ) + HAS_FLASHINFER_CHECKPOINTING_SSU = True except ImportError: HAS_FLASHINFER_CHECKPOINTING_SSU = False try: from flashinfer.mamba.replayssm_materialize import replayssm_materialize - HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = callable(replayssm_materialize) and ( + HAS_FLASHINFER_REPLAYSSM_MATERIALIZE = ( "active_request_indices" in signature(replayssm_materialize).parameters ) except ImportError: diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index ffbcc7a43d44..33a166524739 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -84,7 +84,6 @@ def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): 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 group_tracker.dtype == torch.int32 assert torch.count_nonzero(group_tracker) == 0 From 7bb8f001a2717da00c8fb74e6d7bfb3a8cc00df4 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 10:36:33 +0200 Subject: [PATCH 39/53] test: reduce ReplaySSM suite to core coverage Rely on retained layout and ownership contracts plus real-model E2E coverage instead of duplicate runner, configuration, and lifecycle mocks. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 195 ------------------ tests/test_config.py | 13 +- .../core/test_single_type_kv_cache_manager.py | 5 +- tests/v1/e2e/test_replayssm_decode.py | 46 +---- .../worker/test_gpu_model_runner_v2_eplb.py | 77 ------- .../worker/test_mamba_hybrid_model_state.py | 41 ---- 6 files changed, 9 insertions(+), 368 deletions(-) diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 22112d5a7aa0..23b5bcff62e9 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -1,18 +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 -import vllm.model_executor.layers.mamba.ops.ssu_dispatch as ssu_dispatch from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( FlashInferSSUBackend, - ReplaySSMModelContext, TritonSSUBackend, get_mamba_ssu_backend, initialize_mamba_ssu_backend, @@ -230,195 +227,3 @@ def test_replayssm_physical_ring_shape( (8, expected_ring_len), (2, expected_ring_len, 16), ) - - -def _materialize_mixer(device: str = "cpu") -> Mock: - mixer = Mock() - mixer.kv_cache = [ - torch.empty(0, device=device), - torch.empty(8, 4, 3, 5, device=device), - ] - mixer.replayssm_cache = [ - torch.empty(8, 4, 20, 3, device=device), - torch.empty(8, 4, 20, device=device), - torch.empty(8, 2, 20, 5, device=device), - ] - mixer.A = torch.empty(4, 3, 5, device=device) - mixer._replayssm_ring_start = torch.arange(8, dtype=torch.int32, device=device) - mixer._replayssm_prev_num_accepted = torch.zeros( - 8, dtype=torch.int32, device=device - ) - mixer.replayssm_buffer_len = 16 - mixer.mamba_config = SimpleNamespace( - backend=MambaBackendEnum.FLASHINFER, - enable_stochastic_rounding=True, - stochastic_rounding_philox_rounds=6, - ) - return mixer - - -def _modelwide_replayssm_fixture(cache_mode: str = "align"): - groups = [ - [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], - [_materialize_mixer(device="cuda"), _materialize_mixer(device="cuda")], - ] - layer_names: list[list[str]] = [] - forward_context = {} - for group_idx, mixers in enumerate(groups): - # Layers in one cache group share the physical tracker namespace. - for mixer in mixers[1:]: - mixer._replayssm_ring_start = mixers[0]._replayssm_ring_start - mixer._replayssm_prev_num_accepted = mixers[0]._replayssm_prev_num_accepted - names = [f"group{group_idx}.layer{layer_idx}" for layer_idx in range(2)] - layer_names.append(names) - for name, mixer in zip(names, mixers): - mixer.use_replayssm = True - forward_context[name] = mixer - - config = Mock() - specs = [ - MambaSpec( - block_size=1024 if cache_mode == "none" else 4, - shapes=((4, 3, 5),), - dtypes=(torch.float32,), - mamba_cache_mode=cache_mode, - ) - for _ in layer_names - ] - config.kv_cache_groups = [ - Mock(layer_names=names, kv_cache_spec=spec) - for names, spec in zip(layer_names, specs) - ] - block_tables = [ - torch.tensor([[1, 2, 3], [0, 0, 0]], dtype=torch.int32, device="cuda"), - torch.tensor([[4, 5, 6], [0, 0, 0]], dtype=torch.int32, device="cuda"), - ] - return groups, config, forward_context, block_tables - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -def test_modelwide_replayssm_none_commits_trackers_without_materialization( - monkeypatch, -): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture( - cache_mode="none" - ) - materializer = Mock() - monkeypatch.setattr( - ssu_dispatch, "_load_replayssm_materialize", lambda: materializer - ) - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - assert not ctx.materialize_prefixes - - query_len = torch.zeros(2, dtype=torch.int32, device="cuda") - num_computed = torch.zeros(2, dtype=torch.int32, device="cuda") - accepted = torch.ones(2, dtype=torch.int32, device="cuda") - is_prefilling = torch.zeros(2, dtype=torch.bool, device="cuda") - - def step(*, scheduled: int, num_accepted: int, prefilling: bool = False) -> None: - query_len[0] = scheduled - accepted[0] = num_accepted - is_prefilling[0] = prefilling - ctx.postprocess( - idx_mapping=None, - query_metadata=query_len, - query_metadata_is_cumulative=False, - num_computed_tokens=num_computed, - num_computed_is_post_step=False, - num_accepted_tokens=accepted, - is_prefilling=is_prefilling, - live_cols=None, - num_reqs=1, - ) - num_computed[0] += num_accepted - - # Prefill canonicalizes the one live slot in mode none. - for mixers, slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[slot] = 17 - mixers[0]._replayssm_prev_num_accepted[slot] = 9 - step(scheduled=8, num_accepted=1, prefilling=True) - - # Mix STP and MTP-shaped updates. Several scheduled lengths exceed the - # accepted count; repeated crossings eventually wrap the 20-row ring. - for scheduled, num_accepted in ( - (4, 2), - (4, 1), - (14, 3), - (18, 2), - (16, 16), - (1, 1), - ): - step(scheduled=scheduled, num_accepted=num_accepted) - torch.accelerator.synchronize() - - assert materializer.call_count == 0 - assert all(group.plan_flush_count.tolist() == [-1, -1] for group in ctx.groups) - for mixers, live_slot in zip(groups, (1, 4)): - # Both layers share this group tracker. The expected single transition - # per step would differ if either layer committed it independently. - assert mixers[0]._replayssm_ring_start[live_slot].item() == 4 - assert mixers[0]._replayssm_prev_num_accepted[live_slot].item() == 1 - - # A later prefill resets the exact same live physical slot again. - step(scheduled=8, num_accepted=1, prefilling=True) - torch.accelerator.synchronize() - for mixers, live_slot in zip(groups, (1, 4)): - assert mixers[0]._replayssm_ring_start[live_slot].item() == 0 - assert mixers[0]._replayssm_prev_num_accepted[live_slot].item() == 0 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA") -@pytest.mark.parametrize( - ( - "query_metadata", - "query_metadata_is_cumulative", - "num_computed_is_post_step", - "num_computed", - ), - [([1, 0], False, False, 7), ([0, 1], True, True, 8)], - ids=["v1", "v2"], -) -def test_modelwide_replayssm_single_token_final_prefill_commits_as_decode( - query_metadata: list[int], - query_metadata_is_cumulative: bool, - num_computed_is_post_step: bool, - num_computed: int, -): - groups, config, forward_context, block_tables = _modelwide_replayssm_fixture() - for mixers, source_slot in zip(groups, (1, 4)): - mixers[0]._replayssm_ring_start[source_slot] = 7 - mixers[0]._replayssm_prev_num_accepted[source_slot] = 4 - - ctx = ReplaySSMModelContext.create( - config, - [0, 1], - forward_context, - block_tables, - max_num_reqs=2, - ) - assert ctx is not None - ctx.postprocess( - idx_mapping=None, - query_metadata=torch.tensor(query_metadata, dtype=torch.int32, device="cuda"), - query_metadata_is_cumulative=query_metadata_is_cumulative, - num_computed_tokens=torch.tensor( - [num_computed, 0], dtype=torch.int32, device="cuda" - ), - num_computed_is_post_step=num_computed_is_post_step, - num_accepted_tokens=torch.tensor([2, 1], dtype=torch.int32, device="cuda"), - is_prefilling=torch.tensor([True, False], device="cuda"), - live_cols=torch.tensor([0, -1], dtype=torch.int32, device="cuda"), - num_reqs=1, - ) - torch.accelerator.synchronize() - - for mixers, source_slot in zip(groups, (1, 4)): - assert mixers[0]._replayssm_ring_start[source_slot].item() == 7 - assert mixers[0]._replayssm_prev_num_accepted[source_slot].item() == 6 diff --git a/tests/test_config.py b/tests/test_config.py index 4b6634bbd064..2f01dd459721 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -192,7 +192,6 @@ def _replayssm_config( ), [ (MambaBackendEnum.TRITON, True, "none", 0, 16, "requires Model Runner V1"), - (MambaBackendEnum.FLASHINFER, True, "align", 3, 16, None), ( MambaBackendEnum.FLASHINFER, True, @@ -209,7 +208,6 @@ def _replayssm_config( 16, "requires --mamba-backend flashinfer", ), - (MambaBackendEnum.FLASHINFER, False, "all", 0, 16, None), ( MambaBackendEnum.TRITON, False, @@ -229,10 +227,8 @@ def _replayssm_config( ], ids=[ "triton-v2-rejected", - "flashinfer-align-spec-v2", "flashinfer-spec-buffer-too-short", "triton-spec-rejected", - "flashinfer-all", "triton-all-rejected", "flashinfer-buffer-too-long", ], @@ -243,7 +239,7 @@ def test_replayssm_config_matrix( mamba_cache_mode: str, num_speculative_tokens: int, replayssm_buffer_len: int, - error_match: str | None, + error_match: str, ): config = _replayssm_config( backend=backend, @@ -253,11 +249,8 @@ def test_replayssm_config_matrix( config.cache_config.replayssm_buffer_len = replayssm_buffer_len config.num_speculative_tokens = num_speculative_tokens - if error_match is None: - assert VllmConfig.validate_mamba_cached_kernel(config) is config - else: - with pytest.raises(ValueError, match=error_match): - VllmConfig.validate_mamba_cached_kernel(config) + with pytest.raises(ValueError, match=error_match): + VllmConfig.validate_mamba_cached_kernel(config) def test_replayssm_rejects_pipeline_parallelism(): diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 4b985ccac98b..e9ff875cd9ec 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -103,15 +103,14 @@ def test_mamba_speculative_block_relocation_requires_exclusive_ownership(): manager._relocate_speculative_block([pinned_block], 0) -@pytest.mark.parametrize("mamba_cache_mode", ["align", "all"]) -def test_replayssm_queues_live_copy_for_new_state_block(mamba_cache_mode: str): +def test_replayssm_queues_live_copy_for_new_state_block(): spec = MambaSpec( block_size=4, shapes=((2,), (3,)), dtypes=(torch.float32, torch.float32), replayssm_shapes=((4,), (5,), (6,)), replayssm_dtypes=(torch.float32,) * 3, - mamba_cache_mode=mamba_cache_mode, + mamba_cache_mode="all", ) block_pool = BlockPool(num_gpu_blocks=6, enable_caching=True, hash_block_size=4) manager = MambaManager( diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index dad52ba866af..d1e304f6e78c 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -121,17 +121,6 @@ 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(vllm_runner, model_name): - pytest.importorskip("flashinfer.mamba.checkpointing_ssu") - _check_replayssm_parity( - vllm_runner, - model_name, - mamba_backend="flashinfer", - name_1="replayssm_flashinfer", - ) - - @pytest.mark.skipif( not HAS_FLASHINFER_CHECKPOINTING_SSU, reason="flashinfer.mamba.checkpointing_ssu not available", @@ -150,22 +139,6 @@ def test_replayssm_flashinfer_decode_matches_baseline_v2( ) -@multi_gpu_test(num_gpus=2) -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="flashinfer.mamba.checkpointing_ssu not available", -) -@pytest.mark.parametrize("model_name", [MAMBA2_MODEL]) -def test_replayssm_flashinfer_decode_matches_baseline_tp2(vllm_runner, model_name): - _check_replayssm_parity( - vllm_runner, - model_name, - tensor_parallel_size=2, - mamba_backend="flashinfer", - name_1="replayssm_flashinfer_tp2", - ) - - @pytest.mark.skipif( not HAS_FLASHINFER_CHECKPOINTING_SSU, reason="flashinfer.mamba.checkpointing_ssu not available", @@ -363,29 +336,18 @@ def run() -> None: @requires_flashinfer_replayssm_materialization @pytest.mark.parametrize("model_name", MODELS) -@pytest.mark.parametrize( - ("mamba_cache_mode", "use_v2", "use_ngram"), - [ - pytest.param("align", False, False, id="align-v1-stp"), - pytest.param("align", False, True, id="align-v1-ngram-t4"), - pytest.param("align", True, False, id="align-v2-stp"), - ], -) -def test_flashinfer_replayssm_prefix_cache_tp1( +def test_flashinfer_replayssm_align_prefix_cache_v1_ngram( vllm_runner, model_name, monkeypatch: pytest.MonkeyPatch, - mamba_cache_mode: str, - use_v2: bool, - use_ngram: bool, ): _check_flashinfer_replayssm_prefix_caching( vllm_runner, model_name, monkeypatch, - mamba_cache_mode=mamba_cache_mode, - use_ngram=use_ngram, - use_v2=use_v2, + mamba_cache_mode="align", + use_ngram=True, + use_v2=False, tensor_parallel_size=1, ) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 4d3457f6e781..559af5572de0 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import nullcontext from types import SimpleNamespace from typing import Any @@ -203,79 +202,3 @@ def fake_receive(*args, **kwargs): output = mrv2.GPUModelRunner.sample_tokens(runner, None) assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] - - -def test_v2_sample_tokens_postprocesses_state_before_drafting(monkeypatch): - events: list[Any] = [] - runner = _make_runner() - input_batch = SimpleNamespace( - req_ids=["request"], - idx_mapping=torch.tensor([0], dtype=torch.int64), - query_start_loc=torch.tensor([0, 1], dtype=torch.int32), - ) - hidden_states = torch.zeros(1, 1) - runner.execute_model_state = SimpleNamespace( - input_batch=input_batch, - attn_metadata=None, - slot_mappings_by_layer=None, - hidden_states=hidden_states, - aux_hidden_states=None, - dp_sync=None, - finished_req_ids=set(), - ec_connector_output=None, - routed_experts=None, - ) - sampled_token_ids = torch.tensor([[1]]) - num_sampled = torch.tensor([1], dtype=torch.int32) - num_rejected = torch.tensor([0], dtype=torch.int32) - runner.sample = lambda *_: ( - SimpleNamespace(sampled_token_ids=sampled_token_ids), - num_sampled, - num_rejected, - ) - runner.pp_handler = None - runner.prompt_logprobs_worker = SimpleNamespace( - compute_prompt_logprobs=lambda *_: {} - ) - runner.model = SimpleNamespace(compute_logits=None) - runner.main_stream = None - runner.output_copy_stream = None - runner.check_ep_fault = None - runner.pcp_manager = None - runner._draft_workspace_lane = None - runner.adaptive_verification = None - runner.sampler = SimpleNamespace( - penalties_state=SimpleNamespace(output_bin_counts=None), - sampling_states=SimpleNamespace( - temperature=SimpleNamespace(gpu=None), - seeds=SimpleNamespace(gpu=None), - ), - ) - runner.req_states = SimpleNamespace( - all_token_ids=SimpleNamespace(gpu=None), - num_computed_tokens=SimpleNamespace(gpu=torch.zeros(1, dtype=torch.int32)), - prompt_len=SimpleNamespace(np=None), - last_sampled_tokens=None, - next_prefill_tokens=None, - total_len=SimpleNamespace(gpu=None), - draft_tokens=torch.zeros((1, 1), dtype=torch.int64), - ) - - def postprocess_state(*_): - events.append("postprocess") - - def propose(*_, **__): - events.append("draft") - return torch.tensor([[2]]) - - runner.speculator = SimpleNamespace(supports_mm_inputs=False, propose=propose) - runner.model_state = SimpleNamespace( - postprocess_state=postprocess_state, - ) - monkeypatch.setattr(mrv2, "AsyncOutput", lambda **_: object()) - monkeypatch.setattr(mrv2, "post_update", lambda *_: None) - monkeypatch.setattr(mrv2, "use_workspace_lane", lambda _: nullcontext()) - - mrv2.GPUModelRunner.sample_tokens(runner, None) - - assert events == ["postprocess", "draft"] diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 214f795da8aa..a9b9ce27b015 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -59,47 +59,6 @@ def test_postprocess_state_scalar_with_int32_mapping( torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") -def test_flashinfer_replayssm_prefix_uses_original_accepted_counts() -> None: - state = object.__new__(MambaHybridModelState) - state._align_mode = True - state._needs_prefix_state_migration = True - state._use_flashinfer_replayssm = True - state.recoverssm = None - state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") - state._mamba_state_idx_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") - state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") - state._replayssm_query_start_loc = torch.tensor( - [0, 4], dtype=torch.int32, device="cuda" - ) - replayssm = Mock(materialize_prefixes=True) - accepted_snapshot = torch.zeros(4, dtype=torch.int32, device="cuda") - ctx = Mock( - is_initialized=True, - replayssm=replayssm, - num_accepted_tokens_snapshot=accepted_snapshot, - ) - - def normalize_live(*_args) -> None: - accepted_snapshot.copy_(state.num_accepted_tokens_gpu) - state.num_accepted_tokens_gpu[2] = 1 - - ctx.run_fused_postprocess_align.side_effect = normalize_live - state._mamba_ctx = ctx - - state.postprocess_state( - torch.tensor([2], dtype=torch.int32, device="cuda"), - torch.tensor([3], dtype=torch.int32, device="cuda"), - num_computed_tokens=torch.tensor([0, 0, 8, 0], device="cuda"), - ) - - kwargs = replayssm.postprocess.call_args.kwargs - assert kwargs["num_accepted_tokens"] is accepted_snapshot - assert accepted_snapshot[2].item() == 3 - assert state.num_accepted_tokens_gpu[2].item() == 1 - assert state._replayssm_query_start_loc is None - - def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: state = RecoverSSMState() metadata = Mock(spec=RecoverSSMMetadata) From 487b9cc34da6d550276c4a1d1f6acc596b26f98c Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 11:04:47 +0200 Subject: [PATCH 40/53] test: restore load-bearing ReplaySSM coverage Restore the focused ordering, snapshot, config, allocation-path, and align-V2 guards that protect behavior not observable through the remaining parity matrix. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/test_config.py | 13 +++- .../core/test_single_type_kv_cache_manager.py | 5 +- tests/v1/e2e/test_replayssm_decode.py | 15 +++- .../worker/test_gpu_model_runner_v2_eplb.py | 77 +++++++++++++++++++ .../worker/test_mamba_hybrid_model_state.py | 41 ++++++++++ 5 files changed, 143 insertions(+), 8 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 2f01dd459721..4b6634bbd064 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -192,6 +192,7 @@ def _replayssm_config( ), [ (MambaBackendEnum.TRITON, True, "none", 0, 16, "requires Model Runner V1"), + (MambaBackendEnum.FLASHINFER, True, "align", 3, 16, None), ( MambaBackendEnum.FLASHINFER, True, @@ -208,6 +209,7 @@ def _replayssm_config( 16, "requires --mamba-backend flashinfer", ), + (MambaBackendEnum.FLASHINFER, False, "all", 0, 16, None), ( MambaBackendEnum.TRITON, False, @@ -227,8 +229,10 @@ def _replayssm_config( ], ids=[ "triton-v2-rejected", + "flashinfer-align-spec-v2", "flashinfer-spec-buffer-too-short", "triton-spec-rejected", + "flashinfer-all", "triton-all-rejected", "flashinfer-buffer-too-long", ], @@ -239,7 +243,7 @@ def test_replayssm_config_matrix( mamba_cache_mode: str, num_speculative_tokens: int, replayssm_buffer_len: int, - error_match: str, + error_match: str | None, ): config = _replayssm_config( backend=backend, @@ -249,8 +253,11 @@ def test_replayssm_config_matrix( config.cache_config.replayssm_buffer_len = replayssm_buffer_len config.num_speculative_tokens = num_speculative_tokens - with pytest.raises(ValueError, match=error_match): - VllmConfig.validate_mamba_cached_kernel(config) + if error_match is None: + assert VllmConfig.validate_mamba_cached_kernel(config) is config + else: + with pytest.raises(ValueError, match=error_match): + VllmConfig.validate_mamba_cached_kernel(config) def test_replayssm_rejects_pipeline_parallelism(): diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index e9ff875cd9ec..4b985ccac98b 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -103,14 +103,15 @@ def test_mamba_speculative_block_relocation_requires_exclusive_ownership(): manager._relocate_speculative_block([pinned_block], 0) -def test_replayssm_queues_live_copy_for_new_state_block(): +@pytest.mark.parametrize("mamba_cache_mode", ["align", "all"]) +def test_replayssm_queues_live_copy_for_new_state_block(mamba_cache_mode: str): spec = MambaSpec( block_size=4, shapes=((2,), (3,)), dtypes=(torch.float32, torch.float32), replayssm_shapes=((4,), (5,), (6,)), replayssm_dtypes=(torch.float32,) * 3, - mamba_cache_mode="all", + mamba_cache_mode=mamba_cache_mode, ) block_pool = BlockPool(num_gpu_blocks=6, enable_caching=True, hash_block_size=4) manager = MambaManager( diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index d1e304f6e78c..6db6635ab752 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -336,18 +336,27 @@ def run() -> None: @requires_flashinfer_replayssm_materialization @pytest.mark.parametrize("model_name", MODELS) -def test_flashinfer_replayssm_align_prefix_cache_v1_ngram( +@pytest.mark.parametrize( + ("use_v2", "use_ngram"), + [ + pytest.param(False, True, id="align-v1-ngram-t4"), + pytest.param(True, False, id="align-v2-stp"), + ], +) +def test_flashinfer_replayssm_prefix_cache_tp1( vllm_runner, model_name, monkeypatch: pytest.MonkeyPatch, + use_v2: bool, + use_ngram: bool, ): _check_flashinfer_replayssm_prefix_caching( vllm_runner, model_name, monkeypatch, mamba_cache_mode="align", - use_ngram=True, - use_v2=False, + use_ngram=use_ngram, + use_v2=use_v2, tensor_parallel_size=1, ) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index 559af5572de0..4d3457f6e781 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from contextlib import nullcontext from types import SimpleNamespace from typing import Any @@ -202,3 +203,79 @@ def fake_receive(*args, **kwargs): output = mrv2.GPUModelRunner.sample_tokens(runner, None) assert output in (EMPTY_MODEL_RUNNER_OUTPUT, None) assert events == ["receive", "postprocess_num_computed_tokens", "eplb"] + + +def test_v2_sample_tokens_postprocesses_state_before_drafting(monkeypatch): + events: list[Any] = [] + runner = _make_runner() + input_batch = SimpleNamespace( + req_ids=["request"], + idx_mapping=torch.tensor([0], dtype=torch.int64), + query_start_loc=torch.tensor([0, 1], dtype=torch.int32), + ) + hidden_states = torch.zeros(1, 1) + runner.execute_model_state = SimpleNamespace( + input_batch=input_batch, + attn_metadata=None, + slot_mappings_by_layer=None, + hidden_states=hidden_states, + aux_hidden_states=None, + dp_sync=None, + finished_req_ids=set(), + ec_connector_output=None, + routed_experts=None, + ) + sampled_token_ids = torch.tensor([[1]]) + num_sampled = torch.tensor([1], dtype=torch.int32) + num_rejected = torch.tensor([0], dtype=torch.int32) + runner.sample = lambda *_: ( + SimpleNamespace(sampled_token_ids=sampled_token_ids), + num_sampled, + num_rejected, + ) + runner.pp_handler = None + runner.prompt_logprobs_worker = SimpleNamespace( + compute_prompt_logprobs=lambda *_: {} + ) + runner.model = SimpleNamespace(compute_logits=None) + runner.main_stream = None + runner.output_copy_stream = None + runner.check_ep_fault = None + runner.pcp_manager = None + runner._draft_workspace_lane = None + runner.adaptive_verification = None + runner.sampler = SimpleNamespace( + penalties_state=SimpleNamespace(output_bin_counts=None), + sampling_states=SimpleNamespace( + temperature=SimpleNamespace(gpu=None), + seeds=SimpleNamespace(gpu=None), + ), + ) + runner.req_states = SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=None), + num_computed_tokens=SimpleNamespace(gpu=torch.zeros(1, dtype=torch.int32)), + prompt_len=SimpleNamespace(np=None), + last_sampled_tokens=None, + next_prefill_tokens=None, + total_len=SimpleNamespace(gpu=None), + draft_tokens=torch.zeros((1, 1), dtype=torch.int64), + ) + + def postprocess_state(*_): + events.append("postprocess") + + def propose(*_, **__): + events.append("draft") + return torch.tensor([[2]]) + + runner.speculator = SimpleNamespace(supports_mm_inputs=False, propose=propose) + runner.model_state = SimpleNamespace( + postprocess_state=postprocess_state, + ) + monkeypatch.setattr(mrv2, "AsyncOutput", lambda **_: object()) + monkeypatch.setattr(mrv2, "post_update", lambda *_: None) + monkeypatch.setattr(mrv2, "use_workspace_lane", lambda _: nullcontext()) + + mrv2.GPUModelRunner.sample_tokens(runner, None) + + assert events == ["postprocess", "draft"] diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index a9b9ce27b015..214f795da8aa 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -59,6 +59,47 @@ def test_postprocess_state_scalar_with_int32_mapping( torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +def test_flashinfer_replayssm_prefix_uses_original_accepted_counts() -> None: + state = object.__new__(MambaHybridModelState) + state._align_mode = True + state._needs_prefix_state_migration = True + state._use_flashinfer_replayssm = True + state.recoverssm = None + state.num_accepted_tokens_gpu = torch.ones(4, dtype=torch.int32, device="cuda") + state._mamba_state_idx_gpu = torch.zeros(4, dtype=torch.int32, device="cuda") + state._is_prefilling_gpu = torch.zeros(4, dtype=torch.bool, device="cuda") + state._replayssm_query_start_loc = torch.tensor( + [0, 4], dtype=torch.int32, device="cuda" + ) + replayssm = Mock(materialize_prefixes=True) + accepted_snapshot = torch.zeros(4, dtype=torch.int32, device="cuda") + ctx = Mock( + is_initialized=True, + replayssm=replayssm, + num_accepted_tokens_snapshot=accepted_snapshot, + ) + + def normalize_live(*_args) -> None: + accepted_snapshot.copy_(state.num_accepted_tokens_gpu) + state.num_accepted_tokens_gpu[2] = 1 + + ctx.run_fused_postprocess_align.side_effect = normalize_live + state._mamba_ctx = ctx + + state.postprocess_state( + torch.tensor([2], dtype=torch.int32, device="cuda"), + torch.tensor([3], dtype=torch.int32, device="cuda"), + num_computed_tokens=torch.tensor([0, 0, 8, 0], device="cuda"), + ) + + kwargs = replayssm.postprocess.call_args.kwargs + assert kwargs["num_accepted_tokens"] is accepted_snapshot + assert accepted_snapshot[2].item() == 3 + assert state.num_accepted_tokens_gpu[2].item() == 1 + assert state._replayssm_query_start_loc is None + + def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: state = RecoverSSMState() metadata = Mock(spec=RecoverSSMMetadata) From 9d29c8b1dcf575f813b7c02540d1d60d900ecc78 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 11:58:08 +0200 Subject: [PATCH 41/53] [Mamba] Scope hybrid fallback to ReplaySSM Restore baseline hybrid grouping, pipeline projection, and V2 align seeding behavior while retaining the non-Mamba fallback and Mamba block-size seeding only for ReplaySSM. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/v1/core/test_kv_cache_utils.py | 51 +++++++++-- .../worker/test_mamba_hybrid_model_state.py | 10 ++- vllm/v1/core/kv_cache_utils.py | 88 ++++++++++++++----- .../worker/gpu/model_states/mamba_hybrid.py | 7 +- 4 files changed, 123 insertions(+), 33 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 9c98e04d4920..5bdcb005f658 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -1284,7 +1284,7 @@ def test_project_kv_cache_groups_to_worker(): spec_b = new_kv_cache_spec(num_kv_heads=4) global_groups = [ - KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a, is_eagle_group=True), + KVCacheGroupSpec(["layer1", "layer2", "layer3"], spec_a), ] worker_spec = {"layer1": spec_a, "layer2": spec_a} projected = kv_cache_utils._project_kv_cache_groups_to_worker( @@ -1300,7 +1300,6 @@ def test_project_kv_cache_groups_to_worker(): assert len(projected) == 1 assert projected[0].layer_names == [] assert projected[0].kv_cache_spec is spec_a - assert projected[0].is_eagle_group uniform_spec = UniformTypeKVCacheSpecs( block_size=16, @@ -3284,14 +3283,21 @@ def test_iter_layer_specs_returns_group_members(): assert list(iter_layer_specs(wrapped)) == [full, mla] -def _spec_decode_grouping_config(method="dspark", model_type=None): +def _spec_decode_grouping_config( + method="dspark", + model_type=None, + use_replayssm=False, + use_kda_recoverssm=False, +): """Grouping config with an EAGLE-family speculative method enabled.""" return SimpleNamespace( scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), cache_config=SimpleNamespace( + use_replayssm=use_replayssm, + use_kda_recoverssm=use_kda_recoverssm, get_resolved_kv_cache_layout=lambda: SimpleNamespace( is_block_outermost=True - ) + ), ), model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), speculative_config=SimpleNamespace( @@ -3363,11 +3369,32 @@ def test_draft_group_not_annotated_without_spec_decode(): assert not any(g.is_eagle_group for g in groups) -def test_unidentifiable_draft_flags_only_non_mamba_groups(): - # When no group carries a draft marker, retain the conservative fallback - # for attention while excluding Mamba state from the widened lookup window. +@pytest.mark.parametrize( + ("use_replayssm", "use_kda_recoverssm"), [(False, False), (True, True)] +) +def test_unidentifiable_draft_with_mamba_warns( + caplog_vllm, use_replayssm, use_kda_recoverssm +): + # Baseline behavior remains unchanged when ReplaySSM is not enabled. groups = get_kv_cache_groups( - _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) + _spec_decode_grouping_config( + use_replayssm=use_replayssm, + use_kda_recoverssm=use_kda_recoverssm, + ), + _hybrid_specs_with_draft(draft=False), + ) + + assert not any(g.is_eagle_group for g in groups) + assert "no KV cache group could be identified as the draft model's" in ( + caplog_vllm.text + ) + assert "Mamba groups" in caplog_vllm.text + + +def test_replayssm_unidentifiable_draft_flags_only_non_mamba_groups(): + groups = get_kv_cache_groups( + _spec_decode_grouping_config(use_replayssm=True), + _hybrid_specs_with_draft(draft=False), ) for group in groups: @@ -3378,6 +3405,14 @@ def test_unidentifiable_draft_flags_only_non_mamba_groups(): assert group.is_eagle_group is not contains_mamba +def test_no_warning_when_draft_group_is_identified(caplog_vllm): + get_kv_cache_groups( + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) + ) + + assert "could be identified as the draft model's" not in caplog_vllm.text + + def _deepseek_v4_specs(model_version="deepseek_v4"): """DeepseekV4-shaped specs: full MLA layers plus sliding-window MLA layers at differing page sizes, with the MTP draft layer registered last.""" diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 214f795da8aa..14f5ce78f6fe 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -16,7 +16,12 @@ from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState -def test_add_request_seeds_state_with_mamba_block_size() -> None: +@pytest.mark.parametrize( + ("use_flashinfer_replayssm", "expected_state_idx"), [(False, 1), (True, 2)] +) +def test_add_request_seeds_state_with_scoped_block_size( + use_flashinfer_replayssm: bool, expected_state_idx: int +) -> None: state = object.__new__(MambaHybridModelState) state.rope_state = None state.prompt_embeds_state = None @@ -26,13 +31,14 @@ def test_add_request_seeds_state_with_mamba_block_size() -> None: mamba_cache_mode="align", ) state._needs_prefix_state_migration = True + state._use_flashinfer_replayssm = use_flashinfer_replayssm state.num_accepted_tokens_gpu = torch.full((2,), 9, dtype=torch.int32) state._mamba_state_idx_gpu = torch.full((2,), -1, dtype=torch.int32) state.add_request(1, Mock(num_computed_tokens=17)) assert state.num_accepted_tokens_gpu.tolist() == [9, 1] - assert state._mamba_state_idx_gpu.tolist() == [-1, 2] + assert state._mamba_state_idx_gpu.tolist() == [-1, expected_state_idx] @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index ee35d95b3074..05f7a9c6b868 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1817,6 +1817,7 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: groups, use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config), ) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups @@ -1856,10 +1857,10 @@ def _annotate_eagle_groups( ``use_deepseek_v4_fallback`` False. The caller gates this fallback on the configured model type. FIXME(yifan): avoid/generalize this hacky check. - 3. Hybrid fallback. If neither rule identifies a draft group and Mamba - groups are present, conservatively flag every non-Mamba group. This - preserves the existing all-groups fallback for attention caches without - applying its widened lookup window to Mamba state. + 3. ReplaySSM hybrid fallback. If neither rule identifies a draft group, + flag every non-Mamba group only for Mamba ReplaySSM. This avoids applying + the downstream widened lookup window to ReplaySSM state without changing + the baseline hybrid-cache fallback. Args: vllm_config: Config supplying the speculative method, if any. @@ -1886,18 +1887,66 @@ def _annotate_eagle_groups( group.is_eagle_group = True break - if not any(group.is_eagle_group for group in kv_cache_groups): - non_mamba_groups = [ - group - for group in kv_cache_groups - if not any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if len(non_mamba_groups) < len(kv_cache_groups): - for group in non_mamba_groups: - group.is_eagle_group = True + cache_config = vllm_config.cache_config + if ( + any(group.is_eagle_group for group in kv_cache_groups) + or not cache_config.use_replayssm + or cache_config.use_kda_recoverssm + ): + return + non_mamba_groups = [ + group + for group in kv_cache_groups + if not any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if len(non_mamba_groups) < len(kv_cache_groups): + for group in non_mamba_groups: + group.is_eagle_group = True + + +def _warn_if_unannotated_eagle_mamba( + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> None: + """Warn when the flag-all eagle fallback will silently disable reuse. + + With no group annotated, consumers flag every group as a draft group. That + widens a Mamba group's required lookup window to two consecutive chunks, + which align-mode checkpointing never produces, so reuse drops to zero with + no error and no metric to show it. + + Args: + vllm_config: Config supplying the speculative method, if any. + kv_cache_groups: Groups as they will be handed to consumers. + """ + spec_config = vllm_config.speculative_config + if spec_config is None or not spec_config.use_eagle(): + return + if any(group.is_eagle_group for group in kv_cache_groups): + return + mamba_groups = [ + idx + for idx, group in enumerate(kv_cache_groups) + if any( + isinstance(spec, MambaSpec) + for spec in iter_layer_specs(group.kv_cache_spec) + ) + ] + if not mamba_groups: + return + logger.warning( + "Speculative decoding (method=%s) is enabled but no KV cache group " + "could be identified as the draft model's, so every group -- " + "including Mamba groups %s -- will be treated as a draft group. A " + "Mamba group cannot satisfy the widened lookup window that implies, " + "so prefix-cache reuse across requests will be disabled and any " + "external KV offload tier will store without ever serving a hit.", + spec_config.method, + mamba_groups, + ) def _largest_divisor_at_most(value: int, limit: int) -> int: @@ -1990,6 +2039,7 @@ def get_kv_cache_groups( groups.append(KVCacheGroupSpec([name], aligned)) _annotate_eagle_groups(vllm_config, kv_cache_spec, groups) + _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups @@ -2212,11 +2262,7 @@ def _project_kv_cache_groups_to_worker( KVCacheGroupSpec( worker_layer_names, group_spec, - # Empty projected groups preserve global group identity across - # PP ranks. Keep the annotation too, so a Mamba-only stage does - # not reinterpret "no local draft layers" as "all groups are - # draft groups" in local or external-cache coordinators. - is_eagle_group=group.is_eagle_group, + is_eagle_group=group.is_eagle_group and bool(worker_layer_names), ) ) return projected_groups diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 8f6d64a1f48b..9dc0f06110ee 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -133,8 +133,11 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: self.num_accepted_tokens_gpu[req_index].fill_(1) if self._needs_prefix_state_migration: # Seed the running state block from the resumed/prefilled position. - state_block_size = self.cache_config.mamba_block_size - assert state_block_size is not None + state_block_size = self.cache_config.block_size + if self._use_flashinfer_replayssm: + mamba_block_size = self.cache_config.mamba_block_size + assert mamba_block_size is not None + state_block_size = mamba_block_size self._mamba_state_idx_gpu[req_index].fill_( (new_req_data.num_computed_tokens - 1) // state_block_size ) From b7ec902da4a2ed1d19fe6254783b7636d2b9de95 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 12:02:40 +0200 Subject: [PATCH 42/53] [Mamba] Clarify ReplaySSM copied state Make the shared ReplaySSM SSM slot explicit and document which backend-owned tensors must follow scheduler block copies. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- vllm/model_executor/layers/mamba/mamba_mixer2.py | 8 +++++--- vllm/v1/worker/utils.py | 9 ++++++++- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index ab59edf7f9ba..d0290c2812f2 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1007,11 +1007,13 @@ def conv_ssm_forward( if self.use_replayssm: # The ownership pre-copy seeds the last-scheduled page before # forward. Keep both convolution and ReplaySSM on that private - # live page instead of touching the cached prefix source. + # live page instead of touching the cached prefix source, so + # SSM input and output deliberately use the same slot. assert block_idx_last_scheduled_token_d is not None assert replayssm_state_indices_d is not None - state_indices_tensor_d_input = replayssm_state_indices_d - state_indices_tensor_d_output = replayssm_state_indices_d + state_indices_tensor_d_input = state_indices_tensor_d_output = ( + replayssm_state_indices_d + ) conv_initial_state_idx_d = block_idx_last_scheduled_token_d elif self.num_spec > 0: assert block_idx_last_scheduled_token_prev_step_d is not None diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index b11ad163c99c..7fdb0358dc41 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -740,7 +740,14 @@ def copy_kv_cache_blocks_inplace( def get_replayssm_block_copy_tensors( forward_context: Mapping[str, Any], ) -> list[torch.Tensor]: - """Return ReplaySSM rings and FlashInfer's group-shared cursors.""" + """Collect ReplaySSM-owned state that must follow scheduler block copies. + + The runner's normal KV-cache list already contains canonical convolution + and SSM state. ReplaySSM ring caches use separate allocations on every + backend, so they are added here. FlashInfer additionally keeps its ring + position and accepted-token trackers in separate group-shared tensors; + the block-copy helper deduplicates those aliases by storage. + """ extra_tensors: list[torch.Tensor] = [] for layer in forward_context.values(): if not getattr(layer, "use_replayssm", False): From 9bf580c5ea0b7673558c26d15c03018c34da3ef3 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 12:09:44 +0200 Subject: [PATCH 43/53] [Mamba] Keep ReplaySSM changes feature-scoped Restore general hybrid-cache grouping to the base behavior, trim redundant configuration coverage, and document the remaining ReplaySSM lifecycle branches. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/test_config.py | 107 ------------------ tests/v1/core/test_kv_cache_utils.py | 42 ++----- .../v1/e2e/general/test_mamba_prefix_cache.py | 8 +- vllm/v1/core/kv_cache_utils.py | 34 +----- vllm/v1/worker/mamba_utils.py | 57 ++++++---- 5 files changed, 55 insertions(+), 193 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 4b6634bbd064..d27e4c50282d 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -161,113 +161,6 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected): assert envs.VLLM_USE_V2_MODEL_RUNNER is expected -def _replayssm_config( - *, - backend: MambaBackendEnum, - use_v2_model_runner: bool = False, -) -> SimpleNamespace: - return SimpleNamespace( - cache_config=SimpleNamespace( - use_replayssm=True, - mamba_cache_mode="none", - replayssm_buffer_len=16, - ), - model_config=None, - num_speculative_tokens=0, - mamba_config=SimpleNamespace(backend=backend), - parallel_config=SimpleNamespace(pipeline_parallel_size=1), - use_v2_model_runner=use_v2_model_runner, - kv_transfer_config=None, - ) - - -@pytest.mark.parametrize( - ( - "backend", - "use_v2_model_runner", - "mamba_cache_mode", - "num_speculative_tokens", - "replayssm_buffer_len", - "error_match", - ), - [ - (MambaBackendEnum.TRITON, True, "none", 0, 16, "requires Model Runner V1"), - (MambaBackendEnum.FLASHINFER, True, "align", 3, 16, None), - ( - MambaBackendEnum.FLASHINFER, - True, - "align", - 3, - 3, - r"replayssm-buffer-len >= 1 \+ num_speculative_tokens", - ), - ( - MambaBackendEnum.TRITON, - False, - "align", - 3, - 16, - "requires --mamba-backend flashinfer", - ), - (MambaBackendEnum.FLASHINFER, False, "all", 0, 16, None), - ( - MambaBackendEnum.TRITON, - False, - "all", - 0, - 16, - "all mode requires.*flashinfer", - ), - ( - MambaBackendEnum.FLASHINFER, - False, - "none", - 0, - 17, - "replayssm-buffer-len <= 16", - ), - ], - ids=[ - "triton-v2-rejected", - "flashinfer-align-spec-v2", - "flashinfer-spec-buffer-too-short", - "triton-spec-rejected", - "flashinfer-all", - "triton-all-rejected", - "flashinfer-buffer-too-long", - ], -) -def test_replayssm_config_matrix( - backend: MambaBackendEnum, - use_v2_model_runner: bool, - mamba_cache_mode: str, - num_speculative_tokens: int, - replayssm_buffer_len: int, - error_match: str | None, -): - config = _replayssm_config( - backend=backend, - use_v2_model_runner=use_v2_model_runner, - ) - config.cache_config.mamba_cache_mode = mamba_cache_mode - config.cache_config.replayssm_buffer_len = replayssm_buffer_len - config.num_speculative_tokens = num_speculative_tokens - - if error_match is None: - assert VllmConfig.validate_mamba_cached_kernel(config) is config - else: - with pytest.raises(ValueError, match=error_match): - VllmConfig.validate_mamba_cached_kernel(config) - - -def test_replayssm_rejects_pipeline_parallelism(): - config = _replayssm_config(backend=MambaBackendEnum.FLASHINFER) - config.parallel_config.pipeline_parallel_size = 2 - - with pytest.raises(ValueError, match="pipeline_parallel_size=1"): - VllmConfig.validate_mamba_cached_kernel(config) - - def test_rocm_keeps_compiled_deepseek_defaults(monkeypatch): """ROCm keeps the DSA models (DeepSeek V3.2/V4, GLM-5.2) on their compiled MRV1 paths and off breakable cudagraphs by default.""" diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 5bdcb005f658..c5159e0d1627 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3283,21 +3283,14 @@ def test_iter_layer_specs_returns_group_members(): assert list(iter_layer_specs(wrapped)) == [full, mla] -def _spec_decode_grouping_config( - method="dspark", - model_type=None, - use_replayssm=False, - use_kda_recoverssm=False, -): +def _spec_decode_grouping_config(method="dspark", model_type=None): """Grouping config with an EAGLE-family speculative method enabled.""" return SimpleNamespace( scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False), cache_config=SimpleNamespace( - use_replayssm=use_replayssm, - use_kda_recoverssm=use_kda_recoverssm, get_resolved_kv_cache_layout=lambda: SimpleNamespace( is_block_outermost=True - ), + ) ), model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type=model_type)), speculative_config=SimpleNamespace( @@ -3369,19 +3362,12 @@ def test_draft_group_not_annotated_without_spec_decode(): assert not any(g.is_eagle_group for g in groups) -@pytest.mark.parametrize( - ("use_replayssm", "use_kda_recoverssm"), [(False, False), (True, True)] -) -def test_unidentifiable_draft_with_mamba_warns( - caplog_vllm, use_replayssm, use_kda_recoverssm -): - # Baseline behavior remains unchanged when ReplaySSM is not enabled. +def test_unidentifiable_draft_with_mamba_warns(caplog_vllm): + # No group carries the draft marker, so every consumer falls back to + # flagging all groups -- including Mamba ones, which then can never report + # a hit. That is silent today; it must at least be visible. groups = get_kv_cache_groups( - _spec_decode_grouping_config( - use_replayssm=use_replayssm, - use_kda_recoverssm=use_kda_recoverssm, - ), - _hybrid_specs_with_draft(draft=False), + _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=False) ) assert not any(g.is_eagle_group for g in groups) @@ -3391,20 +3377,6 @@ def test_unidentifiable_draft_with_mamba_warns( assert "Mamba groups" in caplog_vllm.text -def test_replayssm_unidentifiable_draft_flags_only_non_mamba_groups(): - groups = get_kv_cache_groups( - _spec_decode_grouping_config(use_replayssm=True), - _hybrid_specs_with_draft(draft=False), - ) - - for group in groups: - contains_mamba = any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - assert group.is_eagle_group is not contains_mamba - - def test_no_warning_when_draft_group_is_identified(caplog_vllm): get_kv_cache_groups( _spec_decode_grouping_config(), _hybrid_specs_with_draft(draft=True) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index ea77d29cbe92..d8fc041b43f7 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -937,8 +937,9 @@ def _run_mamba_prefix_cache_mrv1_configured( def _run_mamba_prefix_cache_mrv1( monkeypatch: pytest.MonkeyPatch, async_scheduling: bool ): - # This test patches the V1 model runner, so pin V1 explicitly: MoE/hybrid - # models like Qwen3-Next now default to the V2 runner. + # The test patches V1 runner methods, while Qwen3-Next now defaults to V2. + # Scope the V1 override to this call, then clear the cached env value after + # monkeypatch restores it so following tests see the original runner choice. try: with monkeypatch.context() as patch: patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") @@ -1239,6 +1240,9 @@ def fake_sample( def _run_mamba_prefix_cache_mrv2( monkeypatch: pytest.MonkeyPatch, async_scheduling: bool ): + # The test patches V2 runner methods in this process, so disable engine-core + # multiprocessing and select V2 only for this call. Clear the cached env + # values after monkeypatch restores them to avoid leaking either override. try: with monkeypatch.context() as patch: patch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 05f7a9c6b868..ec24e987413a 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1839,7 +1839,7 @@ def _annotate_eagle_groups( ) -> None: """Flag the KV cache groups that hold drafter attention layers. - Three detection rules, in order of preference: + Two detection rules, in order of preference: 1. Spec-driven. ``non_causal_multi_token_decode`` is declared on MLAAttentionSpec and set by drafter attention layers that run a @@ -1857,10 +1857,6 @@ def _annotate_eagle_groups( ``use_deepseek_v4_fallback`` False. The caller gates this fallback on the configured model type. FIXME(yifan): avoid/generalize this hacky check. - 3. ReplaySSM hybrid fallback. If neither rule identifies a draft group, - flag every non-Mamba group only for Mamba ReplaySSM. This avoids applying - the downstream widened lookup window to ReplaySSM state without changing - the baseline hybrid-cache fallback. Args: vllm_config: Config supplying the speculative method, if any. @@ -1880,31 +1876,13 @@ def _annotate_eagle_groups( ): group.is_eagle_group = True - if use_deepseek_v4_fallback: - last_layer = next(reversed(kv_cache_spec)) - for group in kv_cache_groups: - if last_layer in group.layer_names: - group.is_eagle_group = True - break - - cache_config = vllm_config.cache_config - if ( - any(group.is_eagle_group for group in kv_cache_groups) - or not cache_config.use_replayssm - or cache_config.use_kda_recoverssm - ): + if not use_deepseek_v4_fallback: return - non_mamba_groups = [ - group - for group in kv_cache_groups - if not any( - isinstance(spec, MambaSpec) - for spec in iter_layer_specs(group.kv_cache_spec) - ) - ] - if len(non_mamba_groups) < len(kv_cache_groups): - for group in non_mamba_groups: + last_layer = next(reversed(kv_cache_spec)) + for group in kv_cache_groups: + if last_layer in group.layer_names: group.is_eagle_group = True + break def _warn_if_unannotated_eagle_mamba( diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 8e89f58c48c1..b9e2af97d938 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -411,8 +411,10 @@ def postprocess_mamba_fused_kernel( # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds # the post-step new_num_computed value (V2 supplies the advanced count). PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False, - # FlashInfer ReplaySSM owns temporal state, while this kernel continues to - # snapshot convolution state at the same block boundary. + # FlashInfer ReplaySSM migrates temporal state through its ring materializer. + # Skip the generic temporal copy in that case so it cannot overwrite the + # ReplaySSM-owned transition; canonical convolution state still needs the + # generic snapshot at the same block boundary. SKIP_TEMPORAL_STATE_COPY: tl.constexpr = False, # TEMPORAL_TILES: when > 1, the temporal copy body is partitioned across # TEMPORAL_TILES CTAs along the u64 inner range. Callers must launch a @@ -481,8 +483,10 @@ def postprocess_mamba_fused_kernel( return if SKIP_TEMPORAL_STATE_COPY: - # state_conv_widths is also the state-kind metadata: convolution widths - # are positive, while zero explicitly denotes a temporal state. + # The flattened metadata interleaves convolution and temporal states. + # state_conv_widths doubles as the state-kind tag: a positive width is + # canonical convolution state, while zero marks ReplaySSM-owned temporal + # state and must bypass _copy_mamba_state_block. conv_width = tl.load(state_conv_widths_ptr + state_idx) if conv_width == 0: return @@ -533,9 +537,11 @@ def preprocess_mamba_align_fused_kernel( 2. Store the pre-copy src columns for ``precopy_mamba_align_fused_kernel``: - src_col = state_idx (the previous running block column) - src_off = max(num_accepted - 1, 0) (the accepted-token bias) - 3. Advance state_idx to the new running block. Baseline Mamba resets - num_accepted after shifting its state; exact ReplaySSM block copies - preserve it so the copied live slot keeps the same accepted position. + 3. Advance state_idx to the new running block. Unless PRESERVE_ACCEPTED is + set, reset num_accepted to 1 when the block changes: the migrated state + now sits at the start of the new block and must be read with the neutral + accepted-token bias. ReplaySSM block copies preserve the accepted count + because its copied live slot retains the same ring position. """ offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < num_reqs @@ -841,8 +847,9 @@ class MambaSpecDecodeGPUContext: # Flag to track if metadata has been populated is_initialized: bool = False - # True when the model-wide FlashInfer ReplaySSM lifecycle owns temporal - # state. Mixed ReplaySSM/baseline Mamba layers are rejected at populate time. + # False for the ordinary hybrid spec-decode state-copy context. True only + # when model-wide FlashInfer ReplaySSM owns temporal state; mixed ReplaySSM + # and baseline Mamba layers are rejected at populate time. has_flashinfer_replayssm: bool = False # Persistent all-layer ReplaySSM descriptors, populated with the cache # addresses on first real forward. None for non-FlashInfer configurations. @@ -1520,6 +1527,8 @@ def preprocess_mamba( ) fused.src_col.np[:num_reqs] = -1 + # ReplaySSM reset consumes only source/destination columns. token_bias + # belongs to the generic pre-copy kernel, which ReplaySSM does not call. if fused.ctx.replayssm is None: fused.token_bias.np[:num_reqs] = 0 @@ -1555,6 +1564,9 @@ def preprocess_mamba( # not mistaken for fresh slot ownership. fused.src_col.np[i] = prev_state_idx + # Baseline Mamba migrates canonical state when ownership crosses a + # column. ReplaySSM receives that ownership move through scheduler block + # copies and uses src/dst only to reset genuinely fresh ring slots. if ( prev_state_idx != -1 and prev_state_idx != curr_state_idx @@ -1661,13 +1673,15 @@ def postprocess_mamba_gpu( mamba_state_copy_funcs: MambaStateCopyFuncsByType, run_prefix_state_migration: bool, ) -> None: - """Run model-wide Mamba state maintenance after token acceptance. - - Lazily binds the fused-kernel context to the persistent block tables and - forward-context state pointers on the first call. Prefix modes run the - generic state-copy planner before committing ReplaySSM trackers; mode none - commits only the trackers. The accepted counts are then copied back for any - CPU-side consumer on the next iteration. + """Publish accepted-token results to model-wide Mamba state. + + The first call binds persistent block tables and state pointers. When prefix + migration is enabled, the fused kernel snapshots the original accepted + counts, copies canonical state at crossed boundaries, and normalizes the live + counts for the next step. ReplaySSM then commits its group-shared trackers + from the original counts and materializes any planned prefix snapshots; mode + ``none`` skips the copy phase and commits directly. Finally, the normalized + live counts are copied back only when the next iteration has a CPU consumer. """ ctx = bufs.postprocess_align # The caller enables this context for spec-decode hybrid state copies or @@ -1735,12 +1749,13 @@ def stage_postprocess_inputs_to_gpu( mamba_state_idx: dict[str, int], run_prefix_state_migration: bool, ) -> None: - """Stage all per-request inputs the fused mamba postprocess kernel reads. + """Stage the per-request decisions consumed after token acceptance. - Walks ``req_ids[:num_reqs]`` once, writing each request's mamba block - index and scheduled/computed/draft token counts into the matching pinned - numpy views, then issues non-blocking H→D copies. Mode ``none`` does not - stage a live column because ReplaySSM's live state is always column zero. + One host pass writes scheduled, computed, draft, and prefill values into + pinned views shared by generic Mamba postprocess and ReplaySSM tracker + publication, then launches non-blocking H→D copies. Prefix migration also + stages the live Mamba column used by the copy planner and ReplaySSM; mode + ``none`` omits it because ReplaySSM live state remains in column zero. """ assert ctx.num_scheduled_tokens_buf is not None assert ctx.num_computed_tokens_buf is not None From ce01679f89ddb4ffae5f0173a5fac40d5755cb0b Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 12:16:33 +0200 Subject: [PATCH 44/53] [Mamba] Explain ReplaySSM state transitions Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- vllm/v1/core/single_type_kv_cache_manager.py | 25 +++++++++++++++++++- vllm/v1/worker/mamba_utils.py | 7 ++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 717e3f8322ea..85adfab2daec 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1372,8 +1372,16 @@ def __init__( self.block_size = kv_cache_spec.block_size self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks + # ``replayssm_shapes`` means that this group's temporal state includes + # ReplaySSM rings indexed by the scheduler's physical block IDs. When + # allocation moves a request's live block, those rings and their + # trackers must move with it; otherwise the next forward reads the new + # block ID with the previous live history left behind. Both Triton and + # FlashInfer ReplaySSM use this scheduler-level migration path. self._copy_replayssm_live_state = bool(kv_cache_spec.replayssm_shapes) if self._copy_replayssm_live_state: + # ReplaySSM stores speculative history inside its ring. It must not + # also reserve the baseline Mamba speculative scratch blocks. assert self.num_speculative_blocks == 0, ( "ReplaySSM keeps speculative state in its replay rings, not " "separate scheduler blocks" @@ -1676,6 +1684,11 @@ def allocate_new_blocks( return super().allocate_new_blocks( request_id, num_tokens, num_tokens_main_model ) + # The base allocator may append a new final block. Remember the + # current live owner before it mutates req_blocks so the worker can + # copy the complete ReplaySSM state into that new write slot. + # For a partial cache hit, the hit block—not the request's old + # tail—is the state source. req_blocks = self.req_to_blocks[request_id] prev_block_len = len(req_blocks) partial_hit = self._partial_hit_reqs.get(request_id) @@ -1709,6 +1722,9 @@ def allocate_new_blocks( has_partial_hit = partial_hit is not None live_source = None if self._copy_replayssm_live_state: + # Capture the live state owner before align-mode allocation + # relocates/nulls table entries. The eventual destination is + # the last new block, where the next forward writes state. if partial_hit is not None: live_source = partial_hit[1] elif prev_block_len > 0: @@ -1811,7 +1827,14 @@ def _queue_replayssm_live_copy( source_block: KVCacheBlock | None, destination_block: KVCacheBlock, ) -> None: - """Queue and retain one complete live ReplaySSM slot migration.""" + """Queue and retain one complete live ReplaySSM slot migration. + + Reuse the existing CoW copy channel because the worker already applies + each queued physical-block pair to every cache tensor registered for + this group, including ReplaySSM rings and trackers. The extra references + keep both slots alive until the coordinator drains and executes the + copy; the scheduler releases those retained references afterward. + """ if ( source_block is None or source_block.is_null diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index b9e2af97d938..6566fbe1d48b 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -1718,6 +1718,13 @@ def postprocess_mamba_gpu( ) accepted_tokens_for_postprocess = ctx.num_accepted_tokens_snapshot if ctx.replayssm is not None: + # Keep this separate from run_fused_postprocess: that kernel is a + # per-(request, canonical-state, tile) copy and is skipped when prefix + # migration is disabled. ReplaySSM must instead publish its shared + # per-(request, cache-group) ring trackers in every cache mode, then + # optionally compact/materialize one plan per group. Folding the two + # would either duplicate tracker writes for every state tensor or make + # the generic Mamba kernel depend on FlashInfer-only group descriptors. ctx.replayssm.postprocess( idx_mapping=None, query_metadata=ctx.num_scheduled_tokens_buf.gpu, From b0185c6592512a39b4e33d70993a8917013037d9 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 12:18:51 +0200 Subject: [PATCH 45/53] [Mamba] Document runner state staging Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- vllm/v1/worker/gpu_model_runner.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 93dd09ae811b..56fe5b58586a 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4469,6 +4469,10 @@ def execute_model( if mamba_bufs is not None and mamba_bufs.postprocess_align is not None: mamba_ctx = mamba_bufs.postprocess_align if not mamba_ctx.is_initialized: + # First real batch only: model loading creates the buffers, + # but the physical cache tensors and block tables are not + # bound until KV-cache initialization. Capture their stable + # addresses now for later graph-safe GPU postprocessing. mamba_ctx.initialize_from_forward_context( self.kv_cache_config, self.compilation_config.static_forward_context, @@ -4480,6 +4484,11 @@ def execute_model( for gid in mamba_ctx.mamba_group_ids ], ) + # Every forward: copy this batch's scheduler decisions into + # persistent GPU buffers before model execution. After sampling + # reveals the accepted-token counts, postprocess_mamba_gpu uses + # the same snapshot to migrate canonical Mamba state and/or + # publish ReplaySSM's model-wide ring trackers. mamba_utils.stage_postprocess_inputs_to_gpu( mamba_ctx, scheduler_output, From 7d8c5750fccf27d29f6560097b8e98fd86836f59 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 12:57:17 +0200 Subject: [PATCH 46/53] test(mamba): pin Triton ReplaySSM to runner V1 Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 269734bf3646..9f0481cfc1e9 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -30,6 +30,16 @@ HAS_FLASHINFER_CHECKPOINTING_SSU = False +@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, From 669da5b0a95d284fa9ec922a530842c7c33e24a2 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 13:05:16 +0200 Subject: [PATCH 47/53] test(mamba): cover FlashInfer ReplaySSM on V1 Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 9f0481cfc1e9..14aba0b1cbfe 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -47,7 +47,7 @@ def _check_replayssm_parity( tensor_parallel_size=1, mamba_backend: str = "triton", name_1: str = "replayssm", - require_v2: bool = False, + 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 @@ -61,14 +61,14 @@ 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 ) 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( @@ -97,19 +97,20 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): reason="flashinfer.mamba.checkpointing_ssu not available", ) @pytest.mark.parametrize("model_name", MODELS) -def test_replayssm_flashinfer_decode_matches_baseline_v2( - vllm_runner, model_name, monkeypatch +@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 ): try: with monkeypatch.context() as patch: - patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + 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_v2", - require_v2=True, + name_1="replayssm_flashinfer", + expected_v2=use_v2_model_runner, ) finally: # The context restores the environment before the final cache reset. From 19fe1606a14f6c6187c52c4e4b478790f364da33 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 13:08:28 +0200 Subject: [PATCH 48/53] test(mamba): cover FlashInfer MTP on V1 Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 14aba0b1cbfe..57f194df2835 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -156,7 +156,8 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na reason="flashinfer.mamba.checkpointing_ssu not available", ) @large_gpu_mark(min_gb=40) -def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): +@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, @@ -169,7 +170,7 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): ) try: with monkeypatch.context() as patch: - patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner))) envs.disable_envs_cache() with vllm_runner( MAMBA2_MTP_MODEL, @@ -177,7 +178,10 @@ def test_replayssm_flashinfer_mtp_v2(vllm_runner, monkeypatch): replayssm_buffer_len=16, **common, ) as llm: - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner + 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 From 53fe7ed2e7a84359802ea65e5dfd2a0dd6fda267 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 13:09:55 +0200 Subject: [PATCH 49/53] test(mamba): require FlashInfer ReplaySSM coverage Signed-off-by: Andrii Skliar --- tests/v1/e2e/test_replayssm_decode.py | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 57f194df2835..9533ee230e4e 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -22,13 +22,6 @@ "Once upon a time, in a small village,", ] -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 _use_v1_model_runner_by_default(monkeypatch): @@ -92,10 +85,6 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2) -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="flashinfer.mamba.checkpointing_ssu not available", -) @pytest.mark.parametrize("model_name", MODELS) @pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"]) def test_replayssm_flashinfer_decode_matches_baseline( @@ -117,10 +106,6 @@ def test_replayssm_flashinfer_decode_matches_baseline( envs.disable_envs_cache() -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="flashinfer.mamba.checkpointing_ssu not available", -) @pytest.mark.parametrize("model_name", MODELS) def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_name): common = dict( @@ -151,10 +136,6 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na @multi_gpu_test(num_gpus=2) -@pytest.mark.skipif( - not HAS_FLASHINFER_CHECKPOINTING_SSU, - reason="flashinfer.mamba.checkpointing_ssu not available", -) @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): From 4c223a7da579ee5fb14add6292c99b9c317cfcb3 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Fri, 4 Sep 2026 13:25:23 +0200 Subject: [PATCH 50/53] [Mamba] Tighten ReplaySSM integration boundaries Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- .../test_mamba_update_block_table.py | 41 ++++++++++++++++++- tests/v1/worker/test_mamba_utils.py | 2 +- .../layers/mamba/mamba_utils.py | 5 --- .../layers/mamba/ops/ssu_dispatch.py | 30 ++++++++------ vllm/v1/attention/backends/mamba_attn.py | 14 +++---- vllm/v1/worker/gpu_model_runner.py | 8 ++-- vllm/v1/worker/mamba_utils.py | 17 +++++--- 7 files changed, 81 insertions(+), 36 deletions(-) diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index 4ec138270203..a5dbfe97b480 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -16,7 +16,11 @@ import torch -from tests.v1.attention.utils import MockMambaBuilder +from tests.v1.attention.utils import ( + BatchSpec, + MockMambaBuilder, + create_common_attn_metadata, +) from vllm.config.compilation import CUDAGraphMode from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadata from vllm.v1.kv_cache_interface import MambaSpec @@ -321,6 +325,41 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): assert builder.block_idx_last_scheduled_token_prev_step.dtype == torch.int32 +def test_all_spec_decode_without_previous_anchor_leaves_metadata_unset(): + """A step without a previous-step anchor keeps the optional field unset. + + The mixer already falls back to the last computed block when this metadata + is absent; populating it here would change the baseline all-mode path. + """ + block_size = 16 + seq_lens = [33, 49] + config = _make_vllm_config( + max_model_len=256, + max_num_seqs=len(seq_lens), + num_speculative_tokens=3, + ) + config.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + spec = MambaSpec( + block_size=block_size, + shapes=((1,), (1,)), + dtypes=(torch.float32,), + mamba_cache_mode="all", + num_speculative_blocks=2, + ) + builder = MockMambaBuilder(spec, ["layer0"], config, torch.device("cpu")) + common = create_common_attn_metadata( + BatchSpec(seq_lens=seq_lens, query_lens=[1, 1]), + block_size, + torch.device("cpu"), + arange_block_indices=True, + ).replace(is_prefilling=torch.zeros(len(seq_lens), dtype=torch.bool)) + + metadata = builder.build(0, common) + + assert metadata.num_decodes == len(seq_lens) + assert metadata.block_idx_last_scheduled_token_prev_step is None + + def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode(): """Without spec decode, the prev-step buffer is unused and must not be allocated — the input anchor reduces to last_computed_token.""" diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a40df81db0ec..4998cc079f1e 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -11,7 +11,6 @@ from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFunc, MambaStateCopyFuncsByType, - _reinterpret_u64_as_i64, get_conv_copy_spec, get_temporal_copy_spec, ) @@ -27,6 +26,7 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, + _reinterpret_u64_as_i64, batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 734823b9fb21..4acfe82c0ddd 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -50,11 +50,6 @@ def is_conv_state_dim_first() -> bool: return get_conv_state_layout() == "DS" -def _reinterpret_u64_as_i64(value: int) -> int: - """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" - return value if value < (1 << 63) else value - (1 << 64) - - class MambaStateDtypeCalculator: @classmethod def linear_attention_state_dtype( diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 4d912e99192d..0fa1077338cd 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -19,7 +19,6 @@ from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.mamba_utils import _reinterpret_u64_as_i64 from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.utils import NULL_BLOCK_ID @@ -28,18 +27,23 @@ logger = init_logger(__name__) +def _reinterpret_u64_as_i64(value: int) -> int: + """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" + return value if value < (1 << 63) else value - (1 << 64) + + @triton.jit -def _mamba_state_copy_boundary( +def mamba_state_copy_boundary( num_tokens_running_state, new_num_computed, block_size: tl.constexpr, ): - """Return the canonical aligned Mamba state-copy decision.""" + """Return the aligned Mamba state-copy decision and destination.""" aligned_new_computed = (new_num_computed // block_size) * block_size needs_copy = aligned_new_computed >= num_tokens_running_state accept_token_bias = aligned_new_computed - num_tokens_running_state - dest_col = aligned_new_computed // block_size - 1 - return needs_copy, accept_token_bias, dest_col + dest_block_idx = aligned_new_computed // block_size - 1 + return needs_copy, accept_token_bias, dest_block_idx @triton.jit @@ -145,9 +149,7 @@ def _postprocess_replayssm_kernel( prefilling = prefilling & ((query_len != 1) | (computed_before <= 0)) accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) - # Derive this request's pre/post-step positions from ReplaySSM metadata, - # then share the canonical copy-boundary calculation with the generic - # Mamba state-copy kernel. + # Derive this request's pre/post-step positions from ReplaySSM metadata. computed_after = tl.where( prefilling, computed_before + query_len, @@ -156,7 +158,7 @@ def _postprocess_replayssm_kernel( running_state_pos = tl.where( prefilling, computed_after, computed_after - accepted + 1 ) - boundary, accept_token_bias, dst_col = _mamba_state_copy_boundary( + boundary, accept_token_bias, dst_col = mamba_state_copy_boundary( running_state_pos, computed_after, MAMBA_BLOCK_SIZE, @@ -442,7 +444,7 @@ def postprocess( BLOCK_SIZE=triton.next_power_of_2(self.max_num_reqs), ) - def materialize(self) -> None: + def materialize(self, materialize_fn: Callable[..., None]) -> None: """Publish the canonical prefix snapshots prepared by ``postprocess``.""" first = self.mixers[0] mamba_config = first.mamba_config @@ -453,7 +455,7 @@ def materialize(self) -> None: 0, 2**32, (1,), device=self.src_slots.device, dtype=torch.int64 ) philox_rounds = mamba_config.stochastic_rounding_philox_rounds or 10 - _load_replayssm_materialize()( + materialize_fn( *self.materialize_tables, self.src_slots, self.dst_slots, @@ -545,8 +547,12 @@ def postprocess(self, **kwargs: Any) -> None: group.postprocess(**kwargs) def materialize(self) -> None: + # Keep the optional FlashInfer dependency lazy: mode ``none`` never + # reaches this path. Resolve the cached callable once for this model + # operation, then invoke it once per physical cache-slot namespace. + materialize_fn = _load_replayssm_materialize() for group in self.groups: - group.materialize() + group.materialize(materialize_fn) class MambaSSUBackend(ABC): diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 4401c0d94689..8bf25c732d35 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -548,14 +548,12 @@ def _compute_common_metadata( ) = self._compute_prefix_caching_block_indices( common_attn_metadata, mamba_block_size ) - if self.use_spec_decode: - block_idx_last_scheduled_token_prev_step = block_idx_last_computed_token - if prev_last_scheduled_idx is not None: - block_idx_last_scheduled_token_prev_step = torch.where( - prev_last_scheduled_idx >= 0, - prev_last_scheduled_idx, - block_idx_last_computed_token, - ) + if self.use_spec_decode and prev_last_scheduled_idx is not None: + block_idx_last_scheduled_token_prev_step = torch.where( + prev_last_scheduled_idx >= 0, + prev_last_scheduled_idx, + block_idx_last_computed_token, + ) else: state_indices_tensor = mamba_get_block_table_tensor( common_attn_metadata.block_table_tensor, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 56fe5b58586a..402875b86063 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -1643,9 +1643,11 @@ def _update_states_after_model_execute( run_prefix_state_migration=self._needs_prefix_state_migration, ) - if self.num_accepted_tokens_event is not None: - # STP ReplaySSM has no accepted-count D2H copy and therefore - # does not allocate an event. + if self._use_flashinfer_replayssm and not self.num_spec_tokens: + # STP ReplaySSM has no accepted-count D2H copy or event. + assert self.num_accepted_tokens_event is None + else: + assert self.num_accepted_tokens_event is not None self.num_accepted_tokens_event.record() else: self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_( diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 6566fbe1d48b..95d59cc1fa30 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -12,14 +12,13 @@ from vllm.logger import init_logger from vllm.model_executor.layers.mamba.mamba_utils import ( MambaStateCopyFuncsByType, - _reinterpret_u64_as_i64, get_conv_copy_spec, get_temporal_copy_spec, is_conv_state_dim_first, ) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( ReplaySSMModelContext, - _mamba_state_copy_boundary, + mamba_state_copy_boundary, ) from vllm.triton_utils import tl, triton from vllm.utils.gpu_sync_debug import gpu_sync_allowed @@ -189,6 +188,11 @@ def _memcpy_u64_tiled( tl.store(dst_u8 + i + offsets, data, mask=mask) +def _reinterpret_u64_as_i64(value: int) -> int: + """Preserve a uint64 pointer bit pattern in a torch.int64 tensor.""" + return value if value < (1 << 63) else value - (1 << 64) + + @triton.jit def _copy_mamba_state_block( state_idx, @@ -463,7 +467,7 @@ def postprocess_mamba_fused_kernel( num_tokens_running_state = num_computed + num_scheduled - num_draft new_num_computed = num_tokens_running_state + num_accepted - 1 - needs_copy, accept_token_bias, dest_block_idx = _mamba_state_copy_boundary( + needs_copy, accept_token_bias, dest_block_idx = mamba_state_copy_boundary( num_tokens_running_state, new_num_computed, block_size, @@ -625,6 +629,7 @@ def precopy_mamba_align_fused_kernel( # so there is nothing to copy. if src_col < 0 or src_col == dst_col: return + token_bias = tl.load(token_bias_ptr + req_idx) _copy_mamba_state_block( state_idx, @@ -730,9 +735,9 @@ def validate_mamba_state_copy_funcs( ) state_copy_funcs = copy_funcs[mamba_spec.mamba_type] assert 0 < len(state_copy_funcs) <= len(mamba_spec.shapes), ( - f"{mamba_spec.mamba_type} expects {len(mamba_spec.shapes)} state copy " - "funcs for its canonical state tensors, but provides " - f"{len(state_copy_funcs)}; expected a non-empty copyable prefix" + f"{mamba_spec.mamba_type} declares {len(mamba_spec.shapes)} states, " + f"but provides {len(state_copy_funcs)} state copy funcs; expected " + "a non-empty copyable prefix" ) From fc4313dcc0e762701af55f672c51cb01bffbfd39 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Sun, 6 Sep 2026 00:02:32 +0200 Subject: [PATCH 51/53] [Mamba] Fix ReplaySSM cache paths and MTP grouping Keep Triton on its packed cache path while scoping external ReplaySSM trackers to FlashInfer. Preserve V2 all-mode state and retain hybrid Mamba boundaries across native-MTP prefix reuse. Identify native-MTP draft groups and cover the supported layouts and model variants. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- tests/model_executor/test_nemotron_h_mtp.py | 76 +++++++++++++ .../test_nemotron_h_quantization.py | 40 +++++++ tests/test_config.py | 7 ++ .../test_mamba_update_block_table.py | 6 +- .../test_replayssm_metadata_builder.py | 10 +- tests/v1/core/test_kv_cache_utils.py | 71 +++++++++++++ tests/v1/core/test_prefix_caching.py | 63 +++++++++++ tests/v1/e2e/test_replayssm_decode.py | 37 +++++-- .../worker/test_mamba_hybrid_model_state.py | 66 +++++++++++- tests/v1/worker/test_utils.py | 100 ++++++++++++++++++ vllm/config/vllm.py | 10 +- .../layers/mamba/mamba_mixer2.py | 64 ++++++++--- vllm/model_executor/models/nemotron_h.py | 36 ++++++- vllm/v1/attention/backends/mamba_attn.py | 9 +- vllm/v1/core/kv_cache_coordinator.py | 11 +- vllm/v1/core/kv_cache_utils.py | 61 ++++++++--- vllm/v1/core/single_type_kv_cache_manager.py | 39 ++++--- vllm/v1/worker/gpu/model_runner.py | 4 +- .../worker/gpu/model_states/mamba_hybrid.py | 65 +++++++++++- vllm/v1/worker/utils.py | 28 +++-- 20 files changed, 712 insertions(+), 91 deletions(-) create mode 100644 tests/model_executor/test_nemotron_h_mtp.py diff --git a/tests/model_executor/test_nemotron_h_mtp.py b/tests/model_executor/test_nemotron_h_mtp.py new file mode 100644 index 000000000000..e38378f0f54d --- /dev/null +++ b/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", + ] diff --git a/tests/model_executor/test_nemotron_h_quantization.py b/tests/model_executor/test_nemotron_h_quantization.py index 81a3fb5b13fc..8adcac3f0da4 100644 --- a/tests/model_executor/test_nemotron_h_quantization.py +++ b/tests/model_executor/test_nemotron_h_quantization.py @@ -3,6 +3,11 @@ from unittest.mock import Mock, patch +import pytest +import torch + +from vllm.config.mamba import MambaBackendEnum + def test_nemotron_h_lm_head_receives_quant_config(): from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM @@ -32,3 +37,38 @@ def test_nemotron_h_lm_head_receives_quant_config(): MockLMHead.assert_called_once() call_kwargs = MockLMHead.call_args.kwargs assert call_kwargs["quant_config"] is mock_quant_config + + +@pytest.mark.parametrize( + ("backend", "expected_num_states"), + [ + (MambaBackendEnum.TRITON, 5), + (MambaBackendEnum.FLASHINFER, 2), + ], +) +def test_nemotron_h_replayssm_platform_sizing_is_backend_scoped( + backend: MambaBackendEnum, + expected_num_states: int, +): + from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM + + config = Mock() + config.cache_config.use_replayssm = True + config.cache_config.replayssm_buffer_len = 16 + config.cache_config.mamba_cache_dtype = "auto" + config.cache_config.mamba_ssm_cache_dtype = "float32" + config.mamba_config.backend = backend + config.model_config.dtype = torch.bfloat16 + config.model_config.hf_config.mamba_num_heads = 32 + config.model_config.hf_config.mamba_head_dim = 64 + config.model_config.hf_config.n_groups = 8 + config.model_config.hf_config.ssm_state_size = 128 + config.model_config.hf_config.conv_kernel = 4 + config.parallel_config.tensor_parallel_size = 1 + config.num_speculative_tokens = 0 + + shapes = NemotronHForCausalLM.get_mamba_state_shape_from_config(config) + dtypes = NemotronHForCausalLM.get_mamba_state_dtype_from_config(config) + + assert len(shapes) == expected_num_states + assert len(dtypes) == expected_num_states diff --git a/tests/test_config.py b/tests/test_config.py index d27e4c50282d..2d4becd37cc5 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -92,6 +92,13 @@ def test_kda_recoverssm_derivation_is_revalidated(): with pytest.raises(ValueError, match="pipeline_parallel_size=1"): VllmConfig.validate_mamba_cached_kernel(config) + # Ordinary Triton ReplaySSM keeps its pre-existing PP support surface. + config.model_config.architecture = "NemotronHForCausalLM" + config.num_speculative_tokens = 0 + config.cache_config.use_kda_recoverssm = False + config.use_v2_model_runner = False + VllmConfig.validate_mamba_cached_kernel(config) + def test_per_request_spec_decode_metrics_requires_spec_decode(): # The flag only makes sense with speculative decoding configured; enabling diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index a5dbfe97b480..81dbebfd1cc0 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -326,11 +326,7 @@ def test_block_idx_prev_step_persistent_buffer_allocated(): def test_all_spec_decode_without_previous_anchor_leaves_metadata_unset(): - """A step without a previous-step anchor keeps the optional field unset. - - The mixer already falls back to the last computed block when this metadata - is absent; populating it here would change the baseline all-mode path. - """ + """A step without a previous-step anchor keeps the optional field unset.""" block_size = 16 seq_lens = [33, 49] config = _make_vllm_config( diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index 740988c81352..211a7a3afd76 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -206,12 +206,14 @@ def _make_mamba_spec( (1, ring_buffer_len), (1, ring_buffer_len, 1), ) + base_shapes = ((1, 1), (1, 1, 1)) + flashinfer = mamba_backend == MambaBackendEnum.FLASHINFER return MambaSpec( block_size=BLOCK_SIZE, - shapes=((1, 1), (1, 1, 1)), - dtypes=(torch.float32,) * 2, - replayssm_shapes=replayssm_shapes, - replayssm_dtypes=(torch.float32,) * 3, + shapes=base_shapes if flashinfer else (*base_shapes, *replayssm_shapes), + dtypes=(torch.float32,) * (2 if flashinfer else 5), + replayssm_shapes=replayssm_shapes if flashinfer else (), + replayssm_dtypes=(torch.float32,) * 3 if flashinfer else (), mamba_cache_mode=mamba_cache_mode, ) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c5159e0d1627..e69887a23af1 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3426,3 +3426,74 @@ def test_deepseek_v4_annotation_requires_model_type(): ) assert not any(g.is_eagle_group for g in groups) + + +@pytest.mark.parametrize("model_type", ["nemotron_h", "nemotron_h_puzzle", "qwen3_5"]) +def test_native_mtp_draft_group_annotated_on_general_path(model_type: str): + specs = { + "language_model.model.layers.0.self_attn.attn": new_kv_cache_spec(), + "language_model.model.layers.1.linear_attn": new_mamba_spec( + block_size=64, mamba_cache_mode="align" + ), + # These models load native MTP after the target. The ordinary + # full-attention spec carries no drafter marker, so registration order + # is the model-scoped discriminator. + "model.layers.0.self_attn.attn": new_kv_cache_spec(), + } + assert len({spec.page_size_bytes for spec in specs.values()}) == 1 + + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type=model_type), specs + ) + + flagged = [group for group in groups if group.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.0.self_attn.attn" in flagged[0].layer_names + assert not any( + group.is_eagle_group + and "language_model.model.layers.1.linear_attn" in group.layer_names + for group in groups + ) + + +@pytest.mark.parametrize("model_type", ["nemotron_h", "nemotron_h_puzzle", "qwen3_5"]) +def test_native_mtp_draft_group_annotated_on_packed_path(model_type: str): + specs = { + "language_model.model.layers.0.self_attn.attn": new_kv_cache_spec( + num_kv_heads=1 + ), + "language_model.model.layers.1.linear_attn": new_mamba_spec( + block_size=64, mamba_cache_mode="align" + ), + "model.layers.0.self_attn.attn": new_kv_cache_spec(num_kv_heads=1), + } + assert len({spec.page_size_bytes for spec in specs.values()}) > 1 + + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="mtp", model_type=model_type), specs + ) + + flagged = [group for group in groups if group.is_eagle_group] + assert len(flagged) == 1 + assert "model.layers.0.self_attn.attn" in flagged[0].layer_names + assert not any( + group.is_eagle_group + and "language_model.model.layers.1.linear_attn" in group.layer_names + for group in groups + ) + + +@pytest.mark.parametrize("model_type", ["nemotron_h", "nemotron_h_puzzle", "qwen3_5"]) +def test_last_layer_fallback_requires_native_mtp(model_type: str): + groups = get_kv_cache_groups( + _spec_decode_grouping_config(method="eagle", model_type=model_type), + { + "language_model.model.layers.0.self_attn.attn": new_kv_cache_spec(), + "language_model.model.layers.1.linear_attn": new_mamba_spec( + block_size=64, mamba_cache_mode="align" + ), + "model.layers.0.self_attn.attn": new_kv_cache_spec(), + }, + ) + + assert not any(group.is_eagle_group for group in groups) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index ec8515a373a6..4ea34e10acec 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -4226,6 +4226,69 @@ def retained(retention_interval, shared_prefix_boundary, end_block=16): assert retained(0, None) == {14} +@pytest.mark.parametrize( + ("prompt_length", "hash_block_size"), + [(40, 16), (32, 16), (32, 8)], + ids=("non-exact-coarse-hash", "exact-coarse-hash", "exact-fine-hash"), +) +def test_mamba_sparse_retention_keeps_hybrid_eagle_replay_boundary( + prompt_length, hash_block_size +): + """The first prompt must retain the state that a sibling MTP group can hit. + + The full-attention draft group drops its last matching block. The Mamba + group is not itself an EAGLE group, but its sparse retention boundary must + still shift by the same scheduler block or the first repeated prompt misses. + """ + block_size = 16 + config = KVCacheConfig( + num_blocks=100, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_mtp"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + is_eagle_group=True, + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=block_size, + shapes=((1, 1),), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + retention_interval=0, + use_eagle=True, + ) + + prompt = [i // hash_block_size for i in range(prompt_length)] + prime = make_request("prime", prompt, hash_block_size, sha256) + assert manager.allocate_slots(prime, block_size) is not None + prime.num_computed_tokens = block_size + assert manager.allocate_slots(prime, len(prompt) - block_size) is not None + prime.num_computed_tokens = len(prompt) + manager.free(prime) + + repeated = make_request("repeated", [*prompt, 999], hash_block_size, sha256) + blocks, num_computed_tokens, _ = manager.get_computed_blocks(repeated) + assert num_computed_tokens == block_size + assert all(len(group_blocks) == 1 for group_blocks in blocks.blocks) + + def test_mamba_shared_prefix_survives_zero_retention(): """Manager-level check of the full wiring: a pinned shared-prefix boundary (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 6db6635ab752..0c683e65e8b2 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -254,7 +254,7 @@ def _prefix_cache_hits(llm) -> int: ) -def _check_flashinfer_replayssm_prefix_caching( +def _check_replayssm_prefix_caching( vllm_runner, model_name, monkeypatch: pytest.MonkeyPatch, @@ -264,6 +264,7 @@ def _check_flashinfer_replayssm_prefix_caching( use_ngram: bool, use_v2: bool, tensor_parallel_size: int, + mamba_backend: str = "flashinfer", ): def run() -> None: # ReplaySSM materializes the exact SSM state at each cacheable block @@ -274,7 +275,7 @@ def run() -> None: enable_prefix_caching=True, enable_chunked_prefill=True, mamba_cache_mode=mamba_cache_mode, - mamba_backend="flashinfer", + mamba_backend=mamba_backend, disable_log_stats=False, # required for llm.get_metrics() tensor_parallel_size=tensor_parallel_size, ) @@ -303,7 +304,13 @@ def run() -> None: ) as llm: assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is use_v2 replay_block_size = llm.llm.llm_engine.vllm_config.cache_config.block_size - assert replay_block_size == baseline_block_size + if mamba_backend == "flashinfer": + # FlashInfer rings are auxiliary and cannot affect the shared page. + assert replay_block_size == baseline_block_size + else: + # Triton retains the original packed five-state page. Its rings may + # increase the attention block size needed to match that page. + assert replay_block_size >= baseline_block_size llm.generate_greedy_logprobs( PREFIX_CACHING_PROMPTS, max_tokens=32, num_logprobs=5 ) @@ -321,8 +328,8 @@ def run() -> None: check_logprobs_close( outputs_0_lst=baseline, outputs_1_lst=replay, - name_0=f"flashinfer_baseline_{mamba_cache_mode}_pc", - name_1=f"flashinfer_replayssm_{mamba_cache_mode}_pc", + name_0=f"{mamba_backend}_baseline_{mamba_cache_mode}_pc", + name_1=f"{mamba_backend}_replayssm_{mamba_cache_mode}_pc", ) try: @@ -350,7 +357,7 @@ def test_flashinfer_replayssm_prefix_cache_tp1( use_v2: bool, use_ngram: bool, ): - _check_flashinfer_replayssm_prefix_caching( + _check_replayssm_prefix_caching( vllm_runner, model_name, monkeypatch, @@ -361,11 +368,27 @@ def test_flashinfer_replayssm_prefix_cache_tp1( ) +@pytest.mark.parametrize("model_name", MODELS) +def test_triton_replayssm_align_prefix_cache_matches_baseline_v1( + vllm_runner, model_name, monkeypatch: pytest.MonkeyPatch +): + _check_replayssm_prefix_caching( + vllm_runner, + model_name, + monkeypatch, + mamba_cache_mode="align", + use_ngram=False, + use_v2=False, + tensor_parallel_size=1, + mamba_backend="triton", + ) + + @requires_flashinfer_replayssm_materialization @large_gpu_mark(min_gb=40) @pytest.mark.parametrize("use_v2", [False, True], ids=["v1", "v2"]) def test_flashinfer_replayssm_all_prefix_cache(vllm_runner, monkeypatch, use_v2: bool): - _check_flashinfer_replayssm_prefix_caching( + _check_replayssm_prefix_caching( vllm_runner, MAMBA2_PREFIX_MODEL, monkeypatch, diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 14f5ce78f6fe..5fa2043789c1 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -8,11 +8,16 @@ import torch from vllm.platforms import current_platform +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder 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.mamba_hybrid import ( + MambaHybridAttnMetadata, + MambaHybridModelState, +) from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState @@ -34,11 +39,70 @@ def test_add_request_seeds_state_with_scoped_block_size( state._use_flashinfer_replayssm = use_flashinfer_replayssm state.num_accepted_tokens_gpu = torch.full((2,), 9, dtype=torch.int32) state._mamba_state_idx_gpu = torch.full((2,), -1, dtype=torch.int32) + state._mamba_prev_last_scheduled_idx_gpu = torch.full((2,), 9, dtype=torch.int32) state.add_request(1, Mock(num_computed_tokens=17)) assert state.num_accepted_tokens_gpu.tolist() == [9, 1] assert state._mamba_state_idx_gpu.tolist() == [-1, expected_state_idx] + assert state._mamba_prev_last_scheduled_idx_gpu.tolist() == [9, -1] + + +def test_all_spec_tracks_previous_scheduled_page_by_request() -> None: + state = object.__new__(MambaHybridModelState) + state.cache_config = SimpleNamespace(mamba_block_size=16) + state.vllm_config = SimpleNamespace(num_speculative_tokens=3) + state._mamba_prev_last_scheduled_idx_gpu = torch.full((4,), -1, dtype=torch.int32) + + first_batch = SimpleNamespace( + num_reqs=2, + idx_mapping=torch.tensor([2, 0], dtype=torch.int32), + seq_lens=torch.tensor([18, 33], dtype=torch.int32), + query_start_loc=torch.tensor([0, 4, 8], dtype=torch.int32), + ) + first_prev = state._stage_prev_last_scheduled_idx(first_batch, num_reqs=3) + + assert first_prev is not None + assert first_prev.tolist() == [-1, -1, -1] + assert state._mamba_prev_last_scheduled_idx_gpu.tolist() == [2, -1, 1, -1] + + # Request 2 previously started with N=14 tokens and scheduled q=4, so its + # state window is anchored at page 1. Accepting one token leaves N=15 and + # logical last-computed page 0, which must not replace that physical anchor. + logical_last_computed = (15 - 1) // 16 + state._mamba_prev_last_scheduled_idx_gpu[1] = 7 + second_batch = SimpleNamespace( + num_reqs=3, + idx_mapping=torch.tensor([0, 2, 1], dtype=torch.int32), + seq_lens=torch.tensor([34, 19, 5], dtype=torch.int32), + query_start_loc=torch.tensor([0, 4, 8, 9], dtype=torch.int32), + ) + second_prev = state._stage_prev_last_scheduled_idx(second_batch, num_reqs=4) + + assert second_prev is not None + assert second_prev.tolist() == [2, 1, 7, -1] + assert second_prev[1].item() != logical_last_computed + # The one-token row may consume its old anchor in this step, but must clear + # it rather than advertising a speculative window to the following step. + assert state._mamba_prev_last_scheduled_idx_gpu.tolist() == [2, -1, 1, -1] + + +def test_previous_scheduled_page_is_passed_only_to_mamba2() -> None: + prev_last_scheduled_idx = torch.tensor([3, 5], dtype=torch.int32) + metadata = MambaHybridAttnMetadata( + is_prefilling=torch.zeros(2, dtype=torch.bool), + prev_last_scheduled_idx=prev_last_scheduled_idx, + ) + + mamba2_args = metadata.get_extra_attn_kwargs( + Mock(spec=Mamba2AttentionMetadataBuilder), 2 + ) + gdn_args = metadata.get_extra_attn_kwargs(Mock(spec=GDNAttentionMetadataBuilder), 2) + + mamba2_prev = mamba2_args["prev_last_scheduled_idx"] + assert torch.equal(mamba2_prev, prev_last_scheduled_idx) + assert mamba2_prev.data_ptr() == prev_last_scheduled_idx.data_ptr() + assert "prev_last_scheduled_idx" not in gdn_args @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 33a166524739..97eac4168c3b 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch from vllm.config.mamba import MambaBackendEnum, MambaConfig @@ -19,6 +20,7 @@ class _TestReplaySSMMixer(MambaMixer2): def __init__(self) -> None: torch.nn.Module.__init__(self) self.use_replayssm = True + self.use_flashinfer_replayssm = True self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER) self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) @@ -36,6 +38,25 @@ def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: return (torch.float32,) * 3 +class _TestTritonReplaySSMMixer(_TestReplaySSMMixer): + def __init__(self) -> None: + super().__init__() + self.use_flashinfer_replayssm = False + self.mamba_config = MambaConfig(backend=MambaBackendEnum.TRITON) + + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + return ((1,),) * 5 + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: + return (torch.float32,) * 5 + + def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: + return () + + def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: + return () + + def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: return torch.zeros((num_blocks, 1, 1, 20), dtype=torch.int8) @@ -140,6 +161,85 @@ def test_replayssm_block_copy_includes_rings_and_group_trackers(monkeypatch): assert mixers[1]._replayssm_prev_num_accepted[dst].item() == 31 +def test_replayssm_block_copy_excludes_triton_auxiliary_state(): + mixer = _TestReplaySSMMixer() + mixer.use_flashinfer_replayssm = False + mixer.mamba_config = MambaConfig(backend=MambaBackendEnum.TRITON) + mixer.replayssm_cache = tuple(torch.zeros((2, 1)) for _ in range(3)) + + assert get_replayssm_block_copy_tensors({"mixer": mixer}) == [] + + +def test_triton_replayssm_raw_page_copy_includes_all_five_states(monkeypatch): + monkeypatch.setattr( + "vllm.v1.worker.utils.async_tensor_h2d", + lambda array, *, device, **_: torch.from_numpy(array).to(device), + ) + mixer = _TestTritonReplaySSMMixer() + layer_name = "layers.0.mixer" + raw_cache = _packed_replayssm_cache(2) + runner_kv_caches: list[torch.Tensor] = [] + bind_kv_cache( + {layer_name: raw_cache}, + {layer_name: mixer}, + runner_kv_caches, + ) + + src, dst = 0, 1 + for state_idx, state in enumerate(mixer.kv_cache): + state[src].fill_(state_idx + 1) + state[dst].fill_(-1) + + copy_kv_cache_blocks_inplace( + runner_kv_caches, + 2, + [KVCacheBlockCopy(src, dst)], + ) + + assert len(mixer.kv_cache) == 5 + for state_idx, state in enumerate(mixer.kv_cache): + assert state[dst].item() == state_idx + 1 + + +@pytest.mark.parametrize( + ("backend", "state_count", "auxiliary_count"), + [ + (MambaBackendEnum.TRITON, 5, 0), + (MambaBackendEnum.FLASHINFER, 2, 3), + ], +) +def test_replayssm_cache_layout_is_backend_scoped( + monkeypatch, backend: MambaBackendEnum, state_count: int, auxiliary_count: int +): + monkeypatch.setattr( + "vllm.model_executor.layers.mamba.mamba_mixer2." + "get_tensor_model_parallel_world_size", + lambda: 1, + ) + mixer = MambaMixer2.__new__(MambaMixer2) + torch.nn.Module.__init__(mixer) + mixer.use_replayssm = True + mixer.use_flashinfer_replayssm = backend == MambaBackendEnum.FLASHINFER + mixer.mamba_config = MambaConfig(backend=backend) + mixer.model_config = SimpleNamespace(dtype=torch.bfloat16) + mixer.cache_config = SimpleNamespace( + mamba_cache_dtype="auto", mamba_ssm_cache_dtype="auto" + ) + mixer.intermediate_size = 32 + mixer.n_groups = 2 + mixer.num_heads = 4 + mixer.head_dim = 8 + mixer.ssm_state_size = 16 + mixer.conv_kernel_size = 4 + mixer.num_spec = 0 + mixer.replayssm_buffer_len = 16 + + assert len(mixer.get_state_shape()) == state_count + assert len(mixer.get_state_dtype()) == state_count + assert len(mixer.get_replayssm_state_shape()) == auxiliary_count + assert len(mixer.get_replayssm_state_dtype()) == auxiliary_count + + def test_bind_kv_cache(default_vllm_config): from vllm.model_executor.layers.attention import Attention diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index d7d04a191f35..484e2d616908 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2876,8 +2876,6 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": "--use-replayssm is not supported for architecture " f"{self.model_config.architecture!r}" ) - if self.parallel_config.pipeline_parallel_size > 1: - raise ValueError("ReplaySSM currently requires pipeline_parallel_size=1") if self.mamba_config.backend == MambaBackendEnum.TRITON: if self.cache_config.use_kda_recoverssm: if self.mamba_config.enable_stochastic_rounding: @@ -2897,6 +2895,10 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": 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" + ) else: if use_mamba_replayssm_spec: raise ValueError( @@ -2916,6 +2918,10 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: if self.cache_config.use_kda_recoverssm: raise ValueError("RecoverSSM requires --mamba-backend triton") + if self.parallel_config.pipeline_parallel_size > 1: + raise ValueError( + "FlashInfer ReplaySSM currently requires pipeline_parallel_size=1" + ) if self.cache_config.replayssm_buffer_len > 16: raise ValueError( "FlashInfer ReplaySSM requires --replayssm-buffer-len <= 16" diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index d0290c2812f2..98432cf0483e 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -519,13 +519,22 @@ def __init__( else None ) self.mamba_config = vllm_config.mamba_config + self.use_flashinfer_replayssm = ( + self.use_replayssm + and self.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) if self.use_replayssm and self.num_heads % self.tp_size != 0: raise ValueError( "--use-replayssm requires tensor-parallel heads to divide evenly" ) - self.kv_cache = tuple(torch.tensor([]) for _ in range(2)) + # Keep Triton's established five-state packed page. FlashInfer's rings + # are auxiliary because its materializer addresses them independently. + num_states = 2 if not self.use_replayssm or self.use_flashinfer_replayssm else 5 + self.kv_cache = tuple(torch.tensor([]) for _ in range(num_states)) self.replayssm_cache = ( - tuple(torch.tensor([]) for _ in range(3)) if self.use_replayssm else () + tuple(torch.tensor([]) for _ in range(3)) + if self.use_flashinfer_replayssm + else () ) self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) @@ -734,10 +743,12 @@ def conv_ssm_forward( ) ssm_state = self.kv_cache[1] if self.use_replayssm: - x_cache, dt_cache, B_cache = self.replayssm_cache - if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: + if self.use_flashinfer_replayssm: + x_cache, dt_cache, B_cache = self.replayssm_cache ring_start = self._replayssm_ring_start prev_num_accepted = self._replayssm_prev_num_accepted + else: + x_cache, dt_cache, B_cache = self.kv_cache[2:5] else: x_cache = dt_cache = B_cache = None has_initial_states_p = attn_metadata.has_initial_states_p @@ -1196,36 +1207,55 @@ def conv_ssm_forward( def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None - return MambaStateDtypeCalculator.mamba2_state_dtype( + dtypes = MambaStateDtypeCalculator.mamba2_state_dtype( self.model_config.dtype, self.cache_config.mamba_cache_dtype, self.cache_config.mamba_ssm_cache_dtype, ) + if self.use_replayssm and not self.use_flashinfer_replayssm: + dtypes = ( + *dtypes, + *MambaStateDtypeCalculator.replayssm_ring_dtypes( + self.model_config.dtype + ), + ) + return dtypes def get_state_shape(self) -> tuple[tuple[int, ...], ...]: tp_world_size = get_tensor_model_parallel_world_size() - return MambaStateShapeCalculator.mamba2_state_shape( - intermediate_size=self.intermediate_size, - tp_world_size=tp_world_size, - n_groups=self.n_groups, - num_heads=self.num_heads, - head_dim=self.head_dim, - state_size=self.ssm_state_size, - conv_kernel=self.conv_kernel_size, - num_spec=self.num_spec, + shapes: tuple[tuple[int, ...], ...] = ( + MambaStateShapeCalculator.mamba2_state_shape( + intermediate_size=self.intermediate_size, + tp_world_size=tp_world_size, + n_groups=self.n_groups, + num_heads=self.num_heads, + head_dim=self.head_dim, + state_size=self.ssm_state_size, + conv_kernel=self.conv_kernel_size, + num_spec=self.num_spec, + ) ) + if self.use_replayssm and not self.use_flashinfer_replayssm: + shapes = (*shapes, *self._get_replayssm_ring_shapes(tp_world_size)) + return shapes def get_replayssm_state_dtype(self) -> tuple[torch.dtype, ...]: - if not self.use_replayssm: + if not self.use_flashinfer_replayssm: return () assert self.model_config is not None return MambaStateDtypeCalculator.replayssm_ring_dtypes(self.model_config.dtype) def get_replayssm_state_shape(self) -> tuple[tuple[int, ...], ...]: - if not self.use_replayssm: + if not self.use_flashinfer_replayssm: return () assert self.replayssm_buffer_len is not None tp_world_size = get_tensor_model_parallel_world_size() + return self._get_replayssm_ring_shapes(tp_world_size) + + def _get_replayssm_ring_shapes( + self, tp_world_size: int + ) -> tuple[tuple[int, ...], ...]: + assert self.replayssm_buffer_len is not None return MambaStateShapeCalculator.replayssm_ring_shapes( num_heads=self.num_heads, head_dim=self.head_dim, @@ -1262,7 +1292,7 @@ def share_replayssm_ring_trackers( if ( isinstance(layer, MambaMixer2) and layer.use_replayssm - and layer.mamba_config.backend == MambaBackendEnum.FLASHINFER + and layer.use_flashinfer_replayssm ): replayssm_mixers[layer_name] = layer diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index a875c5c1da02..234f314c8981 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -26,6 +26,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, ModelConfig, VllmConfig +from vllm.config.mamba import MambaBackendEnum from vllm.config.parallel import ParallelConfig from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size from vllm.distributed.communication_op import tensor_model_parallel_all_gather @@ -740,11 +741,22 @@ def get_mamba_state_dtype_from_config( vllm_config: "VllmConfig", ) -> tuple[torch.dtype, ...]: cache_config = vllm_config.cache_config - return MambaStateDtypeCalculator.mamba2_state_dtype( + dtypes = MambaStateDtypeCalculator.mamba2_state_dtype( vllm_config.model_config.dtype, cache_config.mamba_cache_dtype, cache_config.mamba_ssm_cache_dtype, ) + if ( + cache_config.use_replayssm + and vllm_config.mamba_config.backend == MambaBackendEnum.TRITON + ): + dtypes = ( + *dtypes, + *MambaStateDtypeCalculator.replayssm_ring_dtypes( + vllm_config.model_config.dtype + ), + ) + return dtypes @classmethod def get_mamba_state_shape_from_config( @@ -760,12 +772,14 @@ def get_mamba_state_shape_from_config( Tuple containing: - conv_state_shape: Shape for convolutional state cache - temporal_state_shape: Shape for state space model cache + - packed x/dt/B ReplaySSM rings for the Triton backend """ parallel_config = vllm_config.parallel_config + cache_config = vllm_config.cache_config hf_config = vllm_config.model_config.hf_config intermediate_size = hf_config.mamba_num_heads * hf_config.mamba_head_dim - return MambaStateShapeCalculator.mamba2_state_shape( + shapes = MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=intermediate_size, tp_world_size=parallel_config.tensor_parallel_size, n_groups=hf_config.n_groups, @@ -775,6 +789,24 @@ def get_mamba_state_shape_from_config( conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, ) + if ( + cache_config.use_replayssm + and vllm_config.mamba_config.backend == MambaBackendEnum.TRITON + ): + shapes = ( + *shapes, + *MambaStateShapeCalculator.replayssm_ring_shapes( + num_heads=hf_config.mamba_num_heads, + head_dim=hf_config.mamba_head_dim, + state_size=hf_config.ssm_state_size, + n_groups=hf_config.n_groups, + tp_world_size=parallel_config.tensor_parallel_size, + logical_window=cache_config.replayssm_buffer_len, + backend=vllm_config.mamba_config.backend, + num_speculative_tokens=vllm_config.num_speculative_tokens, + ), + ) + return shapes @classmethod def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]: diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 8bf25c732d35..f9a535eb8440 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -181,9 +181,8 @@ def __init__( tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None ) = None self.decode_replayssm_state_indices_d: torch.Tensor | None = None - # Canonical state is (conv, ssm). ReplaySSM auxiliaries are ordered as - # x=(nheads, ring, head_dim), dt=(nheads, ring), - # B=(ngroups, ring, dstate). + # FlashInfer keeps x/dt/B as auxiliary state. Triton retains its + # established packed (conv, ssm, x, dt, B) page layout. if self.use_replayssm and not self.use_flashinfer_replayssm: self.decode_write_pos_d: torch.Tensor = torch.empty( (self.decode_cudagraph_max_bs,), @@ -195,7 +194,9 @@ def __init__( dtype=torch.int8, device=device, ) - bc_ngroups = kv_cache_spec.replayssm_shapes[2][0] + triton_replayssm_shapes = kv_cache_spec.shapes[2:5] + assert len(triton_replayssm_shapes) == 3 + bc_ngroups = triton_replayssm_shapes[2][0] bc_scratch_bs = max( self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs ) diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index b3a9b4c03210..79636553559c 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -658,8 +658,17 @@ def __init__( ", ".join(sorted(unsupported_partial_hit_managers)), ) cache_hit_alignment_tokens = self._cache_hit_alignment_tokens - for manager in self.single_type_managers: + for manager, group in zip( + self.single_type_managers, + kv_cache_config.kv_cache_groups, + strict=True, + ): manager.cache_hit_alignment_tokens = cache_hit_alignment_tokens + if self.eagle_group_ids and isinstance(group.kv_cache_spec, MambaSpec): + # The draft attention group drops one scheduler block from a + # cache hit. Preserve the matching state in sparse Mamba + # retention without reclassifying the Mamba group as EAGLE. + manager._sparse_replay_boundary_shift = self.scheduler_block_size self.verify_and_split_kv_cache_groups() @property diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index ec24e987413a..a4f67f4140de 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -1815,19 +1815,40 @@ def widest_group_bytes(page_size_layers: dict[int, list[str]], n: int) -> int: vllm_config, kv_cache_spec, groups, - use_deepseek_v4_fallback=_is_deepseek_v4_eagle(vllm_config), + use_last_registered_draft_fallback=( + _uses_last_registered_draft_layer(vllm_config) + ), ) _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups -def _is_deepseek_v4_eagle(vllm_config: VllmConfig) -> bool: +def _uses_last_registered_draft_layer(vllm_config: VllmConfig) -> bool: spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle(): return False - model_config = vllm_config.model_config + model_config = getattr(vllm_config, "model_config", None) + if model_config is None: + return False + model_type = model_config.hf_config.model_type + # DeepSeekV4 uses this ordering for its EAGLE-family heads. Some native MTP + # models have the same invariant, but only for their MTP path. + return model_type == "deepseek_v4" or _uses_native_mtp_draft_layer_fallback( + vllm_config + ) + + +def _uses_native_mtp_draft_layer_fallback(vllm_config: VllmConfig) -> bool: + spec_config = vllm_config.speculative_config + model_config = getattr(vllm_config, "model_config", None) return ( - model_config is not None and model_config.hf_config.model_type == "deepseek_v4" + spec_config is not None + and spec_config.method == "mtp" + and model_config is not None + # These native MTP implementations register their ordinary attention + # draft layer after every target cache layer. + and model_config.hf_config.model_type + in ("nemotron_h", "nemotron_h_puzzle", "qwen3_5") ) @@ -1835,7 +1856,7 @@ def _annotate_eagle_groups( vllm_config: VllmConfig, kv_cache_spec: dict[str, KVCacheSpec], kv_cache_groups: list[KVCacheGroupSpec], - use_deepseek_v4_fallback: bool = False, + use_last_registered_draft_fallback: bool = False, ) -> None: """Flag the KV cache groups that hold drafter attention layers. @@ -1848,14 +1869,13 @@ def _annotate_eagle_groups( spec merging, wherever grouping happens to land. It is sufficient but not necessary: a drafter whose spec is indistinguishable from the target's cannot be found this way. - 2. Model-scoped positional fallback for DeepseekV4, whose MTP block reuses - the target's own decoder layer and so carries no spec marker. Its draft - attention layer is always the last registered layer, so flag whichever - group holds it. This rule is only valid where the groups partition - exactly the layers of ``kv_cache_spec``, which is true on the packed - grouping path and not in general; other callers must leave - ``use_deepseek_v4_fallback`` False. The caller gates this fallback on - the configured model type. + 2. Model-scoped positional fallback for models whose MTP attention carries + no spec marker but is known to register last. Flag whichever group holds + that layer. This rule requires ``kv_cache_groups`` to partition exactly + the layers of ``kv_cache_spec``. Both the packed and general grouping + callers satisfy that invariant; other callers must leave + ``use_last_registered_draft_fallback`` False. The caller also gates this + fallback on the configured model and speculative method. FIXME(yifan): avoid/generalize this hacky check. Args: @@ -1863,7 +1883,9 @@ def _annotate_eagle_groups( kv_cache_spec: The kv cache spec of each attention layer, in layer registration order. Only read by rule 2. kv_cache_groups: Groups to annotate in place. - use_deepseek_v4_fallback: Enable rule 2 for a DeepseekV4 packed group. + use_last_registered_draft_fallback: Enable rule 2 for a model with a + known last-registered draft layer when the groups exactly partition + ``kv_cache_spec``. """ spec_config = vllm_config.speculative_config if spec_config is None or not spec_config.use_eagle_block_drop(): @@ -1876,7 +1898,7 @@ def _annotate_eagle_groups( ): group.is_eagle_group = True - if not use_deepseek_v4_fallback: + if not use_last_registered_draft_fallback: return last_layer = next(reversed(kv_cache_spec)) for group in kv_cache_groups: @@ -2016,7 +2038,14 @@ def get_kv_cache_groups( aligned = replace(spec, block_size=new_bs, page_size_padded=common_page) groups.append(KVCacheGroupSpec([name], aligned)) - _annotate_eagle_groups(vllm_config, kv_cache_spec, groups) + _annotate_eagle_groups( + vllm_config, + kv_cache_spec, + groups, + use_last_registered_draft_fallback=( + _uses_native_mtp_draft_layer_fallback(vllm_config) + ), + ) _warn_if_unannotated_eagle_mamba(vllm_config, groups) return groups diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 85adfab2daec..a94a6eaed802 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -113,6 +113,11 @@ def __init__( # aligned segment (SWA). Initialized lazily by the coordinator after # determining the attention groups. self.use_eagle = False + # Hybrid EAGLE/MTP drops one scheduler block from the common reusable + # prefix. Sparse Mamba retention must keep that shifted replay state, + # even though the Mamba group itself is not an EAGLE group. The hybrid + # coordinator sets this after it identifies the draft cache group. + self._sparse_replay_boundary_shift = 0 # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained # managers (full attention, mamba "align"); harmlessly empty elsewhere. @@ -454,9 +459,16 @@ def cache_blocks( return # Token boundaries whose reachable tail must be retained under sparse - # retention: the replay boundary (``num_prompt - 1``, capped by - # ``get_computed_blocks``) and any detected shared-prefix junction. - reachable_boundaries = [request.num_prompt_tokens - 1] + # retention: the replay boundary, capped by ``get_computed_blocks``, and + # any detected shared-prefix junction. The ordinary boundary is the + # prompt's last token; a shifted boundary is exclusive so an exact + # ``num_prompt - shift`` boundary remains representable. + replay_boundary = ( + request.num_prompt_tokens - self._sparse_replay_boundary_shift + if self._sparse_replay_boundary_shift + else request.num_prompt_tokens - 1 + ) + reachable_boundaries = [max(replay_boundary, 0)] if request.shared_prefix_boundary: reachable_boundaries.append(request.shared_prefix_boundary) @@ -1372,14 +1384,13 @@ def __init__( self.block_size = kv_cache_spec.block_size self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks - # ``replayssm_shapes`` means that this group's temporal state includes - # ReplaySSM rings indexed by the scheduler's physical block IDs. When - # allocation moves a request's live block, those rings and their - # trackers must move with it; otherwise the next forward reads the new - # block ID with the previous live history left behind. Both Triton and - # FlashInfer ReplaySSM use this scheduler-level migration path. - self._copy_replayssm_live_state = bool(kv_cache_spec.replayssm_shapes) - if self._copy_replayssm_live_state: + # FlashInfer ReplaySSM rings and trackers live outside the canonical + # cache page, so explicitly migrate them when the live block moves. + # Triton's packed five-state page follows the normal copy path. + self._copy_flashinfer_replayssm_live_state = bool( + kv_cache_spec.replayssm_shapes + ) + if self._copy_flashinfer_replayssm_live_state: # ReplaySSM stores speculative history inside its ring. It must not # also reserve the baseline Mamba speculative scratch blocks. assert self.num_speculative_blocks == 0, ( @@ -1680,7 +1691,7 @@ def allocate_new_blocks( # speculative decoding (MTP/EAGLE) with linear attention. if self.num_speculative_blocks > 0: num_tokens += self.block_size * self.num_speculative_blocks - if not self._copy_replayssm_live_state: + if not self._copy_flashinfer_replayssm_live_state: return super().allocate_new_blocks( request_id, num_tokens, num_tokens_main_model ) @@ -1721,7 +1732,7 @@ def allocate_new_blocks( partial_hit = self._partial_hit_reqs.get(request_id) has_partial_hit = partial_hit is not None live_source = None - if self._copy_replayssm_live_state: + if self._copy_flashinfer_replayssm_live_state: # Capture the live state owner before align-mode allocation # relocates/nulls table entries. The eventual destination is # the last new block, where the next forward writes state. @@ -1815,7 +1826,7 @@ def allocate_new_blocks( self._apply_cow(request_id, block_idx, source_block, cow_block) returned_blocks = [cow_block] + returned_blocks req_blocks.extend(new_blocks) - if self._copy_replayssm_live_state and new_blocks: + if self._copy_flashinfer_replayssm_live_state and new_blocks: self._queue_replayssm_live_copy(live_source, req_blocks[-1]) self._allocated_block_reqs.add(request_id) self._partial_hit_reqs.pop(request_id, None) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 46e6cab3b58b..59221a761e57 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1674,9 +1674,11 @@ def execute_model( ) assert block_tables is not None attn_groups = self.attn_groups - if dummy_run and is_profile: + if dummy_run and is_profile and not valid_dummy_state_slots: # Mamba layers take a cheap warmup path with no metadata; # attention metadata is still built so those kernels tune. + # ReplaySSM autotuning supplies valid dummy state slots and + # needs Mamba metadata so checkpointing_ssu actually runs. attn_groups = [ [g for g in groups if not isinstance(g.kv_cache_spec, MambaSpec)] for groups in attn_groups diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 9dc0f06110ee..5d33c50236ad 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -42,6 +42,7 @@ class MambaHybridAttnMetadata(ModelSpecificAttnMetadata): is_prefilling: torch.Tensor num_accepted_tokens: torch.Tensor | None = None num_decode_draft_tokens_cpu: torch.Tensor | None = None + prev_last_scheduled_idx: torch.Tensor | None = None def get_extra_common_attn_kwargs( self, @@ -65,7 +66,7 @@ def get_extra_attn_kwargs( ), ): return {} - return { + extra_args = { "num_accepted_tokens": None if self.num_accepted_tokens is None else self.num_accepted_tokens[:num_reqs], @@ -73,6 +74,14 @@ def get_extra_attn_kwargs( if self.num_decode_draft_tokens_cpu is None else self.num_decode_draft_tokens_cpu[:num_reqs], } + if ( + isinstance(attn_metadata_builder, Mamba2AttentionMetadataBuilder) + and self.prev_last_scheduled_idx is not None + ): + extra_args["prev_last_scheduled_idx"] = self.prev_last_scheduled_idx[ + :num_reqs + ] + return extra_args class MambaHybridModelState(DefaultModelState): @@ -90,6 +99,15 @@ def __init__( self.num_accepted_tokens_gpu = torch.ones( self.max_num_reqs, dtype=torch.int32, device=self.device ) + self._mamba_prev_last_scheduled_idx_gpu: torch.Tensor | None = None + if ( + self.cache_config.mamba_cache_mode == "all" + and vllm_config.num_speculative_tokens > 0 + and not self.cache_config.use_replayssm + ): + self._mamba_prev_last_scheduled_idx_gpu = torch.full( + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device + ) self._replayssm_query_start_loc: torch.Tensor | None = None # Pre-copy prefix-cache state (V2). The migration of each request's # mamba state across block boundaries runs as a fused GPU kernel reusing @@ -131,6 +149,8 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: super().add_request(req_index, new_req_data) # Must reset the speculative acceptance count in this idx which could be stale. self.num_accepted_tokens_gpu[req_index].fill_(1) + if self._mamba_prev_last_scheduled_idx_gpu is not None: + self._mamba_prev_last_scheduled_idx_gpu[req_index].fill_(-1) if self._needs_prefix_state_migration: # Seed the running state block from the resumed/prefilled position. state_block_size = self.cache_config.block_size @@ -288,6 +308,12 @@ def prepare_attn( is_prefilling[: input_batch.num_reqs] = torch.from_numpy( input_batch.is_prefilling_np ) + prev_last_scheduled_idx = None + if not for_capture: + prev_last_scheduled_idx = self._stage_prev_last_scheduled_idx( + input_batch, num_reqs + ) + if self._use_flashinfer_replayssm: self._is_prefilling_gpu[:num_reqs].copy_(is_prefilling, non_blocking=True) # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens @@ -346,6 +372,7 @@ def prepare_attn( is_prefilling=is_prefilling, num_accepted_tokens=num_accepted_tokens, num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu, + prev_last_scheduled_idx=prev_last_scheduled_idx, ) attn_metadata = build_attn_metadata( attn_groups=attn_groups, @@ -373,6 +400,42 @@ def prepare_attn( ) return attn_metadata + def _stage_prev_last_scheduled_idx( + self, input_batch: InputBatch, num_reqs: int + ) -> torch.Tensor | None: + state = self._mamba_prev_last_scheduled_idx_gpu + if state is None: + return None + + num_actual_reqs = input_batch.num_reqs + prev_last_scheduled_idx = state.new_full((num_reqs,), -1) + if num_actual_reqs == 0: + return prev_last_scheduled_idx + + idx_mapping = input_batch.idx_mapping[:num_actual_reqs] + # A non-full step can still consume the preceding speculative window. + # Gather that anchor before replacing it with this step's window state. + prev_last_scheduled_idx[:num_actual_reqs] = state[idx_mapping] + + mamba_block_size = self.cache_config.mamba_block_size + assert mamba_block_size is not None + current_last_scheduled_idx = torch.div( + input_batch.seq_lens[:num_actual_reqs] - 1, + mamba_block_size, + rounding_mode="floor", + ).clamp_(min=0) + query_lens = ( + input_batch.query_start_loc[1 : num_actual_reqs + 1] + - input_batch.query_start_loc[:num_actual_reqs] + ) + full_spec_query_len = self.vllm_config.num_speculative_tokens + 1 + state[idx_mapping] = torch.where( + query_lens == full_spec_query_len, + current_last_scheduled_idx, + -1, + ) + return prev_last_scheduled_idx + def postprocess_state( self, idx_mapping: torch.Tensor, diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index 7fdb0358dc41..6d0e1728ae65 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -11,7 +11,6 @@ import torch from vllm.config import CacheConfig, VllmConfig -from vllm.config.mamba import MambaBackendEnum from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.mamba.mamba_mixer2 import share_replayssm_ring_trackers @@ -445,7 +444,7 @@ def allocate_replayssm_caches( kv_cache_config: KVCacheConfig, device: torch.device, ) -> dict[str, tuple[torch.Tensor, ...]]: - """Allocate ReplaySSM ring state separately from canonical Mamba pages.""" + """Allocate FlashInfer ReplaySSM rings outside canonical Mamba pages.""" caches: dict[str, tuple[torch.Tensor, ...]] = {} for group in kv_cache_config.kv_cache_groups: group_spec = group.kv_cache_spec @@ -740,27 +739,24 @@ def copy_kv_cache_blocks_inplace( def get_replayssm_block_copy_tensors( forward_context: Mapping[str, Any], ) -> list[torch.Tensor]: - """Collect ReplaySSM-owned state that must follow scheduler block copies. + """Collect FlashInfer ReplaySSM state for scheduler block copies. The runner's normal KV-cache list already contains canonical convolution - and SSM state. ReplaySSM ring caches use separate allocations on every - backend, so they are added here. FlashInfer additionally keeps its ring - position and accepted-token trackers in separate group-shared tensors; - the block-copy helper deduplicates those aliases by storage. + and SSM state. Triton ReplaySSM retains its packed five-state cache page, so + normal block copies already include its rings. FlashInfer keeps both rings + and group-shared trackers in separate allocations; the block-copy helper + deduplicates layer aliases by storage. """ extra_tensors: list[torch.Tensor] = [] for layer in forward_context.values(): - if not getattr(layer, "use_replayssm", False): + if not getattr(layer, "use_flashinfer_replayssm", False): continue extra_tensors.extend(layer.replayssm_cache) - mamba_config = getattr(layer, "mamba_config", None) - backend = getattr(mamba_config, "backend", None) - if backend == MambaBackendEnum.FLASHINFER: - # Group-shared trackers appear once per layer; the block-copy helper - # deduplicates them by (device, data_ptr()). - extra_tensors.extend( - (layer._replayssm_ring_start, layer._replayssm_prev_num_accepted) - ) + # Group-shared trackers appear once per layer; the block-copy helper + # deduplicates them by (device, data_ptr()). + extra_tensors.extend( + (layer._replayssm_ring_start, layer._replayssm_prev_num_accepted) + ) return extra_tensors From 6cececc639fdfe8c6607f84cdb432e6363c597e4 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Sun, 6 Sep 2026 07:17:16 +0200 Subject: [PATCH 52/53] [Mamba] Commit RecoverSSM state in Model Runner V1 Record KDA RecoverSSM metadata in V1 and commit the accepted recurrent state immediately after sampling. Keep none mode independent of align-only slot mappings, and reject unsupported V1 microbatching. Signed-off-by: Andrii Skliar --- tests/test_config.py | 12 +++++++- tests/v1/worker/test_gpu_model_runner.py | 29 +++++++++++++++++++ .../worker/test_mamba_hybrid_model_state.py | 7 ++--- vllm/config/vllm.py | 4 +++ vllm/v1/worker/gpu/model_states/recoverssm.py | 3 +- vllm/v1/worker/gpu_model_runner.py | 23 ++++++++++++++- 6 files changed, 71 insertions(+), 7 deletions(-) diff --git a/tests/test_config.py b/tests/test_config.py index 2d4becd37cc5..b4996cccb4ce 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -61,7 +61,10 @@ def test_kda_recoverssm_derivation_is_revalidated(): backend=MambaBackendEnum.TRITON, enable_stochastic_rounding=False, ), - parallel_config=SimpleNamespace(pipeline_parallel_size=1), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + use_ubatching=False, + ), kv_transfer_config=None, use_v2_model_runner=True, ) @@ -81,6 +84,13 @@ def test_kda_recoverssm_derivation_is_revalidated(): VllmConfig.validate_mamba_cached_kernel(config) config.cache_config.mamba_cache_mode = "none" + config.use_v2_model_runner = False + config.parallel_config.use_ubatching = True + with pytest.raises(ValueError, match="does not support microbatching"): + VllmConfig.validate_mamba_cached_kernel(config) + config.parallel_config.use_ubatching = False + config.use_v2_model_runner = True + config.model_config.architecture = "NemotronHForCausalLM" config.mamba_config.backend = MambaBackendEnum.FLASHINFER VllmConfig.validate_mamba_cached_kernel(config) diff --git a/tests/v1/worker/test_gpu_model_runner.py b/tests/v1/worker/test_gpu_model_runner.py index 99331e7b354e..176e24325a19 100644 --- a/tests/v1/worker/test_gpu_model_runner.py +++ b/tests/v1/worker/test_gpu_model_runner.py @@ -717,6 +717,35 @@ def test_update_states_request_unscheduled(model_runner, dist_init): assert not _is_req_scheduled(model_runner, req_ids[1]) +def test_v1_recoverssm_commits_post_sampling_acceptance() -> None: + runner = object.__new__(GPUModelRunner) + runner._use_flashinfer_replayssm = False + runner._needs_prefix_state_migration = False + runner.speculative_config = SimpleNamespace() + runner.model_config = SimpleNamespace(is_hybrid=True) + runner.num_spec_tokens = 3 + runner.num_accepted_tokens = SimpleNamespace(gpu=torch.ones(4, dtype=torch.int32)) + runner.recoverssm = Mock() + runner.input_batch = SimpleNamespace( + num_accepted_tokens_cpu_tensor=torch.zeros(4, dtype=torch.int32) + ) + runner.num_accepted_tokens_event = Mock() + runner.cache_config = SimpleNamespace(mamba_cache_mode="none") + + output_token_ids = torch.tensor( + [[11, 12, -1, -1], [21, 22, 23, -1]], dtype=torch.int64 + ) + runner._update_states_after_model_execute(output_token_ids, Mock()) + + expected = torch.tensor([2, 3], dtype=torch.int32) + torch.testing.assert_close(runner.num_accepted_tokens.gpu[:2], expected) + args, kwargs = runner.recoverssm.commit_step.call_args + torch.testing.assert_close(args[0], expected) + assert args[1] is None + assert kwargs["state_indices"] is None + assert kwargs["num_accepted_tokens"] is runner.num_accepted_tokens.gpu + + def test_update_states_pp_non_async_multi_request_keeps_token_buffers_consistent( model_runner, model_runner_2, dist_init, monkeypatch ): diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 5fa2043789c1..63d255bc6cd4 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -170,25 +170,24 @@ def normalize_live(*_args) -> None: assert state._replayssm_query_start_loc is None -def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None: +def test_recoverssm_commits_accepted_window_without_align_mapping() -> 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, + None, state_indices=None, num_accepted_tokens=num_accepted_tokens, ) state.commit_step( num_sampled, - idx_mapping, + None, state_indices=None, num_accepted_tokens=num_accepted_tokens, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 484e2d616908..5e33fac99fa0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2895,6 +2895,10 @@ def validate_mamba_cached_kernel(self) -> "VllmConfig": raise ValueError( "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" ) + if not self.use_v2_model_runner and self.parallel_config.use_ubatching: + raise ValueError( + "RecoverSSM with Model Runner V1 does not support microbatching" + ) if self.parallel_config.pipeline_parallel_size > 1: raise ValueError( "RecoverSSM currently requires pipeline_parallel_size=1" diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py index cfbe1a4d45d1..b336b52222be 100644 --- a/vllm/v1/worker/gpu/model_states/recoverssm.py +++ b/vllm/v1/worker/gpu/model_states/recoverssm.py @@ -38,7 +38,7 @@ def record_step( def commit_step( self, num_sampled: torch.Tensor | int, - idx_mapping: torch.Tensor, + idx_mapping: torch.Tensor | None, *, state_indices: torch.Tensor | None, num_accepted_tokens: torch.Tensor, @@ -52,6 +52,7 @@ def commit_step( postprocess_meta = metadata.commit_recoverssm_state(num_sampled) if postprocess_meta is None: continue + assert idx_mapping is not None 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. diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 402875b86063..2f57259ae456 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -224,6 +224,7 @@ ) from vllm.v1.worker.dp_utils import coordinate_batch_across_dp from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin @@ -1016,6 +1017,9 @@ def __init__( and self.cache_config.mamba_cache_mode == "all" ) ) + self.recoverssm = ( + RecoverSSMState() if self.cache_config.use_kda_recoverssm else None + ) self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None if ( self.cache_config.mamba_cache_mode == "all" @@ -1608,7 +1612,8 @@ def _update_states_after_model_execute( based on the number of accepted tokens. """ if not ( - self._use_flashinfer_replayssm + self.recoverssm is not None + or self._use_flashinfer_replayssm or (self.speculative_config and self.model_config.is_hybrid) ): return @@ -1622,6 +1627,14 @@ def _update_states_after_model_execute( dim=1 ) + if self.recoverssm is not None: + self.recoverssm.commit_step( + self.num_accepted_tokens.gpu[:num_reqs], + None, + state_indices=None, + num_accepted_tokens=self.num_accepted_tokens.gpu, + ) + if self._needs_prefix_state_migration or self._use_flashinfer_replayssm: # Fused GPU postprocess: state copies + per-request accepted-token # update without CPU-GPU sync. The metadata @@ -2685,6 +2698,14 @@ def _build_attn_group_metadata( spec_decode_common_attn_metadata.unpadded(num_tokens, num_reqs) ) + if self.recoverssm is not None: + assert isinstance(attn_metadata, dict) + self.recoverssm.record_step( + attn_metadata, + self.attn_groups, + for_capture=for_cudagraph_capture, + ) + return attn_metadata, spec_decode_common_attn_metadata def _compute_cascade_attn_prefix_lens( From e4be47bebbaffe682a96720d697130b9f79ab265 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Wed, 9 Sep 2026 02:45:44 +0200 Subject: [PATCH 53/53] [Mamba] Preserve ReplaySSM state for padded prompt tails Stage the attention-consistent transition before speculative acceptance so cached prompt tails retain their accepted replay history in both runners. Repair the V2 test wrapper and cache-spec test double, add regression coverage, and include the lifecycle helper in ReplaySSM E2E dependencies. Assisted-by: OpenAI Codex Signed-off-by: Andrii Skliar --- .buildkite/test_areas/engine.yaml | 1 + tests/kernels/mamba/test_ssu_dispatch.py | 60 +++++++++++++++++++ .../test_attention_backends_selection.py | 2 + .../v1/e2e/general/test_mamba_prefix_cache.py | 6 -- .../worker/test_mamba_hybrid_model_state.py | 49 +++++++++++---- tests/v1/worker/test_mamba_utils.py | 36 +++++++++++ .../layers/mamba/ops/ssu_dispatch.py | 4 +- .../worker/gpu/model_states/mamba_hybrid.py | 16 ++++- vllm/v1/worker/mamba_utils.py | 6 +- 9 files changed, 156 insertions(+), 24 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 536b343e3c16..8d46527a51c9 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -109,6 +109,7 @@ steps: - 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 diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index 23b5bcff62e9..fb0dde95f282 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( FlashInferSSUBackend, TritonSSUBackend, + _postprocess_replayssm_kernel, get_mamba_ssu_backend, initialize_mamba_ssu_backend, selective_state_update, @@ -227,3 +228,62 @@ def test_replayssm_physical_ring_shape( (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) diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index 2695517e1421..63c4605f263c 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/tests/v1/attention/test_attention_backends_selection.py @@ -28,6 +28,8 @@ def test_replayssm_does_not_reserve_speculative_state_blocks( layer = SimpleNamespace( get_state_shape=lambda: ((2,),), get_state_dtype=lambda: (torch.float32,), + get_replayssm_state_shape=lambda: (), + get_replayssm_state_dtype=lambda: (), mamba_type=MambaAttentionBackendEnum.MAMBA2, is_kv_cache_tp_replicated=False, ) diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index d8fc041b43f7..2ed4ec918229 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -1045,8 +1045,6 @@ def wrapped_postprocess_state( idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int, num_computed_tokens: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - is_prefilling: torch.Tensor | None = None, ) -> None: action = cur_step_action block_tables = captured.get("block_tables") @@ -1064,8 +1062,6 @@ def wrapped_postprocess_state( idx_mapping, num_sampled, num_computed_tokens, - query_start_loc, - is_prefilling, ) expected = action.postprocess_copy_idx snapshots = [ @@ -1077,8 +1073,6 @@ def wrapped_postprocess_state( idx_mapping, num_sampled, num_computed_tokens, - query_start_loc, - is_prefilling, ) # Comparing device tensors for the assertion is a deliberate D2H. with gpu_sync_allowed(): diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index f5dc7738e914..df4a06a291c3 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -107,27 +107,51 @@ def test_previous_scheduled_page_is_passed_only_to_mamba2() -> None: assert "prev_last_scheduled_idx" not in gdn_args -def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize( + ("computed", "scheduled", "drafts", "prefilling", "expected_prefilling"), + [ + (256, 4, 3, True, False), # Cached prompt tail plus placeholders. + (1, 4, 3, True, False), # Rejection can exceed the cached prefix length. + (256, 1, 0, True, False), + (0, 4, 3, True, True), # No prior state: must stay a prefill. + (256, 4, 0, True, True), + (256, 4, 3, False, False), + ], +) +def test_prepare_attn_forwards_positions_and_stages_replayssm_prefill( + monkeypatch: pytest.MonkeyPatch, + computed, + scheduled, + drafts, + prefilling, + expected_prefilling, +) -> None: state = object.__new__(MambaHybridModelState) - state.vllm_config = SimpleNamespace(num_speculative_tokens=0) + state.vllm_config = SimpleNamespace(num_speculative_tokens=3) state.max_model_len = 8192 state._align_mode = False - state._use_flashinfer_replayssm = False + state._use_flashinfer_replayssm = True + state._is_prefilling_gpu = torch.zeros(1, dtype=torch.bool) + state.num_accepted_tokens_gpu = torch.ones(1, dtype=torch.int32) + state._get_mamba_group_info = Mock(return_value=([], None)) + state._ensure_mamba_postprocess_ctx = Mock() state._mamba_prev_last_scheduled_idx_gpu = None state.recoverssm = None - positions = torch.tensor([1536], dtype=torch.int64) + positions = torch.arange(computed, computed + scheduled, dtype=torch.int64) input_batch = SimpleNamespace( num_reqs=1, - num_tokens=1, + num_tokens=scheduled, num_reqs_after_padding=1, - num_tokens_after_padding=1, - query_start_loc_np=torch.tensor([0, 1], dtype=torch.int32).numpy(), - query_start_loc=torch.tensor([0, 1], dtype=torch.int32), - num_scheduled_tokens=torch.tensor([1], dtype=torch.int32), - seq_lens_cpu_upper_bound=torch.tensor([1537], dtype=torch.int32), - seq_lens=torch.tensor([1537], dtype=torch.int32), - is_prefilling_np=torch.tensor([False]).numpy(), + num_tokens_after_padding=scheduled, + query_start_loc_np=torch.tensor([0, scheduled], dtype=torch.int32).numpy(), + query_start_loc=torch.tensor([0, scheduled], dtype=torch.int32), + num_scheduled_tokens=torch.tensor([scheduled], dtype=torch.int32).numpy(), + num_draft_tokens_per_req=torch.tensor([drafts], dtype=torch.int32).numpy(), + idx_mapping=torch.tensor([0], dtype=torch.int32), + seq_lens_cpu_upper_bound=torch.tensor([computed + scheduled]), + seq_lens=torch.tensor([computed + scheduled]), + is_prefilling_np=torch.tensor([prefilling]).numpy(), dcp_local_seq_lens=None, positions=positions, prompt_lens=torch.tensor([1024], dtype=torch.int32), @@ -147,6 +171,7 @@ def test_prepare_attn_forwards_positions(monkeypatch: pytest.MonkeyPatch) -> Non assert metadata is expected_metadata assert build_attn_metadata.call_args.kwargs["positions"] is positions + assert state._is_prefilling_gpu.item() == expected_prefilling @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index 4998cc079f1e..e8503f71df79 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -751,6 +751,42 @@ def test_stage_postprocess_inputs_to_gpu_fills_pinned_views(): ctx.replayssm.materialize.assert_not_called() +@pytest.mark.parametrize( + ("computed", "scheduled", "drafts", "prompt_len", "expected_prefilling"), + [ + (256, 4, 3, 257, False), + (1, 4, 3, 2, False), + (256, 1, 0, 257, False), + (0, 4, 3, 1, True), + (256, 4, 0, 260, True), + (256, 4, 3, 256, False), + ], +) +def test_stage_replayssm_prefill_classification( + computed, scheduled, drafts, prompt_len, expected_prefilling +): + ctx = _make_staging_ctx(1, torch.device("cpu")) + stage_postprocess_inputs_to_gpu( + ctx, + _make_postprocess_scheduler_output( + req_ids=["req"], + num_scheduled_tokens={"req": scheduled}, + scheduled_spec_decode_tokens={"req": [1] * drafts}, + ), + ["req"], + 1, + _make_requests( + req_ids=["req"], + num_computed_tokens=[computed], + block_ids_per_req=[[0]], + num_prompt_tokens=[prompt_len], + ), + {}, + run_prefix_state_migration=False, + ) + assert ctx.is_prefilling_buf.gpu.item() == expected_prefilling + + def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx(): """If preprocess_mamba didn't populate mamba_state_idx for a req in the batch, staging must fail loudly rather than silently writing a stale index.""" diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 0fa1077338cd..14415df0b27f 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -143,10 +143,8 @@ def _postprocess_replayssm_kernel( computed_before = tl.where( NUM_COMPUTED_IS_POST_STEP, computed - query_len, computed ) - # Mamba attention runs a one-token final prefill chunk with prior state as - # decode. Commit the same transition here instead of resetting its cursors. + # Staged from forward metadata using Mamba attention's classification. prefilling = tl.load(is_prefilling + batch_idx) - prefilling = prefilling & ((query_len != 1) | (computed_before <= 0)) accepted = tl.maximum(tl.load(num_accepted_tokens + req_idx), 1) # Derive this request's pre/post-step positions from ReplaySSM metadata. diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index 2a9365752eff..5d05fb9eaca1 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -316,8 +316,6 @@ def prepare_attn( input_batch, num_reqs ) - if self._use_flashinfer_replayssm: - self._is_prefilling_gpu[:num_reqs].copy_(is_prefilling, non_blocking=True) # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only # compute them during actual (non-capture) forward execution. @@ -347,6 +345,20 @@ def prepare_attn( if self._use_flashinfer_replayssm: self._replayssm_query_start_loc = input_batch.query_start_loc + # Match Mamba attention using pre-step metadata: after acceptance, + # subtracting the scheduled length also subtracts rejected drafts. + query_lens = torch.diff(query_start_loc_cpu) + decode_rows = query_lens == 1 + if num_decode_draft_tokens_cpu is not None: + decode_rows |= (num_decode_draft_tokens_cpu >= 0) & ( + query_lens == num_decode_draft_tokens_cpu + 1 + ) + replayssm_prefilling = is_prefilling & ~( + (seq_lens_cpu_upper_bound[:num_reqs] > query_lens) & decode_rows + ) + self._is_prefilling_gpu[:num_reqs].copy_( + replayssm_prefilling, non_blocking=True + ) mamba_group_ids, _ = self._get_mamba_group_info(kv_cache_config) self._ensure_mamba_postprocess_ctx( kv_cache_config, mamba_group_ids, block_tables diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 95d59cc1fa30..40586db37378 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -1801,7 +1801,11 @@ def stage_postprocess_inputs_to_gpu( scheduled_np[i] = scheduled computed_np[i] = computed draft_np[i] = num_draft - prefill_np[i] = computed < req_state.num_prompt_tokens + # Match Mamba attention: stateful one-token prompt tails, including + # those padded with speculative placeholders, run the decode kernels. + prefill_np[i] = computed < req_state.num_prompt_tokens and not ( + computed > 0 and (scheduled == 1 or scheduled == num_draft + 1) + ) if run_prefix_state_migration: assert ctx.mamba_state_idx_buf is not None ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs)