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

[Core] Remove deprecated mamba_cache_mode "all" (#58997)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Thomas Parnell <tpa@zurich.ibm.com>

Jiangyun Zhu committed Oct 1, 2026 at 17:59 UTC cb6458c5ce61632524c0447d49790725c1e1d9a2
45 files changed +201 -1476
tests/kernels/mamba/test_cpu_short_conv.py
-10
@@ -88,11 +88,6 @@ def test_short_conv_forward_native_prefill(vllm_config):
88 state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32),
89 num_accepted_tokens=None,
90 query_start_loc_d=None,
91 - block_idx_last_scheduled_token=None,
92 - block_idx_first_scheduled_token_p=None,
93 - block_idx_last_computed_token=None,
94 - block_idx_last_scheduled_token_prev_step=None,
95 - num_computed_tokens_p=None,
91 seq_lens=torch.tensor([5]),
92 )
93
@@ -145,11 +140,6 @@ def test_short_conv_forward_native_decode(vllm_config):
140 state_indices_tensor_d=state_indices_tensor_d,
141 num_accepted_tokens=None,
142 query_start_loc_d=torch.tensor([0, 1, 2], dtype=torch.int32),
148 - block_idx_last_scheduled_token=None,
149 - block_idx_first_scheduled_token_p=None,
150 - block_idx_last_computed_token=None,
151 - block_idx_last_scheduled_token_prev_step=None,
152 - num_computed_tokens_p=None,
143 seq_lens=torch.tensor([1, 1]),
144 )
145
tests/kernels/mamba/test_ssu_dispatch.py
-5
@@ -568,15 +568,10 @@ def test_replayssm_mixer_selects_mtp_layout(
568 num_reqs=2,
569 has_initial_states_p=None,
570 query_start_loc_p=None,
571 - num_computed_tokens_p=None,
571 state_indices_tensor_p=None,
572 state_indices_tensor_d=torch.tensor([[1], [2]], dtype=torch.int32),
573 query_start_loc_d=query_start_loc_d,
574 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,
575 seq_lens=torch.tensor([104, 102], dtype=torch.int32),
576 replayssm_scratch=(torch.empty(0), torch.empty(0), torch.empty(0)),
577 replayssm_state_indices_d=torch.tensor([1, 2], dtype=torch.int32),
tests/test_config.py
+6 -3
@@ -270,9 +270,6 @@ def test_kda_recoverssm_derivation_is_revalidated():
270 with pytest.raises(ValueError, match="VLLM_USE_V2_MODEL_RUNNER=1"):
271 VllmConfig.validate_mamba_cached_kernel(config)
272 config.use_v2_model_runner = True
273 - config.cache_config.mamba_cache_mode = "all"
274 - with pytest.raises(ValueError, match="only none and align"):
275 - VllmConfig.validate_mamba_cached_kernel(config)
273 config.cache_config.mamba_cache_mode = "none"
274
275 config.model_config.architecture = "NemotronHForCausalLM"
@@ -286,6 +283,12 @@ def test_kda_recoverssm_derivation_is_revalidated():
283 VllmConfig.validate_mamba_cached_kernel(config)
284
285
286 +def test_mamba_cache_mode_all_is_rejected():
287 + """The removed 'all' mode must fail validation instead of being ignored."""
288 + with pytest.raises(ValidationError, match="mamba_cache_mode"):
289 + CacheConfig(mamba_cache_mode="all")
290 +
291 +
292 def test_per_request_spec_decode_metrics_requires_spec_decode():
293 # The flag only makes sense with speculative decoding configured; enabling
294 # it without --speculative-config should fail fast rather than silently
tests/v1/attention/test_mamba_metadata_builder.py new
+58
@@ -0,0 +1,58 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for BaseMambaAttentionMetadataBuilder batch classification."""
4 +
5 +from types import SimpleNamespace
6 +
7 +from tests.v1.attention.utils import MockMambaBuilder
8 +from vllm.config.compilation import CUDAGraphMode
9 +
10 +
11 +def _make_vllm_config(
12 + max_model_len: int,
13 + max_num_seqs: int,
14 + num_speculative_tokens: int = 0,
15 + block_size: int | None = None,
16 +):
17 + """Create a minimal mock VllmConfig with only the fields the builder
18 + accesses, avoiding any model download / HF config inspection."""
19 + speculative_config = (
20 + SimpleNamespace(
21 + num_speculative_tokens=num_speculative_tokens,
22 + parallel_drafting=False,
23 + )
24 + if num_speculative_tokens > 0
25 + else None
26 + )
27 + return SimpleNamespace(
28 + cache_config=SimpleNamespace(
29 + block_size=block_size,
30 + mamba_cache_mode="none",
31 + use_replayssm=False,
32 + replayssm_buffer_len=16,
33 + ),
34 + compilation_config=SimpleNamespace(
35 + cudagraph_mode=CUDAGraphMode.FULL,
36 + max_cudagraph_capture_size=None,
37 + ),
38 + speculative_config=speculative_config,
39 + num_speculative_tokens=num_speculative_tokens,
40 + parallel_config=SimpleNamespace(decode_context_parallel_size=1),
41 + scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
42 + model_config=SimpleNamespace(max_model_len=max_model_len),
43 + )
44 +
45 +
46 +def test_mamba_single_token_prompt_runs_as_prefill():
47 + seq_lens = [8, 9, 1]
48 + config = _make_vllm_config(256, len(seq_lens), block_size=16)
49 + metadata = MockMambaBuilder.build_mamba_metadata(
50 + config,
51 + seq_lens=seq_lens,
52 + query_lens=[1] * len(seq_lens),
53 + is_prefilling=[False, False, True],
54 + )
55 +
56 + assert metadata.num_decodes == 2
57 + assert metadata.num_prefills == 1
58 + assert metadata.has_initial_states_p.tolist() == [False]
tests/v1/attention/test_mamba_update_block_table.py deleted
-466
@@ -1,466 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Regression test for https://github.com/vllm-project/vllm/issues/34865.
4 -
5 -When multiple KV cache groups share the same MambaSpec (as in Nemotron
6 -hybrid models), the metadata caching optimization reuses metadata from
7 -an earlier group via update_block_table(). In 'all' mode with CUDA graphs,
8 -update_block_table() must copy block_idx_last_scheduled_token and
9 -block_idx_last_computed_token to the *current* builder's persistent
10 -buffers, otherwise CUDA graph replay reads stale values from uninitialized
11 -buffers.
12 -"""
13 -
14 -from types import SimpleNamespace
15 -
16 -import torch
17 -
18 -from tests.v1.attention.utils import MockMambaBuilder
19 -from vllm.config.compilation import CUDAGraphMode
20 -from vllm.v1.attention.backends.mamba_attn import BaseMambaAttentionMetadata
21 -from vllm.v1.kv_cache_interface import MambaSpec
22 -
23 -
24 -def _make_vllm_config(
25 - max_model_len: int,
26 - max_num_seqs: int,
27 - num_speculative_tokens: int = 0,
28 - block_size: int | None = None,
29 -):
30 - """Create a minimal mock VllmConfig with only the fields the builder
31 - accesses, avoiding any model download / HF config inspection."""
32 - speculative_config = (
33 - SimpleNamespace(
34 - num_speculative_tokens=num_speculative_tokens,
35 - parallel_drafting=False,
36 - )
37 - if num_speculative_tokens > 0
38 - else None
39 - )
40 - return SimpleNamespace(
41 - cache_config=SimpleNamespace(
42 - block_size=block_size,
43 - mamba_cache_mode="all",
44 - use_replayssm=False,
45 - replayssm_buffer_len=16,
46 - ),
47 - compilation_config=SimpleNamespace(
48 - cudagraph_mode=CUDAGraphMode.FULL,
49 - max_cudagraph_capture_size=None,
50 - ),
51 - speculative_config=speculative_config,
52 - num_speculative_tokens=num_speculative_tokens,
53 - parallel_config=SimpleNamespace(decode_context_parallel_size=1),
54 - scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
55 - model_config=SimpleNamespace(max_model_len=max_model_len),
56 - )
57 -
58 -
59 -def test_mamba_single_token_prompt_runs_as_prefill():
60 - seq_lens = [8, 9, 1]
61 - config = _make_vllm_config(256, len(seq_lens), block_size=16)
62 - metadata = MockMambaBuilder.build_mamba_metadata(
63 - config,
64 - seq_lens=seq_lens,
65 - query_lens=[1] * len(seq_lens),
66 - is_prefilling=[False, False, True],
67 - )
68 -
69 - assert metadata.num_decodes == 2
70 - assert metadata.num_prefills == 1
71 - assert metadata.has_initial_states_p.tolist() == [False]
72 -
73 -
74 -def test_update_block_table_copies_block_idx_to_persistent_buffers():
75 - """update_block_table() must write block_idx tensors to the current
76 - builder's persistent buffers, not leave them pointing to a different
77 - builder's buffers."""
78 - block_size = 16
79 - max_model_len = 256
80 - num_reqs = 4
81 - device = torch.device("cpu")
82 -
83 - vllm_config = _make_vllm_config(max_model_len, num_reqs)
84 -
85 - spec = MambaSpec(
86 - block_size=block_size,
87 - shapes=((1,), (1,)),
88 - dtypes=(torch.float32,),
89 - mamba_cache_mode="all",
90 - )
91 -
92 - # Two builders simulating two KV cache groups with the same MambaSpec.
93 - builder_a = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
94 - builder_b = MockMambaBuilder(spec, ["layer1"], vllm_config, device)
95 -
96 - # Sanity: each builder has its own persistent buffer.
97 - assert (
98 - builder_a.block_idx_last_scheduled_token.data_ptr()
99 - != builder_b.block_idx_last_scheduled_token.data_ptr()
100 - )
101 -
102 - # Construct decode-only metadata as if builder_a.build() produced it.
103 - max_blocks = max_model_len // block_size
104 - seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
105 - block_idx_vals = (seq_lens - 1) // block_size # [3, 3, 3, 3]
106 -
107 - builder_a.block_idx_last_scheduled_token[:num_reqs].copy_(block_idx_vals)
108 - builder_a.block_idx_last_computed_token[:num_reqs].copy_(block_idx_vals)
109 -
110 - metadata_a = BaseMambaAttentionMetadata(
111 - num_prefills=0,
112 - num_prefill_tokens=0,
113 - num_decodes=num_reqs,
114 - num_decode_tokens=num_reqs,
115 - num_reqs=num_reqs,
116 - has_initial_states_p=None,
117 - query_start_loc_p=None,
118 - num_computed_tokens_p=None,
119 - state_indices_tensor_p=None,
120 - query_start_loc_d=None,
121 - num_accepted_tokens=None,
122 - state_indices_tensor_d=builder_a.state_indices_tensor_d[:num_reqs],
123 - block_idx_last_scheduled_token=(
124 - builder_a.block_idx_last_scheduled_token[:num_reqs]
125 - ),
126 - block_idx_first_scheduled_token_p=None,
127 - block_idx_last_computed_token=(
128 - builder_a.block_idx_last_computed_token[:num_reqs]
129 - ),
130 - block_idx_last_scheduled_token_prev_step=None,
131 - seq_lens=seq_lens,
132 - )
133 -
134 - # Call update_block_table on builder_b (simulates the metadata caching
135 - # optimization reusing metadata from builder_a's group).
136 - blk_table = torch.randint(
137 - 0, 100, (num_reqs, max_blocks), dtype=torch.int32, device=device
138 - )
139 - slot_mapping = torch.zeros(num_reqs, dtype=torch.int64, device=device)
140 -
141 - metadata_b = builder_b.update_block_table(metadata_a, blk_table, slot_mapping)
142 -
143 - # block_idx tensors must live in builder_b's persistent buffers.
144 - def shares_storage(tensor, buffer):
145 - return (
146 - tensor.untyped_storage().data_ptr() == buffer.untyped_storage().data_ptr()
147 - )
148 -
149 - assert shares_storage(
150 - metadata_b.block_idx_last_scheduled_token,
151 - builder_b.block_idx_last_scheduled_token,
152 - ), "block_idx_last_scheduled_token not in builder_b's persistent buffer"
153 -
154 - assert shares_storage(
155 - metadata_b.block_idx_last_computed_token,
156 - builder_b.block_idx_last_computed_token,
157 - ), "block_idx_last_computed_token not in builder_b's persistent buffer"
158 -
159 - # Must NOT point to builder_a's buffers.
160 - assert not shares_storage(
161 - metadata_b.block_idx_last_scheduled_token,
162 - builder_a.block_idx_last_scheduled_token,
163 - ), "block_idx_last_scheduled_token still points to builder_a's buffer"
164 -
165 - # Values must be correct (copied from metadata_a).
166 - torch.testing.assert_close(
167 - metadata_b.block_idx_last_scheduled_token,
168 - block_idx_vals,
169 - )
170 - torch.testing.assert_close(
171 - metadata_b.block_idx_last_computed_token,
172 - block_idx_vals,
173 - )
174 -
175 -
176 -def test_state_indices_tensor_d_includes_num_speculative_blocks():
177 - """Regression test for https://github.com/vllm-project/vllm/issues/39809
178 - bug 1: with mamba_cache_mode='all' and speculative decoding enabled,
179 - the cudagraph buffer for state_indices_tensor_d must allocate the same
180 - per-request column count as the runtime block table, which includes
181 - num_speculative_blocks trailing scratch columns."""
182 - block_size = 16
183 - max_model_len = 256
184 - max_num_seqs = 4
185 - num_speculative_tokens = 1
186 - num_speculative_blocks = 2
187 - device = torch.device("cpu")
188 -
189 - vllm_config = _make_vllm_config(
190 - max_model_len,
191 - max_num_seqs,
192 - num_speculative_tokens=num_speculative_tokens,
193 - )
194 -
195 - spec = MambaSpec(
196 - block_size=block_size,
197 - shapes=((1,), (1,)),
198 - dtypes=(torch.float32,),
199 - mamba_cache_mode="all",
200 - num_speculative_blocks=num_speculative_blocks,
201 - )
202 -
203 - builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
204 -
205 - expected_cols = (max_model_len // block_size) + num_speculative_blocks
206 - assert builder.state_indices_tensor_d.shape == (max_num_seqs, expected_cols)
207 -
208 -
209 -def test_block_idx_cudagraph_capture_padded_by_num_reqs():
210 - """Regression test for https://github.com/vllm-project/vllm/issues/39809
211 - bug 2: with mamba_cache_mode='all' and spec decode, _update_metadata_for
212 - _cudagraph_capture must slice block_idx_last_{scheduled,computed}_token
213 - by the request count (padded_bs == num_reqs), not by num_decode_tokens.
214 - Past num_decodes, the slice must be zero-filled."""
215 - block_size = 16
216 - max_model_len = 256
217 - max_num_seqs = 8
218 - num_speculative_tokens = 1
219 - device = torch.device("cpu")
220 -
221 - vllm_config = _make_vllm_config(
222 - max_model_len,
223 - max_num_seqs,
224 - num_speculative_tokens=num_speculative_tokens,
225 - )
226 -
227 - spec = MambaSpec(
228 - block_size=block_size,
229 - shapes=((1,), (1,)),
230 - dtypes=(torch.float32,),
231 - mamba_cache_mode="all",
232 - num_speculative_blocks=2,
233 - )
234 -
235 - builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
236 -
237 - builder.block_idx_last_scheduled_token.fill_(-1)
238 - builder.block_idx_last_computed_token.fill_(-1)
239 -
240 - num_decodes = 2
241 - num_reqs = 3
242 - num_decode_tokens = num_decodes * (1 + num_speculative_tokens)
243 - seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
244 - block_idx_vals = torch.tensor([3, 5], dtype=torch.int32, device=device)
245 - state_indices_d = torch.zeros(
246 - (num_decodes, builder.state_indices_tensor_d.shape[1]),
247 - dtype=torch.int32,
248 - device=device,
249 - )
250 - query_start_loc_d = torch.arange(
251 - num_decodes + 1, dtype=torch.int32, device=device
252 - ) * (1 + num_speculative_tokens)
253 - num_accepted_tokens = torch.ones(num_decodes, dtype=torch.int32, device=device)
254 -
255 - metadata = BaseMambaAttentionMetadata(
256 - num_prefills=0,
257 - num_prefill_tokens=0,
258 - num_decodes=num_decodes,
259 - num_decode_tokens=num_decode_tokens,
260 - num_reqs=num_reqs,
261 - has_initial_states_p=None,
262 - query_start_loc_p=None,
263 - num_computed_tokens_p=None,
264 - state_indices_tensor_p=None,
265 - state_indices_tensor_d=state_indices_d,
266 - query_start_loc_d=query_start_loc_d,
267 - num_accepted_tokens=num_accepted_tokens,
268 - block_idx_last_scheduled_token=block_idx_vals,
269 - block_idx_first_scheduled_token_p=None,
270 - block_idx_last_computed_token=block_idx_vals,
271 - block_idx_last_scheduled_token_prev_step=None,
272 - seq_lens=seq_lens,
273 - )
274 -
275 - out = builder._update_metadata_for_cudagraph_capture(metadata)
276 -
277 - assert out.block_idx_last_scheduled_token.shape == (num_reqs,)
278 - assert out.block_idx_last_computed_token.shape == (num_reqs,)
279 - torch.testing.assert_close(
280 - out.block_idx_last_scheduled_token[:num_decodes], block_idx_vals
281 - )
282 - torch.testing.assert_close(
283 - out.block_idx_last_computed_token[:num_decodes], block_idx_vals
284 - )
285 - assert torch.all(out.block_idx_last_scheduled_token[num_decodes:] == 0)
286 - assert torch.all(out.block_idx_last_computed_token[num_decodes:] == 0)
287 -
288 -
289 -def test_block_idx_prev_step_persistent_buffer_allocated():
290 - """With mamba_cache_mode='all' + spec decode, the builder must allocate
291 - block_idx_last_scheduled_token_prev_step as a persistent buffer with the
292 - same shape as the existing block_idx_last_{scheduled,computed}_token
293 - buffers, so cudagraph capture records a stable pointer for the prev-step
294 - input anchor consumed by mamba_mixer2's input gather."""
295 - block_size = 16
296 - max_model_len = 256
297 - max_num_seqs = 8
298 - num_speculative_tokens = 1
299 - device = torch.device("cpu")
300 -
301 - vllm_config = _make_vllm_config(
302 - max_model_len,
303 - max_num_seqs,
304 - num_speculative_tokens=num_speculative_tokens,
305 - )
306 - spec = MambaSpec(
307 - block_size=block_size,
308 - shapes=((1,), (1,)),
309 - dtypes=(torch.float32,),
310 - mamba_cache_mode="all",
311 - num_speculative_blocks=2,
312 - )
313 - builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
314 -
315 - assert hasattr(builder, "block_idx_last_scheduled_token_prev_step")
316 - assert builder.block_idx_last_scheduled_token_prev_step.shape == (max_num_seqs,)
317 - assert builder.block_idx_last_scheduled_token_prev_step.dtype == torch.int32
318 -
319 -
320 -def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode():
321 - """Without spec decode, the prev-step buffer is unused and must not be
322 - allocated — the input anchor reduces to last_computed_token."""
323 - block_size = 16
324 - max_model_len = 256
325 - max_num_seqs = 8
326 - device = torch.device("cpu")
327 -
328 - vllm_config = _make_vllm_config(
329 - max_model_len, max_num_seqs, num_speculative_tokens=0
330 - )
331 - spec = MambaSpec(
332 - block_size=block_size,
333 - shapes=((1,), (1,)),
334 - dtypes=(torch.float32,),
335 - mamba_cache_mode="all",
336 - )
337 - builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
338 -
339 - assert not hasattr(builder, "block_idx_last_scheduled_token_prev_step")
340 -
341 -
342 -def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer():
343 - """_update_metadata_for_cudagraph_capture must copy the prev-step anchor
344 - into the builder's persistent buffer (so cudagraph replay reads from the
345 - same underlying memory), pad past num_decodes with zero, and return a
346 - slice of the persistent buffer in the metadata."""
347 - block_size = 16
348 - max_model_len = 256
349 - max_num_seqs = 8
350 - num_speculative_tokens = 1
351 - device = torch.device("cpu")
352 -
353 - vllm_config = _make_vllm_config(
354 - max_model_len,
355 - max_num_seqs,
356 - num_speculative_tokens=num_speculative_tokens,
357 - )
358 - spec = MambaSpec(
359 - block_size=block_size,
360 - shapes=((1,), (1,)),
361 - dtypes=(torch.float32,),
362 - mamba_cache_mode="all",
363 - num_speculative_blocks=2,
364 - )
365 - builder = MockMambaBuilder(spec, ["layer0"], vllm_config, device)
366 - builder.block_idx_last_scheduled_token.fill_(-1)
367 - builder.block_idx_last_computed_token.fill_(-1)
368 - builder.block_idx_last_scheduled_token_prev_step.fill_(-1)
369 -
370 - num_decodes = 2
371 - num_reqs = 3
372 - num_decode_tokens = num_decodes * (1 + num_speculative_tokens)
373 - seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
374 - block_idx_vals = torch.tensor([3, 5], dtype=torch.int32, device=device)
375 - prev_step_vals = torch.tensor([2, 4], dtype=torch.int32, device=device)
376 - state_indices_d = torch.zeros(
377 - (num_decodes, builder.state_indices_tensor_d.shape[1]),
378 - dtype=torch.int32,
379 - device=device,
380 - )
381 - query_start_loc_d = torch.arange(
382 - num_decodes + 1, dtype=torch.int32, device=device
383 - ) * (1 + num_speculative_tokens)
384 - num_accepted_tokens = torch.ones(num_decodes, dtype=torch.int32, device=device)
385 -
386 - metadata = BaseMambaAttentionMetadata(
387 - num_prefills=0,
388 - num_prefill_tokens=0,
389 - num_decodes=num_decodes,
390 - num_decode_tokens=num_decode_tokens,
391 - num_reqs=num_reqs,
392 - has_initial_states_p=None,
393 - query_start_loc_p=None,
394 - num_computed_tokens_p=None,
395 - state_indices_tensor_p=None,
396 - state_indices_tensor_d=state_indices_d,
397 - query_start_loc_d=query_start_loc_d,
398 - num_accepted_tokens=num_accepted_tokens,
399 - block_idx_last_scheduled_token=block_idx_vals,
400 - block_idx_first_scheduled_token_p=None,
401 - block_idx_last_computed_token=block_idx_vals,
402 - block_idx_last_scheduled_token_prev_step=prev_step_vals,
403 - seq_lens=seq_lens,
404 - )
405 -
406 - out = builder._update_metadata_for_cudagraph_capture(metadata)
407 -
408 - # Output field exists and is identity-shared with the persistent buffer.
409 - assert out.block_idx_last_scheduled_token_prev_step is not None
410 - assert (
411 - out.block_idx_last_scheduled_token_prev_step.untyped_storage().data_ptr()
412 - == builder.block_idx_last_scheduled_token_prev_step.untyped_storage().data_ptr()
413 - ), (
414 - "prev-step buffer must live in the builder's persistent buffer, not "
415 - "in the caller-provided tensor"
416 - )
417 -
418 - # Padded by num_reqs (not num_decode_tokens) — same fix as bug 2 for the
419 - # other block_idx_* fields.
420 - assert out.block_idx_last_scheduled_token_prev_step.shape == (num_reqs,)
421 -
422 - # First num_decodes values: input values copied through.
423 - torch.testing.assert_close(
424 - out.block_idx_last_scheduled_token_prev_step[:num_decodes],
425 - prev_step_vals,
426 - )
427 -
428 - # Tail values past num_decodes: zero-filled padding for cudagraph capture.
429 - assert torch.all(out.block_idx_last_scheduled_token_prev_step[num_decodes:] == 0)
430 -
431 -
432 -def test_num_state_writes_p_matches_block_idx_tensors():
433 - """The prefill state save sizes its batched copy from
434 - num_state_writes_p instead of reading the block index tensors back from
435 - the GPU, so the host count must equal what those tensors say."""
436 - # (block_size, seq_lens, query_lens, is_prefilling)
437 - cases = [
438 - # block-aligned and unaligned computed tokens, and a chunked prefill
439 - (16, [64, 33, 17], [48, 17, 1], [True, True, True]),
440 - (32, [200, 100, 64], [200, 36, 64], [True, True, True]),
441 - # decode rows ahead of the prefills
442 - (16, [40, 41, 100, 200], [1, 1, 60, 200], [False, False, True, True]),
443 - # prefills that stay inside their first block: nothing to write
444 - (16, [65, 66, 67], [1, 2, 3], [True, True, True]),
445 - ]
446 -
447 - for block_size, seq_lens, query_lens, is_prefilling in cases:
448 - metadata = MockMambaBuilder.build_mamba_metadata(
449 - _make_vllm_config(4096, len(seq_lens), block_size=block_size),
450 - seq_lens=seq_lens,
451 - query_lens=query_lens,
452 - is_prefilling=is_prefilling,
453 - )
454 - num_reqs, num_prefills = metadata.num_reqs, metadata.num_prefills
455 - block_idx_last_p = metadata.block_idx_last_scheduled_token[
456 - num_reqs - num_prefills : num_reqs
457 - ]
458 - expected = int(
459 - (block_idx_last_p - metadata.block_idx_first_scheduled_token_p)
460 - .clamp(min=0)
461 - .sum()
462 - )
463 -
464 - assert metadata.num_state_writes_p == expected, (
465 - f"seq_lens={seq_lens} query_lens={query_lens} block_size={block_size}"
466 - )
tests/v1/core/test_scheduler.py
+38 -38
@@ -39,6 +39,7 @@ from vllm.sampling_params import SamplingParams, StructuredOutputsParams
39 from vllm.utils.hashing import sha256
40 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager
41 from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator
42 +from vllm.v1.core.kv_cache_manager import KVCacheManager
43 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash
44 from vllm.v1.core.sched.diffusion_scheduler import (
45 DiffusionAsyncScheduler,
@@ -6751,7 +6752,7 @@ def _create_hybrid_mamba_connector_scheduler(
6752 num_blocks: int = 100,
6753 supports_divergent_hits: bool = True,
6754 ) -> Scheduler:
6754 - """FA + Mamba ("all" cache mode) scheduler with a MockKVConnector."""
6755 + """FA + Mamba ("align" cache mode) scheduler with a MockKVConnector."""
6756 model_config = ModelConfig(
6757 model="facebook/opt-125m",
6758 trust_remote_code=True,
@@ -6772,7 +6773,7 @@ def _create_hybrid_mamba_connector_scheduler(
6773 cache_config=CacheConfig(
6774 block_size=block_size,
6775 enable_prefix_caching=True,
6775 - mamba_cache_mode="all",
6776 + mamba_cache_mode="align",
6777 ),
6778 kv_transfer_config=KVTransferConfig(
6779 kv_connector="MockKVConnector",
@@ -6804,7 +6805,7 @@ def _create_hybrid_mamba_connector_scheduler(
6805 block_size=block_size,
6806 shapes=((1, 1),),
6807 dtypes=(torch.float32,),
6807 - mamba_cache_mode="all",
6808 + mamba_cache_mode="align",
6809 ),
6810 ),
6811 ],
@@ -6820,6 +6821,33 @@ def _create_hybrid_mamba_connector_scheduler(
6821 )
6822
6823
6824 +def _seed_hybrid_prefix(
6825 + manager: KVCacheManager, num_blocks: int, block_size: int
6826 +) -> tuple[list[int], list[int]]:
6827 + """Prefill a prefix one block per step so that align-mode Mamba caches
6828 + the state at every block boundary. Returns the FA and Mamba block ids."""
6829 + [fill] = create_requests(
6830 + num_requests=1,
6831 + num_tokens=num_blocks * block_size,
6832 + max_tokens=1,
6833 + same_prompt=True,
6834 + block_size=block_size,
6835 + req_ids=["fill"],
6836 + )
6837 + computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
6838 + assert num_computed == 0
6839 + mamba_ids = []
6840 + for i in range(num_blocks):
6841 + manager.allocate_slots(
6842 + fill, block_size, new_computed_blocks=computed_blocks if i == 0 else None
6843 + )
6844 + fill.num_computed_tokens += block_size
6845 + fa_ids, group_mamba_ids = manager.get_block_ids(fill.request_id)
6846 + mamba_ids.append(group_mamba_ids[i])
6847 + manager.free(fill)
6848 + return fa_ids, mamba_ids
6849 +
6850 +
6851 @pytest.mark.parametrize(
6852 "matched_tokens,expected_num_computed",
6853 [
@@ -6845,23 +6873,8 @@ def test_hybrid_per_group_hit_divergence_with_connector(
6873 manager = scheduler.kv_cache_manager
6874 assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
6875
6848 - # Seed a 4-block prefix so both groups cache all four boundaries
6849 - # (mamba cache mode "all" caches every block's state densely).
6850 - [fill] = create_requests(
6851 - num_requests=1,
6852 - num_tokens=4 * block_size,
6853 - max_tokens=1,
6854 - same_prompt=True,
6855 - block_size=block_size,
6856 - req_ids=["fill"],
6857 - )
6858 - computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
6859 - blocks = manager.allocate_slots(
6860 - fill, fill.num_tokens, num_computed, computed_blocks
6861 - )
6862 - fa_ids = [b.block_id for b in blocks.blocks[0]]
6863 - mamba_ids = [b.block_id for b in blocks.blocks[1]]
6864 - manager.free(fill)
6876 + # Seed a 4-block prefix so both groups cache all four boundaries.
6877 + fa_ids, mamba_ids = _seed_hybrid_prefix(manager, 4, block_size)
6878
6879 # Evict the FA tail and the middle mamba states; block 0 (both groups)
6880 # and the deep mamba state at block 3 survive.
@@ -6885,8 +6898,8 @@ def test_hybrid_per_group_hit_divergence_with_connector(
6898
6899 scheduler.add_request(replay)
6900 output = scheduler.schedule()
6888 - num_scheduled = output.num_scheduled_tokens[replay.request_id]
6889 - assert replay.num_tokens - num_scheduled == expected_num_computed
6901 + [new_req] = output.scheduled_new_reqs
6902 + assert new_req.num_computed_tokens == expected_num_computed
6903
6904
6905 @pytest.mark.parametrize(
@@ -6921,20 +6934,7 @@ def test_hybrid_fa_deeper_hit_respects_connector_lookup_policy(
6934 assert isinstance(manager.coordinator, HybridKVCacheCoordinator)
6935
6936 # Seed a 4-block prefix in both groups.
6924 - [fill] = create_requests(
6925 - num_requests=1,
6926 - num_tokens=4 * block_size,
6927 - max_tokens=1,
6928 - same_prompt=True,
6929 - block_size=block_size,
6930 - req_ids=["fill"],
6931 - )
6932 - computed_blocks, num_computed, _ = manager.get_computed_blocks(fill)
6933 - blocks = manager.allocate_slots(
6934 - fill, fill.num_tokens, num_computed, computed_blocks
6935 - )
6936 - mamba_ids = [b.block_id for b in blocks.blocks[1]]
6937 - manager.free(fill)
6937 + _, mamba_ids = _seed_hybrid_prefix(manager, 4, block_size)
6938
6939 # Keep all FA blocks; evict every mamba state but block 0. FA reaches 4
6940 # blocks, the mamba hit only reaches 1 -> diverged (FA > Mamba).
@@ -6955,8 +6955,8 @@ def test_hybrid_fa_deeper_hit_respects_connector_lookup_policy(
6955
6956 scheduler.add_request(replay)
6957 output = scheduler.schedule()
6958 - num_scheduled = output.num_scheduled_tokens[replay.request_id]
6959 - assert replay.num_tokens - num_scheduled == expected_num_computed
6958 + [new_req] = output.scheduled_new_reqs
6959 + assert new_req.num_computed_tokens == expected_num_computed
6960
6961
6962 def _make_encoder_instance_request(scheduler, text_prefix=8, image_tokens=16):
tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh
+2 -2
@@ -60,7 +60,7 @@ vllm serve "$MODEL" \
60 --block-size 128 \
61 --trust-remote-code \
62 --enable-prefix-caching \
63 - --mamba-cache-mode all \
63 + --mamba-cache-mode align \
64 --kv-transfer-config "$KV_CONFIG" \
65 "${EXTRA_ARGS[@]}" &
66
@@ -78,7 +78,7 @@ vllm serve "$MODEL" \
78 --block-size 128 \
79 --trust-remote-code \
80 --enable-prefix-caching \
81 - --mamba-cache-mode all \
81 + --mamba-cache-mode align \
82 --kv-transfer-config "$KV_CONFIG" \
83 "${EXTRA_ARGS[@]}" &
84
tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py
+13 -48
@@ -67,7 +67,6 @@ class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore
67 self._mamba_group_ids: list[int] = []
68 self._attn_group_ids: list[int] = [0]
69 self._num_ssm_scratch_blocks = 0
70 - self._ssm_state_slots_are_positional = False
70 self._is_hma_required = False
71 self.kv_cache_config = SimpleNamespace(
72 kv_cache_groups=[None, None],
@@ -130,7 +129,7 @@ def _mamba_spec(num_states: int = 2, *, num_speculative_blocks: int = 0):
129 block_size=16,
130 shapes=shapes,
131 dtypes=(torch.float32,) * num_states,
133 - mamba_cache_mode="all",
132 + mamba_cache_mode="align",
133 num_speculative_blocks=num_speculative_blocks,
134 )
135
@@ -146,7 +145,7 @@ def _gate_vllm_config(
145 block_size: int = 16,
146 num_lookahead_tokens: int = 0,
147 num_speculative_tokens: int = 0,
149 - mamba_cache_mode: str = "all",
148 + mamba_cache_mode: str = "align",
149 dcp_size: int = 1,
150 ):
151 return SimpleNamespace(
@@ -193,7 +192,7 @@ def test_split_block_groups_ignores_transfer_disabled_group():
192 block_size=16,
193 shapes=((1, 1),),
194 dtypes=(torch.float32,),
196 - mamba_cache_mode="all",
195 + mamba_cache_mode="align",
196 )
197 config = SimpleNamespace(
198 kv_cache_groups=[
@@ -361,7 +360,6 @@ def test_scheduler_accepts_dspark_with_wrapped_full_attention(
360 block_size=4,
361 num_lookahead_tokens=7,
362 num_speculative_tokens=7,
364 - mamba_cache_mode="align",
363 dcp_size=dcp_size,
364 )
365 scheduler = moriio_connector.MoRIIOConnectorScheduler(vllm_config, "engine", config)
@@ -398,21 +396,6 @@ def test_scheduler_accepts_dspark_with_wrapped_full_attention(
396 )
397
398
401 -def test_scheduler_rejects_dspark_with_positional_mamba_cache():
402 - """Positional mode counts lookahead in the Mamba allocation, so the
403 - one-block recurrent-state tail bound no longer holds."""
404 - with pytest.raises(moriio_common.MoRIIOError, match="mamba_cache_mode"):
405 - moriio_connector.MoRIIOConnectorScheduler(
406 - _gate_vllm_config(
407 - speculative_config=_spec_config("dspark"),
408 - block_size=4,
409 - mamba_cache_mode="all",
410 - ),
411 - "engine",
412 - _dspark_wrapped_attention_config(),
413 - )
414 -
415 -
399 def test_scheduler_rejects_hybrid_write():
400 config = SimpleNamespace(
401 kv_cache_groups=[
@@ -441,20 +424,7 @@ def test_split_block_groups_accepts_empty_abort_payload():
424 assert sched.split_block_groups(()) == ([], [])
425
426
444 -def test_split_block_groups_keeps_positional_slots_in_all_mode():
445 - # mamba_cache_mode="all" keeps a state per block position.
446 - sched = _FakeScheduler(
447 - _has_mamba=True,
448 - _attn_group_ids=[0],
449 - _mamba_group_ids=[1],
450 - _ssm_state_slots_are_positional=True,
451 - )
452 - attn, mamba = sched.split_block_groups(([1], [40, 41, 42]))
453 - assert attn == [1]
454 - assert mamba == [[40, 41, 42]]
455 -
456 -
457 -def test_split_block_groups_keeps_only_running_state_outside_all_mode():
427 +def test_split_block_groups_keeps_only_running_state():
428 sched = _FakeScheduler(
429 _has_mamba=True,
430 _attn_group_ids=[0],
@@ -464,24 +434,21 @@ def test_split_block_groups_keeps_only_running_state_outside_all_mode():
434
435
436 @pytest.mark.parametrize(
467 - "positional,blocks,expected",
437 + "blocks,expected",
438 [
469 - # Trailing scratch slots go in either cache mode.
470 - (True, [40, 41, 42], [40]),
471 - # Outside "all" the running state is picked only after they are gone.
472 - (False, [40, 41, 42, 43], [41]),
439 + # The running state is picked only after the scratch slots are gone.
440 + ([40, 41, 42, 43], [41]),
441 # Never empty: a list shorter than the scratch count keeps one state.
474 - (True, [40], [40]),
475 - (False, [], []),
442 + ([40], [40]),
443 + ([], []),
444 ],
445 )
478 -def test_split_block_groups_strips_dspark_scratch_slots(positional, blocks, expected):
446 +def test_split_block_groups_strips_dspark_scratch_slots(blocks, expected):
447 sched = _FakeScheduler(
448 _has_mamba=True,
449 _attn_group_ids=[0],
450 _mamba_group_ids=[1],
451 _num_ssm_scratch_blocks=2,
484 - _ssm_state_slots_are_positional=positional,
452 )
453
454 assert sched.split_block_groups(([1], blocks)) == ([1], [expected])
@@ -555,7 +522,6 @@ def _make_read_scheduler():
522 _has_mamba=True,
523 _attn_group_ids=[0],
524 _mamba_group_ids=[1],
558 - _ssm_state_slots_are_positional=True,
525 request_id_to_transfer_id={},
526 transfer_id_to_request_id={},
527 _reqs_need_recv={},
@@ -578,17 +544,17 @@ def _make_read_request(remote_block_ids):
544
545 def test_update_state_drops_decode_recompute_tail_block():
546 sched = _make_read_scheduler()
581 - request = _make_read_request([[10, 11], [90, 91]])
547 + request = _make_read_request([[10, 11], [91]])
548 blocks = _FakeBlocks(
549 all_groups=([100, 101, 102], [200, 201, 202]),
550 )
551
552 sched.update_state_after_alloc(request, blocks, num_external_tokens=256)
553
588 - assert sched._reqs_need_recv["req"][1] == [[100, 101], [200, 201]]
554 + assert sched._reqs_need_recv["req"][1] == [[100, 101], [202]]
555 assert sched._req_kv_params["req"]["remote_block_ids"] == [
556 [10, 11],
591 - [90, 91],
557 + [91],
558 ]
559
560
@@ -609,7 +575,6 @@ def test_update_state_pairs_trimmed_mamba_state_with_remote_state():
575 """The DSpark shape end to end: scratch slots are stripped, the running
576 state is then selected, and the attention list loses its lookahead tail."""
577 sched = _make_read_scheduler()
612 - sched._ssm_state_slots_are_positional = False
578 sched._num_ssm_scratch_blocks = 2
579 sched._max_decode_tail_blocks = 2
580 request = _make_read_request([[10, 11], [90]])
tests/v1/kv_connector/unit/test_nixl_connector_hma.py
+3 -39
@@ -856,28 +856,6 @@ def test_apply_prefix_caching_mamba_hybrid(
856 [[6, 7, 8, 9], [99]],
857 id="fa_prefix_hit_and_ssm_trim",
858 ),
859 - # Multi-slot SSM ("all" mode): a local prefix hit leaves fewer local
860 - # slots; the earlier remote slots are covered locally → remote tail.
861 - pytest.param(
862 - 10,
863 - 10,
864 - [list(range(10)), [5, 6]],
865 - [list(range(10)), [1, 2, 3]],
866 - [list(range(10)), [5, 6]],
867 - [list(range(10)), [2, 3]],
868 - id="ssm_multi_block_local_hit_tail",
869 - ),
870 - # Multi-slot SSM ("all" mode): the one trailing local position holds
871 - # the token D recomputes itself → local head-clip.
872 - pytest.param(
873 - 10,
874 - 10,
875 - [list(range(10)), [4, 5, 6]],
876 - [list(range(10)), [8, 9]],
877 - [list(range(10)), [4, 5]],
878 - [list(range(10)), [8, 9]],
879 - id="ssm_multi_block_local_extra_head_clip",
880 - ),
859 ],
860 )
861 def test_apply_prefix_caching_ssm_prefix_cache_hit(
@@ -2034,13 +2012,12 @@ def test_logical_to_kernel_block_ids_with_remote_ratio(
2012
2013 @pytest.mark.cpu_test
2014 def test_exchange_clipped_blocks_ssm_single_state():
2037 - """In single-state cache modes, SSM lists are reduced to the running
2038 - state slot: speculative scratch slots, null placeholders and the previous
2039 - step's state carry nothing. Attention groups pass through untouched."""
2015 + """SSM lists are reduced to the running state slot: speculative scratch
2016 + slots, null placeholders and the previous step's state carry nothing.
2017 + Attention groups pass through untouched."""
2018 sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
2019 sched.blocks_per_sw = [0, 0]
2020 sched._ssm_spec_blocks = [None, 2]
2043 - sched._ssm_state_slots_are_positional = False
2021
2022 # Align-mode list: null placeholders, state block, 2 speculative slots.
2023 clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 0, 7, 8, 9]))
@@ -2069,19 +2046,6 @@ def test_exchange_clipped_blocks_ssm_single_state():
2046 assert fa_sched.get_exchange_clipped_blocks(([1, 2],)) == ([1, 2],)
2047
2048
2072 -@pytest.mark.cpu_test
2073 -def test_exchange_clipped_blocks_ssm_positional_states():
2074 - """In "all" mode every position holds a state, so only the speculative
2075 - slots go; placeholders stay to keep the list position-indexed."""
2076 - sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)
2077 - sched.blocks_per_sw = [0, 0]
2078 - sched._ssm_spec_blocks = [None, 2]
2079 - sched._ssm_state_slots_are_positional = True
2080 -
2081 - clipped = sched.get_exchange_clipped_blocks(([1, 2, 3], [0, 5, 6, 7, 8, 9]))
2082 - assert clipped == ([1, 2, 3], [0, 5, 6, 7])
2083 -
2084 -
2049 # ── Hybrid MLA+SSM (KimiLinear-shaped KDA+MLA) tests ─────────────────────
2050
2051
tests/v1/kv_connector/unit/test_offloading_connector.py
+2 -5
@@ -566,10 +566,7 @@ def test_fs_tiering_offloading(tmp_path) -> None:
566 ("state-spaces/mamba-1.4b-hf", 16, 1),
567 ],
568 )
569 -@pytest.mark.parametrize("mamba_cache_mode", ["align", "all"])
570 -def test_mamba_cpu_offload_boundary(
571 - model: str, block_size: int, tp_size: int, mamba_cache_mode: str
572 -):
569 +def test_mamba_cpu_offload_boundary(model: str, block_size: int, tp_size: int):
570 kv_transfer_config = KVTransferConfig(
571 kv_connector="OffloadingConnector",
572 kv_role="kv_both",
@@ -586,7 +583,7 @@ def test_mamba_cpu_offload_boundary(
583 kv_transfer_config=kv_transfer_config,
584 language_model_only=True,
585 enable_prefix_caching=True,
589 - mamba_cache_mode=mamba_cache_mode,
586 + mamba_cache_mode="align",
587 # Use lossless state storage so this exact-equality regression isolates
588 # offload boundary selection from low-precision Mamba checkpointing.
589 mamba_cache_dtype="float32",
tests/v1/kv_connector/unit/utils.py
+1 -1
@@ -474,7 +474,7 @@ def make_kv_cache_config(
474 mamba_enabled: bool = False,
475 sw_size: int = 128,
476 num_blocks: int = 100,
477 - mamba_cache_mode: Literal["all", "align", "none"] = "none",
477 + mamba_cache_mode: Literal["align", "none"] = "none",
478 ) -> KVCacheConfig:
479 kv_cache_groups = [
480 KVCacheGroupSpec(
tests/v1/simple_kv_offload/test_scheduler.py
+1 -50
@@ -2335,7 +2335,7 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None:
2335 # of the hash block, so it cannot exceed the smallest of them.
2336 hash_block_size=BLOCK_SIZE,
2337 dcp_world_size=dcp_world_size,
2338 - mamba_cache_mode="all",
2338 + mamba_cache_mode="none",
2339 )
2340 sched = fix.scheduler
2341 gpu_pool = fix.gpu_block_pool
@@ -2747,55 +2747,6 @@ def test_eager_store_does_not_scan_mamba_blocks_positionally() -> None:
2747 assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 0]
2748
2749
2750 -def test_eager_store_scans_mamba_all_blocks_positionally() -> None:
2751 - """Mamba all-mode retains positional identity and needs no handoff."""
2752 - block_size = 4 * BLOCK_SIZE
2753 - fix = _make_hybrid_attention_mamba_scheduler(
2754 - num_cpu_blocks=16,
2755 - num_gpu_blocks=24,
2756 - block_size=block_size,
2757 - mamba_cache_mode="all",
2758 - )
2759 - sched = fix.scheduler
2760 - gpu_pool = fix.gpu_block_pool
2761 - attention_manager, mamba_manager = sched.cpu_coordinator.single_type_managers
2762 - assert attention_manager.has_positionally_stable_blocks
2763 - assert mamba_manager.has_positionally_stable_blocks
2764 -
2765 - req = _make_cp_request(num_blocks=2, virtual_block_size=block_size)
2766 - attention_blocks = _allocate_cp_gpu_blocks(
2767 - gpu_pool, req, 2, virtual_block_size=block_size, group_id=0
2768 - )
2769 - mamba_blocks = _allocate_cp_gpu_blocks(
2770 - gpu_pool, req, 2, virtual_block_size=block_size, group_id=1
2771 - )
2772 - kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks))
2773 - sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0)
2774 -
2775 - # build_connector_meta runs before the scheduler commits this step's
2776 - # tokens, so the positional scan picks the blocks up on the next step.
2777 - first = sched.build_connector_meta(
2778 - make_scheduler_output(
2779 - {req.request_id: 2 * block_size},
2780 - new_reqs={req.request_id: kv_blocks.get_block_ids()},
2781 - )
2782 - )
2783 - assert first.store_gpu_blocks == []
2784 - req.num_computed_tokens = 2 * block_size
2785 - meta = sched.build_connector_meta(
2786 - make_scheduler_output(
2787 - {req.request_id: 1},
2788 - cached_req_new_blocks={req.request_id: None},
2789 - )
2790 - )
2791 -
2792 - assert set(meta.store_gpu_blocks) == {
2793 - *(block.block_id for block in attention_blocks),
2794 - *(block.block_id for block in mamba_blocks),
2795 - }
2796 - assert sched._reqs_to_store[req.request_id].num_stored_blocks == [2, 2]
2797 -
2798 -
2750 @pytest.mark.parametrize("gone_via", ["preempted", "finished", "unregistered"])
2751 def test_boundary_handoff_dropped_for_departing_request(gone_via: str) -> None:
2752 """Handoffs for a request leaving this step must not be read.
tests/v1/worker/test_gpu_warmup_blocks.py
+2 -10
@@ -189,7 +189,7 @@ def test_mixed_warmup_reserves_lookahead_blocks():
189 _assert_covers_lookahead(recorder.steps, num_lookahead_tokens)
190
191
192 -@pytest.mark.parametrize("mamba_cache_mode", ["none", "all", "align"])
192 +@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
193 def test_warmup_reserves_mamba_speculative_blocks(mamba_cache_mode):
194 """Mamba groups hold the running-state block plus the speculative tail.
195
@@ -227,14 +227,7 @@ def test_warmup_reserves_mamba_speculative_blocks(mamba_cache_mode):
227
228
229 def _hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
230 - """Attention, circular, and Mamba groups exercise every reservation branch.
231 -
232 - "none" and "all" reach the same branch of both `_reserved_block_count` and
233 - `MambaManager`, which tests only for "align". They are still both listed:
234 - the mode is a spec-level input, and having the real manager confirm the
235 - prediction for each is what keeps a future divergence between them from
236 - landing unnoticed.
237 - """
230 + """Attention, circular, and Mamba groups exercise every reservation branch."""
231 return KVCacheConfig(
232 num_blocks=num_blocks,
233 kv_cache_tensors=[],
@@ -242,7 +235,6 @@ def _hybrid_kv_cache_config(num_blocks: int) -> KVCacheConfig:
235 _attention_group(),
236 _circular_group(),
237 _mamba_group("none"),
245 - _mamba_group("all"),
238 _mamba_group("align"),
239 ],
240 )
tests/v1/worker/test_mamba_utils.py
+1 -1
@@ -398,7 +398,7 @@ def _make_kv_cache_config(cfg: _TestConfig, layer_names: list[str]) -> KVCacheCo
398 (cfg.temporal_state_dim,),
399 ),
400 dtypes=(cfg.dtype, cfg.dtype),
401 - mamba_cache_mode="all",
401 + mamba_cache_mode="align",
402 )
403 group = KVCacheGroupSpec(
404 layer_names=layer_names,
vllm/config/cache.py
+1 -18
@@ -59,7 +59,7 @@ CacheDType = Literal[
59
60
61 MambaDType = Literal["auto", "float32", "float16", "bfloat16"]
62 -MambaCacheMode = Literal["all", "align", "none"]
62 +MambaCacheMode = Literal["align", "none"]
63 PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"]
64 KVOffloadingBackend = Literal["native", "lmcache"]
65
@@ -75,8 +75,6 @@ class CacheConfig:
75 Accepts None (meaning "use default"). After construction, always int."""
76 user_specified_block_size: bool = field(default=False, init=False)
77 """Whether block_size was explicitly provided. Derived automatically."""
78 - user_specified_mamba_block_size: bool = field(default=False, init=False)
79 - """Whether mamba_block_size was explicitly provided. Derived automatically."""
78 kv_cache_layout: str | None = field(default=None, init=False)
79 """Resolved physical KV cache layout name (a ``KVCacheLayout`` member).
80
@@ -191,7 +189,6 @@ class CacheConfig:
189 """The cache strategy for Mamba layers:
190
191 - "none": set when prefix caching is disabled.
194 - - "all": cache the mamba state of all tokens at position i * block_size.
192 - "align": only cache the mamba state of the last token of each scheduler step and
193 when the token is at position i * block_size. This is the default when prefix
194 caching is enabled.
@@ -293,7 +290,6 @@ class CacheConfig:
290 "mamba_page_size_padded",
291 "skip_page_size_padded",
292 "user_specified_block_size",
296 - "user_specified_mamba_block_size",
293 "_block_size_resolved",
294 # Post-init/derived counters
295 "num_gpu_blocks",
@@ -337,21 +333,8 @@ class CacheConfig:
333 self.block_size = self.DEFAULT_BLOCK_SIZE
334 else:
335 self.user_specified_block_size = True
340 - if self.mamba_block_size is not None:
341 - self.user_specified_mamba_block_size = True
336 return self
337
344 - @field_validator("mamba_cache_mode", mode="after")
345 - @classmethod
346 - def _validate_mamba_cache_mode(cls, mode: MambaCacheMode) -> MambaCacheMode:
347 - if mode == "all":
348 - logger.warning_once(
349 - "Mamba cache mode 'all' is deprecated and will be removed in an "
350 - "upcoming release. If this is a problem, please open an issue "
351 - "at https://github.com/vllm-project/vllm/issues."
352 - )
353 - return mode
354 -
338 @field_validator("cache_dtype", mode="after")
339 @classmethod
340 def _validate_cache_dtype(cls, cache_dtype: CacheDType) -> CacheDType:
vllm/config/vllm.py
-12
@@ -3091,9 +3091,6 @@ class VllmConfig:
3091 if self.parallel_config.use_ubatching:
3092 unsupported.extend(self._get_dbo_unsupported_features())
3093
3094 - if self.cache_config.mamba_cache_mode == "all":
3095 - unsupported.append("mamba cache mode 'all'")
3096 -
3094 return unsupported
3095
3096 def _get_v1_model_runner_unsupported_features(self) -> list[str]:
@@ -3499,10 +3496,6 @@ class VllmConfig:
3496 "SSM state caches, not --enable-mamba-cache-stochastic-"
3497 "rounding, which requires an explicit float16 cache"
3498 )
3502 - if self.cache_config.mamba_cache_mode not in ("none", "align"):
3503 - raise ValueError(
3504 - "RecoverSSM supports only none and align Mamba cache modes"
3505 - )
3499 if (
3500 self.cache_config.mamba_cache_mode == "align"
3501 and not self.use_v2_model_runner
@@ -3535,11 +3528,6 @@ class VllmConfig:
3528 "Mamba2 ReplaySSM speculative decoding requires "
3529 "--mamba-backend flashinfer"
3530 )
3538 - elif self.cache_config.mamba_cache_mode == "all":
3539 - raise ValueError(
3540 - "--use-replayssm supports prefix caching only in align mode; "
3541 - "pass --mamba-cache-mode align"
3542 - )
3531 elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER:
3532 if self.cache_config.mamba_cache_mode == "align":
3533 raise ValueError(
vllm/distributed/kv_transfer/kv_connector/utils.py
+3 -5
@@ -33,16 +33,14 @@ EngineId = str
33 BlockIds = tuple[list[int], ...] | list[list[int]]
34
35
36 -def clip_ssm_state_blocks(
37 - blocks: list[int], num_spec_blocks: int, positional: bool
38 -) -> list[int]:
39 - """Drop speculative scratch slots, retaining position-indexed or running state."""
36 +def clip_ssm_state_blocks(blocks: list[int], num_spec_blocks: int) -> list[int]:
37 + """Drop speculative scratch slots and keep only the running-state slot."""
38 if not blocks:
39 return blocks
40 # Keep at least the running-state slot, even before scratch allocation.
41 if num_scratch := min(num_spec_blocks, len(blocks) - 1):
42 blocks = blocks[:-num_scratch]
45 - return blocks if positional else blocks[-1:]
43 + return blocks[-1:]
44
45
46 def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None):
vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py
+2 -15
@@ -450,13 +450,6 @@ def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None:
450 "MoRIIO hybrid READ supports DSpark speculative decoding only, got "
451 f"method={speculative_config.method!r}"
452 )
453 - if vllm_config.cache_config.mamba_cache_mode == "all":
454 - # Positional allocation counts lookahead, so its local tail can exceed
455 - # the single running-state block supported by hybrid READ.
456 - raise MoRIIOError(
457 - "MoRIIO hybrid READ does not support DSpark with "
458 - "mamba_cache_mode='all'; use 'align'"
459 - )
453
454
455 def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None:
@@ -521,11 +514,6 @@ class MoRIIOConnectorScheduler:
514 "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "
515 "mode only; set kv_connector_extra_config.read_mode=true"
516 )
524 - # Only "all" mode keeps a state per block position; the other modes keep
525 - # a single running state in the last slot.
526 - self._ssm_state_slots_are_positional = (
527 - vllm_config.cache_config.mamba_cache_mode == "all"
528 - )
517 self._num_ssm_scratch_blocks = (
518 mamba_spec.num_speculative_blocks if mamba_spec is not None else 0
519 )
@@ -1288,8 +1276,8 @@ class MoRIIOConnectorScheduler:
1276 """Select transferable attention and Mamba block groups.
1277
1278 The wire payload stores the attention group first, followed by every
1291 - Mamba group in transfer-group order. Outside ``mamba_cache_mode="all"``,
1292 - only each group's running-state slot is transferred.
1279 + Mamba group in transfer-group order. Only each group's running-state
1280 + slot is transferred.
1281 """
1282 if not block_ids:
1283 return [], []
@@ -1301,7 +1289,6 @@ class MoRIIOConnectorScheduler:
1289 clip_ssm_state_blocks(
1290 list(transfer_block_ids[group_id]),
1291 self._num_ssm_scratch_blocks,
1304 - self._ssm_state_slots_are_positional,
1292 )
1293 for group_id in self._mamba_group_ids
1294 ]
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py
+4 -13
@@ -154,11 +154,6 @@ class NixlBaseConnectorScheduler:
154 else None
155 for g in kv_cache_config.transfer_groups
156 ]
157 - # Only "all" mode keeps a state per block position; the other modes
158 - # keep a single running state in the last non-speculative slot.
159 - self._ssm_state_slots_are_positional = (
160 - vllm_config.cache_config.mamba_cache_mode == "all"
161 - )
157
158 # Threshold to decide whether to compute kv cache locally
159 # or pull from a remote node: minimum number of remote
@@ -258,11 +253,9 @@ class NixlBaseConnectorScheduler:
253 out-of-window blocks only prior to the `request_finished_all_groups`
254 hook.
255
261 - SSM groups keep only their state-bearing slots: the trailing
262 - speculative scratch slots always go, and in single-state cache modes
263 - so does everything before the running state (null placeholders and
264 - the previous step's superseded state). "all" mode keeps its remaining
265 - slots, which the worker pairs position-wise.
256 + SSM groups keep only their state-bearing slot: the trailing
257 + speculative scratch slots and everything before the running state
258 + (null placeholders and the previous step's superseded state) go.
259
260 Use this at every block-id exchange point. Pass ``clip_ssm=False``
261 for per-step partial lists (host-buffer save), where the SSM strip
@@ -289,9 +282,7 @@ class NixlBaseConnectorScheduler:
282 and blocks
283 and (n_spec_blocks := self._ssm_spec_blocks[i]) is not None
284 ):
292 - blocks = clip_ssm_state_blocks(
293 - blocks, n_spec_blocks, self._ssm_state_slots_are_positional
294 - )
285 + blocks = clip_ssm_state_blocks(blocks, n_spec_blocks)
286 clipped.append(blocks)
287 return tuple(clipped)
288
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+6 -7
@@ -3510,13 +3510,12 @@ class NixlBaseConnectorWorker:
3510 if _is_ssm_spec(self._group_spec_types[i]):
3511 if num_decode_blocks == num_prefill_blocks:
3512 continue
3513 - # Only state-bearing slots reach here, single-state modes
3514 - # just one (see get_exchange_clipped_blocks), so differing
3515 - # counts mean position-indexed "all"-mode lists. A longer
3516 - # prefill list carries earlier positions the decode side
3517 - # already has (prefix hit) -> take its tail; a longer decode
3518 - # list holds the position D recomputes itself, which gets
3519 - # no prefill state.
3513 + # Only state-bearing slots reach here, normally just one
3514 + # (see get_exchange_clipped_blocks). A longer prefill list
3515 + # carries earlier positions the decode side already has
3516 + # (prefix hit) -> take its tail; a longer decode list holds
3517 + # the position D recomputes itself, which gets no prefill
3518 + # state.
3519 assert num_decode_blocks - num_prefill_blocks <= 1, (
3520 f"Group {i}: unpairable SSM state slots, "
3521 f"decode={num_decode_blocks} prefill={num_prefill_blocks}"
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+2 -5
@@ -156,7 +156,7 @@ def resolve_mamba_align_size(
156 """Scan all KV cache groups in *spec* and return the single mamba alignment
157 size, or None if no group requires mamba alignment.
158
159 - For MambaSpec groups in "align" or "all" cache mode the hit window must be
159 + For MambaSpec groups in "align" cache mode the hit window must be
160 rounded down to a multiple of the offloaded chunk size. Asserts that all
161 such groups agree on the same value.
162 """
@@ -164,10 +164,7 @@ def resolve_mamba_align_size(
164 for group in spec.config.groups:
165 tokens_per_block = group.tokens_per_block
166 kv_spec = kv_cache_config.kv_cache_groups[group.group_id].kv_cache_spec
167 - if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode in (
168 - "align",
169 - "all",
170 - ):
167 + if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align":
168 tokens_per_chunk = tokens_per_block * spec.blocks_per_chunk
169 assert mamba_align_size is None or mamba_align_size == tokens_per_chunk
170 mamba_align_size = tokens_per_chunk
vllm/model_executor/layers/mamba/mamba_mixer.py
+2 -55
@@ -260,8 +260,6 @@ class MambaMixer(MambaBase, PluggableLayer):
260 attn_metadata_raw = forward_context.attn_metadata
261
262 assert self.cache_config is not None
263 - mamba_block_size = self.cache_config.mamba_block_size
264 - is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
263
264 attn_metadata: AttentionMetadata | None = None
265 if attn_metadata_raw is not None:
@@ -297,8 +295,6 @@ class MambaMixer(MambaBase, PluggableLayer):
295
296 num_prefill_tokens = attn_metadata.num_prefill_tokens # token count
297 num_decode_tokens = attn_metadata.num_decode_tokens
300 - num_prefills = attn_metadata.num_prefills # request count
301 - num_decodes = attn_metadata.num_decode_tokens # token count (=request)
298 has_prefill = num_prefill_tokens > 0
299 has_decode = num_decode_tokens > 0
300 num_actual_tokens = num_prefill_tokens + num_decode_tokens
@@ -314,34 +310,6 @@ class MambaMixer(MambaBase, PluggableLayer):
310 gate_p = prefill_decode_split.gate_p
311 gate_d = prefill_decode_split.gate_d
312
317 - if is_mamba_cache_all:
318 - block_idx_last_computed_token_d, block_idx_last_computed_token_p = (
319 - torch.split(
320 - attn_metadata.block_idx_last_computed_token,
321 - [num_decodes, num_prefills],
322 - dim=0,
323 - )
324 - )
325 - block_idx_last_scheduled_token_d, block_idx_last_scheduled_token_p = (
326 - torch.split(
327 - attn_metadata.block_idx_last_scheduled_token,
328 - [num_decodes, num_prefills],
329 - dim=0,
330 - )
331 - )
332 -
333 - block_idx_first_scheduled_token_p = (
334 - attn_metadata.block_idx_first_scheduled_token_p
335 - )
336 - num_computed_tokens_p = attn_metadata.num_computed_tokens_p
337 - else:
338 - block_idx_last_computed_token_d = None
339 - block_idx_last_computed_token_p = None
340 - block_idx_last_scheduled_token_d = None
341 - block_idx_last_scheduled_token_p = None
342 - block_idx_first_scheduled_token_p = None
343 - num_computed_tokens_p = None
344 -
313 ssm_outputs = []
314
315 if has_prefill:
@@ -355,11 +323,6 @@ class MambaMixer(MambaBase, PluggableLayer):
323 has_initial_state=has_initial_states_p,
324 cache_indices=state_indices_tensor_p,
325 query_start_loc=query_start_loc_p,
358 - block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,
359 - block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,
360 - initial_state_idx=block_idx_last_computed_token_p,
361 - num_computed_tokens=num_computed_tokens_p,
362 - block_size_to_align=mamba_block_size,
326 metadata=attn_metadata,
327 )
328 # 3. State Space Model sequence transformations.
@@ -383,10 +346,6 @@ class MambaMixer(MambaBase, PluggableLayer):
346 cache_indices=state_indices_tensor_p,
347 has_initial_state=has_initial_states_p,
348 query_start_loc=query_start_loc_p,
386 - block_size=mamba_block_size,
387 - block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,
388 - block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,
389 - initial_state_idx=block_idx_last_computed_token_p,
349 cu_chunk_seqlen=cu_chunk_seqlen_p,
350 last_chunk_indices=last_chunk_indices_p,
351 )
@@ -396,16 +355,6 @@ class MambaMixer(MambaBase, PluggableLayer):
355 # state_indices_tensor_d is assigned when attn_metadata is not None,
356 # and has_decode is only True when attn_metadata is not None
357 assert state_indices_tensor_d is not None
399 - if is_mamba_cache_all:
400 - state_indices_tensor_d_input = state_indices_tensor_d.gather(
401 - 1, block_idx_last_computed_token_d.unsqueeze(1)
402 - ).squeeze(1)
403 - state_indices_tensor_d_output = state_indices_tensor_d.gather(
404 - 1, block_idx_last_scheduled_token_d.unsqueeze(1)
405 - ).squeeze(1)
406 - else:
407 - state_indices_tensor_d_input = state_indices_tensor_d
408 - state_indices_tensor_d_output = state_indices_tensor_d
358 # 2. Convolution sequence transformation
359 conv_out_d = causal_conv1d_update(
360 hidden_states_BC_d.transpose(0, 1),
@@ -414,8 +363,6 @@ class MambaMixer(MambaBase, PluggableLayer):
363 self.conv1d.bias,
364 self.activation,
365 conv_state_indices=state_indices_tensor_d,
417 - block_idx_last_scheduled_token=block_idx_last_scheduled_token_d,
418 - initial_state_idx=block_idx_last_computed_token_d,
366 ).transpose(0, 1)
367
368 # 3. State Space Model sequence transformation.
@@ -437,8 +384,8 @@ class MambaMixer(MambaBase, PluggableLayer):
384 time_proj_bias,
385 z=gate_d.transpose(0, 1),
386 dt_softplus=True,
440 - state_batch_indices=state_indices_tensor_d_input,
441 - dst_state_batch_indices=state_indices_tensor_d_output,
387 + state_batch_indices=state_indices_tensor_d,
388 + dst_state_batch_indices=state_indices_tensor_d,
389 out=scan_outputs_d,
390 )
391 scan_outputs_d = scan_outputs_d.transpose(0, 1)
vllm/model_executor/layers/mamba/mamba_mixer2.py
+17 -182
@@ -549,12 +549,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
549 self._updates_replayssm_trackers = True
550
551 self.num_spec = vllm_config.num_speculative_tokens
552 - if self.num_spec > 0:
553 - self.register_buffer(
554 - "_decode_state_offsets",
555 - torch.arange(1 + self.num_spec, dtype=torch.int32).unsqueeze(0),
556 - persistent=False,
557 - )
552
553 # Pre-compute sizes for forward pass
554 self.tped_intermediate_size = self.intermediate_size // self.tp_size
@@ -732,8 +726,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
726 attn_metadata_raw = forward_context.attn_metadata
727
728 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"
729 use_spec_decode = self.num_spec > 0
730 ring_start = prev_num_accepted = prev_query_len = None
731
@@ -803,45 +795,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
795 dim=0,
796 )
797
806 - if is_mamba_cache_all:
807 - # If prefix caching is enabled, retrieve the relevant variables
808 - # for prefill and decode
809 - block_idx_last_computed_token_d, block_idx_last_computed_token_p = (
810 - torch.split(
811 - attn_metadata.block_idx_last_computed_token,
812 - [num_decodes, num_prefills],
813 - dim=0,
814 - )
815 - )
816 - block_idx_last_scheduled_token_d, block_idx_last_scheduled_token_p = (
817 - torch.split(
818 - attn_metadata.block_idx_last_scheduled_token,
819 - [num_decodes, num_prefills],
820 - dim=0,
821 - )
822 - )
823 - if attn_metadata.block_idx_last_scheduled_token_prev_step is not None:
824 - block_idx_last_scheduled_token_prev_step_d, _ = torch.split(
825 - attn_metadata.block_idx_last_scheduled_token_prev_step,
826 - [num_decodes, num_prefills],
827 - dim=0,
828 - )
829 - else:
830 - block_idx_last_scheduled_token_prev_step_d = None
831 - # Prefill-only variables:
832 - block_idx_first_scheduled_token_p = (
833 - attn_metadata.block_idx_first_scheduled_token_p
834 - )
835 - num_computed_tokens_p = attn_metadata.num_computed_tokens_p
836 - else:
837 - block_idx_last_computed_token_p = None
838 - block_idx_last_scheduled_token_p = None
839 - block_idx_first_scheduled_token_p = None
840 - block_idx_last_scheduled_token_d = None
841 - block_idx_last_computed_token_d = None
842 - block_idx_last_scheduled_token_prev_step_d = None
843 - num_computed_tokens_p = None
844 -
798 preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
799 output[:num_actual_tokens],
800 [num_decode_tokens, num_prefill_tokens],
@@ -858,11 +811,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
811 # to by "state_indices_tensor_p".
812 # In particular, it will always write the state at the
813 # sequence end.
861 - # In addition, "block_idx_first_scheduled_token_p" and
862 - # "block_idx_last_scheduled_token_p"
863 - # are provided (which are pointers into
864 - # "state_indices_tensor_p"), it will write additional cache
865 - # states aligned at "block_size_to_align".
814 x = hidden_states_B_C_p.transpose(
815 0, 1
816 ) # this is the form that causal-conv see
@@ -874,11 +822,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
822 conv_states=conv_state,
823 has_initial_state=has_initial_states_p,
824 cache_indices=state_indices_tensor_p,
877 - block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,
878 - block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,
879 - initial_state_idx=block_idx_last_computed_token_p,
880 - num_computed_tokens=num_computed_tokens_p,
881 - block_size_to_align=mamba_block_size,
825 metadata=attn_metadata,
826 query_start_loc=query_start_loc_p,
827 ).transpose(0, 1)[:num_prefill_tokens]
@@ -891,14 +834,9 @@ class MambaMixer2(MambaBase, PluggableLayer):
834 initial_states = None
835 if has_initial_states_p is not None and prep_initial_states:
836 assert state_indices_tensor_p is not None
894 - kernel_ssm_indices = state_indices_tensor_p
895 - if is_mamba_cache_all:
896 - kernel_ssm_indices = state_indices_tensor_p.gather(
897 - 1, block_idx_last_computed_token_p.unsqueeze(1)
898 - ).squeeze(1)
837 initial_states = torch.where(
838 has_initial_states_p[:, None, None, None],
901 - ssm_state[kernel_ssm_indices],
839 + ssm_state[state_indices_tensor_p],
840 0,
841 )
842
@@ -921,130 +859,29 @@ class MambaMixer2(MambaBase, PluggableLayer):
859 cu_chunk_seqlens=cu_chunk_seqlen_p,
860 last_chunk_indices=last_chunk_indices_p,
861 initial_states=initial_states,
924 - return_intermediate_states=is_mamba_cache_all,
862 dt_softplus=True,
863 dt_limit=(0.0, float("inf")),
864 out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim),
865 state_dtype=ssm_state.dtype,
866 )
867
931 - if is_mamba_cache_all:
932 - assert mamba_block_size is not None
933 - assert state_indices_tensor_p is not None
934 - assert block_idx_first_scheduled_token_p is not None
935 - assert block_idx_last_scheduled_token_p is not None
936 - assert last_chunk_indices_p is not None
937 - assert num_computed_tokens_p is not None
938 -
939 - # The chunk_stride is the number of chunks per mamba block
940 - # e.g., if mamba_block_size = 512 and chunk_size = 256,
941 - # then chunk_stride = 2
942 - chunk_stride = mamba_block_size // chunk_size
943 -
944 - # Save state for sequences with more than just final state,
945 - # batched over sequences so that the number of kernel launches
946 - # does not scale with the number of prefills.
947 - n_blocks_to_fill = (
948 - block_idx_last_scheduled_token_p - block_idx_first_scheduled_token_p
949 - ).clamp(min=0)
950 - total_fill = attn_metadata.num_state_writes_p
951 - if total_fill > 0:
952 - first_chunk = torch.zeros_like(last_chunk_indices_p)
953 - first_chunk[1:] = last_chunk_indices_p[:-1] + 1
954 -
955 - # The chunk that completes this sequence's first mamba block.
956 - # A block spans chunk_stride chunks. When the sequence resumed
957 - # mid-block, the chunks before this step are not in varlen_states,
958 - # so subtract them: the boundary sits earlier by their count.
959 - num_unaligned_computed_tokens = (
960 - num_computed_tokens_p % mamba_block_size
961 - )
962 - unaligned_chunk_shift = num_unaligned_computed_tokens // chunk_size
963 - first_aligned_chunk = (
964 - first_chunk + (chunk_stride - 1) - unaligned_chunk_shift
965 - )
966 -
967 - # Flatten all writes into one list: fill_counts=[2,1,2]
968 - # -> seq_ids=[0,0,1,2,2], write_offsets=[0,1,0,0,1]
969 - fill_counts = n_blocks_to_fill.long()
970 - # output_size, or repeat_interleave reads the size off the device
971 - seq_ids = torch.repeat_interleave(
972 - torch.arange(num_prefills, device=ssm_state.device),
973 - fill_counts,
974 - output_size=total_fill,
975 - )
976 - write_starts = torch.cumsum(fill_counts, dim=0) - fill_counts
977 - write_offsets = torch.arange(
978 - total_fill, device=ssm_state.device
979 - ) - torch.repeat_interleave(
980 - write_starts, fill_counts, output_size=total_fill
981 - )
982 -
983 - cache_blocks_to_fill = state_indices_tensor_p[
984 - seq_ids,
985 - block_idx_first_scheduled_token_p[seq_ids] + write_offsets,
986 - ]
987 - # Write the states. Each row goes to its own cache block,
988 - # so the indexed copy has no duplicate destinations.
989 - ssm_state[cache_blocks_to_fill] = varlen_states[
990 - first_aligned_chunk[seq_ids] + write_offsets * chunk_stride
991 - ]
992 -
993 - # For all seqs, store the last state (note: might be partial):
994 - assert state_indices_tensor_p is not None
995 - ssm_state[
996 - state_indices_tensor_p.gather(
997 - 1, block_idx_last_scheduled_token_p.unsqueeze(1)
998 - ).squeeze(1)
999 - ] = varlen_states[last_chunk_indices_p]
1000 -
1001 - else:
1002 - # update ssm states
1003 - # - varlen state is a (num_prefills, nheads, headdim, dstate)
1004 - # tensor
1005 - assert state_indices_tensor_p is not None
1006 - ssm_state[state_indices_tensor_p] = varlen_states
1007 - if ring_start is not None and self._updates_replayssm_trackers:
1008 - assert prev_num_accepted is not None
1009 - reset_replayssm_ring_trackers(
1010 - ring_start,
1011 - prev_num_accepted,
1012 - prev_query_len,
1013 - state_indices_tensor_p,
1014 - )
868 + # update ssm states
869 + # - varlen state is a (num_prefills, nheads, headdim, dstate)
870 + # tensor
871 + assert state_indices_tensor_p is not None
872 + ssm_state[state_indices_tensor_p] = varlen_states
873 + if ring_start is not None and self._updates_replayssm_trackers:
874 + assert prev_num_accepted is not None
875 + reset_replayssm_ring_trackers(
876 + ring_start,
877 + prev_num_accepted,
878 + prev_query_len,
879 + state_indices_tensor_p,
880 + )
881
882 # Process decode requests
883 if has_decode:
884 assert state_indices_tensor_d is not None
1019 - if is_mamba_cache_all:
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)
1024 - + self._decode_state_offsets
1025 - )
1026 - output_indices = (
1027 - block_idx_last_scheduled_token_d.unsqueeze(1)
1028 - + self._decode_state_offsets
1029 - )
1030 - state_indices_tensor_d_input = state_indices_tensor_d.gather(
1031 - 1, input_indices
1032 - )
1033 - state_indices_tensor_d_output = state_indices_tensor_d.gather(
1034 - 1, output_indices
1035 - )
1036 - else:
1037 - state_indices_tensor_d_input = state_indices_tensor_d.gather(
1038 - 1, block_idx_last_computed_token_d.unsqueeze(1)
1039 - ).squeeze(1)
1040 - state_indices_tensor_d_output = state_indices_tensor_d.gather(
1041 - 1, block_idx_last_scheduled_token_d.unsqueeze(1)
1042 - ).squeeze(1)
1043 - else:
1044 - # Without caching, read and write in-place to the same blocks:
1045 - state_indices_tensor_d_input = state_indices_tensor_d
1046 - state_indices_tensor_d_output = state_indices_tensor_d
1047 -
885 if (
886 self.use_replayssm
887 and self.mamba_config.backend == MambaBackendEnum.FLASHINFER
@@ -1081,8 +918,6 @@ class MambaMixer2(MambaBase, PluggableLayer):
918 self.conv1d.bias,
919 self.activation,
920 conv_state_indices=state_indices_tensor_d,
1084 - block_idx_last_scheduled_token=block_idx_last_scheduled_token_d,
1085 - initial_state_idx=block_idx_last_computed_token_d,
921 num_accepted_tokens=num_accepted_tokens,
922 query_start_loc=query_start_loc_d,
923 # ReplaySSM keeps one physical state block while a speculative
@@ -1207,7 +1042,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
1042 write_pos=attn_metadata.write_pos_d,
1043 is_flush=attn_metadata.is_flush_d,
1044 max_cache_len=self.replayssm_buffer_len,
1210 - state_batch_indices=state_indices_tensor_d_input,
1045 + state_batch_indices=state_indices_tensor_d,
1046 out=preallocated_ssm_out_d,
1047 enable_stochastic_rounding=(
1048 self.mamba_config.enable_stochastic_rounding
@@ -1227,8 +1062,8 @@ class MambaMixer2(MambaBase, PluggableLayer):
1062 D_d,
1063 dt_bias,
1064 dt_softplus=True,
1230 - state_batch_indices=state_indices_tensor_d_input,
1231 - dst_state_batch_indices=state_indices_tensor_d_output,
1065 + state_batch_indices=state_indices_tensor_d,
1066 + dst_state_batch_indices=state_indices_tensor_d,
1067 out=preallocated_ssm_out_d,
1068 num_accepted_tokens=num_accepted_tokens,
1069 cu_seqlens=query_start_loc_d,
vllm/model_executor/models/config.py
+3 -14
@@ -646,20 +646,9 @@ class MambaModelConfig(VerifyAndUpdateConfig):
646 cache_config.mamba_cache_mode,
647 model_config.architecture,
648 )
649 - if (
650 - cache_config.mamba_cache_mode == "all"
651 - and not model_config.supports_mamba_prefix_caching
652 - ):
653 - cache_config.mamba_cache_mode = "align"
654 - logger.warning(
655 - "Hybrid or mamba-based model detected without support "
656 - "for prefix caching with Mamba cache 'all' mode: "
657 - "falling back to 'align' mode."
658 - )
659 - if cache_config.mamba_cache_mode == "align":
660 - assert vllm_config.scheduler_config.enable_chunked_prefill, (
661 - "Chunked prefill is required for mamba cache mode 'align'."
662 - )
649 + assert vllm_config.scheduler_config.enable_chunked_prefill, (
650 + "Chunked prefill is required for mamba cache mode 'align'."
651 + )
652 # By default, mamba block size will be set to max_model_len (see
653 # below). When enabling prefix caching, we align mamba block size
654 # to the block size as the basic granularity for prefix caching.
vllm/model_executor/models/interns2_mobius.py
-6
@@ -418,12 +418,6 @@ class InternS2MobiusForCausalLM(Qwen3_5ForCausalLMBase):
418 self.quant_config = vllm_config.quant_config
419 self.scheduler_config = vllm_config.scheduler_config
420
421 - if vllm_config.cache_config.mamba_cache_mode == "all":
422 - raise NotImplementedError(
423 - "Intern-S2-Mobius does not support 'all' prefix caching; "
424 - "use '--mamba-cache-mode=align' instead."
425 - )
426 -
421 nn.Module.__init__(self)
422 self.config = config
423 self.model = InternS2MobiusModel(
vllm/model_executor/models/lfm2.py
-7
@@ -458,13 +458,6 @@ class Lfm2ForCausalLM(
458 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
459 config = vllm_config.model_config.hf_config
460 quant_config = vllm_config.quant_config
461 - cache_config = vllm_config.cache_config
462 - if cache_config.mamba_cache_mode == "all":
463 - raise NotImplementedError(
464 - "Lfm2 currently does not support 'all' prefix caching, "
465 - "please use '--mamba-cache-mode=align' instead"
466 - )
467 -
461 super().__init__()
462 self.config = config
463 self.model = Lfm2Model(
vllm/model_executor/models/lfm2_moe.py
-7
@@ -573,13 +573,6 @@ class Lfm2MoeForCausalLM(
573 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
574 config = vllm_config.model_config.hf_config
575 quant_config = vllm_config.quant_config
576 - cache_config = vllm_config.cache_config
577 -
578 - if cache_config.mamba_cache_mode == "all":
579 - raise NotImplementedError(
580 - "Lfm2Moe currently does not support 'all' prefix caching, "
581 - "please use '--mamba-cache-mode=align' instead"
582 - )
576
577 super().__init__()
578 self.config = config
vllm/model_executor/models/qwen3_5.py
-6
@@ -329,14 +329,8 @@ class Qwen3_5ForCausalLMBase(
329 config = vllm_config.model_config.hf_text_config
330 self.vllm_config = vllm_config
331 self.model_config = vllm_config.model_config
332 - cache_config = vllm_config.cache_config
332
333 scheduler_config = vllm_config.scheduler_config
335 - if cache_config.mamba_cache_mode == "all":
336 - raise NotImplementedError(
337 - "Qwen3.5 currently does not support 'all' prefix caching, "
338 - "please use '--mamba-cache-mode=align' instead"
339 - )
334 self.quant_config = vllm_config.quant_config
335
336 super().__init__()
vllm/model_executor/models/qwen3_5_mtp.py
-7
@@ -227,13 +227,6 @@ class Qwen3_5MTP(LocalArgmaxMixin, nn.Module, SupportsMultiModal, SupportsPP):
227 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
228 config = vllm_config.model_config.hf_text_config
229 self.vllm_config = vllm_config
230 - cache_config = vllm_config.cache_config
231 - if cache_config.mamba_cache_mode == "all":
232 - raise NotImplementedError(
233 - "Qwen3_5MTP currently does not support 'all' prefix caching, "
234 - "please use '--mamba-cache-mode=align' instead"
235 - )
236 -
230 self.quant_config = vllm_config.quant_config
231
232 super().__init__()
vllm/model_executor/models/qwen3_next.py
-6
@@ -828,14 +828,8 @@ class Qwen3NextForCausalLM(
828 config = vllm_config.model_config.hf_text_config
829 self.vllm_config = vllm_config
830 self.model_config = vllm_config.model_config
831 - cache_config = vllm_config.cache_config
831
832 scheduler_config = vllm_config.scheduler_config
834 - if cache_config.mamba_cache_mode == "all":
835 - raise NotImplementedError(
836 - "Qwen3Next currently does not support 'all' prefix caching, "
837 - "please use '--mamba-cache-mode=align' instead"
838 - )
833 self.quant_config = vllm_config.quant_config
834
835 super().__init__()
vllm/model_executor/models/qwen3_next_mtp.py
-7
@@ -187,13 +187,6 @@ class Qwen3NextMTP(nn.Module, QwenNextMixtureOfExperts):
187 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
188 config = vllm_config.model_config.hf_config
189 self.vllm_config = vllm_config
190 - cache_config = vllm_config.cache_config
191 - if cache_config.mamba_cache_mode == "all":
192 - raise NotImplementedError(
193 - "Qwen3NextMTP currently does not support 'all' prefix caching, "
194 - "please use '--mamba-cache-mode=align' instead"
195 - )
196 -
190 self.quant_config = vllm_config.quant_config
191
192 super().__init__()
vllm/models/kimi_k3/nvidia/kda_metadata.py
+1 -1
@@ -104,7 +104,7 @@ def _mamba_get_block_table_tensor(
104 kv_cache_spec: MambaSpec,
105 mamba_cache_mode: str,
106 ) -> torch.Tensor:
107 - if mamba_cache_mode in ("all", "none"):
107 + if mamba_cache_mode == "none":
108 return block_table
109
110 assert block_table.is_cuda and seq_lens.is_cuda
vllm/models/qwen4_exp/amd/model.py
-5
@@ -638,11 +638,6 @@ class Qwen4ExpForCausalLM(
638 self.quant_config = vllm_config.quant_config
639 self.config = config
640 self.scheduler_config = vllm_config.scheduler_config
641 - if vllm_config.cache_config.mamba_cache_mode == "all":
642 - raise NotImplementedError(
643 - "Qwen4Exp currently does not support 'all' prefix caching, "
644 - "please use '--mamba-cache-mode=align' instead"
645 - )
641 self.model = Qwen4ExpModel(
642 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
643 )
vllm/models/qwen4_exp/amd/mtp.py
-7
@@ -395,13 +395,6 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts):
395 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
396 config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
397 self.vllm_config = vllm_config
398 - cache_config = vllm_config.cache_config
399 - if cache_config.mamba_cache_mode == "all":
400 - raise NotImplementedError(
401 - "Qwen4ExpMTP currently does not support 'all' prefix caching, "
402 - "please use '--mamba-cache-mode=align' instead"
403 - )
404 -
398 self.quant_config = vllm_config.quant_config
399
400 super().__init__()
vllm/models/qwen4_exp/nvidia/model.py
-5
@@ -669,11 +669,6 @@ class Qwen4ExpForCausalLM(
669 self.quant_config = vllm_config.quant_config
670 self.config = config
671 self.scheduler_config = vllm_config.scheduler_config
672 - if vllm_config.cache_config.mamba_cache_mode == "all":
673 - raise NotImplementedError(
674 - "Qwen4Exp currently does not support 'all' prefix caching, "
675 - "please use '--mamba-cache-mode=align' instead"
676 - )
672 self.model = Qwen4ExpModel(
673 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
674 )
vllm/models/qwen4_exp/nvidia/mtp.py
-7
@@ -375,13 +375,6 @@ class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts):
375 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
376 config: Qwen4ExpTextConfig = vllm_config.model_config.hf_text_config
377 self.vllm_config = vllm_config
378 - cache_config = vllm_config.cache_config
379 - if cache_config.mamba_cache_mode == "all":
380 - raise NotImplementedError(
381 - "Qwen4ExpMTP currently does not support 'all' prefix caching, "
382 - "please use '--mamba-cache-mode=align' instead"
383 - )
384 -
378 self.quant_config = vllm_config.quant_config
379
380 super().__init__()
vllm/platforms/interface.py
+9 -28
@@ -939,13 +939,6 @@ class Platform:
939 if mamba_page_size == 0:
940 return
941
942 - # mamba_block_size here should either be user specified value or None
943 - mamba_block_size = (
944 - cache_config.mamba_block_size
945 - if cache_config.user_specified_mamba_block_size
946 - else None
947 - )
948 -
942 # Get kernel block alignment from the backend's supported sizes
943 with set_current_vllm_config(vllm_config):
944 kernel_block_alignment_size = max(
@@ -962,27 +955,15 @@ class Platform:
955 # multiple of 128 so split kernel blocks keep that invariant.
956 kernel_block_alignment_size = max(kernel_block_alignment_size, 128)
957
965 - if cache_config.mamba_cache_mode == "all":
966 - # With prefix caching, align to mamba chunk size for kernel perf
967 - # TODO(tdoublep): this constraint can be relaxed fairly
968 - # easily by changing the way we layout chunks in the
969 - # mamba2 kernels.
970 - base_chunk_size = mamba_block_size or model_config.get_mamba_chunk_size()
971 - assert base_chunk_size is not None
972 - attn_tokens_per_mamba_state = cdiv(mamba_page_size, attn_page_size_1_token)
973 - chunk_size = lcm(base_chunk_size, kernel_block_alignment_size)
974 - attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
975 - cache_config.mamba_block_size = attn_block_size
976 - else:
977 - # Without prefix caching, use minimum block size that satisfies
978 - # both backend alignment and mamba page size compatibility
979 - attn_block_size = kernel_block_alignment_size * cdiv(
980 - mamba_page_size,
981 - kernel_block_alignment_size * attn_page_size_1_token,
982 - )
983 - indexer_align = cls._get_indexer_block_alignment(vllm_config)
984 - if indexer_align:
985 - attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align)
958 + # Use minimum block size that satisfies both backend alignment and
959 + # mamba page size compatibility
960 + attn_block_size = kernel_block_alignment_size * cdiv(
961 + mamba_page_size,
962 + kernel_block_alignment_size * attn_page_size_1_token,
963 + )
964 + indexer_align = cls._get_indexer_block_alignment(vllm_config)
965 + if indexer_align:
966 + attn_block_size = indexer_align * cdiv(attn_block_size, indexer_align)
967
968 if cache_config.block_size < attn_block_size:
969 cache_config.block_size = attn_block_size
vllm/v1/attention/backends/mamba1_attn.py
+2 -21
@@ -1,7 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from dataclasses import dataclass, replace
4 +from dataclasses import dataclass
5 from typing import Any
6
7 from vllm.v1.attention.backend import AttentionBackend, CommonAttentionMetadata
@@ -42,23 +42,4 @@ class Mamba1AttentionMetadataBuilder(
42 fast_build: bool = False,
43 **kwargs: Any,
44 ) -> Mamba1AttentionMetadata:
45 - common = self._compute_common_metadata(common_attn_metadata)
46 -
47 - if (
48 - common.num_prefills > 0
49 - and self.vllm_config.cache_config.mamba_cache_mode == "all"
50 - ):
51 - cu_chunk_seqlen_p, _, last_chunk_indices_p = (
52 - self._build_chunk_metadata_tensors(
53 - self.kv_cache_spec.block_size,
54 - common,
55 - common_attn_metadata,
56 - )
57 - )
58 - return replace(
59 - common,
60 - cu_chunk_seqlen_p=cu_chunk_seqlen_p,
61 - last_chunk_indices_p=last_chunk_indices_p,
62 - )
63 -
64 - return common
45 + return self._compute_common_metadata(common_attn_metadata)
vllm/v1/attention/backends/mamba2_attn.py
-1
@@ -139,7 +139,6 @@ class Mamba2AttentionMetadataBuilder(
139 common = self._compute_common_metadata(
140 common_attn_metadata,
141 num_accepted_tokens=kwargs.get("num_accepted_tokens"),
142 - prev_last_scheduled_idx=kwargs.get("prev_last_scheduled_idx"),
142 num_decode_draft_tokens_cpu=kwargs.get("num_decode_draft_tokens_cpu"),
143 )
144
vllm/v1/attention/backends/mamba_attn.py
+17 -250
@@ -39,7 +39,6 @@ class BaseMambaAttentionMetadata:
39 # the batch has no prefill requests.
40 has_initial_states_p: torch.Tensor | None
41 query_start_loc_p: torch.Tensor | None
42 - num_computed_tokens_p: torch.Tensor | None
42 state_indices_tensor_p: torch.Tensor | None
43
44 # The following tensors are used for decode requests and
@@ -52,13 +51,6 @@ class BaseMambaAttentionMetadata:
51 # Includes the bonus token (so minimum is 1)
52 num_accepted_tokens: torch.Tensor | None # shape: [batch,]
53
55 - # The following tensors are only used for prefix caching in all mode and
56 - # are None if disabled
57 - block_idx_last_scheduled_token: torch.Tensor | None
58 - block_idx_first_scheduled_token_p: torch.Tensor | None
59 - block_idx_last_computed_token: torch.Tensor | None
60 - block_idx_last_scheduled_token_prev_step: torch.Tensor | None
61 -
54 # The following tensor is only used for prefix caching in align mode
55 seq_lens: torch.Tensor
56
@@ -70,9 +62,6 @@ class BaseMambaAttentionMetadata:
62 # last_chunk_indices_p is a tensor of shape (batch,) that contains the
63 # index of the last chunk for every sequence in the (prefill) batch.
64 last_chunk_indices_p: torch.Tensor | None = None
73 - # Number of block-aligned SSM states the prefill path writes, i.e.
74 - # sum(block_idx_last_scheduled_token_p - block_idx_first_scheduled_token_p).
75 - num_state_writes_p: int = 0
65
66 # The following attributes are for triton implementation of causal_conv1d
67 nums_dict: dict | None = None
@@ -130,47 +119,11 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
119 self.compilation_config.max_cudagraph_capture_size,
120 )
121
133 - if self.vllm_config.cache_config.mamba_cache_mode == "all":
134 - max_num_blocks = (
135 - cdiv(
136 - self.vllm_config.model_config.max_model_len,
137 - kv_cache_spec.block_size,
138 - )
139 - + kv_cache_spec.num_speculative_blocks
140 - )
141 - # TODO: reduce this size as needed for decode-only cudagraph capture
142 - self.state_indices_tensor_d: torch.Tensor = torch.empty(
143 - (
144 - self.decode_cudagraph_max_bs,
145 - max_num_blocks,
146 - ),
147 - dtype=torch.int32,
148 - device=device,
149 - )
150 - self.block_idx_last_scheduled_token: torch.Tensor = torch.empty(
151 - (self.decode_cudagraph_max_bs,),
152 - dtype=torch.int32,
153 - device=device,
154 - )
155 - self.block_idx_last_computed_token: torch.Tensor = torch.empty(
156 - (self.decode_cudagraph_max_bs,),
157 - dtype=torch.int32,
158 - device=device,
159 - )
160 - if self.use_spec_decode:
161 - self.block_idx_last_scheduled_token_prev_step: torch.Tensor = (
162 - torch.empty(
163 - (self.decode_cudagraph_max_bs,),
164 - dtype=torch.int32,
165 - device=device,
166 - )
167 - )
168 - else:
169 - self.state_indices_tensor_d = torch.empty(
170 - (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens),
171 - dtype=torch.int32,
172 - device=device,
173 - )
122 + self.state_indices_tensor_d = torch.empty(
123 + (self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens),
124 + dtype=torch.int32,
125 + device=device,
126 + )
127
128 # For speculative decoding, we need to store the following buffers
129 # for CUDA graph capture during decode
@@ -259,23 +212,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
212 if self.num_spec_tokens > 0:
213 num_accepted_tokens = torch.diff(m.query_start_loc)
214
262 - prev_last_scheduled_idx = None
263 - if (
264 - self.use_spec_decode
265 - and self.vllm_config.cache_config.mamba_cache_mode == "all"
266 - ):
267 - prev_last_scheduled_idx = torch.zeros(
268 - (m.num_reqs,),
269 - dtype=torch.int32,
270 - device=m.query_start_loc.device,
271 - )
272 -
273 - return self.build(
274 - 0,
275 - m,
276 - num_accepted_tokens=num_accepted_tokens,
277 - prev_last_scheduled_idx=prev_last_scheduled_idx,
278 - )
215 + return self.build(0, m, num_accepted_tokens=num_accepted_tokens)
216
217 def build(
218 self,
@@ -284,7 +221,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
221 fast_build: bool = False,
222 *,
223 num_accepted_tokens: torch.Tensor | None = None,
287 - prev_last_scheduled_idx: torch.Tensor | None = None,
224 num_decode_draft_tokens_cpu: torch.Tensor | None = None,
225 **kwargs: Any,
226 ) -> M:
@@ -294,7 +230,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
230 return self._compute_common_metadata(
231 common_attn_metadata,
232 num_accepted_tokens=num_accepted_tokens,
297 - prev_last_scheduled_idx=prev_last_scheduled_idx,
233 num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
234 )
235
@@ -423,56 +358,11 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
358 )
359 return cu_chunk_seqlen_p, seq_idx_p, last_chunk_indices_p
360
426 - def _compute_prefix_caching_block_indices(
427 - self,
428 - common_attn_metadata: CommonAttentionMetadata,
429 - mamba_block_size: int,
430 - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
431 - num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
432 - # Block index of the last computed token
433 - block_idx_last_computed_token = (
434 - torch.div(
435 - num_computed_tokens + mamba_block_size - 1,
436 - mamba_block_size,
437 - rounding_mode="floor",
438 - )
439 - - 1
440 - )
441 - # which is <= block index for the first scheduled token
442 - block_idx_first_scheduled_token = (
443 - torch.div(
444 - num_computed_tokens + mamba_block_size,
445 - mamba_block_size,
446 - rounding_mode="floor",
447 - )
448 - - 1
449 - )
450 - # which is <= block index of the last scheduled token
451 - block_idx_last_scheduled_token = (
452 - torch.div(
453 - common_attn_metadata.seq_lens + mamba_block_size - 1,
454 - mamba_block_size,
455 - rounding_mode="floor",
456 - )
457 - - 1
458 - )
459 - # -1 in case it's non-computed and causes later issues with indexing
460 - block_idx_last_computed_token.clamp_(min=0)
461 - # -1 in the case we have a padded request (0 seq-len)
462 - block_idx_last_scheduled_token.clamp_(min=0)
463 -
464 - return (
465 - block_idx_last_computed_token,
466 - block_idx_first_scheduled_token,
467 - block_idx_last_scheduled_token,
468 - )
469 -
361 def _compute_common_metadata(
362 self,
363 common_attn_metadata: CommonAttentionMetadata,
364 *,
365 num_accepted_tokens: torch.Tensor | None = None,
475 - prev_last_scheduled_idx: torch.Tensor | None = None,
366 num_decode_draft_tokens_cpu: torch.Tensor | None = None,
367 ) -> M:
368 """Compute metadata common to both Mamba1 and Mamba2."""
@@ -540,16 +430,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
430 has_initial_states_p = None
431 query_start_loc_p = None
432 query_start_loc_d = None
543 - num_computed_tokens = None
544 - num_computed_tokens_p = None
545 -
546 - # for prefix caching
547 - block_idx_first_scheduled_token = None
548 - block_idx_first_scheduled_token_p = None
549 - block_idx_last_computed_token = None
550 - block_idx_last_scheduled_token = None
551 - block_idx_last_scheduled_token_prev_step = None
552 - num_state_writes_p = 0
433
434 # for causal_conv1d
435 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
@@ -557,35 +437,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
437 is_flush_d = None
438 replayssm_scratch = None
439
560 - if self.vllm_config.cache_config.mamba_cache_mode == "all":
561 - num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
562 -
563 - # Return a tensor of shape (#requests, #max blocks)
564 - state_indices_tensor = common_attn_metadata.block_table_tensor
565 - # Additional cache-related variables:
566 - mamba_block_size = self.kv_cache_spec.block_size
567 - (
568 - block_idx_last_computed_token,
569 - block_idx_first_scheduled_token,
570 - block_idx_last_scheduled_token,
571 - ) = self._compute_prefix_caching_block_indices(
572 - common_attn_metadata, mamba_block_size
573 - )
574 - if self.use_spec_decode and prev_last_scheduled_idx is not None:
575 - fallback = (num_computed_tokens - 1) // mamba_block_size
576 - fallback.clamp_(min=0)
577 - block_idx_last_scheduled_token_prev_step = torch.where(
578 - prev_last_scheduled_idx >= 0,
579 - prev_last_scheduled_idx,
580 - fallback,
581 - )
582 - else:
583 - state_indices_tensor = mamba_get_block_table_tensor(
584 - common_attn_metadata.block_table_tensor,
585 - common_attn_metadata.seq_lens,
586 - self.kv_cache_spec,
587 - self.vllm_config.cache_config.mamba_cache_mode,
588 - )
440 + state_indices_tensor = mamba_get_block_table_tensor(
441 + common_attn_metadata.block_table_tensor,
442 + common_attn_metadata.seq_lens,
443 + self.kv_cache_spec,
444 + self.vllm_config.cache_config.mamba_cache_mode,
445 + )
446
447 if state_indices_tensor.dim() == 1:
448 state_indices_tensor = state_indices_tensor.unsqueeze(-1)
@@ -595,11 +452,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
452 [num_decodes, num_prefills],
453 dim=0,
454 )
598 - if self.vllm_config.cache_config.mamba_cache_mode != "all":
599 - state_indices_tensor_d = state_indices_tensor_d[
600 - :, : 1 + self.num_spec_tokens
601 - ]
602 - state_indices_tensor_p = state_indices_tensor_p[:, 0]
455 + state_indices_tensor_d = state_indices_tensor_d[:, : 1 + self.num_spec_tokens]
456 + state_indices_tensor_p = state_indices_tensor_p[:, 0]
457
458 if num_decodes > 0 and self.use_spec_decode:
459 query_start_loc_d = common_attn_metadata.query_start_loc[: num_decodes + 1]
@@ -613,8 +467,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
467 num_accepted_tokens = num_accepted_tokens[:num_decodes]
468
469 if num_prefills > 0:
616 - if num_computed_tokens is None:
617 - num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
470 + num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
471
472 query_start_loc_p = (
473 common_attn_metadata.query_start_loc[-num_prefills - 1 :]
@@ -636,33 +489,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
489 )
490 )
491
639 - if self.vllm_config.cache_config.mamba_cache_mode == "all":
640 - assert num_computed_tokens is not None
641 - num_computed_tokens_p = num_computed_tokens[
642 - num_reqs - num_prefills : num_reqs
643 - ]
644 - assert block_idx_first_scheduled_token is not None
645 - block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[
646 - num_reqs - num_prefills : num_reqs
647 - ]
648 -
649 - # How many block-aligned states prefill will write: the same
650 - # (block_idx_last - block_idx_first) as above, from CPU data so
651 - # the mixer needs no D2H read. The -1 in both indices cancels.
652 - block_size = self.kv_cache_spec.block_size
653 - seq_lens_p_cpu = seq_lens_cpu[num_reqs - num_prefills : num_reqs]
654 - num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata(
655 - common_attn_metadata, num_reqs, num_prefills, num_decode_tokens
656 - )
657 - num_state_writes_p = int(
658 - (
659 - (seq_lens_p_cpu + block_size - 1) // block_size
660 - - (num_computed_tokens_p_cpu + block_size) // block_size
661 - )
662 - .clamp(min=0)
663 - .sum()
664 - )
665 -
492 if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0:
493 decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu
494 seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
@@ -764,14 +590,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
590 replayssm_scratch=replayssm_scratch,
591 num_accepted_tokens=num_accepted_tokens,
592 query_start_loc_d=query_start_loc_d,
767 - block_idx_last_scheduled_token=block_idx_last_scheduled_token,
768 - block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
769 - num_state_writes_p=num_state_writes_p,
770 - block_idx_last_computed_token=block_idx_last_computed_token,
771 - block_idx_last_scheduled_token_prev_step=(
772 - block_idx_last_scheduled_token_prev_step
773 - ),
774 - num_computed_tokens_p=num_computed_tokens_p,
593 num_reqs=num_reqs,
594 seq_lens=common_attn_metadata.seq_lens,
595 nums_dict=nums_dict,
@@ -791,11 +609,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
609 state_indices_tensor_d = metadata.state_indices_tensor_d
610 query_start_loc_d = metadata.query_start_loc_d
611 num_accepted_tokens = metadata.num_accepted_tokens
794 - block_idx_last_scheduled_token = metadata.block_idx_last_scheduled_token
795 - block_idx_last_computed_token = metadata.block_idx_last_computed_token
796 - block_idx_last_scheduled_token_prev_step = (
797 - metadata.block_idx_last_scheduled_token_prev_step
798 - )
612 write_pos_d = metadata.write_pos_d
613 is_flush_d = metadata.is_flush_d
614 bc_pre_scratch = metadata.bc_pre_scratch
@@ -824,44 +637,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
637 1 # pad with 1st slot index
638 )
639
827 - if self.vllm_config.cache_config.mamba_cache_mode == "all":
828 - assert block_idx_last_scheduled_token is not None
829 - assert block_idx_last_computed_token is not None
830 - self.block_idx_last_scheduled_token[: metadata.num_decodes].copy_(
831 - block_idx_last_scheduled_token[: metadata.num_decodes],
832 - non_blocking=True,
833 - )
834 - block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
835 - :padded_bs
836 - ]
837 - block_idx_last_scheduled_token[metadata.num_decodes :] = 0
838 -
839 - self.block_idx_last_computed_token[: metadata.num_decodes].copy_(
840 - block_idx_last_computed_token[: metadata.num_decodes],
841 - non_blocking=True,
842 - )
843 - block_idx_last_computed_token = self.block_idx_last_computed_token[
844 - :padded_bs
845 - ]
846 - block_idx_last_computed_token[metadata.num_decodes :] = 0
847 -
848 - if (
849 - self.use_spec_decode
850 - and block_idx_last_scheduled_token_prev_step is not None
851 - ):
852 - self.block_idx_last_scheduled_token_prev_step[
853 - : metadata.num_decodes
854 - ].copy_(
855 - block_idx_last_scheduled_token_prev_step[
856 - : metadata.num_decodes
857 - ],
858 - non_blocking=True,
859 - )
860 - block_idx_last_scheduled_token_prev_step = (
861 - self.block_idx_last_scheduled_token_prev_step[:padded_bs]
862 - )
863 - block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0
864 -
640 if self.use_replayssm and not self.use_flashinfer_replayssm:
641 assert write_pos_d is not None
642 assert is_flush_d is not None
@@ -914,11 +689,6 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
689 bc_pre_scratch=bc_pre_scratch,
690 replayssm_scratch=replayssm_scratch,
691 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=(
920 - block_idx_last_scheduled_token_prev_step
921 - ),
692 )
693
694 def update_block_table(
@@ -950,11 +720,8 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
720 [metadata.num_decodes, metadata.num_prefills],
721 dim=0,
722 )
953 - if self.vllm_config.cache_config.mamba_cache_mode != "all":
954 - state_indices_tensor_d = state_indices_tensor_d[
955 - :, : 1 + self.num_spec_tokens
956 - ]
957 - state_indices_tensor_p = state_indices_tensor_p[:, 0]
723 + state_indices_tensor_d = state_indices_tensor_d[:, : 1 + self.num_spec_tokens]
724 + state_indices_tensor_p = state_indices_tensor_p[:, 0]
725
726 new_metadata = replace(
727 metadata,
vllm/v1/attention/backends/short_conv_attn.py
+3 -10
@@ -283,7 +283,6 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
283
284 has_initial_states_p = None
285 has_initial_states_d = None
286 - num_computed_tokens_p = None
286
287 # Per-request classification by mask, NOT by position. With
288 # spec-decode, the front decode group can contain both spec-decode
@@ -402,10 +401,9 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
401 assert non_spec_query_start_loc is not None
402 query_start_loc_d = non_spec_query_start_loc[: num_decodes + 1]
403 if num_prefills > 0:
405 - num_computed_tokens_p = num_computed_tokens[
406 - num_decodes : num_decodes + num_prefills
407 - ]
408 - has_initial_states_p = num_computed_tokens_p > 0
404 + has_initial_states_p = (
405 + num_computed_tokens[num_decodes : num_decodes + num_prefills] > 0
406 + )
407 assert non_spec_query_start_loc is not None
408 query_start_loc_p = (
409 non_spec_query_start_loc[num_decodes:] - num_decode_tokens
@@ -478,11 +476,6 @@ class PleShortConvAttentionMetadataBuilder(ShortConvAttentionMetadataBuilder):
476 query_start_loc_d=query_start_loc_d,
477 state_indices_tensor_p=state_indices_tensor_p,
478 state_indices_tensor_d=state_indices_tensor_d,
481 - num_computed_tokens_p=num_computed_tokens_p,
482 - block_idx_last_scheduled_token=None,
483 - block_idx_first_scheduled_token_p=None,
484 - block_idx_last_computed_token=None,
485 - block_idx_last_scheduled_token_prev_step=None,
479 seq_lens=m.seq_lens,
480 )
481
vllm/v1/attention/backends/utils.py
+1 -6
@@ -1162,11 +1162,6 @@ def mamba_get_block_table_tensor(
1162 """Get the block table tensor for mamba kernels from the input
1163 common_attn_metadata.block_table_tensor given different mamba cache modes.
1164
1165 - - "all": input (#requests, cdiv(max_model_len, block_size)
1166 - + num_speculative_blocks);
1167 - output (#requests, cdiv(max_model_len, block_size)
1168 - + num_speculative_blocks).
1169 -
1165 - "none": input (#requests, 1 + num_speculative_blocks);
1166 output (#requests, 1 + num_speculative_blocks).
1167
@@ -1174,7 +1169,7 @@ def mamba_get_block_table_tensor(
1169 output (#requests, 1 + num_speculative_blocks), which are the last
1170 1 + num_speculative_blocks of each request.
1171 """
1177 - if mamba_cache_mode in ("all", "none"):
1172 + if mamba_cache_mode == "none":
1173 return block_table
1174 else:
1175 assert isinstance(kv_cache_spec, MambaSpec)
vllm/v1/kv_cache_interface.py
+1 -6
@@ -1079,12 +1079,7 @@ class MambaSpec(KVCacheSpec):
1079 return None
1080
1081 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
1082 - if vllm_config.cache_config.mamba_cache_mode == "all":
1083 - max_model_len = vllm_config.model_config.max_model_len
1084 - return (
1085 - cdiv(max_model_len, self.block_size) + self.num_speculative_blocks
1086 - ) * self.page_size_bytes
1087 - elif vllm_config.cache_config.mamba_cache_mode == "align":
1082 + if vllm_config.cache_config.mamba_cache_mode == "align":
1083 return self.page_size_bytes * (
1084 2 + self.num_speculative_blocks + self.num_prefill_checkpoint_blocks
1085 )
vllm/v1/worker/gpu_model_runner.py
-32
@@ -967,11 +967,6 @@ class GPUModelRunner(
967 self.mamba_state_idx: dict[str, int] = {}
968 self._mamba_bufs: mamba_utils.MambaBuffers | None = None
969 self._mamba_state_copy_funcs: MambaStateCopyFuncsByType | None = None
970 - self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None
971 - if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0:
972 - self.mamba_prev_last_scheduled_idx = self._make_buffer(
973 - self.max_num_reqs, dtype=torch.int32
974 - )
970 self.layerwise_nvtx_hooks_registered = False
971
972 def update_max_model_len(self, max_model_len: int) -> None:
@@ -1608,17 +1603,6 @@ class GPUModelRunner(
1603 assert self.num_accepted_tokens_event is not None
1604 self.num_accepted_tokens_event.record()
1605
1611 - if self.cache_config.mamba_cache_mode == "all":
1612 - mamba_utils.postprocess_mamba_all(
1613 - scheduler_output,
1614 - self.kv_cache_config,
1615 - self.input_batch,
1616 - self.requests,
1617 - self.mamba_state_idx,
1618 - self.num_spec_tokens,
1619 - num_reqs,
1620 - )
1621 -
1606 def _update_streaming_request(
1607 self, req_id: str, new_req_data: NewRequestData
1608 ) -> CachedRequestState:
@@ -2127,15 +2111,6 @@ class GPUModelRunner(
2111 self.num_accepted_tokens.np.fill(1)
2112 self.num_accepted_tokens.gpu.fill_(1)
2113
2130 - if self.mamba_prev_last_scheduled_idx is not None:
2131 - mamba_utils.preprocess_mamba_all_specdec(
2132 - scheduler_output,
2133 - self.input_batch,
2134 - self.mamba_state_idx,
2135 - num_reqs,
2136 - self.mamba_prev_last_scheduled_idx,
2137 - )
2138 -
2114 # Update num_computed_tokens on GPU. In async spec decode,
2115 # CPU values are optimistic (all drafts accepted). The kernel
2116 # corrects on GPU using the previous step's
@@ -2511,13 +2486,6 @@ class GPUModelRunner(
2486 :num_reqs_padded
2487 ],
2488 )
2514 - if (
2515 - isinstance(builder, Mamba2AttentionMetadataBuilder)
2516 - and self.mamba_prev_last_scheduled_idx is not None
2517 - ):
2518 - extra_attn_metadata_args["prev_last_scheduled_idx"] = (
2519 - self.mamba_prev_last_scheduled_idx.gpu[:num_reqs_padded]
2520 - )
2489
2490 if for_cudagraph_capture:
2491 attn_metadata_i = builder.build_for_cudagraph_capture(
vllm/v1/worker/mamba_utils.py
-47
@@ -1535,53 +1535,6 @@ def preprocess_mamba(
1535 do_mamba_copy_block(copy_bufs)
1536
1537
1538 -def postprocess_mamba_all(
1539 - scheduler_output: SchedulerOutput,
1540 - kv_cache_config: KVCacheConfig,
1541 - input_batch: InputBatch,
1542 - requests: dict[str, CachedRequestState],
1543 - mamba_state_idx: dict[str, int],
1544 - num_spec_tokens: int,
1545 - num_reqs: int,
1546 -):
1547 - """All-mode postprocess (only meaningful with num_spec_tokens > 0):
1548 - record per-request the block index of the last token scheduled this
1549 - step, so the next step can anchor its in-place writes when accepted
1550 - drafts leave the sequence at a non-block-aligned position.
1551 - """
1552 - if num_spec_tokens <= 0:
1553 - return
1554 - mamba_groups = get_mamba_groups(kv_cache_config)
1555 - block_sizes = {mamba_spec.block_size for mamba_spec in mamba_groups}
1556 - assert len(block_sizes) == 1, "all mamba groups must share block_size"
1557 - block_size = next(iter(block_sizes))
1558 - full_decode_len = 1 + num_spec_tokens
1559 - scheduled = scheduler_output.num_scheduled_tokens
1560 - for req_id in input_batch.req_ids[:num_reqs]:
1561 - num_query = scheduled.get(req_id, 0)
1562 - if num_query == full_decode_len:
1563 - req = requests[req_id]
1564 - seq_len = req.num_computed_tokens + num_query
1565 - mamba_state_idx[req_id] = max(0, (seq_len - 1) // block_size)
1566 - else:
1567 - mamba_state_idx.pop(req_id, None)
1568 -
1569 -
1570 -def preprocess_mamba_all_specdec(
1571 - scheduler_output: SchedulerOutput,
1572 - input_batch: InputBatch,
1573 - mamba_state_idx: dict[str, int],
1574 - num_reqs: int,
1575 - prev_last_scheduled_idx_buf: CpuGpuBuffer,
1576 -) -> None:
1577 - cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)
1578 - np_view = prev_last_scheduled_idx_buf.np
1579 - for i, req_id in enumerate(input_batch.req_ids[:num_reqs]):
1580 - np_view[i] = mamba_state_idx.get(req_id, -1)
1581 - np_view[num_reqs:].fill(-1)
1582 - prev_last_scheduled_idx_buf.copy_to_gpu()
1583 -
1584 -
1538 def postprocess_mamba_align_gpu(
1539 *,
1540 bufs: "MambaBuffers",