From adebc41b7e9f1085d3f73434e23beb76883b9eb4 Mon Sep 17 00:00:00 2001 From: Andrii Skliar Date: Tue, 1 Sep 2026 16:26:43 +0200 Subject: [PATCH] [Mamba] Add FlashInfer ReplaySSM backend (#52506) Signed-off-by: Andrii Skliar Signed-off-by: Andrii Skliar Co-authored-by: Andrii Skliar --- tests/kernels/mamba/test_ssu_dispatch.py | 72 +++++- tests/model_executor/test_replayssm_warmup.py | 150 +++++++++++ .../test_replayssm_metadata_builder.py | 58 ++++- tests/v1/e2e/test_replayssm_decode.py | 39 ++- .../worker/test_kv_cache_allocation_scope.py | 4 +- tests/v1/worker/test_utils.py | 57 +++++ vllm/config/cache.py | 11 +- vllm/config/vllm.py | 19 +- .../layers/mamba/mamba_mixer2.py | 185 +++++++++++--- .../layers/mamba/mamba_utils.py | 22 +- .../layers/mamba/ops/ssu_dispatch.py | 235 ++++++++++++++++-- vllm/model_executor/models/nemotron_h.py | 12 +- vllm/model_executor/warmup/kernel_warmup.py | 4 + .../model_executor/warmup/replayssm_warmup.py | 159 ++++++++++++ vllm/v1/attention/backends/mamba_attn.py | 66 ++++- vllm/v1/worker/gpu/attn_utils.py | 8 +- vllm/v1/worker/gpu/model_runner.py | 22 +- vllm/v1/worker/gpu_model_runner.py | 5 +- vllm/v1/worker/utils.py | 6 + 19 files changed, 1029 insertions(+), 105 deletions(-) create mode 100644 tests/model_executor/test_replayssm_warmup.py create mode 100644 vllm/model_executor/warmup/replayssm_warmup.py diff --git a/tests/kernels/mamba/test_ssu_dispatch.py b/tests/kernels/mamba/test_ssu_dispatch.py index c37f74eb50..6226340a99 100644 --- a/tests/kernels/mamba/test_ssu_dispatch.py +++ b/tests/kernels/mamba/test_ssu_dispatch.py @@ -7,12 +7,15 @@ import pytest import torch 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, TritonSSUBackend, get_mamba_ssu_backend, initialize_mamba_ssu_backend, + reset_replayssm_ring_trackers, selective_state_update, + update_replayssm_ring_trackers, ) from vllm.utils.torch_utils import set_random_seed from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum @@ -30,6 +33,47 @@ except ImportError: HAS_FLASHINFER = False +@pytest.fixture(autouse=True) +def restore_backend_state(): + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod + + old_backend = mod._mamba_ssu_backend + old_replayssm_kernel = mod._flashinfer_replayssm_kernel + yield + mod._mamba_ssu_backend = old_backend + mod._flashinfer_replayssm_kernel = old_replayssm_kernel + + +def test_flashinfer_replayssm_ring_tracker_lifecycle(): + ring_start = torch.zeros(2, dtype=torch.int32, device="cuda") + prev_num_accepted = torch.zeros(2, dtype=torch.int32, device="cuda") + state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda") + + observed = [] + for _ in range(33): + update_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + state_batch_indices, + logical_window=16, + ring_buffer_len=17, + ) + observed.append((int(ring_start[1]), int(prev_num_accepted[1]))) + + assert observed[4] == (0, 5) + assert observed[15] == (0, 16) + assert observed[16] == (16, 1) + assert observed[31] == (16, 16) + assert observed[32] == (15, 1) + + reset_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + state_batch_indices, + ) + assert (ring_start[1].item(), prev_num_accepted[1].item()) == (0, 0) + + def _kv_cache_config_with_ssu( mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2, ) -> KVCacheConfig: @@ -116,11 +160,10 @@ def test_flashinfer_forwards_ssu_algorithm( def test_uninitialized_backend_raises(): import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod - old = mod._mamba_ssu_backend + # restore_backend_state (autouse) puts the global back afterwards. mod._mamba_ssu_backend = None with pytest.raises(RuntimeError, match="not been initialized"): get_mamba_ssu_backend() - mod._mamba_ssu_backend = old @pytest.mark.parametrize( @@ -186,3 +229,28 @@ def test_triton_basic_call(): out=out, ) assert not torch.isnan(out).any() + + +@pytest.mark.parametrize( + ("backend", "expected_ring_len"), + [ + (MambaBackendEnum.TRITON, 16), + (MambaBackendEnum.FLASHINFER, 17), + ], +) +def test_replayssm_physical_ring_shape(backend, expected_ring_len): + base_shapes = ((64, 3), (8, 4, 16)) + + shapes = MambaStateShapeCalculator.append_replayssm_ring( + base_shapes, + n_groups=4, + tp_world_size=2, + logical_window=16, + backend=backend, + ) + + assert shapes[2:] == ( + (8, expected_ring_len, 4), + (8, expected_ring_len), + (2, expected_ring_len, 16), + ) diff --git a/tests/model_executor/test_replayssm_warmup.py b/tests/model_executor/test_replayssm_warmup.py new file mode 100644 index 0000000000..272fa2cf85 --- /dev/null +++ b/tests/model_executor/test_replayssm_warmup.py @@ -0,0 +1,150 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np +import pytest +import torch + +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, + "skip_eplb": True, + "is_profile": True, + "randomize_inputs": True, +} + + +def _autotune_runner( + *, + use_v2_model_runner: bool = False, + query_len: int = 6, + max_num_seqs: int = 32, + num_blocks: int = 17, + max_num_tokens: int = 100, + use_replayssm: bool = True, + backend: MambaBackendEnum = MambaBackendEnum.FLASHINFER, +) -> SimpleNamespace: + return SimpleNamespace( + vllm_config=SimpleNamespace( + cache_config=SimpleNamespace(use_replayssm=use_replayssm), + mamba_config=SimpleNamespace(backend=backend), + use_v2_model_runner=use_v2_model_runner, + ), + uniform_decode_query_len=query_len, + decode_query_len=query_len, + max_num_tokens=max_num_tokens, + scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs), + kv_cache_config=SimpleNamespace(num_blocks=num_blocks), + ) + + +@pytest.mark.parametrize( + ("runner_kwargs", "expected_num_reqs"), + [ + # max_num_seqs (32) vs max_num_tokens // query_len (16) vs blocks-1 (16). + (dict(query_len=6, use_v2_model_runner=False), 16), + (dict(query_len=6, use_v2_model_runner=True), 16), + # num_blocks - 1 is the binding constraint. + (dict(query_len=1, max_num_tokens=128, max_num_seqs=64, num_blocks=5), 4), + ], + ids=["v1", "v2", "clamped_to_state_capacity"], +) +def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs): + query_len = runner_kwargs["query_len"] + with patch.object( + warmup, "flashinfer_replayssm_autotune_supported", return_value=True + ): + result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs)) + + expected_kwargs = { + **PREFILL_KWARGS, + "num_tokens": expected_num_reqs * query_len, + "uniform_decode": True, + } + if runner_kwargs.get("use_v2_model_runner"): + expected_kwargs["valid_dummy_state_slots"] = True + else: + expected_kwargs.update( + allow_microbatching=False, + force_attention=True, + profile_seq_lens=query_len + 1, + ) + assert result == (expected_num_reqs, expected_kwargs) + + +@pytest.mark.parametrize( + ("runner_kwargs", "flashinfer_supported"), + [ + (dict(use_replayssm=False), True), + (dict(backend=MambaBackendEnum.TRITON), True), + ({}, False), + ], + ids=["replayssm_disabled", "non_flashinfer_backend", "kernel_unavailable"], +) +def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported): + with patch.object( + warmup, + "flashinfer_replayssm_autotune_supported", + return_value=flashinfer_supported, + ): + result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs)) + assert result is None + + +def test_replayssm_autotune_slots_restore_state_and_trackers(): + mixer = MambaMixer2.__new__(MambaMixer2) + torch.nn.Module.__init__(mixer) + mixer.use_replayssm = True + mixer.replayssm_buffer_len = 16 + 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_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_ring_start, + mixer._replayssm_prev_num_accepted, + ) + + block_ids = np.arange(10, 14, dtype=np.int32).reshape(4, 1) + original_block_ids = block_ids.copy() + block_table = SimpleNamespace(block_table=SimpleNamespace(np=block_ids)) + multi_group_block_table = SimpleNamespace( + block_tables=[block_table], commit_block_table=Mock() + ) + runner = SimpleNamespace( + vllm_config=SimpleNamespace(use_v2_model_runner=False), + input_batch=SimpleNamespace(block_table=multi_group_block_table), + get_model=lambda: SimpleNamespace(modules=lambda: (mixer,)), + ) + + with warmup._temporary_replayssm_autotune_state(runner, 2): + assert block_ids[:2, 0].tolist() == [1, 2] + for tensor in tracked: + tensor[1:3].fill_(9) + + assert np.array_equal(block_ids, original_block_ids) + assert multi_group_block_table.commit_block_table.call_args_list == [ + ((2,), {}), + ((2,), {}), + ] + for tensor in tracked: + assert torch.count_nonzero(tensor[1:3]) == 0 + assert torch.all(tensor[0] == 3) + assert torch.all(tensor[3] == 3) diff --git a/tests/v1/attention/test_replayssm_metadata_builder.py b/tests/v1/attention/test_replayssm_metadata_builder.py index cbbfedee96..89cb0f5664 100644 --- a/tests/v1/attention/test_replayssm_metadata_builder.py +++ b/tests/v1/attention/test_replayssm_metadata_builder.py @@ -16,6 +16,7 @@ from tests.v1.attention.utils import ( create_common_attn_metadata, create_vllm_config, ) +from vllm.config.mamba import MambaBackendEnum from vllm.v1.kv_cache_interface import MambaSpec BLOCK_SIZE = 16 @@ -187,34 +188,47 @@ REPLAYSSM_BUILD_CASES = { } -def _make_mamba_spec(buffer_len: int) -> MambaSpec: - # Five-tensor ReplaySSM page; the builder only reads shapes[4][0] (bc groups). +def _make_mamba_spec( + buffer_len: int, + mamba_backend: MambaBackendEnum, +) -> MambaSpec: + ring_buffer_len = buffer_len + ( + 1 if mamba_backend == MambaBackendEnum.FLASHINFER else 0 + ) + shapes = ( + (1, 1), + (1, 1, 1), + (1, ring_buffer_len, 1), + (1, ring_buffer_len), + (1, ring_buffer_len, 1), + ) return MambaSpec( block_size=BLOCK_SIZE, - shapes=( - (1, 1), - (1, 1, 1), - (1, buffer_len, 1), - (1, buffer_len), - (1, buffer_len, 1), - ), + shapes=shapes, dtypes=(torch.float32,), ) def _create_replayssm_builder( - buffer_len: int, mamba_cache_mode: str = "none" + buffer_len: int, + mamba_cache_mode: str = "none", + *, + mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON, ) -> MockMambaBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE ) # Set the flags after construction to skip validate_mamba_cached_kernel - # (it requires a Triton backend) on the mock model. + # (it requires a real SupportsReplaySSM model) on the mock model. vllm_config.cache_config.use_replayssm = True 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 return MockMambaBuilder( - _make_mamba_spec(buffer_len), ["layer0"], vllm_config, DEVICE + _make_mamba_spec(buffer_len, mamba_backend), + ["layer0"], + vllm_config, + DEVICE, ) @@ -254,3 +268,23 @@ def test_resumed_request_differs_from_fresh(): assert meta.write_pos_d.tolist()[:2] == [5, 0] 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") + + 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/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 4fc6768e74..2e66a46689 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -4,6 +4,7 @@ import pytest +import vllm.envs as envs from vllm.v1.metrics.reader import Counter from ...models.utils import check_logprobs_close @@ -21,29 +22,48 @@ PROMPTS = [ ] -def _check_replayssm_parity(vllm_runner, model_name, *, tensor_parallel_size=1): +def _check_replayssm_parity( + vllm_runner, + model_name, + *, + tensor_parallel_size=1, + 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, 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="replayssm", + name_1=name_1, ) @@ -60,6 +80,21 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name): _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2) +@pytest.mark.parametrize("model_name", MODELS) +def test_replayssm_flashinfer_decode_matches_baseline_v2( + vllm_runner, model_name, monkeypatch +): + pytest.importorskip("flashinfer.mamba.checkpointing_ssu") + _check_replayssm_parity( + vllm_runner, + model_name, + mamba_backend="flashinfer", + name_1="replayssm_flashinfer_v2", + require_v2=True, + monkeypatch=monkeypatch, + ) + + # Prefix spans several mamba blocks; prefix caching only reuses full blocks. _PC_SENTENCE = ( "In a detailed survey of state space models, the authors compared many " diff --git a/tests/v1/worker/test_kv_cache_allocation_scope.py b/tests/v1/worker/test_kv_cache_allocation_scope.py index 1818f173d9..a9579f1af6 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 test_mrv2_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: result = attn_utils.init_kv_cache( [], {}, - object(), + SimpleNamespace(kv_cache_groups=[]), torch.device("cpu"), [], config, @@ -83,7 +83,7 @@ def test_mrv1_kv_pool_only_wraps_backing_allocation(monkeypatch) -> None: ) result = gpu_model_runner.GPUModelRunner.initialize_kv_cache_tensors( runner, - object(), + SimpleNamespace(kv_cache_groups=[]), [], kv_cache_allocation_context=scope, ) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 016b8aa5a6..a4474c8a58 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1,11 +1,68 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + 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 +class _TestReplaySSMMixer(MambaMixer2): + 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._updates_replayssm_trackers = True + + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: + return ((2,), (3,), (4,), (5,), (6,)) + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: + return (torch.float32,) * 5 + + +def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: + return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8) + + +def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(): + mixers = [_TestReplaySSMMixer() for _ in range(3)] + layer_names = [f"layers.{i}.mixer" for i in range(3)] + ctx = dict(zip(layer_names, mixers)) + # Reverse insertion order: updater must follow layer index, not dict order. + kv_cache = { + layer_names[2]: _packed_replayssm_cache(4), + layer_names[1]: _packed_replayssm_cache(4), + layer_names[0]: _packed_replayssm_cache(4), + } + kv_cache_groups = [ + SimpleNamespace(layer_names=[layer_names[0], layer_names[2]]), + SimpleNamespace(layer_names=[layer_names[1]]), + ] + + 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() + ) + # Group {0, 2} shares trackers; layer 2 (not 0) updates after both run. + assert [m._updates_replayssm_trackers for m in mixers] == [False, True, 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 3a853000c3..0ffd27c281 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -196,15 +196,16 @@ class CacheConfig: caching is enabled. """ replayssm_buffer_len: int = Field(default=16, gt=0) - """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3 - speculative decoding does not use B. Default 16.""" + """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.""" 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 - 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; standard (non-speculative) decode only. In align + mode flushes are most efficient when mamba_block_size is a multiple of + replayssm_buffer_len, but this is not required.""" use_kda_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 3fcbb83685..7fb53cdc24 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -2889,13 +2889,28 @@ class VllmConfig: raise ValueError( "RecoverSSM currently requires pipeline_parallel_size=1" ) + if self.mamba_config.backend != MambaBackendEnum.TRITON: + raise ValueError("RecoverSSM requires --mamba-backend triton") elif self.cache_config.mamba_cache_mode == "all": raise ValueError( "--use-replayssm supports prefix caching only in align mode; " "pass --mamba-cache-mode align" ) - if self.mamba_config.backend != MambaBackendEnum.TRITON: - raise ValueError("--use-replayssm requires --mamba-backend triton") + 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: + raise ValueError( + "--use-replayssm requires --mamba-backend triton or flashinfer" + ) + elif 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 513402a0b1..0de22ead83 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -2,10 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence + import torch from torch import nn from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config +from vllm.config.mamba import MambaBackendEnum from vllm.distributed import ( divide, get_tensor_model_parallel_rank, @@ -16,6 +19,7 @@ from vllm.distributed import ( from vllm.forward_context import ForwardContext, get_forward_context from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp, PluggableLayer +from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, @@ -38,7 +42,11 @@ from vllm.model_executor.layers.mamba.ops.selective_state_update_replayssm_outpu from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) -from vllm.model_executor.layers.mamba.ops.ssu_dispatch import selective_state_update +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + reset_replayssm_ring_trackers, + selective_state_update, + selective_state_update_replayssm_flashinfer, +) from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import ( LoaderFunction, @@ -58,6 +66,7 @@ from vllm.utils.torch_utils import ( from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import KVCacheGroupSpec logger = init_logger(__name__) @@ -515,10 +524,12 @@ class MambaMixer2(MambaBase, PluggableLayer): raise ValueError( "--use-replayssm requires tensor-parallel heads to divide evenly" ) - # The tuple is (conv_state, ssm_state); with the cached (ReplaySSM) decode - # kernel enabled it is (conv_state, ssm_state, x_cache, dt_cache, B_cache). + # 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._replayssm_ring_start = torch.empty(0, dtype=torch.int32) + self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) + self._updates_replayssm_trackers = True self.num_spec = vllm_config.num_speculative_tokens if self.num_spec > 0: @@ -706,6 +717,7 @@ class MambaMixer2(MambaBase, PluggableLayer): assert self.cache_config is not None mamba_block_size = self.cache_config.mamba_block_size is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all" + ring_start = prev_num_accepted = None attn_metadata: AttentionMetadata | None = None if attn_metadata_raw is not None: @@ -723,7 +735,10 @@ class MambaMixer2(MambaBase, PluggableLayer): ) ssm_state = self.kv_cache[1] if self.use_replayssm: - x_cache, dt_cache, B_cache = self.kv_cache[2:] + x_cache, dt_cache, B_cache = self.kv_cache[2:5] + if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: + ring_start = self._replayssm_ring_start + prev_num_accepted = self._replayssm_prev_num_accepted else: x_cache = dt_cache = B_cache = None has_initial_states_p = attn_metadata.has_initial_states_p @@ -982,6 +997,13 @@ class MambaMixer2(MambaBase, PluggableLayer): # 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, + state_indices_tensor_p, + ) # Process decode requests if has_decode: @@ -1060,35 +1082,64 @@ class MambaMixer2(MambaBase, PluggableLayer): ) if self.use_replayssm: assert self.replayssm_buffer_len is not None - selective_state_update_replayssm_output_only( - ssm_state, - hidden_states_d, - dt_d, - A_d, - B_d, - C_d, - D_d, - dt_bias, - dt_softplus=True, - x_cache=x_cache, - dt_cache=dt_cache, - B_cache=B_cache, - bc_pre=attn_metadata.bc_pre_scratch, - write_pos=attn_metadata.write_pos_d, - is_flush=attn_metadata.is_flush_d, - max_cache_len=self.replayssm_buffer_len, - state_batch_indices=state_indices_tensor_d_input, - out=preallocated_ssm_out_d, - # Stochastic Rounding for the vanilla decode path is read - # from mamba_config inside ssu_dispatch; the replay kernel - # isn't a dispatch backend, so pass it here. - enable_stochastic_rounding=( - self.mamba_config.enable_stochastic_rounding - ), - cache_philox_rounds=( - self.mamba_config.stochastic_rounding_philox_rounds - ), - ) + if self.mamba_config.backend == MambaBackendEnum.FLASHINFER: + assert ring_start is not None + assert prev_num_accepted is not None + assert attn_metadata.replayssm_scratch is not None + selective_state_update_replayssm_flashinfer( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + preallocated_ssm_out_d, + x_cache, + B_cache, + dt_cache, + ring_start, + prev_num_accepted, + logical_window=self.replayssm_buffer_len, + D=D_d, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_indices_tensor_d_input, + scratch=attn_metadata.replayssm_scratch, + update_trackers=self._updates_replayssm_trackers, + enable_stochastic_rounding=( + self.mamba_config.enable_stochastic_rounding + ), + stochastic_rounding_philox_rounds=( + self.mamba_config.stochastic_rounding_philox_rounds + ), + ) + else: + selective_state_update_replayssm_output_only( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + dt_bias, + dt_softplus=True, + x_cache=x_cache, + dt_cache=dt_cache, + B_cache=B_cache, + bc_pre=attn_metadata.bc_pre_scratch, + write_pos=attn_metadata.write_pos_d, + is_flush=attn_metadata.is_flush_d, + max_cache_len=self.replayssm_buffer_len, + state_batch_indices=state_indices_tensor_d_input, + out=preallocated_ssm_out_d, + enable_stochastic_rounding=( + self.mamba_config.enable_stochastic_rounding + ), + cache_philox_rounds=( + self.mamba_config.stochastic_rounding_philox_rounds + ), + ) else: selective_state_update( ssm_state, @@ -1137,10 +1188,11 @@ class MambaMixer2(MambaBase, PluggableLayer): if self.use_replayssm: assert self.replayssm_buffer_len is not None return MambaStateShapeCalculator.append_replayssm_ring( - base_shape, - self.n_groups, - tp_world_size, - self.replayssm_buffer_len, + 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, ) return base_shape @@ -1149,6 +1201,65 @@ class MambaMixer2(MambaBase, PluggableLayer): return MambaAttentionBackendEnum.MAMBA2 +def share_replayssm_ring_trackers( + ordered_layer_names: list[str], + forward_context: dict[str, Attention], + kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, +) -> None: + """Share ring cursors within each cache-slot index namespace. + + 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. + """ + + replayssm_mixers: dict[str, MambaMixer2] = {} + for layer_name in ordered_layer_names: + layer = forward_context[layer_name] + if ( + isinstance(layer, MambaMixer2) + and layer.use_replayssm + and layer.mamba_config.backend == MambaBackendEnum.FLASHINFER + ): + replayssm_mixers[layer_name] = layer + + layer_to_group: dict[str, int] = {} + if kv_cache_groups: + for group_idx, group in enumerate(kv_cache_groups): + for layer_name in group.layer_names: + layer_to_group[layer_name] = group_idx + + groups_by_namespace: dict[int | str, list[str]] = {} + for layer_name in ordered_layer_names: + if layer_name not in replayssm_mixers: + continue + namespace = layer_to_group.get(layer_name, layer_name) + groups_by_namespace.setdefault(namespace, []).append(layer_name) + + for group_layer_names in groups_by_namespace.values(): + 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 + for layer_name in group_layer_names: + state = replayssm_mixers[layer_name].kv_cache[1] + if (state.shape[0], state.device) != (num_blocks, device): + raise ValueError( + "ReplaySSM layers in one cache group must share cache capacity" + ) + + ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device) + prev_num_accepted = 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 + + def mamba_mixer2( projected_states: torch.Tensor, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index 99059c6b54..73b1d934e8 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -10,6 +10,7 @@ import torch import vllm.envs as envs from vllm.config.cache import MambaDType +from vllm.config.mamba import MambaBackendEnum from vllm.config.model import ModelDType from vllm.distributed import divide from vllm.logger import init_logger @@ -214,20 +215,25 @@ class MambaStateShapeCalculator: base_shapes: tuple[tuple[int, ...], ...], n_groups: int, tp_world_size: int, - replayssm_buffer_len: int, + logical_window: int, + backend: MambaBackendEnum, ) -> tuple[tuple[int, ...], ...]: - """Append the ReplaySSM ring shapes (x_cache, dt_cache, B_cache) to a - base ``(conv, ssm)`` tuple. ``base_shapes[1]`` is the ssm shape - ``(nheads // tp, head_dim, state_size)``; B_cache uses the un-extended - ``n_groups``. + """Append the physical ReplaySSM ring shapes. + + ``base_shapes[1]`` is ``(nheads // tp, head_dim, state_size)``; + B_cache uses the un-extended ``n_groups``. """ + ring_buffer_len = logical_window + if backend == MambaBackendEnum.FLASHINFER: + # FlashInfer keeps the live window and appended token together. + ring_buffer_len += 1 local_nheads, head_dim, state_size = base_shapes[1] local_ngroups = divide(n_groups, tp_world_size) return ( *base_shapes, - (local_nheads, replayssm_buffer_len, head_dim), - (local_nheads, replayssm_buffer_len), - (local_ngroups, replayssm_buffer_len, state_size), + (local_nheads, ring_buffer_len, head_dim), + (local_nheads, ring_buffer_len), + (local_ngroups, ring_buffer_len, state_size), ) @classmethod diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 24bdba688e..62e4307b37 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -10,11 +10,14 @@ the backend defaults to 'cpu'. """ from abc import ABC, abstractmethod +from collections.abc import Callable +from functools import cache import torch from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm from vllm.logger import init_logger +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 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec @@ -22,6 +25,99 @@ from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec 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, + 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, +) -> 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) + if RESET: + tl.store(ring_start + slots, 0, mask=valid) + tl.store(prev_num_accepted + 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) + must_checkpoint = prev + 1 > logical_window + next_start = tl.where( + must_checkpoint, + (start + prev) % ring_buffer_len, + start, + ) + 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) + + +def update_replayssm_ring_trackers( + ring_start: torch.Tensor, + prev_num_accepted: 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, + state_batch_indices, + state_batch_indices.stride(0), + n_slots, + min(ring_start.numel(), prev_num_accepted.numel()), + logical_window, + ring_buffer_len, + pad_slot_id, + RESET=reset, + BLOCK=block, + ) + + +def reset_replayssm_ring_trackers( + ring_start: torch.Tensor, + prev_num_accepted: 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, + state_batch_indices, + pad_slot_id=pad_slot_id, + ) + + class MambaSSUBackend(ABC): """Abstract base class for Mamba SSU backends.""" @@ -262,15 +358,115 @@ _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { _mamba_ssu_backend: MambaSSUBackend | None = None +_flashinfer_replayssm_kernel: Callable[..., torch.Tensor] | None = None + + +@cache +def flashinfer_replayssm_autotune_supported() -> bool: + """Return True when FlashInfer exposes ReplaySSM autotuning.""" + try: + from flashinfer.mamba.checkpointing_ssu import ( # noqa: F401 + CheckpointingSSURunner, + ) + except ImportError: + return False + return True + + +def selective_state_update_replayssm_flashinfer( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor, + x_cache: torch.Tensor, + B_cache: torch.Tensor, + dt_cache: torch.Tensor, + ring_start: torch.Tensor, + prev_num_accepted_tokens: 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, +) -> torch.Tensor: + """Run FlashInfer checkpointing SSU and optionally advance shared trackers.""" + if _flashinfer_replayssm_kernel is None: + raise RuntimeError( + "FlashInfer ReplaySSM has not been initialized. " + "Call initialize_mamba_ssu_backend() with use_replayssm=True." + ) + + if x.dim() == 3: + x = x.unsqueeze(1) + dt = dt.unsqueeze(1) + B = B.unsqueeze(1) + C = C.unsqueeze(1) + out = out.unsqueeze(1) + + indices = state_batch_indices + if indices is not None and indices.dim() > 1: + indices = indices[:, 0] + + cb_scaled = cumAdt_vec = cb_old = None + if scratch is not None: + cb_scaled, cumAdt_vec, cb_old = scratch + + rand_seed = ( + torch.randint(0, 2**32, (1,), device=state.device, dtype=torch.int64) + if enable_stochastic_rounding + else None + ) + result = _flashinfer_replayssm_kernel( + state, + x_cache, + B_cache, + dt_cache, + ring_start, + prev_num_accepted_tokens, + x, + dt, + A, + B, + C, + out, + D=D, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=indices, + pad_slot_id=null_block_id, + rand_seed=rand_seed, + philox_rounds=stochastic_rounding_philox_rounds or 10, + cb_scaled=cb_scaled, + 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, + indices, + logical_window=logical_window, + ring_buffer_len=x_cache.size(2), + pad_slot_id=null_block_id, + ) + return result + + def initialize_mamba_ssu_backend( mamba_config: MambaConfig, kv_cache_config: KVCacheConfig, + *, + use_replayssm: bool = False, ) -> None: - """Initialize the global Mamba SSU backend. - - No-op if `kv_cache_config` contains no specs that call - selective_state_update. - """ + """Initialize the Mamba SSU backend and optional FlashInfer ReplaySSM.""" if not any( isinstance(g.kv_cache_spec, MambaSpec) and g.kv_cache_spec.mamba_type @@ -279,13 +475,9 @@ def initialize_mamba_ssu_backend( ): return - global _mamba_ssu_backend - + global _flashinfer_replayssm_kernel, _mamba_ssu_backend backend = mamba_config.backend - # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is - # unstable or unavailable. Silently fall back to the CPU - # backend unless the user explicitly chose something other than "triton". if backend == MambaBackendEnum.TRITON: from vllm.platforms import current_platform @@ -301,13 +493,28 @@ def initialize_mamba_ssu_backend( f"Unknown Mamba SSU backend: {backend}. " f"Valid options: {list(_BACKEND_REGISTRY.keys())}" ) + if use_replayssm and backend not in ( + MambaBackendEnum.TRITON, + MambaBackendEnum.FLASHINFER, + ): + raise ValueError(f"ReplaySSM does not support mamba backend {backend.value!r}") backend_cls = _BACKEND_REGISTRY[backend] - if isinstance(_mamba_ssu_backend, backend_cls): - return + if not isinstance(_mamba_ssu_backend, backend_cls): + _mamba_ssu_backend = backend_cls(mamba_config) + logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name) - _mamba_ssu_backend = backend_cls(mamba_config) - logger.info("Using %s Mamba SSU backend.", _mamba_ssu_backend.name) + _flashinfer_replayssm_kernel = None + if use_replayssm and backend == MambaBackendEnum.FLASHINFER: + try: + from flashinfer.mamba.checkpointing_ssu import checkpointing_ssu + except ImportError as e: + raise ImportError( + "FlashInfer ReplaySSM requires a compatible flashinfer-python package" + ) from e + _flashinfer_replayssm_kernel = checkpointing_ssu + if use_replayssm: + logger.info("Using %s ReplaySSM backend.", backend.value) def get_mamba_ssu_backend() -> MambaSSUBackend: diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index b62e46d44a..e65e51f3bb 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -747,7 +747,8 @@ class NemotronHForCausalLM( ) if cache_config.use_replayssm: return MambaStateDtypeCalculator.append_replayssm_ring( - base_dtype, vllm_config.model_config.dtype + base_dtype, + vllm_config.model_config.dtype, ) return base_dtype @@ -784,10 +785,11 @@ class NemotronHForCausalLM( ) if cache_config.use_replayssm: return MambaStateShapeCalculator.append_replayssm_ring( - base_shape, - hf_config.n_groups, - parallel_config.tensor_parallel_size, - cache_config.replayssm_buffer_len, + 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, ) return base_shape diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b406ec3873..ca816cdb44 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -38,6 +38,9 @@ from vllm.model_executor.warmup.kimi_k3_triton_warmup import ( kimi_k3_triton_warmup, ) from vllm.model_executor.warmup.qwen_triton_warmup import qwen_triton_warmup +from vllm.model_executor.warmup.replayssm_warmup import ( + replayssm_autotune_warmup, +) from vllm.platforms import current_platform from vllm.utils.deep_gemm import is_deep_gemm_supported from vllm.utils.flashinfer import has_flashinfer @@ -335,6 +338,7 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None: fi_utils.autotune(tune_mode=True, **autotune_kwargs), ): _run_flashinfer_autotune_dummy_runs(runner) + replayssm_autotune_warmup(runner) finally: set_autotune_process_group(None) diff --git a/vllm/model_executor/warmup/replayssm_warmup.py b/vllm/model_executor/warmup/replayssm_warmup.py new file mode 100644 index 0000000000..60b3eced20 --- /dev/null +++ b/vllm/model_executor/warmup/replayssm_warmup.py @@ -0,0 +1,159 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any + +import torch + +from vllm.config.mamba import MambaBackendEnum +from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( + flashinfer_replayssm_autotune_supported, +) + +if TYPE_CHECKING: + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +logger = init_logger(__name__) + + +def _replayssm_autotune_kwargs( + runner: "GPUModelRunner", +) -> tuple[int, dict[str, Any]] | None: + config = runner.vllm_config + if not ( + config.cache_config.use_replayssm + and config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ): + return None + if not flashinfer_replayssm_autotune_supported(): + logger.info_once( + "Skipping FlashInfer ReplaySSM autotuning because " + "flashinfer.mamba.checkpointing_ssu.CheckpointingSSURunner " + "is unavailable." + ) + return None + v2_runner: Any = runner + query_len = ( + v2_runner.decode_query_len + if config.use_v2_model_runner + else runner.uniform_decode_query_len + ) + max_num_reqs = min( + runner.scheduler_config.max_num_seqs, + runner.max_num_tokens // query_len, + runner.kv_cache_config.num_blocks - 1, + ) + decode_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: + decode_kwargs["valid_dummy_state_slots"] = True + else: + decode_kwargs.update( + allow_microbatching=False, + force_attention=True, + profile_seq_lens=query_len + 1, + ) + return max_num_reqs, decode_kwargs + + +@contextmanager +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, 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 + tracker_specs.setdefault( + ring_start.data_ptr(), + ( + ring_start, + prev_num_accepted, + module.replayssm_buffer_len, + module.kv_cache[2].size(2), + ), + ) + tensors = ( + *module.kv_cache, + ring_start, + prev_num_accepted, + ) + for tensor in tensors: + if tensor.numel(): + reset_tensors.setdefault(tensor.data_ptr(), tensor) + + v2_runner: Any = runner + block_tables = saved_block_ids = None + if not runner.vllm_config.use_v2_model_runner: + block_tables = runner.input_batch.block_table.block_tables + saved_block_ids = tuple( + block_table.block_table.np[:max_num_reqs, 0].copy() + for block_table in block_tables + ) + dummy_block_ids = range(1, max_num_reqs + 1) + for block_table in block_tables: + 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, + 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) + update_replayssm_ring_trackers( + ring_start, + prev_num_accepted, + state_slots, + logical_window, + ring_buffer_len, + ) + reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots) + + try: + yield + finally: + if runner.vllm_config.use_v2_model_runner: + v2_runner.block_tables.get_dummy_block_tables(max_num_reqs) + else: + assert block_tables is not None and saved_block_ids is not None + for block_table, block_ids in zip(block_tables, saved_block_ids): + block_table.block_table.np[:max_num_reqs, 0] = block_ids + runner.input_batch.block_table.commit_block_table(max_num_reqs) + for tensor in reset_tensors.values(): + tensor[1 : max_num_reqs + 1].zero_() + + +def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None: + autotune = _replayssm_autotune_kwargs(runner) + if autotune is None: + return + max_num_reqs, decode_kwargs = autotune + with _temporary_replayssm_autotune_state(runner, max_num_reqs): + runner._dummy_run(**decode_kwargs) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 6b7ac71ae8..b1a0c2d0ef 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -8,6 +8,7 @@ from typing import Any, ClassVar, TypeVar import torch from vllm.config import VllmConfig +from vllm.config.mamba import MambaBackendEnum from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( @@ -74,12 +75,15 @@ class BaseMambaAttentionMetadata: nums_dict: dict | None = None batch_ptr: torch.Tensor | None = None token_chunk_offset_ptr: torch.Tensor | None = None - # ReplaySSM standard decode: per-row ring cursor and flush flag, plus the - # per-step (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for the - # precomputed k^T q products. All None when use_replayssm is disabled. + # ReplaySSM standard decode — Triton: per-row ring cursor and flush flag, + # plus (decode_rows, ngroups, replayssm_buffer_len) fp32 scratch for + # precomputed B·C products. All None when use_replayssm is disabled or + # the FlashInfer ReplaySSM backend is selected. write_pos_d: torch.Tensor | None = None is_flush_d: torch.Tensor | None = None 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 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): @@ -107,6 +111,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): self.use_spec_decode = self.num_spec_tokens > 0 self.use_replayssm = vllm_config.cache_config.use_replayssm self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len + self.use_flashinfer_replayssm = ( + self.use_replayssm + and vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + ) scheduler_config = vllm_config.scheduler_config self.decode_cudagraph_max_bs: int = scheduler_config.max_num_seqs @@ -166,9 +174,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): dtype=torch.int32, device=device, ) - # ReplaySSM standard-decode CUDA-graph buffers: per-row ring cursor, - # flush flag, and the k^T q precompute scratch. - if self.use_replayssm: + self.decode_bc_pre_scratch: torch.Tensor | None = None + self.decode_replayssm_scratch: ( + tuple[torch.Tensor, torch.Tensor, 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( (self.decode_cudagraph_max_bs,), dtype=torch.int32, @@ -185,7 +196,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): bc_scratch_bs = max( self.decode_cudagraph_max_bs, scheduler_config.max_num_seqs ) - self.decode_bc_pre_scratch: torch.Tensor = torch.empty( + self.decode_bc_pre_scratch = torch.empty( ( bc_scratch_bs, bc_ngroups, @@ -194,8 +205,20 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): dtype=torch.float32, device=device, ) - else: - self.decode_bc_pre_scratch = None + elif self.use_flashinfer_replayssm: + from flashinfer.mamba.checkpointing_ssu import ( + allocate_checkpointing_ssu_scratch, + ) + + nheads = kv_cache_spec.shapes[2][0] + self.decode_replayssm_scratch = allocate_checkpointing_ssu_scratch( + batch_size=scheduler_config.max_num_seqs, + num_heads=nheads, + num_predicted_tokens=1, + max_window=self.replayssm_buffer_len, + dtype=vllm_config.model_config.dtype, + device=device, + ) self._init_reorder_batch_threshold(1, self.use_spec_decode) if self.use_spec_decode: @@ -500,6 +523,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None write_pos_d = None is_flush_d = None + replayssm_scratch = None if self.vllm_config.cache_config.mamba_cache_mode == "all": num_computed_tokens = common_attn_metadata.compute_num_computed_tokens() @@ -586,7 +610,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): num_reqs - num_prefills : num_reqs ] - if self.use_replayssm and num_decodes > 0: + if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0: decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu num_computed_tokens_cpu = common_attn_metadata._num_computed_tokens_cpu if decode_base_cpu is None or num_computed_tokens_cpu is None: @@ -651,6 +675,15 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): device=common_attn_metadata.query_start_loc.device, ) + if self.use_flashinfer_replayssm and num_decodes > 0: + assert self.decode_replayssm_scratch is not None + cb_scaled, cumAdt_vec, cb_old = self.decode_replayssm_scratch + replayssm_scratch = ( + cb_scaled[:num_decodes], + cumAdt_vec[:num_decodes], + cb_old[:num_decodes], + ) + bc_pre_scratch = None if ( self.use_replayssm @@ -671,6 +704,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): write_pos_d=write_pos_d, is_flush_d=is_flush_d, bc_pre_scratch=bc_pre_scratch, + replayssm_scratch=replayssm_scratch, num_accepted_tokens=num_accepted_tokens, query_start_loc_d=query_start_loc_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token, @@ -708,6 +742,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): write_pos_d = metadata.write_pos_d is_flush_d = metadata.is_flush_d bc_pre_scratch = metadata.bc_pre_scratch + replayssm_scratch = metadata.replayssm_scratch if ( metadata.num_prefills == 0 and metadata.num_decodes <= self.decode_cudagraph_max_bs @@ -769,7 +804,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): ) block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0 - if self.use_replayssm: + if self.use_replayssm and not self.use_flashinfer_replayssm: assert write_pos_d is not None assert is_flush_d is not None self.decode_write_pos_d[: metadata.num_decodes].copy_( @@ -788,6 +823,14 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): if self.decode_bc_pre_scratch is not None: bc_pre_scratch = self.decode_bc_pre_scratch[:padded_bs] + elif self.use_flashinfer_replayssm: + assert self.decode_replayssm_scratch is not None + cb_scaled, cumAdt_vec, cb_old = self.decode_replayssm_scratch + replayssm_scratch = ( + cb_scaled[:padded_bs], + cumAdt_vec[:padded_bs], + cb_old[:padded_bs], + ) return replace( metadata, @@ -797,6 +840,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): write_pos_d=write_pos_d, is_flush_d=is_flush_d, bc_pre_scratch=bc_pre_scratch, + replayssm_scratch=replayssm_scratch, 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=( diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 05febec328..8a7043e54a 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -229,7 +229,13 @@ def init_kv_cache( in ("longcat_flash", "longcat_flash_ngram") else 1 ) - bind_kv_cache(kv_caches, forward_context, runner_kv_caches, num_attn_module) + bind_kv_cache( + kv_caches, + forward_context, + runner_kv_caches, + num_attn_module, + kv_cache_groups=kv_cache_config.kv_cache_groups, + ) return kv_caches diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3886b28a7a..24e38d47a7 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -633,7 +633,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): cls=self.pcp_manager_cls, ) initialize_mamba_ssu_backend( - self.vllm_config.mamba_config, self.kv_cache_config + self.vllm_config.mamba_config, + self.kv_cache_config, + use_replayssm=self.vllm_config.cache_config.use_replayssm, ) if self.adaptive_verification is not None: self.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE @@ -707,6 +709,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): context_len: int = 0, skip_eplb: bool = False, is_profile: bool = False, + valid_dummy_state_slots: bool = False, **kwargs, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: if skip_attn and not is_profile: @@ -764,6 +767,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): skip_attn_for_dummy_run=skip_attn, is_profile=is_profile, context_len=context_len, + valid_dummy_state_slots=valid_dummy_state_slots, ) self.kv_connector.set_disabled(False) @@ -1390,9 +1394,18 @@ class GPUModelRunner(LoRAModelRunnerMixin): return block_tables, slot_mappings def prepare_dummy_attn( - self, input_batch: InputBatch + self, input_batch: InputBatch, valid_state_slots: bool = False ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) + if valid_state_slots: + state_slots = torch.arange( + 1, + input_batch.num_reqs + 1, + dtype=torch.int32, + device=self.device, + ) + for block_table in block_tables: + block_table[:, 0].copy_(state_slots) slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings( self.pcp_manager, self.block_tables, input_batch.num_tokens ) @@ -1525,6 +1538,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): skip_attn_for_dummy_run: bool = False, is_profile: bool = False, context_len: int = 0, + valid_dummy_state_slots: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: if not dummy_run: # Update the request states. @@ -1624,7 +1638,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): max_query_len=batch_desc.max_query_len, ) if not skip_attn_for_dummy_run: - block_tables, slot_mappings = self.prepare_dummy_attn(input_batch) + block_tables, slot_mappings = self.prepare_dummy_attn( + input_batch, valid_dummy_state_slots + ) if context_len: set_dummy_context( input_batch, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 84c87fb79c..1f79e6ef44 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -7434,6 +7434,7 @@ class GPUModelRunner( self.compilation_config.static_forward_context, self.kv_caches, num_attn_module, + kv_cache_groups=kv_cache_config.kv_cache_groups, ) return kv_caches @@ -7485,7 +7486,9 @@ class GPUModelRunner( self.maybe_add_kv_sharing_layers_to_kv_cache_groups(kv_cache_config) self.initialize_attn_backend(kv_cache_config, is_profiling=is_profiling) initialize_mamba_ssu_backend( - self.vllm_config.mamba_config, self.kv_cache_config + self.vllm_config.mamba_config, + self.kv_cache_config, + use_replayssm=self.vllm_config.cache_config.use_replayssm, ) # The kernel block size for all KV cache groups. For example, if # kv_cache_manager uses block_size 256 for a given group, but the attention diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index adf0127dbd..d767a8fffd 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -13,6 +13,7 @@ import torch from vllm.config import CacheConfig, VllmConfig 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 from vllm.model_executor.layers.utils import warmup_rocm_skinny_gemm_workspaces from vllm.model_executor.models.interfaces import MultiModalEmbeddings from vllm.model_executor.models.utils import extract_layer_index @@ -575,6 +576,7 @@ def bind_kv_cache( forward_context: dict[str, Attention], runner_kv_caches: list[torch.Tensor], num_attn_module: int = 1, + kv_cache_groups: Sequence[KVCacheGroupSpec] | None = None, ) -> None: """ Bind the allocated KV cache to both ModelRunner and forward context so @@ -600,6 +602,7 @@ def bind_kv_cache( for layer_name in kv_caches: index2name[extract_layer_index(layer_name, num_attn_module)].append(layer_name) + ordered_layer_names: list[str] = [] for layer_index in sorted(index2name.keys()): layer_names = index2name[layer_index] if len(layer_names) > 1: @@ -613,6 +616,7 @@ def bind_kv_cache( current_platform.check_runner_kv_caches_multi_layer() for layer_name in layer_names: runner_kv_caches.append(kv_caches[layer_name]) + ordered_layer_names.append(layer_name) # Bind kv_caches to forward context. Each layer's bind_kv_cache unpacks # its raw allocation into the per-layer view(s) it needs (e.g. Mamba @@ -621,6 +625,8 @@ def bind_kv_cache( for layer_name, kv_cache in kv_caches.items(): forward_context[layer_name].bind_kv_cache(kv_cache) + share_replayssm_ring_trackers(ordered_layer_names, forward_context, kv_cache_groups) + def clear_layer_kv_caches(layers: Iterable[Any]) -> None: """Detach the KV/state cache tensors installed by bind_kv_cache().