@tduncan / Nvidia-Optimization-Stack-VLLM / commits / deb148e730

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

Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>

Andrii Skliar committed Sep 30, 2026 at 23:41 UTC deb148e730225c64427f8ccf1f313cb4b33325b0
18 files changed +1026 -110
.buildkite/test_areas/engine.yaml
+21
@@ -133,6 +133,27 @@ steps:
133 - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
134 - 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
135
136 +- label: ":nvidia: (H100) ReplaySSM E2E"
137 + device: h100
138 + num_devices: 2
139 + key: replayssm-e2e
140 + timeout_in_minutes: 90
141 + working_dir: "/vllm-workspace/tests"
142 + source_file_dependencies:
143 + - tests/v1/e2e/test_replayssm_decode.py
144 + - vllm/config/cache.py
145 + - vllm/config/vllm.py
146 + - vllm/model_executor/layers/mamba/
147 + - vllm/model_executor/models/nemotron_h.py
148 + - vllm/model_executor/warmup/replayssm_warmup.py
149 + - vllm/v1/attention/backends/mamba_attn.py
150 + - vllm/v1/spec_decode/
151 + - vllm/v1/worker/gpu/
152 + - vllm/v1/worker/gpu_model_runner.py
153 + - vllm/v1/worker/gpu_worker.py
154 + commands:
155 + - pytest -v -s v1/e2e/test_replayssm_decode.py
156 +
157 - label: ":nvidia: (L4) V1 E2E"
158 key: v1-e2e-2-gpus
159 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
@@ -1,24 +1,30 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from types import SimpleNamespace
5 from unittest.mock import Mock
6
7 import pytest
8 import torch
9
10 from vllm.config.mamba import MambaBackendEnum, MambaConfig, MambaSSUAlgorithm
11 +from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
12 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator
13 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
14 FlashInferSSUBackend,
15 TritonSSUBackend,
16 + commit_replayssm_ring_trackers,
17 get_mamba_ssu_backend,
18 initialize_mamba_ssu_backend,
19 reset_replayssm_ring_trackers,
20 selective_state_update,
21 + selective_state_update_replayssm_flashinfer,
22 update_replayssm_ring_trackers,
23 )
24 from vllm.utils.torch_utils import set_random_seed
25 +from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
26 from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
27 +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
28 from vllm.v1.kv_cache_interface import (
29 KVCacheConfig,
30 KVCacheGroupSpec,
@@ -32,6 +38,13 @@ try:
38 except ImportError:
39 HAS_FLASHINFER = False
40
41 +try:
42 + from flashinfer.mamba.checkpointing_ssu import CheckpointingSSURunner
43 +
44 + HAS_FLASHINFER_CHECKPOINTING_SSU = CheckpointingSSURunner is not None
45 +except ImportError:
46 + HAS_FLASHINFER_CHECKPOINTING_SSU = False
47 +
48
49 @pytest.fixture(autouse=True)
50 def restore_backend_state():
@@ -47,6 +60,7 @@ def restore_backend_state():
60 def test_flashinfer_replayssm_ring_tracker_lifecycle():
61 ring_start = torch.zeros(2, dtype=torch.int32, device="cuda")
62 prev_num_accepted = torch.zeros(2, dtype=torch.int32, device="cuda")
63 + prev_query_len = torch.zeros(2, dtype=torch.int32, device="cuda")
64 state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda")
65
66 observed = []
@@ -54,6 +68,7 @@ def test_flashinfer_replayssm_ring_tracker_lifecycle():
68 update_replayssm_ring_trackers(
69 ring_start,
70 prev_num_accepted,
71 + prev_query_len,
72 state_batch_indices,
73 logical_window=16,
74 ring_buffer_len=17,
@@ -69,9 +84,205 @@ def test_flashinfer_replayssm_ring_tracker_lifecycle():
84 reset_replayssm_ring_trackers(
85 ring_start,
86 prev_num_accepted,
87 + prev_query_len,
88 state_batch_indices,
89 )
74 - assert (ring_start[1].item(), prev_num_accepted[1].item()) == (0, 0)
90 + assert (
91 + ring_start[1].item(),
92 + prev_num_accepted[1].item(),
93 + prev_query_len[1].item(),
94 + ) == (0, 0, 0)
95 +
96 +
97 +@pytest.mark.parametrize(
98 + ("accepted_sequence", "expected"),
99 + [
100 + pytest.param(
101 + [4] * 22,
102 + [
103 + (0, 0, 4),
104 + (0, 4, 4),
105 + (0, 8, 4),
106 + (0, 12, 4),
107 + (0, 16, 4),
108 + (16, 4, 4),
109 + (16, 8, 4),
110 + (16, 12, 4),
111 + (16, 16, 4),
112 + (12, 4, 4),
113 + (12, 8, 4),
114 + (12, 12, 4),
115 + (12, 16, 4),
116 + (8, 4, 4),
117 + (8, 8, 4),
118 + (8, 12, 4),
119 + (8, 16, 4),
120 + (4, 4, 4),
121 + (4, 8, 4),
122 + (4, 12, 4),
123 + (4, 16, 4),
124 + (0, 4, 4),
125 + ],
126 + id="all-accepted",
127 + ),
128 + pytest.param(
129 + [4, 4, 0, 3, 4, 2, 4, 1],
130 + [
131 + (0, 0, 4),
132 + (0, 4, 4),
133 + (0, 4, 4),
134 + (0, 7, 4),
135 + (0, 11, 4),
136 + (0, 13, 4),
137 + (13, 4, 4),
138 + (13, 5, 4),
139 + ],
140 + id="mixed",
141 + ),
142 + ],
143 +)
144 +def test_replayssm_commit_tracker_acceptance_sequence(accepted_sequence, expected):
145 + logical_window = 16
146 + num_speculative_tokens = 3
147 + query_len = 1 + num_speculative_tokens
148 + ring_buffer_len = logical_window + 1 + num_speculative_tokens
149 + ring_start = torch.zeros(2, dtype=torch.int32, device="cuda")
150 + prev_num_accepted = torch.zeros(2, dtype=torch.int32, device="cuda")
151 + prev_query_len = torch.zeros(2, dtype=torch.int32, device="cuda")
152 + state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda")
153 + query_start_loc = torch.tensor([0, query_len], dtype=torch.int32, device="cuda")
154 +
155 + observed = []
156 + for accepted in accepted_sequence:
157 + commit_replayssm_ring_trackers(
158 + ring_start,
159 + prev_num_accepted,
160 + prev_query_len,
161 + state_batch_indices,
162 + torch.tensor([accepted], dtype=torch.int32, device="cuda"),
163 + query_start_loc,
164 + logical_window,
165 + ring_buffer_len,
166 + )
167 + snapshot = (
168 + ring_start[1].item(),
169 + prev_num_accepted[1].item(),
170 + prev_query_len[1].item(),
171 + )
172 + observed.append(snapshot)
173 + assert snapshot[1] + snapshot[2] <= ring_buffer_len
174 +
175 + assert observed == expected
176 +
177 +
178 +def test_replayssm_resume_resets_commit_history():
179 + ring_start = torch.tensor([0, 13], dtype=torch.int32, device="cuda")
180 + prev_num_accepted = torch.tensor([0, 13], dtype=torch.int32, device="cuda")
181 + prev_query_len = torch.tensor([0, 4], dtype=torch.int32, device="cuda")
182 + state_batch_indices = torch.tensor([1], dtype=torch.int32, device="cuda")
183 +
184 + reset_replayssm_ring_trackers(
185 + ring_start,
186 + prev_num_accepted,
187 + prev_query_len,
188 + state_batch_indices,
189 + )
190 + assert (
191 + ring_start[1].item(),
192 + prev_num_accepted[1].item(),
193 + prev_query_len[1].item(),
194 + ) == (0, 0, 0)
195 +
196 + commit_replayssm_ring_trackers(
197 + ring_start,
198 + prev_num_accepted,
199 + prev_query_len,
200 + state_batch_indices,
201 + torch.tensor([3], dtype=torch.int32, device="cuda"),
202 + torch.tensor([0, 4], dtype=torch.int32, device="cuda"),
203 + logical_window=16,
204 + ring_buffer_len=20,
205 + )
206 + assert (
207 + ring_start[1].item(),
208 + prev_num_accepted[1].item(),
209 + prev_query_len[1].item(),
210 + ) == (0, 0, 4)
211 +
212 +
213 +def test_replayssm_commit_tracker_ragged_query_lengths():
214 + ring_start = torch.zeros(3, dtype=torch.int32, device="cuda")
215 + prev_num_accepted = torch.zeros(3, dtype=torch.int32, device="cuda")
216 + prev_query_len = torch.zeros(3, dtype=torch.int32, device="cuda")
217 + state_batch_indices = torch.tensor([1, 2], dtype=torch.int32, device="cuda")
218 + query_start_loc = torch.tensor([0, 4, 6], dtype=torch.int32, device="cuda")
219 +
220 + observed = []
221 + for accepted in ([4, 2], [3, 1]):
222 + commit_replayssm_ring_trackers(
223 + ring_start,
224 + prev_num_accepted,
225 + prev_query_len,
226 + state_batch_indices,
227 + torch.tensor(accepted, dtype=torch.int32, device="cuda"),
228 + query_start_loc,
229 + logical_window=16,
230 + ring_buffer_len=20,
231 + )
232 + observed.append(
233 + [
234 + (
235 + ring_start[slot].item(),
236 + prev_num_accepted[slot].item(),
237 + prev_query_len[slot].item(),
238 + )
239 + for slot in (1, 2)
240 + ]
241 + )
242 +
243 + assert observed == [[(0, 0, 4), (0, 0, 2)], [(0, 3, 4), (0, 1, 2)]]
244 +
245 +
246 +@pytest.mark.parametrize("operation", ["commit", "reset"])
247 +def test_replayssm_tracker_kernels_mask_invalid_slots(operation):
248 + num_states = 3
249 + ring_start = torch.tensor([11, 2, 33], dtype=torch.int32, device="cuda")
250 + prev_num_accepted = torch.tensor([11, 3, 33], dtype=torch.int32, device="cuda")
251 + prev_query_len = torch.tensor([11, 4, 33], dtype=torch.int32, device="cuda")
252 + state_batch_indices = torch.tensor(
253 + [-1, num_states, NULL_BLOCK_ID, 1], dtype=torch.int32, device="cuda"
254 + )
255 +
256 + if operation == "commit":
257 + commit_replayssm_ring_trackers(
258 + ring_start,
259 + prev_num_accepted,
260 + prev_query_len,
261 + state_batch_indices,
262 + torch.tensor([4, 4, 4, 2], dtype=torch.int32, device="cuda"),
263 + torch.tensor([0, 4, 8, 12, 16], dtype=torch.int32, device="cuda"),
264 + logical_window=16,
265 + ring_buffer_len=20,
266 + )
267 + expected_valid = (2, 5, 4)
268 + else:
269 + reset_replayssm_ring_trackers(
270 + ring_start,
271 + prev_num_accepted,
272 + prev_query_len,
273 + state_batch_indices,
274 + )
275 + expected_valid = (0, 0, 0)
276 +
277 + assert (
278 + ring_start.tolist(),
279 + prev_num_accepted.tolist(),
280 + prev_query_len.tolist(),
281 + ) == (
282 + [11, expected_valid[0], 33],
283 + [11, expected_valid[1], 33],
284 + [11, expected_valid[2], 33],
285 + )
286
287
288 def _kv_cache_config_with_ssu(
@@ -231,14 +442,190 @@ def test_triton_basic_call():
442 assert not torch.isnan(out).any()
443
444
445 +@pytest.mark.parametrize("layout", ["packed", "dense"])
446 +def test_replayssm_flashinfer_call_forwards_mtp_layout(monkeypatch, layout):
447 + import vllm.model_executor.layers.mamba.ops.ssu_dispatch as mod
448 +
449 + kernel = Mock(return_value=torch.empty(0))
450 + monkeypatch.setattr(mod, "_flashinfer_replayssm_kernel", kernel)
451 +
452 + batch, max_seqlen, nheads, dim, dstate, ngroups = 2, 4, 2, 4, 8, 1
453 + state = torch.empty(2, nheads, dim, dstate)
454 + x_shape: tuple[int, ...]
455 + B_shape: tuple[int, ...]
456 + expected_x_shape: tuple[int, ...]
457 + expected_B_shape: tuple[int, ...]
458 + if layout == "packed":
459 + x_shape = (6, nheads, dim)
460 + B_shape = (6, ngroups, dstate)
461 + expected_x_shape = (1, 6, nheads, dim)
462 + expected_B_shape = (1, 6, ngroups, dstate)
463 + cu_seqlens = torch.tensor([0, 4, 6], dtype=torch.int32)
464 + kernel_max_seqlen = max_seqlen
465 + else:
466 + x_shape = (batch, max_seqlen, nheads, dim)
467 + B_shape = (batch, max_seqlen, ngroups, dstate)
468 + expected_x_shape = x_shape
469 + expected_B_shape = B_shape
470 + cu_seqlens = None
471 + kernel_max_seqlen = None
472 + x = torch.empty(x_shape)
473 + dt = torch.empty_like(x)
474 + A = torch.empty(nheads, dim, dstate)
475 + B = torch.empty(B_shape)
476 + C = torch.empty_like(B)
477 + out = torch.empty_like(x)
478 + x_cache = torch.empty(2, nheads, 20, dim)
479 + dt_cache = torch.empty(2, nheads, 20)
480 + B_cache = torch.empty(2, ngroups, 20, dstate)
481 + ring_start = torch.zeros(2, dtype=torch.int32)
482 + prev_num_accepted = torch.zeros(2, dtype=torch.int32)
483 + prev_query_len = torch.zeros(2, dtype=torch.int32)
484 + selective_state_update_replayssm_flashinfer(
485 + state,
486 + x,
487 + dt,
488 + A,
489 + B,
490 + C,
491 + out,
492 + x_cache,
493 + B_cache,
494 + dt_cache,
495 + ring_start,
496 + prev_num_accepted,
497 + prev_query_len,
498 + logical_window=16,
499 + state_batch_indices=torch.tensor([0, 1], dtype=torch.int32),
500 + cu_seqlens=cu_seqlens,
501 + max_seqlen=kernel_max_seqlen,
502 + update_trackers=False,
503 + )
504 +
505 + args = kernel.call_args.args
506 + assert args[6].shape == expected_x_shape
507 + assert args[7].shape == expected_x_shape
508 + assert args[9].shape == expected_B_shape
509 + assert args[10].shape == expected_B_shape
510 + assert args[11].shape == expected_x_shape
511 + assert kernel.call_args.kwargs["cu_seqlens"] is cu_seqlens
512 + assert kernel.call_args.kwargs["max_seqlen"] == kernel_max_seqlen
513 +
514 +
515 +@pytest.mark.parametrize(
516 + ("query_start_loc", "expected_shape", "expected_max_seqlen"),
517 + [
518 + pytest.param([0, 4, 8], (2, 4, 2, 4), None, id="dense"),
519 + pytest.param([0, 4, 6], (6, 2, 4), 4, id="packed"),
520 + ],
521 +)
522 +def test_replayssm_mixer_selects_mtp_layout(
523 + monkeypatch, query_start_loc, expected_shape, expected_max_seqlen
524 +):
525 + import vllm.model_executor.layers.mamba.mamba_mixer2 as mod
526 +
527 + mixer = MambaMixer2.__new__(MambaMixer2)
528 + torch.nn.Module.__init__(mixer)
529 + mixer.prefix = "mixer"
530 + mixer.tped_intermediate_size = 0
531 + mixer.tped_conv_size = 1
532 + mixer.tped_dt_size = 2
533 + mixer.num_heads = 2
534 + mixer.head_dim = 4
535 + mixer.n_groups = mixer.tp_size = 1
536 + mixer.ssm_state_size = 8
537 + mixer.num_spec = 3
538 + mixer.use_replayssm = True
539 + mixer.replayssm_buffer_len = 16
540 + mixer._commits_replayssm_trackers = True
541 + mixer._updates_replayssm_trackers = False
542 + mixer.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER)
543 + mixer.cache_config = SimpleNamespace(mamba_block_size=16, mamba_cache_mode="none")
544 + mixer.conv_weights = torch.empty(0)
545 + mixer.conv1d = SimpleNamespace(bias=None)
546 + mixer.activation = "silu"
547 + mixer.A = torch.empty(2)
548 + mixer.dt_bias = torch.empty(2)
549 + mixer.D = torch.empty(2)
550 + mixer._replayssm_ring_start = torch.zeros(3, dtype=torch.int32)
551 + mixer._replayssm_prev_num_accepted = torch.zeros(3, dtype=torch.int32)
552 + mixer._replayssm_prev_query_len = torch.zeros(3, dtype=torch.int32)
553 + mixer.kv_cache = (
554 + torch.empty(3, 1),
555 + torch.empty(3, 2, 4, 8),
556 + torch.empty(3, 2, 20, 4),
557 + torch.empty(3, 2, 20),
558 + torch.empty(3, 1, 20, 8),
559 + )
560 +
561 + num_decode_tokens = query_start_loc[-1]
562 + query_start_loc_d = torch.tensor(query_start_loc, dtype=torch.int32)
563 + metadata = Mamba2AttentionMetadata(
564 + num_prefills=0,
565 + num_prefill_tokens=0,
566 + num_decodes=2,
567 + num_decode_tokens=num_decode_tokens,
568 + num_reqs=2,
569 + has_initial_states_p=None,
570 + query_start_loc_p=None,
571 + num_computed_tokens_p=None,
572 + state_indices_tensor_p=None,
573 + state_indices_tensor_d=torch.tensor([[1], [2]], dtype=torch.int32),
574 + query_start_loc_d=query_start_loc_d,
575 + num_accepted_tokens=torch.tensor([4, 2], dtype=torch.int32),
576 + block_idx_last_scheduled_token=None,
577 + block_idx_first_scheduled_token_p=None,
578 + block_idx_last_computed_token=None,
579 + block_idx_last_scheduled_token_prev_step=None,
580 + seq_lens=torch.tensor([104, 102], dtype=torch.int32),
581 + replayssm_scratch=(torch.empty(0), torch.empty(0), torch.empty(0)),
582 + replayssm_state_indices_d=torch.tensor([1, 2], dtype=torch.int32),
583 + )
584 +
585 + def split_hidden_states_B_C(values):
586 + tokens = values.size(0)
587 + return (
588 + torch.empty(tokens, 8),
589 + torch.empty(tokens, 8),
590 + torch.empty(tokens, 8),
591 + )
592 +
593 + mixer.split_hidden_states_B_C_fn = split_hidden_states_B_C
594 + kernel = Mock()
595 + monkeypatch.setattr(
596 + mod,
597 + "get_forward_context",
598 + lambda: SimpleNamespace(attn_metadata={mixer.prefix: metadata}),
599 + )
600 + monkeypatch.setattr(mod, "commit_replayssm_ring_trackers", Mock())
601 + monkeypatch.setattr(
602 + mod, "causal_conv1d_update", lambda values, *args, **kwargs: values
603 + )
604 + monkeypatch.setattr(mod, "selective_state_update_replayssm_flashinfer", kernel)
605 +
606 + mixer.conv_ssm_forward(
607 + torch.empty(num_decode_tokens, 3), torch.empty(num_decode_tokens, 8)
608 + )
609 +
610 + assert kernel.call_args.args[1].shape == expected_shape
611 + if expected_max_seqlen is None:
612 + assert kernel.call_args.kwargs["cu_seqlens"] is None
613 + else:
614 + assert kernel.call_args.kwargs["cu_seqlens"] is query_start_loc_d
615 + assert kernel.call_args.kwargs["max_seqlen"] == expected_max_seqlen
616 +
617 +
618 @pytest.mark.parametrize(
235 - ("backend", "expected_ring_len"),
619 + ("backend", "num_speculative_tokens", "expected_ring_len"),
620 [
237 - (MambaBackendEnum.TRITON, 16),
238 - (MambaBackendEnum.FLASHINFER, 17),
621 + (MambaBackendEnum.TRITON, 0, 16),
622 + (MambaBackendEnum.FLASHINFER, 0, 17),
623 + (MambaBackendEnum.FLASHINFER, 3, 20),
624 ],
625 )
241 -def test_replayssm_physical_ring_shape(backend, expected_ring_len):
626 +def test_replayssm_physical_ring_shape(
627 + backend, num_speculative_tokens, expected_ring_len
628 +):
629 base_shapes = ((64, 3), (8, 4, 16))
630
631 shapes = MambaStateShapeCalculator.append_replayssm_ring(
@@ -247,6 +634,7 @@ def test_replayssm_physical_ring_shape(backend, expected_ring_len):
634 tp_world_size=2,
635 logical_window=16,
636 backend=backend,
637 + num_speculative_tokens=num_speculative_tokens,
638 )
639
640 assert shapes[2:] == (
tests/model_executor/test_replayssm_warmup.py
+19 -4
@@ -67,8 +67,13 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
67 with patch.object(
68 warmup, "flashinfer_replayssm_autotune_supported", return_value=True
69 ):
70 - result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs))
70 + result = warmup._replayssm_autotune_kwargs(
71 + _autotune_runner(**runner_kwargs), PREFILL_KWARGS
72 + )
73
74 + assert result is not None
75 + max_num_reqs, decode_kwargs = result
76 + assert max_num_reqs == expected_num_reqs
77 expected_kwargs = {
78 **PREFILL_KWARGS,
79 "num_tokens": expected_num_reqs * query_len,
@@ -82,7 +87,7 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
87 force_attention=True,
88 profile_seq_lens=query_len + 1,
89 )
85 - assert result == (expected_num_reqs, expected_kwargs)
90 + assert decode_kwargs == expected_kwargs
91
92
93 @pytest.mark.parametrize(
@@ -91,8 +96,14 @@ def test_replayssm_autotune_decode_kwargs(runner_kwargs, expected_num_reqs):
96 (dict(use_replayssm=False), True),
97 (dict(backend=MambaBackendEnum.TRITON), True),
98 ({}, False),
99 + (dict(num_blocks=1), True),
100 + ],
101 + ids=[
102 + "replayssm_disabled",
103 + "non_flashinfer_backend",
104 + "kernel_unavailable",
105 + "zero_non_padding_slots",
106 ],
95 - ids=["replayssm_disabled", "non_flashinfer_backend", "kernel_unavailable"],
107 )
108 def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported):
109 with patch.object(
@@ -100,7 +111,9 @@ def test_replayssm_autotune_kwargs_skipped(runner_kwargs, flashinfer_supported):
111 "flashinfer_replayssm_autotune_supported",
112 return_value=flashinfer_supported,
113 ):
103 - result = warmup._replayssm_autotune_kwargs(_autotune_runner(**runner_kwargs))
114 + result = warmup._replayssm_autotune_kwargs(
115 + _autotune_runner(**runner_kwargs), PREFILL_KWARGS
116 + )
117 assert result is None
118
119
@@ -116,10 +129,12 @@ def test_replayssm_autotune_slots_restore_state_and_trackers():
129 )
130 mixer._replayssm_ring_start = torch.full((4,), 3, dtype=torch.int32)
131 mixer._replayssm_prev_num_accepted = torch.full((4,), 3, dtype=torch.int32)
132 + mixer._replayssm_prev_query_len = torch.full((4,), 3, dtype=torch.int32)
133 tracked = (
134 *mixer.kv_cache,
135 mixer._replayssm_ring_start,
136 mixer._replayssm_prev_num_accepted,
137 + mixer._replayssm_prev_query_len,
138 )
139
140 block_ids = np.arange(10, 14, dtype=np.int32).reshape(4, 1)
tests/test_config.py
+4 -2
@@ -153,6 +153,7 @@ def test_kda_recoverssm_derivation_is_revalidated():
153 use_replayssm=True,
154 use_kda_recoverssm=False,
155 mamba_cache_mode="none",
156 + replayssm_buffer_len=16,
157 ),
158 num_speculative_tokens=3,
159 model_config=SimpleNamespace(
@@ -184,8 +185,9 @@ def test_kda_recoverssm_derivation_is_revalidated():
185 config.cache_config.mamba_cache_mode = "none"
186
187 config.model_config.architecture = "NemotronHForCausalLM"
187 - with pytest.raises(ValueError, match="only supported for Kimi-K3 KDA"):
188 - VllmConfig.validate_mamba_cached_kernel(config)
188 + config.mamba_config.backend = MambaBackendEnum.FLASHINFER
189 + VllmConfig.validate_mamba_cached_kernel(config)
190 + assert not config.cache_config.use_kda_recoverssm
191
192 config.model_config.architecture = "KimiLinearForCausalLM"
193 config.parallel_config.pipeline_parallel_size = 2
tests/v1/attention/test_attention_backends_selection.py
+28
@@ -5,7 +5,9 @@
5 from types import SimpleNamespace
6
7 import pytest
8 +import torch
9
10 +from vllm.model_executor.layers.mamba.abstract import MambaBase
11 from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
12 MiniMaxText01LinearAttention,
13 )
@@ -19,6 +21,32 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum
21 from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionBackend
22
23
24 +@pytest.mark.parametrize(("use_replayssm", "expected_blocks"), [(False, 3), (True, 0)])
25 +def test_replayssm_does_not_reserve_speculative_state_blocks(
26 + use_replayssm, expected_blocks
27 +):
28 + layer = SimpleNamespace(
29 + get_state_shape=lambda: ((2,),),
30 + get_state_dtype=lambda: (torch.float32,),
31 + mamba_type=MambaAttentionBackendEnum.MAMBA2,
32 + is_kv_cache_tp_replicated=False,
33 + )
34 + vllm_config = SimpleNamespace(
35 + cache_config=SimpleNamespace(
36 + mamba_block_size=1,
37 + mamba_page_size_padded=None,
38 + mamba_cache_mode="none",
39 + use_replayssm=use_replayssm,
40 + ),
41 + num_speculative_tokens=3,
42 + )
43 +
44 + spec = MambaBase.get_kv_cache_spec(layer, vllm_config)
45 +
46 + assert spec is not None
47 + assert spec.num_speculative_blocks == expected_blocks
48 +
49 +
50 @pytest.mark.parametrize(
51 "layer_class, init_kwargs, expected_backend, expected_mamba_type",
52 [
tests/v1/attention/test_replayssm_metadata_builder.py
+74 -18
@@ -16,6 +16,8 @@ from tests.v1.attention.utils import (
16 create_common_attn_metadata,
17 create_vllm_config,
18 )
19 +from vllm.config import SpeculativeConfig
20 +from vllm.config.compilation import CUDAGraphMode
21 from vllm.config.mamba import MambaBackendEnum
22 from vllm.v1.kv_cache_interface import MambaSpec
23
@@ -191,9 +193,12 @@ REPLAYSSM_BUILD_CASES = {
193 def _make_mamba_spec(
194 buffer_len: int,
195 mamba_backend: MambaBackendEnum,
196 + num_speculative_tokens: int = 0,
197 ) -> MambaSpec:
198 ring_buffer_len = buffer_len + (
196 - 1 if mamba_backend == MambaBackendEnum.FLASHINFER else 0
199 + 1 + num_speculative_tokens
200 + if mamba_backend == MambaBackendEnum.FLASHINFER
201 + else 0
202 )
203 shapes = (
204 (1, 1),
@@ -214,6 +219,7 @@ def _create_replayssm_builder(
219 mamba_cache_mode: str = "none",
220 *,
221 mamba_backend: MambaBackendEnum = MambaBackendEnum.TRITON,
222 + num_speculative_tokens: int = 0,
223 ) -> MockMambaBuilder:
224 vllm_config = create_vllm_config(
225 model_name="Qwen/Qwen3.5-0.8B", block_size=BLOCK_SIZE
@@ -224,21 +230,32 @@ def _create_replayssm_builder(
230 vllm_config.cache_config.replayssm_buffer_len = buffer_len
231 vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
232 vllm_config.mamba_config.backend = mamba_backend
233 + if num_speculative_tokens > 0:
234 + # Triton ReplaySSM requires exact synchronous CPU metadata.
235 + vllm_config.scheduler_config.async_scheduling = False
236 + vllm_config.speculative_config = SpeculativeConfig(
237 + method="ngram",
238 + num_speculative_tokens=num_speculative_tokens,
239 + )
240 return MockMambaBuilder(
228 - _make_mamba_spec(buffer_len, mamba_backend),
241 + _make_mamba_spec(buffer_len, mamba_backend, num_speculative_tokens),
242 ["layer0"],
243 vllm_config,
244 DEVICE,
245 )
246
247
235 -def _build(builder: MockMambaBuilder, case: ReplaySSMBuildCase):
248 +def _build(
249 + builder: MockMambaBuilder,
250 + case: ReplaySSMBuildCase,
251 + num_accepted_tokens: torch.Tensor | None = None,
252 +):
253 batch = BatchSpec(seq_lens=case.seq_lens, query_lens=case.query_lens)
254 common = create_common_attn_metadata(batch, BLOCK_SIZE, DEVICE).replace(
255 is_prefilling=torch.tensor(case.is_prefilling, dtype=torch.bool),
256 replayssm_decode_base_cpu=torch.tensor(case.decode_base, dtype=torch.int32),
257 )
241 - return builder.build(0, common)
258 + return builder.build(0, common, num_accepted_tokens=num_accepted_tokens)
259
260
261 @pytest.mark.parametrize(
@@ -270,21 +287,60 @@ def test_resumed_request_differs_from_fresh():
287 assert meta.is_flush_d.tolist()[:2] == [0, 0]
288
289
273 -def test_flashinfer_replayssm_scratch_metadata_fresh_decode():
290 +def test_spec_decode_single_token_chunk_synthesizes_acceptance_metadata():
291 + builder = _create_replayssm_builder(16, num_speculative_tokens=3)
292 + case = REPLAYSSM_BUILD_CASES["leftover_prompt_one_token_flush"]
293 +
294 + meta = _build(builder, case)
295 +
296 + assert meta.query_start_loc_d is not None
297 + assert meta.query_start_loc_d.tolist() == [0, 1]
298 + assert meta.num_accepted_tokens is not None
299 + assert meta.num_accepted_tokens.tolist() == [1]
300 +
301 +
302 +def test_flashinfer_replayssm_state_indices_are_stable_for_full_cudagraph():
303 checkpointing_ssu = pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
304 if not hasattr(checkpointing_ssu, "allocate_checkpointing_ssu_scratch"):
276 - pytest.skip("FlashInfer does not expose ReplaySSM scratch allocation")
305 + pytest.skip("requires FlashInfer ReplaySSM autotuning support")
306
278 - builder = _create_replayssm_builder(16, mamba_backend=MambaBackendEnum.FLASHINFER)
279 - case = REPLAYSSM_BUILD_CASES["fresh_decode"]
280 - meta = _build(builder, case)
307 + builder = _create_replayssm_builder(
308 + 16,
309 + mamba_backend=MambaBackendEnum.FLASHINFER,
310 + num_speculative_tokens=3,
311 + )
312 + builder.compilation_config.cudagraph_mode = CUDAGraphMode.FULL
313
282 - assert meta.write_pos_d is None
283 - assert meta.is_flush_d is None
284 - assert meta.bc_pre_scratch is None
285 - assert meta.replayssm_scratch is not None
286 - assert [tensor.shape for tensor in meta.replayssm_scratch] == [
287 - (1, 1, 32, 8),
288 - (1, 1, 16),
289 - (1, 1, 32, 8),
290 - ]
314 + first = _build(
315 + builder,
316 + ReplaySSMBuildCase(
317 + seq_lens=[106, 106],
318 + query_lens=[1, 1],
319 + is_prefilling=[False, False],
320 + decode_base=[100, 100],
321 + buffer_len=16,
322 + expected_write_pos=[],
323 + expected_is_flush=[],
324 + ),
325 + )
326 + first_indices = first.replayssm_state_indices_d
327 + assert first_indices is not None
328 + assert first_indices.is_contiguous()
329 + first_ptr = first_indices.data_ptr()
330 +
331 + second = _build(
332 + builder,
333 + ReplaySSMBuildCase(
334 + seq_lens=[122, 122],
335 + query_lens=[1, 1],
336 + is_prefilling=[False, False],
337 + decode_base=[116, 116],
338 + buffer_len=16,
339 + expected_write_pos=[],
340 + expected_is_flush=[],
341 + ),
342 + )
343 + second_indices = second.replayssm_state_indices_d
344 + assert second_indices is not None
345 + assert second_indices.data_ptr() == first_ptr
346 + assert torch.equal(second_indices, second.state_indices_tensor_d[:, 0])
tests/v1/e2e/test_replayssm_decode.py
+118 -24
@@ -10,11 +10,16 @@ from vllm.v1.metrics.reader import Counter
10 from ...models.utils import check_logprobs_close
11 from ...utils import large_gpu_mark, multi_gpu_test
12
13 -# Mamba2 (Nemotron-3) hybrid.
13 +# Mamba2 (Nemotron-3) hybrids.
14 MAMBA2_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"
15 +# FlashInfer 0.7.0's standard SSU cannot compile Nano's DIM=80 configuration.
16 +MAMBA2_FLASHINFER_MODEL = "nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4"
17 MODELS = [
18 pytest.param(MAMBA2_MODEL, marks=large_gpu_mark(min_gb=40)),
19 ]
20 +FLASHINFER_MODELS = [
21 + pytest.param(MAMBA2_FLASHINFER_MODEL, marks=large_gpu_mark(min_gb=40)),
22 +]
23
24 PROMPTS = [
25 "The capital of France is",
@@ -22,6 +27,16 @@ PROMPTS = [
27 ]
28
29
30 +@pytest.fixture(autouse=True)
31 +def _use_v1_model_runner_by_default(monkeypatch):
32 + # Triton ReplaySSM is V1-only. FlashInfer V2 tests override this locally.
33 + with monkeypatch.context() as patch:
34 + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
35 + envs.disable_envs_cache()
36 + yield
37 + envs.disable_envs_cache()
38 +
39 +
40 def _check_replayssm_parity(
41 vllm_runner,
42 model_name,
@@ -29,17 +44,11 @@ def _check_replayssm_parity(
44 tensor_parallel_size=1,
45 mamba_backend: str = "triton",
46 name_1: str = "replayssm",
32 - require_v2: bool = False,
33 - monkeypatch: pytest.MonkeyPatch | None = None,
47 + expected_v2: bool | None = None,
48 ):
49 # Compare logprobs, not greedy ids: ReplaySSM's fp arithmetic can flip a
36 - # near-tie. Baseline and ReplaySSM run at the same TP, so TP numerics are
37 - # common-mode and only ReplaySSM varies.
38 - if require_v2:
39 - assert monkeypatch is not None
40 - monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
41 - envs.disable_envs_cache()
42 -
50 + # near-tie. Baseline and ReplaySSM use the same backend so only ReplaySSM
51 + # behavior varies.
52 common = dict(
53 max_model_len=1024,
54 trust_remote_code=True,
@@ -49,14 +58,17 @@ def _check_replayssm_parity(
58 mamba_backend=mamba_backend,
59 )
60 with vllm_runner(model_name, **common) as llm:
52 - if require_v2:
53 - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner
61 + if expected_v2 is not None:
62 + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2
63 baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
64 with vllm_runner(
56 - model_name, use_replayssm=True, replayssm_buffer_len=16, **common
65 + model_name,
66 + use_replayssm=True,
67 + replayssm_buffer_len=16,
68 + **common,
69 ) as llm:
58 - if require_v2:
59 - assert llm.llm.llm_engine.vllm_config.use_v2_model_runner
70 + if expected_v2 is not None:
71 + assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2
72 replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
73
74 check_logprobs_close(
@@ -80,21 +92,103 @@ def test_replayssm_decode_matches_baseline_tp2(vllm_runner, model_name):
92 _check_replayssm_parity(vllm_runner, model_name, tensor_parallel_size=2)
93
94
83 -@pytest.mark.parametrize("model_name", MODELS)
84 -def test_replayssm_flashinfer_decode_matches_baseline_v2(
85 - vllm_runner, model_name, monkeypatch
95 +@pytest.mark.parametrize("model_name", FLASHINFER_MODELS)
96 +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
97 +def test_replayssm_flashinfer_decode_matches_baseline(
98 + vllm_runner, model_name, monkeypatch, use_v2_model_runner
99 ):
87 - pytest.importorskip("flashinfer.mamba.checkpointing_ssu")
88 - _check_replayssm_parity(
89 - vllm_runner,
100 + try:
101 + with monkeypatch.context() as patch:
102 + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
103 + envs.disable_envs_cache()
104 + _check_replayssm_parity(
105 + vllm_runner,
106 + model_name,
107 + mamba_backend="flashinfer",
108 + name_1="replayssm_flashinfer",
109 + expected_v2=use_v2_model_runner,
110 + )
111 + finally:
112 + # The context restores the environment before the final cache reset.
113 + envs.disable_envs_cache()
114 +
115 +
116 +@pytest.mark.parametrize("model_name", FLASHINFER_MODELS)
117 +def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_name):
118 + common = dict(
119 + max_model_len=1024,
120 + trust_remote_code=True,
121 + enable_prefix_caching=False,
122 + mamba_cache_mode="none",
123 + speculative_config={
124 + "method": "ngram",
125 + "num_speculative_tokens": 3,
126 + "prompt_lookup_max": 3,
127 + },
128 + )
129 + with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm:
130 + baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
131 + with vllm_runner(
132 model_name,
133 + use_replayssm=True,
134 + replayssm_buffer_len=16,
135 mamba_backend="flashinfer",
92 - name_1="replayssm_flashinfer_v2",
93 - require_v2=True,
94 - monkeypatch=monkeypatch,
136 + **common,
137 + ) as llm:
138 + replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
139 +
140 + check_logprobs_close(
141 + outputs_0_lst=baseline,
142 + outputs_1_lst=replay,
143 + name_0="baseline_spec",
144 + name_1="replayssm_flashinfer_spec",
145 )
146
147
148 +@multi_gpu_test(num_gpus=2)
149 +@large_gpu_mark(min_gb=40)
150 +@pytest.mark.parametrize("use_v2_model_runner", [False, True], ids=["v1", "v2"])
151 +def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner):
152 + common = dict(
153 + max_model_len=1024,
154 + trust_remote_code=True,
155 + enable_prefix_caching=False,
156 + mamba_cache_mode="none",
157 + mamba_backend="flashinfer",
158 + tensor_parallel_size=2,
159 + disable_log_stats=False,
160 + speculative_config={"method": "mtp", "num_speculative_tokens": 3},
161 + )
162 + try:
163 + with monkeypatch.context() as patch:
164 + patch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
165 + envs.disable_envs_cache()
166 + with vllm_runner(
167 + MAMBA2_FLASHINFER_MODEL,
168 + use_replayssm=True,
169 + replayssm_buffer_len=16,
170 + **common,
171 + ) as llm:
172 + assert (
173 + llm.llm.llm_engine.vllm_config.use_v2_model_runner
174 + is use_v2_model_runner
175 + )
176 + outputs = llm.generate_greedy(PROMPTS, max_tokens=32)
177 + draft_count = sum(
178 + metric.value
179 + for metric in llm.llm.get_metrics()
180 + if isinstance(metric, Counter)
181 + and metric.name == "vllm:spec_decode_num_drafts"
182 + )
183 + finally:
184 + envs.disable_envs_cache()
185 +
186 + # At least one request must run past the 16-token replay window; another
187 + # may legitimately stop early on EOS.
188 + assert any(len(token_ids) > 16 for token_ids, _ in outputs)
189 + assert draft_count > 0
190 +
191 +
192 # Prefix spans several mamba blocks; prefix caching only reuses full blocks.
193 _PC_SENTENCE = (
194 "In a detailed survey of state space models, the authors compared many "
tests/v1/worker/test_utils.py
+14 -12
@@ -1260,6 +1260,8 @@ class _TestReplaySSMMixer(MambaMixer2):
1260 self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER)
1261 self._replayssm_ring_start = torch.empty(0, dtype=torch.int32)
1262 self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32)
1263 + self._replayssm_prev_query_len = torch.empty(0, dtype=torch.int32)
1264 + self._commits_replayssm_trackers = True
1265 self._updates_replayssm_trackers = True
1266
1267 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):
1296 else:
1297 bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups)
1298
1297 - assert (
1298 - mixers[0]._replayssm_ring_start.data_ptr()
1299 - == mixers[2]._replayssm_ring_start.data_ptr()
1300 - )
1301 - assert (
1302 - mixers[0]._replayssm_prev_num_accepted.data_ptr()
1303 - == mixers[2]._replayssm_prev_num_accepted.data_ptr()
1299 + tracker_names = (
1300 + "_replayssm_ring_start",
1301 + "_replayssm_prev_num_accepted",
1302 + "_replayssm_prev_query_len",
1303 )
1305 - assert (
1306 - mixers[1]._replayssm_ring_start.data_ptr()
1307 - != mixers[0]._replayssm_ring_start.data_ptr()
1308 - )
1309 - # Group {0, 2} shares trackers; layer 2 (not 0) updates after both run.
1304 + for tracker_name in tracker_names:
1305 + group_tracker = getattr(mixers[0], tracker_name)
1306 + assert group_tracker.data_ptr() == getattr(mixers[2], tracker_name).data_ptr()
1307 + assert group_tracker.data_ptr() != getattr(mixers[1], tracker_name).data_ptr()
1308 + assert group_tracker.shape == (4,)
1309 + assert torch.count_nonzero(group_tracker) == 0
1310 +
1311 + assert [m._commits_replayssm_trackers for m in mixers] == [True, True, False]
1312 assert [m._updates_replayssm_trackers for m in mixers] == [False, True, True]
1313
1314
tools/pre_commit/test_tethering_allowlist.txt
-1
@@ -60,7 +60,6 @@ tests/models/test_language_model_cache_is_weak.py # no models_*.yaml job collec
60 tests/models/transformers/test_layer_registry.py # no models_*.yaml job collects this path
61
62 # tests/v1/ - the v1 subtree is otherwise covered by many `pytest v1/...` jobs.
63 -tests/v1/e2e/test_replayssm_decode.py # ReplaySSM decode test (#52506) wired into no job
63 tests/v1/ec_connector/integration/test_epd_correctness.py # run_epd_correctness_test.sh exists but no Buildkite job invokes it
64 tests/v1/ec_connector/integration/test_nixl_failure.py # 2-GPU + NIXL E/PD integration test (#55290); no Buildkite job collects ec_connector/integration/
65 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
@@ -204,15 +204,16 @@ class CacheConfig:
204 Mamba block size."""
205 replayssm_buffer_len: int = Field(default=16, gt=0)
206 """ReplaySSM logical history length B for Mamba2. Triton uses B physical
207 - rows and FlashInfer uses B+1. Kimi-K3 speculative decode does not use B.
208 - Default 16."""
207 + rows and FlashInfer uses B+T, where T is the target verification length.
208 + Kimi-K3 speculative decode does not use B. Default 16."""
209 use_replayssm: bool = False
210 """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip
211 the per-step full-state store, writing the checkpoint back only on flush.
212 Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton
213 - or FlashInfer mamba backend; standard (non-speculative) decode only. In align
214 - mode flushes are most efficient when mamba_block_size is a multiple of
215 - replayssm_buffer_len, but this is not required."""
213 + or FlashInfer mamba backend. Mamba2 speculative decode requires FlashInfer
214 + and mamba_cache_mode 'none'. In align mode flushes are most efficient when
215 + mamba_block_size is a multiple of replayssm_buffer_len, but this is not
216 + required."""
217 use_kda_recoverssm: bool = field(default=False, init=False)
218 """Whether Kimi-K3 KDA uses RecoverSSM speculative decode."""
219
vllm/config/vllm.py
+41 -6
@@ -3458,19 +3458,35 @@ class VllmConfig:
3458 if not self.cache_config.use_replayssm:
3459 self.cache_config.use_kda_recoverssm = False
3460 return self
3461 - self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0
3461 +
3462 + kda_architectures = (
3463 + "KimiLinearForCausalLM",
3464 + "KimiK3ForConditionalGeneration",
3465 + )
3466 + is_kda_model = (
3467 + self.model_config is not None
3468 + and self.model_config.architecture in kda_architectures
3469 + )
3470 + self.cache_config.use_kda_recoverssm = (
3471 + self.num_speculative_tokens > 0 and is_kda_model
3472 + )
3473 + use_mamba_replayssm_spec = (
3474 + self.num_speculative_tokens > 0 and not self.cache_config.use_kda_recoverssm
3475 + )
3476
3477 if self.model_config is not None and not self.model_config.supports_replayssm:
3478 raise ValueError(
3479 "--use-replayssm is not supported for architecture "
3480 f"{self.model_config.architecture!r}"
3481 )
3482 + if (
3483 + self.mamba_config.backend == MambaBackendEnum.FLASHINFER
3484 + and self.cache_config.replayssm_buffer_len > 16
3485 + ):
3486 + raise ValueError(
3487 + "FlashInfer ReplaySSM requires --replayssm-buffer-len <= 16"
3488 + )
3489 if self.cache_config.use_kda_recoverssm:
3469 - if self.model_config is not None and self.model_config.architecture not in (
3470 - "KimiLinearForCausalLM",
3471 - "KimiK3ForConditionalGeneration",
3472 - ):
3473 - raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA")
3490 if self.mamba_config.enable_stochastic_rounding:
3491 raise ValueError(
3492 "RecoverSSM supports bfloat16/float32 "
@@ -3494,6 +3510,25 @@ class VllmConfig:
3510 )
3511 if self.mamba_config.backend != MambaBackendEnum.TRITON:
3512 raise ValueError("RecoverSSM requires --mamba-backend triton")
3513 + elif use_mamba_replayssm_spec:
3514 + if self.cache_config.mamba_cache_mode != "none":
3515 + raise ValueError(
3516 + "FlashInfer ReplaySSM speculative decoding requires "
3517 + "--mamba-cache-mode none"
3518 + )
3519 + query_len = 1 + self.num_speculative_tokens
3520 + if self.cache_config.replayssm_buffer_len < query_len:
3521 + raise ValueError(
3522 + "FlashInfer ReplaySSM speculative decoding requires "
3523 + "--replayssm-buffer-len >= 1 + num_speculative_tokens "
3524 + f"({query_len}); got "
3525 + f"{self.cache_config.replayssm_buffer_len}"
3526 + )
3527 + if self.mamba_config.backend != MambaBackendEnum.FLASHINFER:
3528 + raise ValueError(
3529 + "Mamba2 ReplaySSM speculative decoding requires "
3530 + "--mamba-backend flashinfer"
3531 + )
3532 elif self.cache_config.mamba_cache_mode == "all":
3533 raise ValueError(
3534 "--use-replayssm supports prefix caching only in align mode; "
vllm/model_executor/layers/mamba/abstract.py
+2 -2
@@ -74,11 +74,11 @@ class MambaBase(AttentionLayerBase):
74 mamba_type=self.mamba_type,
75 tp_replicated=self.is_kv_cache_tp_replicated,
76 mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
77 - # RecoverSSM verifies the whole window off one checkpoint, so it
77 + # ReplaySSM verifies the whole window off one checkpoint, so it
78 # never writes the baseline's per-draft-token state slots.
79 num_speculative_blocks=(
80 0
81 - if vllm_config.cache_config.use_kda_recoverssm
81 + if vllm_config.cache_config.use_replayssm
82 else vllm_config.num_speculative_tokens
83 ),
84 )
vllm/model_executor/layers/mamba/mamba_mixer2.py
+100 -9
@@ -43,6 +43,7 @@ from vllm.model_executor.layers.mamba.ops.ssd_combined import (
43 mamba_chunk_scan_combined_varlen,
44 )
45 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
46 + commit_replayssm_ring_trackers,
47 reset_replayssm_ring_trackers,
48 selective_state_update,
49 selective_state_update_replayssm_flashinfer,
@@ -69,6 +70,13 @@ from vllm.v1.kv_cache_interface import KVCacheGroupSpec
70
71 logger = init_logger(__name__)
72
73 +
74 +def _view_mtp_decode_tensor(
75 + tensor: torch.Tensor, decode_batch: int, spec_query_len: int
76 +) -> torch.Tensor:
77 + return tensor.view(decode_batch, spec_query_len, *tensor.shape[1:])
78 +
79 +
80 # Added by the IBM Team, 2024
81
82
@@ -536,6 +544,8 @@ class MambaMixer2(MambaBase, PluggableLayer):
544 self.kv_cache = tuple(torch.tensor([]) for _ in range(_n_state))
545 self._replayssm_ring_start = torch.empty(0, dtype=torch.int32)
546 self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32)
547 + self._replayssm_prev_query_len = torch.empty(0, dtype=torch.int32)
548 + self._commits_replayssm_trackers = True
549 self._updates_replayssm_trackers = True
550
551 self.num_spec = vllm_config.num_speculative_tokens
@@ -724,7 +734,8 @@ class MambaMixer2(MambaBase, PluggableLayer):
734 assert self.cache_config is not None
735 mamba_block_size = self.cache_config.mamba_block_size
736 is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
727 - ring_start = prev_num_accepted = None
737 + use_spec_decode = self.num_spec > 0
738 + ring_start = prev_num_accepted = prev_query_len = None
739
740 attn_metadata: AttentionMetadata | None = None
741 if attn_metadata_raw is not None:
@@ -746,6 +757,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
757 if self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
758 ring_start = self._replayssm_ring_start
759 prev_num_accepted = self._replayssm_prev_num_accepted
760 + prev_query_len = self._replayssm_prev_query_len
761 else:
762 x_cache = dt_cache = B_cache = None
763 has_initial_states_p = attn_metadata.has_initial_states_p
@@ -757,6 +769,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
769 last_chunk_indices_p = attn_metadata.last_chunk_indices_p
770 state_indices_tensor_p = attn_metadata.state_indices_tensor_p
771 state_indices_tensor_d = attn_metadata.state_indices_tensor_d
772 + replayssm_state_indices_d = attn_metadata.replayssm_state_indices_d
773 num_accepted_tokens = attn_metadata.num_accepted_tokens
774 query_start_loc_d = attn_metadata.query_start_loc_d
775 num_decodes = attn_metadata.num_decodes
@@ -996,6 +1009,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
1009 reset_replayssm_ring_trackers(
1010 ring_start,
1011 prev_num_accepted,
1012 + prev_query_len,
1013 state_indices_tensor_p,
1014 )
1015
@@ -1003,7 +1017,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
1017 if has_decode:
1018 assert state_indices_tensor_d is not None
1019 if is_mamba_cache_all:
1006 - if self.num_spec > 0:
1020 + if use_spec_decode:
1021 assert block_idx_last_scheduled_token_prev_step_d is not None
1022 input_indices = (
1023 block_idx_last_scheduled_token_prev_step_d.unsqueeze(1)
@@ -1031,6 +1045,34 @@ class MambaMixer2(MambaBase, PluggableLayer):
1045 state_indices_tensor_d_input = state_indices_tensor_d
1046 state_indices_tensor_d_output = state_indices_tensor_d
1047
1048 + if (
1049 + self.use_replayssm
1050 + and self.mamba_config.backend == MambaBackendEnum.FLASHINFER
1051 + and use_spec_decode
1052 + and self._commits_replayssm_trackers
1053 + ):
1054 + assert ring_start is not None
1055 + assert prev_num_accepted is not None
1056 + assert prev_query_len is not None
1057 + assert replayssm_state_indices_d is not None
1058 + assert num_accepted_tokens is not None
1059 + assert query_start_loc_d is not None
1060 + assert x_cache is not None
1061 + assert self.replayssm_buffer_len is not None
1062 + # The previous MTP query's accepted prefix is known only now.
1063 + # Commit it once, before this cache group's layers evaluate the
1064 + # current speculative query with the shared ring positions.
1065 + commit_replayssm_ring_trackers(
1066 + ring_start,
1067 + prev_num_accepted,
1068 + prev_query_len,
1069 + replayssm_state_indices_d,
1070 + num_accepted_tokens,
1071 + query_start_loc_d,
1072 + logical_window=self.replayssm_buffer_len,
1073 + ring_buffer_len=x_cache.size(2),
1074 + )
1075 +
1076 # 2. Convolution sequence transformation
1077 hidden_states_B_C_d = causal_conv1d_update(
1078 hidden_states_B_C_d,
@@ -1043,7 +1085,13 @@ class MambaMixer2(MambaBase, PluggableLayer):
1085 initial_state_idx=block_idx_last_computed_token_d,
1086 num_accepted_tokens=num_accepted_tokens,
1087 query_start_loc=query_start_loc_d,
1046 - max_query_len=state_indices_tensor_d.size(-1),
1088 + # ReplaySSM keeps one physical state block while a speculative
1089 + # decode call still processes the full target + draft window.
1090 + max_query_len=(
1091 + 1 + self.num_spec
1092 + if self.use_replayssm and use_spec_decode
1093 + else state_indices_tensor_d.size(-1)
1094 + ),
1095 )
1096
1097 hidden_states_d, B_d, C_d = self.split_hidden_states_B_C_fn(
@@ -1079,7 +1127,35 @@ class MambaMixer2(MambaBase, PluggableLayer):
1127 if self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
1128 assert ring_start is not None
1129 assert prev_num_accepted is not None
1130 + assert prev_query_len is not None
1131 assert attn_metadata.replayssm_scratch is not None
1132 + fi_cu_seqlens = query_start_loc_d
1133 + fi_max_seqlen = None
1134 + if use_spec_decode:
1135 + spec_query_len = 1 + self.num_spec
1136 + fi_max_seqlen = spec_query_len
1137 + assert replayssm_state_indices_d is not None
1138 + decode_batch = replayssm_state_indices_d.size(0)
1139 + if num_decode_tokens == decode_batch * spec_query_len:
1140 + hidden_states_d = _view_mtp_decode_tensor(
1141 + hidden_states_d, decode_batch, spec_query_len
1142 + )
1143 + dt_d = _view_mtp_decode_tensor(
1144 + dt_d, decode_batch, spec_query_len
1145 + )
1146 + B_d = _view_mtp_decode_tensor(
1147 + B_d, decode_batch, spec_query_len
1148 + )
1149 + C_d = _view_mtp_decode_tensor(
1150 + C_d, decode_batch, spec_query_len
1151 + )
1152 + preallocated_ssm_out_d = _view_mtp_decode_tensor(
1153 + preallocated_ssm_out_d,
1154 + decode_batch,
1155 + spec_query_len,
1156 + )
1157 + fi_cu_seqlens = None
1158 + fi_max_seqlen = None
1159 selective_state_update_replayssm_flashinfer(
1160 ssm_state,
1161 hidden_states_d,
@@ -1093,19 +1169,25 @@ class MambaMixer2(MambaBase, PluggableLayer):
1169 dt_cache,
1170 ring_start,
1171 prev_num_accepted,
1172 + prev_query_len,
1173 logical_window=self.replayssm_buffer_len,
1174 D=D_d,
1175 dt_bias=dt_bias,
1176 dt_softplus=True,
1100 - state_batch_indices=state_indices_tensor_d_input,
1177 + state_batch_indices=replayssm_state_indices_d,
1178 scratch=attn_metadata.replayssm_scratch,
1102 - update_trackers=self._updates_replayssm_trackers,
1179 + update_trackers=(
1180 + self._updates_replayssm_trackers and not use_spec_decode
1181 + ),
1182 enable_stochastic_rounding=(
1183 self.mamba_config.enable_stochastic_rounding
1184 ),
1185 stochastic_rounding_philox_rounds=(
1186 self.mamba_config.stochastic_rounding_philox_rounds
1187 ),
1188 + cu_seqlens=fi_cu_seqlens,
1189 + max_seqlen=fi_max_seqlen,
1190 + enable_pdl=False,
1191 )
1192 else:
1193 selective_state_update_replayssm_output_only(
@@ -1187,6 +1269,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
1269 tp_world_size=tp_world_size,
1270 logical_window=self.replayssm_buffer_len,
1271 backend=self.mamba_config.backend,
1272 + num_speculative_tokens=self.num_spec,
1273 )
1274 return base_shape
1275
@@ -1204,9 +1287,10 @@ def share_replayssm_ring_trackers(
1287
1288 Layers backed by one KV-cache group use the same physical block indices and
1289 can therefore share cursors. Different KV-cache groups may assign different
1207 - block indices to the same request and must keep separate cursor tensors.
1208 - The final local layer in each group advances its cursors after every layer
1209 - in that group has consumed the previous values.
1290 + block indices to the same request and must keep separate cursor tensors. For
1291 + speculative decode, the first local layer commits the preceding acceptance;
1292 + for standard decode, the final local layer advances after all layers consume
1293 + the previous values.
1294 """
1295 replayssm_mixers: dict[str, MambaMixer2] = {}
1296 for layer_name in ordered_layer_names:
@@ -1232,6 +1316,7 @@ def share_replayssm_ring_trackers(
1316 groups_by_namespace.setdefault(namespace, []).append(layer_name)
1317
1318 for group_layer_names in groups_by_namespace.values():
1319 + first_layer_name = group_layer_names[0]
1320 last_layer_name = group_layer_names[-1]
1321
1322 first_mixer = replayssm_mixers[group_layer_names[0]]
@@ -1246,11 +1331,17 @@ def share_replayssm_ring_trackers(
1331
1332 ring_start = torch.zeros(num_blocks, dtype=torch.int32, device=device)
1333 prev_num_accepted = torch.zeros_like(ring_start)
1334 + prev_query_len = torch.zeros_like(ring_start)
1335 for layer_name in group_layer_names:
1336 mixer = replayssm_mixers[layer_name]
1337 mixer._replayssm_ring_start = ring_start
1338 mixer._replayssm_prev_num_accepted = prev_num_accepted
1253 - mixer._updates_replayssm_trackers = layer_name == last_layer_name
1339 + mixer._replayssm_prev_query_len = prev_query_len
1340 + mixer._commits_replayssm_trackers = False
1341 + mixer._updates_replayssm_trackers = False
1342 +
1343 + replayssm_mixers[first_layer_name]._commits_replayssm_trackers = True
1344 + replayssm_mixers[last_layer_name]._updates_replayssm_trackers = True
1345
1346
1347 def mamba_mixer2(
vllm/model_executor/layers/mamba/mamba_utils.py
+3 -2
@@ -227,6 +227,7 @@ class MambaStateShapeCalculator:
227 tp_world_size: int,
228 logical_window: int,
229 backend: MambaBackendEnum,
230 + num_speculative_tokens: int = 0,
231 ) -> tuple[
232 tuple[int, int],
233 tuple[int, int, int],
@@ -241,8 +242,8 @@ class MambaStateShapeCalculator:
242 """
243 ring_buffer_len = logical_window
244 if backend == MambaBackendEnum.FLASHINFER:
244 - # FlashInfer keeps the live window and appended token together.
245 - ring_buffer_len += 1
245 + # FlashInfer keeps the live window and current verify window together.
246 + ring_buffer_len += 1 + num_speculative_tokens
247 local_nheads, head_dim, state_size = base_shapes[1]
248 local_ngroups = divide(n_groups, tp_world_size)
249 return (
vllm/model_executor/layers/mamba/ops/ssu_dispatch.py
+137 -6
@@ -31,6 +31,7 @@ logger = init_logger(__name__)
31 def _update_replayssm_ring_trackers_kernel(
32 ring_start,
33 prev_num_accepted,
34 + prev_query_len,
35 state_batch_indices,
36 state_batch_indices_stride,
37 n_slots,
@@ -52,6 +53,9 @@ def _update_replayssm_ring_trackers_kernel(
53 if RESET:
54 tl.store(ring_start + slots, 0, mask=valid)
55 tl.store(prev_num_accepted + slots, 0, mask=valid)
56 + # MTP alone reads prev_query_len. Standard decode deliberately leaves
57 + # it unchanged; reset clears the pending MTP window after prefill.
58 + tl.store(prev_query_len + slots, 0, mask=valid)
59 else:
60 prev = tl.load(prev_num_accepted + slots, mask=valid, other=0)
61 start = tl.load(ring_start + slots, mask=valid, other=0)
@@ -66,9 +70,70 @@ def _update_replayssm_ring_trackers_kernel(
70 tl.store(prev_num_accepted + slots, next_prev, mask=valid)
71
72
73 +@triton.jit(
74 + do_not_specialize=["n_slots", "state_batch_indices_stride"],
75 + do_not_specialize_on_alignment=["state_batch_indices"],
76 +)
77 +def _commit_replayssm_ring_trackers_kernel(
78 + ring_start,
79 + prev_num_accepted,
80 + prev_query_len,
81 + state_batch_indices,
82 + num_accepted_tokens,
83 + query_start_loc,
84 + state_batch_indices_stride,
85 + n_slots,
86 + num_states,
87 + logical_window: tl.constexpr,
88 + ring_buffer_len: tl.constexpr,
89 + pad_slot_id: tl.constexpr,
90 + BLOCK: tl.constexpr,
91 +) -> None:
92 + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
93 + mask = offsets < n_slots
94 + slots = tl.load(
95 + state_batch_indices + offsets * state_batch_indices_stride,
96 + mask=mask,
97 + other=pad_slot_id,
98 + )
99 + valid = mask & (slots != pad_slot_id) & (slots >= 0) & (slots < num_states)
100 + accepted_since_checkpoint = tl.load(prev_num_accepted + slots, mask=valid, other=0)
101 + current_ring_start = tl.load(ring_start + slots, mask=valid, other=0)
102 + previous_speculative_query_len = tl.load(
103 + prev_query_len + slots, mask=valid, other=0
104 + )
105 + accepted_from_previous_query = tl.load(
106 + num_accepted_tokens + offsets, mask=mask, other=0
107 + )
108 + must_checkpoint = (previous_speculative_query_len > 0) & (
109 + accepted_since_checkpoint + previous_speculative_query_len > logical_window
110 + )
111 + next_start = tl.where(
112 + must_checkpoint,
113 + (current_ring_start + accepted_since_checkpoint) % ring_buffer_len,
114 + current_ring_start,
115 + )
116 + next_prev = tl.where(
117 + previous_speculative_query_len == 0,
118 + 0,
119 + tl.where(
120 + must_checkpoint,
121 + accepted_from_previous_query,
122 + accepted_since_checkpoint + accepted_from_previous_query,
123 + ),
124 + )
125 + current_speculative_query_len = tl.load(
126 + query_start_loc + offsets + 1, mask=mask, other=0
127 + ) - tl.load(query_start_loc + offsets, mask=mask, other=0)
128 + tl.store(ring_start + slots, next_start, mask=valid)
129 + tl.store(prev_num_accepted + slots, next_prev, mask=valid)
130 + tl.store(prev_query_len + slots, current_speculative_query_len, mask=valid)
131 +
132 +
133 def update_replayssm_ring_trackers(
134 ring_start: torch.Tensor,
135 prev_num_accepted: torch.Tensor,
136 + prev_query_len: torch.Tensor,
137 state_batch_indices: torch.Tensor,
138 logical_window: int | None = None,
139 ring_buffer_len: int | None = None,
@@ -90,10 +155,15 @@ def update_replayssm_ring_trackers(
155 _update_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)](
156 ring_start,
157 prev_num_accepted,
158 + prev_query_len,
159 state_batch_indices,
160 state_batch_indices.stride(0),
161 n_slots,
96 - min(ring_start.numel(), prev_num_accepted.numel()),
162 + min(
163 + ring_start.numel(),
164 + prev_num_accepted.numel(),
165 + prev_query_len.numel(),
166 + ),
167 logical_window,
168 ring_buffer_len,
169 pad_slot_id,
@@ -105,6 +175,7 @@ def update_replayssm_ring_trackers(
175 def reset_replayssm_ring_trackers(
176 ring_start: torch.Tensor,
177 prev_num_accepted: torch.Tensor,
178 + prev_query_len: torch.Tensor,
179 state_batch_indices: torch.Tensor,
180 pad_slot_id: int = NULL_BLOCK_ID,
181 ) -> None:
@@ -112,11 +183,62 @@ def reset_replayssm_ring_trackers(
183 update_replayssm_ring_trackers(
184 ring_start,
185 prev_num_accepted,
186 + prev_query_len,
187 state_batch_indices,
188 pad_slot_id=pad_slot_id,
189 )
190
191
192 +def commit_replayssm_ring_trackers(
193 + ring_start: torch.Tensor,
194 + prev_num_accepted: torch.Tensor,
195 + prev_query_len: torch.Tensor,
196 + state_batch_indices: torch.Tensor,
197 + num_accepted_tokens: torch.Tensor,
198 + query_start_loc: torch.Tensor,
199 + logical_window: int,
200 + ring_buffer_len: int,
201 + pad_slot_id: int = NULL_BLOCK_ID,
202 +) -> None:
203 + """Commit the preceding speculative window and record the current one.
204 +
205 + MTP evaluates a target token and its draft tokens together, but the number
206 + accepted from that query is available only on the next forward pass. This
207 + function then advances the per-request ring by the accepted prefix and
208 + records the current query length for the following pass. A zero
209 + ``prev_query_len`` means that reset/prefill left no prior MTP query to
210 + commit; slot validity is handled independently by the kernel mask.
211 +
212 + Standard single-token decode needs no delayed commit: its ReplaySSM kernel
213 + advances the shared ring trackers directly after every token.
214 + """
215 + if state_batch_indices.dim() > 1:
216 + state_batch_indices = state_batch_indices[:, 0]
217 + n_slots = state_batch_indices.numel()
218 + if n_slots == 0:
219 + return
220 + block = 128
221 + _commit_replayssm_ring_trackers_kernel[(triton.cdiv(n_slots, block),)](
222 + ring_start,
223 + prev_num_accepted,
224 + prev_query_len,
225 + state_batch_indices,
226 + num_accepted_tokens,
227 + query_start_loc,
228 + state_batch_indices.stride(0),
229 + n_slots,
230 + min(
231 + ring_start.numel(),
232 + prev_num_accepted.numel(),
233 + prev_query_len.numel(),
234 + ),
235 + logical_window,
236 + ring_buffer_len,
237 + pad_slot_id,
238 + BLOCK=block,
239 + )
240 +
241 +
242 class MambaSSUBackend(ABC):
243 """Abstract base class for Mamba SSU backends."""
244
@@ -385,6 +507,7 @@ def selective_state_update_replayssm_flashinfer(
507 dt_cache: torch.Tensor,
508 ring_start: torch.Tensor,
509 prev_num_accepted_tokens: torch.Tensor,
510 + prev_query_len: torch.Tensor,
511 logical_window: int,
512 D: torch.Tensor | None = None,
513 dt_bias: torch.Tensor | None = None,
@@ -395,6 +518,9 @@ def selective_state_update_replayssm_flashinfer(
518 update_trackers: bool = True,
519 enable_stochastic_rounding: bool = False,
520 stochastic_rounding_philox_rounds: int = 0,
521 + cu_seqlens: torch.Tensor | None = None,
522 + max_seqlen: int | None = None,
523 + enable_pdl: bool = False,
524 ) -> torch.Tensor:
525 """Run FlashInfer checkpointing SSU and optionally advance shared trackers."""
526 if _flashinfer_replayssm_kernel is None:
@@ -404,11 +530,12 @@ def selective_state_update_replayssm_flashinfer(
530 )
531
532 if x.dim() == 3:
407 - x = x.unsqueeze(1)
408 - dt = dt.unsqueeze(1)
409 - B = B.unsqueeze(1)
410 - C = C.unsqueeze(1)
411 - out = out.unsqueeze(1)
533 + dim = 0 if cu_seqlens is not None else 1
534 + x = x.unsqueeze(dim)
535 + dt = dt.unsqueeze(dim)
536 + B = B.unsqueeze(dim)
537 + C = C.unsqueeze(dim)
538 + out = out.unsqueeze(dim)
539
540 indices = state_batch_indices
541 if indices is not None and indices.dim() > 1:
@@ -443,6 +570,9 @@ def selective_state_update_replayssm_flashinfer(
570 pad_slot_id=null_block_id,
571 rand_seed=rand_seed,
572 philox_rounds=stochastic_rounding_philox_rounds or 10,
573 + cu_seqlens=cu_seqlens,
574 + max_seqlen=max_seqlen,
575 + enable_pdl=enable_pdl,
576 cb_scaled=cb_scaled,
577 cumAdt_vec=cumAdt_vec,
578 cb_old=cb_old,
@@ -451,6 +581,7 @@ def selective_state_update_replayssm_flashinfer(
581 update_replayssm_ring_trackers(
582 ring_start,
583 prev_num_accepted_tokens,
584 + prev_query_len,
585 indices,
586 logical_window=logical_window,
587 ring_buffer_len=x_cache.size(2),
vllm/model_executor/models/nemotron_h.py
+1
@@ -815,6 +815,7 @@ class NemotronHForCausalLM(
815 tp_world_size=parallel_config.tensor_parallel_size,
816 logical_window=cache_config.replayssm_buffer_len,
817 backend=vllm_config.mamba_config.backend,
818 + num_speculative_tokens=vllm_config.num_speculative_tokens,
819 )
820 return base_shape
821
vllm/model_executor/warmup/replayssm_warmup.py
+30 -7
@@ -21,6 +21,7 @@ logger = init_logger(__name__)
21
22 def _replayssm_autotune_kwargs(
23 runner: "GPUModelRunner",
24 + max_token_prefill_kwargs: dict[str, Any],
25 ) -> tuple[int, dict[str, Any]] | None:
26 config = runner.vllm_config
27 if not (
@@ -46,11 +47,16 @@ def _replayssm_autotune_kwargs(
47 runner.max_num_tokens // query_len,
48 runner.kv_cache_config.num_blocks - 1,
49 )
50 + if max_num_reqs <= 0:
51 + logger.warning_once(
52 + "Skipping FlashInfer ReplaySSM autotuning because no non-padding "
53 + "state slot is available."
54 + )
55 + return None
56 +
57 decode_kwargs = {
58 + **max_token_prefill_kwargs,
59 "num_tokens": max_num_reqs * query_len,
51 - "skip_eplb": True,
52 - "is_profile": True,
53 - "randomize_inputs": True,
60 "uniform_decode": True,
61 }
62 if config.use_v2_model_runner:
@@ -75,18 +81,22 @@ def _temporary_replayssm_autotune_state(
81 )
82
83 reset_tensors: dict[int, torch.Tensor] = {}
78 - tracker_specs: dict[int, tuple[torch.Tensor, torch.Tensor, int, int]] = {}
84 + tracker_specs: dict[
85 + int, tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int]
86 + ] = {}
87 for module in runner.get_model().modules():
88 if not isinstance(module, MambaMixer2) or not module.use_replayssm:
89 continue
90 assert module.replayssm_buffer_len is not None
91 ring_start = module._replayssm_ring_start
92 prev_num_accepted = module._replayssm_prev_num_accepted
93 + prev_query_len = module._replayssm_prev_query_len
94 tracker_specs.setdefault(
95 ring_start.data_ptr(),
96 (
97 ring_start,
98 prev_num_accepted,
99 + prev_query_len,
100 module.replayssm_buffer_len,
101 module.kv_cache[2].size(2),
102 ),
@@ -95,6 +105,7 @@ def _temporary_replayssm_autotune_state(
105 *module.kv_cache,
106 ring_start,
107 prev_num_accepted,
108 + prev_query_len,
109 )
110 for tensor in tensors:
111 if tensor.numel():
@@ -121,20 +132,26 @@ def _temporary_replayssm_autotune_state(
132 for (
133 ring_start,
134 prev_num_accepted,
135 + prev_query_len,
136 logical_window,
137 ring_buffer_len,
138 ) in tracker_specs.values():
139 # Compile reset (prefill) and advance (decode) before inference.
140 # The final reset leaves the decode tuning run in a clean state.
129 - reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots)
141 + reset_replayssm_ring_trackers(
142 + ring_start, prev_num_accepted, prev_query_len, state_slots
143 + )
144 update_replayssm_ring_trackers(
145 ring_start,
146 prev_num_accepted,
147 + prev_query_len,
148 state_slots,
149 logical_window,
150 ring_buffer_len,
151 )
137 - reset_replayssm_ring_trackers(ring_start, prev_num_accepted, state_slots)
152 + reset_replayssm_ring_trackers(
153 + ring_start, prev_num_accepted, prev_query_len, state_slots
154 + )
155
156 try:
157 yield
@@ -151,7 +168,13 @@ def _temporary_replayssm_autotune_state(
168
169
170 def replayssm_autotune_warmup(runner: "GPUModelRunner") -> None:
154 - autotune = _replayssm_autotune_kwargs(runner)
171 + max_token_prefill_kwargs = {
172 + "num_tokens": runner.scheduler_config.max_num_batched_tokens,
173 + "skip_eplb": True,
174 + "is_profile": True,
175 + "randomize_inputs": True,
176 + }
177 + autotune = _replayssm_autotune_kwargs(runner, max_token_prefill_kwargs)
178 if autotune is None:
179 return
180 max_num_reqs, decode_kwargs = autotune
vllm/v1/attention/backends/mamba_attn.py
+35 -7
@@ -87,6 +87,8 @@ class BaseMambaAttentionMetadata:
87 bc_pre_scratch: torch.Tensor | None = None
88 # ReplaySSM — FlashInfer checkpointing_ssu two-kernel scratch.
89 replayssm_scratch: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None
90 + # Contiguous cache-slot indices shared by all FlashInfer ReplaySSM layers.
91 + replayssm_state_indices_d: torch.Tensor | None = None
92
93
94 class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
@@ -182,6 +184,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
184 self.decode_replayssm_scratch: (
185 tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None
186 ) = None
187 + self.decode_replayssm_state_indices_d: torch.Tensor | None = None
188 # ReplaySSM CUDA-graph buffers for the selected backend.
189 if self.use_replayssm and not self.use_flashinfer_replayssm:
190 self.decode_write_pos_d: torch.Tensor = torch.empty(
@@ -218,11 +221,17 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
221 self.decode_replayssm_scratch = allocate_checkpointing_ssu_scratch(
222 batch_size=scheduler_config.max_num_seqs,
223 num_heads=nheads,
221 - num_predicted_tokens=1,
224 + num_predicted_tokens=1 + self.num_spec_tokens,
225 max_window=self.replayssm_buffer_len,
226 dtype=vllm_config.model_config.dtype,
227 device=device,
228 )
229 + # Full CUDA graphs retain capture-time tensor addresses. Keep the
230 + # contiguous first-column view used by FlashInfer in a persistent
231 + # buffer and refresh its contents before each replay.
232 + self.decode_replayssm_state_indices_d = torch.empty(
233 + (self.decode_cudagraph_max_bs,), dtype=torch.int32, device=device
234 + )
235
236 self._init_reorder_batch_threshold(1, self.use_spec_decode)
237 if self.use_spec_decode:
@@ -592,13 +601,16 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
601 ]
602 state_indices_tensor_p = state_indices_tensor_p[:, 0]
603
595 - # Sometimes even with specdec enabled we get single-token prefill chunks that
596 - # should be treated as decodes but don't have num_accepted_tokens set.
597 - # These should be fine to process as non-spec decodes since there's only
598 - # one token, so no risk of placing accepted tokens in the wrong slot.
599 - if num_decodes > 0 and self.use_spec_decode and num_accepted_tokens is not None:
604 + if num_decodes > 0 and self.use_spec_decode:
605 query_start_loc_d = common_attn_metadata.query_start_loc[: num_decodes + 1]
601 - num_accepted_tokens = num_accepted_tokens[:num_decodes]
606 + if num_accepted_tokens is None:
607 + # Single-token prefill chunks can be reclassified as decodes before
608 + # speculative decoding has produced acceptance counts. Treat each
609 + # token as accepted so recurrent state and ReplaySSM trackers follow
610 + # the normal speculative-decode path.
611 + num_accepted_tokens = torch.diff(query_start_loc_d)
612 + else:
613 + num_accepted_tokens = num_accepted_tokens[:num_decodes]
614
615 if num_prefills > 0:
616 if num_computed_tokens is None:
@@ -788,6 +800,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
800 is_flush_d = metadata.is_flush_d
801 bc_pre_scratch = metadata.bc_pre_scratch
802 replayssm_scratch = metadata.replayssm_scratch
803 + replayssm_state_indices_d = None
804 if (
805 metadata.num_prefills == 0
806 and metadata.num_decodes <= self.decode_cudagraph_max_bs
@@ -876,6 +889,20 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
889 cumAdt_vec[:padded_bs],
890 cb_old[:padded_bs],
891 )
892 + assert self.decode_replayssm_state_indices_d is not None
893 + self.decode_replayssm_state_indices_d[:padded_bs].copy_(
894 + state_indices_tensor_d[:, 0], non_blocking=True
895 + )
896 + replayssm_state_indices_d = self.decode_replayssm_state_indices_d[
897 + :padded_bs
898 + ]
899 +
900 + if (
901 + self.use_flashinfer_replayssm
902 + and state_indices_tensor_d is not None
903 + and replayssm_state_indices_d is None
904 + ):
905 + replayssm_state_indices_d = state_indices_tensor_d[:, 0].contiguous()
906
907 return replace(
908 metadata,
@@ -886,6 +913,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
913 is_flush_d=is_flush_d,
914 bc_pre_scratch=bc_pre_scratch,
915 replayssm_scratch=replayssm_scratch,
916 + replayssm_state_indices_d=replayssm_state_indices_d,
917 block_idx_last_scheduled_token=block_idx_last_scheduled_token,
918 block_idx_last_computed_token=block_idx_last_computed_token,
919 block_idx_last_scheduled_token_prev_step=(