iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Mamba] Add FlashInfer ReplaySSM support for MTP (#52928)

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