A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 10:30 UTC
[Bugfix][Hybrid][NemotronH] Fix mamba_cache_mode=all + speculative decoding crash (#41233)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
May 18, 2026 at 14:54 UTC
737bfa3a43ce386bd1894792f3302d9f3f9d73fa
10 files changed
+568
-117
tests/v1/attention/test_mamba_update_block_table.py
+271
-4
@@ -32,17 +32,25 @@ class _ConcreteMambaBuilder(
32
metadata_cls = BaseMambaAttentionMetadata
33
34
35
-def _make_vllm_config(block_size, max_model_len, max_num_seqs):
35
+def _make_vllm_config(max_model_len, max_num_seqs, num_speculative_tokens=0):
36
"""Create a minimal mock VllmConfig with only the fields the builder
37
accesses, avoiding any model download / HF config inspection."""
38
+ speculative_config = (
39
+ SimpleNamespace(
40
+ num_speculative_tokens=num_speculative_tokens,
41
+ parallel_drafting=False,
42
+ )
43
+ if num_speculative_tokens > 0
44
+ else None
45
+ )
46
return SimpleNamespace(
47
cache_config=SimpleNamespace(mamba_cache_mode="all"),
48
compilation_config=SimpleNamespace(
49
cudagraph_mode=CUDAGraphMode.FULL,
50
max_cudagraph_capture_size=None,
51
),
44
- speculative_config=None,
45
- num_speculative_tokens=0,
52
+ speculative_config=speculative_config,
53
+ num_speculative_tokens=num_speculative_tokens,
54
parallel_config=SimpleNamespace(decode_context_parallel_size=1),
55
scheduler_config=SimpleNamespace(max_num_seqs=max_num_seqs),
56
model_config=SimpleNamespace(max_model_len=max_model_len),
@@ -59,7 +67,7 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
67
num_reqs = 4
68
device = torch.device("cpu")
69
62
- vllm_config = _make_vllm_config(block_size, max_model_len, num_reqs)
70
+ vllm_config = _make_vllm_config(max_model_len, num_reqs)
71
72
spec = MambaSpec(
73
block_size=block_size,
@@ -106,6 +114,7 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
114
block_idx_last_computed_token=(
115
builder_a.block_idx_last_computed_token[:num_reqs]
116
),
117
+ block_idx_last_scheduled_token_prev_step=None,
118
seq_lens=seq_lens,
119
)
120
@@ -149,3 +158,261 @@ def test_update_block_table_copies_block_idx_to_persistent_buffers():
158
metadata_b.block_idx_last_computed_token,
159
block_idx_vals,
160
)
161
+
162
+
163
+def test_state_indices_tensor_d_includes_num_speculative_blocks():
164
+ """Regression test for https://github.com/vllm-project/vllm/issues/39809
165
+ bug 1: with mamba_cache_mode='all' and speculative decoding enabled,
166
+ the cudagraph buffer for state_indices_tensor_d must allocate the same
167
+ per-request column count as the runtime block table, which includes
168
+ num_speculative_blocks trailing scratch columns."""
169
+
170
+ block_size = 16
171
+ max_model_len = 256
172
+ max_num_seqs = 4
173
+ num_speculative_tokens = 1
174
+ num_speculative_blocks = 2
175
+ device = torch.device("cpu")
176
+
177
+ vllm_config = _make_vllm_config(
178
+ max_model_len,
179
+ max_num_seqs,
180
+ num_speculative_tokens=num_speculative_tokens,
181
+ )
182
+
183
+ spec = MambaSpec(
184
+ block_size=block_size,
185
+ shapes=((1,), (1,)),
186
+ dtypes=(torch.float32,),
187
+ mamba_cache_mode="all",
188
+ num_speculative_blocks=num_speculative_blocks,
189
+ )
190
+
191
+ builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
192
+
193
+ expected_cols = (max_model_len // block_size) + num_speculative_blocks
194
+ assert builder.state_indices_tensor_d.shape == (max_num_seqs, expected_cols)
195
+
196
+
197
+def test_block_idx_cudagraph_capture_padded_by_num_reqs():
198
+ """Regression test for https://github.com/vllm-project/vllm/issues/39809
199
+ bug 2: with mamba_cache_mode='all' and spec decode, _update_metadata_for
200
+ _cudagraph_capture must slice block_idx_last_{scheduled,computed}_token
201
+ by the request count (padded_bs == num_reqs), not by num_decode_tokens.
202
+ Past num_decodes, the slice must be zero-filled."""
203
+
204
+ block_size = 16
205
+ max_model_len = 256
206
+ max_num_seqs = 8
207
+ num_speculative_tokens = 1
208
+ device = torch.device("cpu")
209
+
210
+ vllm_config = _make_vllm_config(
211
+ max_model_len,
212
+ max_num_seqs,
213
+ num_speculative_tokens=num_speculative_tokens,
214
+ )
215
+
216
+ spec = MambaSpec(
217
+ block_size=block_size,
218
+ shapes=((1,), (1,)),
219
+ dtypes=(torch.float32,),
220
+ mamba_cache_mode="all",
221
+ num_speculative_blocks=2,
222
+ )
223
+
224
+ builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
225
+
226
+ builder.block_idx_last_scheduled_token.fill_(-1)
227
+ builder.block_idx_last_computed_token.fill_(-1)
228
+
229
+ num_decodes = 2
230
+ num_reqs = 3
231
+ num_decode_tokens = num_decodes * (1 + num_speculative_tokens)
232
+ seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
233
+ block_idx_vals = torch.tensor([3, 5], dtype=torch.int32, device=device)
234
+ state_indices_d = torch.zeros(
235
+ (num_decodes, builder.state_indices_tensor_d.shape[1]),
236
+ dtype=torch.int32,
237
+ device=device,
238
+ )
239
+ query_start_loc_d = torch.arange(
240
+ num_decodes + 1, dtype=torch.int32, device=device
241
+ ) * (1 + num_speculative_tokens)
242
+ num_accepted_tokens = torch.ones(num_decodes, dtype=torch.int32, device=device)
243
+
244
+ metadata = BaseMambaAttentionMetadata(
245
+ num_prefills=0,
246
+ num_prefill_tokens=0,
247
+ num_decodes=num_decodes,
248
+ num_decode_tokens=num_decode_tokens,
249
+ num_reqs=num_reqs,
250
+ has_initial_states_p=None,
251
+ query_start_loc_p=None,
252
+ num_computed_tokens_p=None,
253
+ state_indices_tensor_p=None,
254
+ state_indices_tensor_d=state_indices_d,
255
+ query_start_loc_d=query_start_loc_d,
256
+ num_accepted_tokens=num_accepted_tokens,
257
+ block_idx_last_scheduled_token=block_idx_vals,
258
+ block_idx_first_scheduled_token_p=None,
259
+ block_idx_last_computed_token=block_idx_vals,
260
+ block_idx_last_scheduled_token_prev_step=None,
261
+ seq_lens=seq_lens,
262
+ )
263
+
264
+ out = builder._update_metadata_for_cudagraph_capture(metadata)
265
+
266
+ assert out.block_idx_last_scheduled_token.shape == (num_reqs,)
267
+ assert out.block_idx_last_computed_token.shape == (num_reqs,)
268
+ torch.testing.assert_close(
269
+ out.block_idx_last_scheduled_token[:num_decodes], block_idx_vals
270
+ )
271
+ torch.testing.assert_close(
272
+ out.block_idx_last_computed_token[:num_decodes], block_idx_vals
273
+ )
274
+ assert torch.all(out.block_idx_last_scheduled_token[num_decodes:] == 0)
275
+ assert torch.all(out.block_idx_last_computed_token[num_decodes:] == 0)
276
+
277
+
278
+def test_block_idx_prev_step_persistent_buffer_allocated():
279
+ """With mamba_cache_mode='all' + spec decode, the builder must allocate
280
+ block_idx_last_scheduled_token_prev_step as a persistent buffer with the
281
+ same shape as the existing block_idx_last_{scheduled,computed}_token
282
+ buffers, so cudagraph capture records a stable pointer for the prev-step
283
+ input anchor consumed by mamba_mixer2's input gather."""
284
+ block_size = 16
285
+ max_model_len = 256
286
+ max_num_seqs = 8
287
+ num_speculative_tokens = 1
288
+ device = torch.device("cpu")
289
+
290
+ vllm_config = _make_vllm_config(
291
+ max_model_len,
292
+ max_num_seqs,
293
+ num_speculative_tokens=num_speculative_tokens,
294
+ )
295
+ spec = MambaSpec(
296
+ block_size=block_size,
297
+ shapes=((1,), (1,)),
298
+ dtypes=(torch.float32,),
299
+ mamba_cache_mode="all",
300
+ num_speculative_blocks=2,
301
+ )
302
+ builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
303
+
304
+ assert hasattr(builder, "block_idx_last_scheduled_token_prev_step")
305
+ assert builder.block_idx_last_scheduled_token_prev_step.shape == (max_num_seqs,)
306
+ assert builder.block_idx_last_scheduled_token_prev_step.dtype == torch.int32
307
+
308
+
309
+def test_block_idx_prev_step_persistent_buffer_skipped_without_spec_decode():
310
+ """Without spec decode, the prev-step buffer is unused and must not be
311
+ allocated — the input anchor reduces to last_computed_token."""
312
+ block_size = 16
313
+ max_model_len = 256
314
+ max_num_seqs = 8
315
+ device = torch.device("cpu")
316
+
317
+ vllm_config = _make_vllm_config(
318
+ max_model_len, max_num_seqs, num_speculative_tokens=0
319
+ )
320
+ spec = MambaSpec(
321
+ block_size=block_size,
322
+ shapes=((1,), (1,)),
323
+ dtypes=(torch.float32,),
324
+ mamba_cache_mode="all",
325
+ )
326
+ builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
327
+
328
+ assert not hasattr(builder, "block_idx_last_scheduled_token_prev_step")
329
+
330
+
331
+def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer():
332
+ """_update_metadata_for_cudagraph_capture must copy the prev-step anchor
333
+ into the builder's persistent buffer (so cudagraph replay reads from the
334
+ same underlying memory), pad past num_decodes with zero, and return a
335
+ slice of the persistent buffer in the metadata."""
336
+ block_size = 16
337
+ max_model_len = 256
338
+ max_num_seqs = 8
339
+ num_speculative_tokens = 1
340
+ device = torch.device("cpu")
341
+
342
+ vllm_config = _make_vllm_config(
343
+ max_model_len,
344
+ max_num_seqs,
345
+ num_speculative_tokens=num_speculative_tokens,
346
+ )
347
+ spec = MambaSpec(
348
+ block_size=block_size,
349
+ shapes=((1,), (1,)),
350
+ dtypes=(torch.float32,),
351
+ mamba_cache_mode="all",
352
+ num_speculative_blocks=2,
353
+ )
354
+ builder = _ConcreteMambaBuilder(spec, ["layer0"], vllm_config, device)
355
+ builder.block_idx_last_scheduled_token.fill_(-1)
356
+ builder.block_idx_last_computed_token.fill_(-1)
357
+ builder.block_idx_last_scheduled_token_prev_step.fill_(-1)
358
+
359
+ num_decodes = 2
360
+ num_reqs = 3
361
+ num_decode_tokens = num_decodes * (1 + num_speculative_tokens)
362
+ seq_lens = torch.full((num_reqs,), 64, dtype=torch.int32, device=device)
363
+ block_idx_vals = torch.tensor([3, 5], dtype=torch.int32, device=device)
364
+ prev_step_vals = torch.tensor([2, 4], dtype=torch.int32, device=device)
365
+ state_indices_d = torch.zeros(
366
+ (num_decodes, builder.state_indices_tensor_d.shape[1]),
367
+ dtype=torch.int32,
368
+ device=device,
369
+ )
370
+ query_start_loc_d = torch.arange(
371
+ num_decodes + 1, dtype=torch.int32, device=device
372
+ ) * (1 + num_speculative_tokens)
373
+ num_accepted_tokens = torch.ones(num_decodes, dtype=torch.int32, device=device)
374
+
375
+ metadata = BaseMambaAttentionMetadata(
376
+ num_prefills=0,
377
+ num_prefill_tokens=0,
378
+ num_decodes=num_decodes,
379
+ num_decode_tokens=num_decode_tokens,
380
+ num_reqs=num_reqs,
381
+ has_initial_states_p=None,
382
+ query_start_loc_p=None,
383
+ num_computed_tokens_p=None,
384
+ state_indices_tensor_p=None,
385
+ state_indices_tensor_d=state_indices_d,
386
+ query_start_loc_d=query_start_loc_d,
387
+ num_accepted_tokens=num_accepted_tokens,
388
+ block_idx_last_scheduled_token=block_idx_vals,
389
+ block_idx_first_scheduled_token_p=None,
390
+ block_idx_last_computed_token=block_idx_vals,
391
+ block_idx_last_scheduled_token_prev_step=prev_step_vals,
392
+ seq_lens=seq_lens,
393
+ )
394
+
395
+ out = builder._update_metadata_for_cudagraph_capture(metadata)
396
+
397
+ # Output field exists and is identity-shared with the persistent buffer.
398
+ assert out.block_idx_last_scheduled_token_prev_step is not None
399
+ assert (
400
+ out.block_idx_last_scheduled_token_prev_step.untyped_storage().data_ptr()
401
+ == builder.block_idx_last_scheduled_token_prev_step.untyped_storage().data_ptr()
402
+ ), (
403
+ "prev-step buffer must live in the builder's persistent buffer, not "
404
+ "in the caller-provided tensor"
405
+ )
406
+
407
+ # Padded by num_reqs (not num_decode_tokens) — same fix as bug 2 for the
408
+ # other block_idx_* fields.
409
+ assert out.block_idx_last_scheduled_token_prev_step.shape == (num_reqs,)
410
+
411
+ # First num_decodes values: input values copied through.
412
+ torch.testing.assert_close(
413
+ out.block_idx_last_scheduled_token_prev_step[:num_decodes],
414
+ prev_step_vals,
415
+ )
416
+
417
+ # Tail values past num_decodes: zero-filled padding for cudagraph capture.
418
+ assert torch.all(out.block_idx_last_scheduled_token_prev_step[num_decodes:] == 0)
tests/v1/e2e/general/test_mamba_prefix_cache.py
+14
-6
@@ -364,26 +364,34 @@ def get_fake_process_mamba_fn(
364
def fake_post_process_mamba_fn(
365
scheduler_output: SchedulerOutput,
366
kv_cache_config: KVCacheConfig,
367
+ cache_config: CacheConfig,
368
input_batch: GPUInputBatch,
369
requests: dict[str, CachedRequestState],
370
mamba_state_idx: dict[str, int],
370
- forward_context: dict[str, Any],
371
- mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
372
- copy_bufs: mamba_utils.MambaCopyBuffers,
371
+ num_spec_tokens: int,
372
+ num_reqs: int,
373
+ *,
374
+ forward_context: dict[str, Any] | None = None,
375
+ mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...] | None = None,
376
+ copy_bufs: mamba_utils.MambaCopyBuffers | None = None,
377
):
378
nonlocal copy_info
379
copy_info = None
380
ret = original_post_process_mamba_fn(
381
scheduler_output,
382
kv_cache_config,
383
+ cache_config,
384
input_batch,
385
requests,
386
mamba_state_idx,
382
- forward_context,
383
- mamba_state_copy_funcs,
384
- copy_bufs,
387
+ num_spec_tokens,
388
+ num_reqs,
389
+ forward_context=forward_context,
390
+ mamba_state_copy_funcs=mamba_state_copy_funcs,
391
+ copy_bufs=copy_bufs,
392
)
393
if cur_step_action is not None:
394
+ assert forward_context is not None
395
check_copy_info(
396
cur_step_action.postprocess_copy_idx,
397
kv_cache_config,
vllm/model_executor/layers/mamba/mamba_mixer2.py
+38
-12
@@ -502,6 +502,12 @@ class MambaMixer2(MambaBase, PluggableLayer):
502
self.prefix = prefix
503
504
self.num_spec = vllm_config.num_speculative_tokens
505
+ if self.num_spec > 0:
506
+ self.register_buffer(
507
+ "_decode_state_offsets",
508
+ torch.arange(1 + self.num_spec, dtype=torch.int32).unsqueeze(0),
509
+ persistent=False,
510
+ )
511
512
# Pre-compute sizes for forward pass
513
self.tped_intermediate_size = self.intermediate_size // self.tp_size
@@ -755,6 +761,14 @@ class MambaMixer2(MambaBase, PluggableLayer):
761
dim=0,
762
)
763
)
764
+ if attn_metadata.block_idx_last_scheduled_token_prev_step is not None:
765
+ block_idx_last_scheduled_token_prev_step_d, _ = torch.split(
766
+ attn_metadata.block_idx_last_scheduled_token_prev_step,
767
+ [num_decodes, num_prefills],
768
+ dim=0,
769
+ )
770
+ else:
771
+ block_idx_last_scheduled_token_prev_step_d = None
772
# Prefill-only variables:
773
block_idx_first_scheduled_token_p = (
774
attn_metadata.block_idx_first_scheduled_token_p
@@ -766,6 +780,7 @@ class MambaMixer2(MambaBase, PluggableLayer):
780
block_idx_first_scheduled_token_p = None
781
block_idx_last_scheduled_token_d = None
782
block_idx_last_computed_token_d = None
783
+ block_idx_last_scheduled_token_prev_step_d = None
784
num_computed_tokens_p = None
785
786
preallocated_ssm_out_d, preallocated_ssm_out_p = torch.split(
@@ -944,18 +959,29 @@ class MambaMixer2(MambaBase, PluggableLayer):
959
if has_decode:
960
assert state_indices_tensor_d is not None
961
if is_mamba_cache_all:
947
- state_indices_tensor_d_input = state_indices_tensor_d.gather(
948
- 1, block_idx_last_computed_token_d.unsqueeze(1)
949
- ).squeeze(1)
950
- state_indices_tensor_d_output = state_indices_tensor_d.gather(
951
- 1, block_idx_last_scheduled_token_d.unsqueeze(1)
952
- ).squeeze(1)
953
- # for decode:
954
- # block_idx_first_scheduled_token_d ==
955
- # block_idx_last_scheduled_token_d
956
- # at block boundaries:
957
- # block_idx_first_scheduled_token_d >
958
- # block_idx_last_computed_token_d
962
+ if self.num_spec > 0:
963
+ assert block_idx_last_scheduled_token_prev_step_d is not None
964
+ input_indices = (
965
+ block_idx_last_scheduled_token_prev_step_d.unsqueeze(1)
966
+ + self._decode_state_offsets
967
+ )
968
+ output_indices = (
969
+ block_idx_last_scheduled_token_d.unsqueeze(1)
970
+ + self._decode_state_offsets
971
+ )
972
+ state_indices_tensor_d_input = state_indices_tensor_d.gather(
973
+ 1, input_indices
974
+ )
975
+ state_indices_tensor_d_output = state_indices_tensor_d.gather(
976
+ 1, output_indices
977
+ )
978
+ else:
979
+ state_indices_tensor_d_input = state_indices_tensor_d.gather(
980
+ 1, block_idx_last_computed_token_d.unsqueeze(1)
981
+ ).squeeze(1)
982
+ state_indices_tensor_d_output = state_indices_tensor_d.gather(
983
+ 1, block_idx_last_scheduled_token_d.unsqueeze(1)
984
+ ).squeeze(1)
985
else:
986
# Without caching, read and write in-place to the same blocks:
987
state_indices_tensor_d_input = state_indices_tensor_d
vllm/model_executor/models/config.py
+9
-20
@@ -350,26 +350,15 @@ class MambaModelConfig(VerifyAndUpdateConfig):
350
351
if cache_config.enable_prefix_caching:
352
if cache_config.mamba_cache_mode == "none":
353
- if (
354
- model_config.supports_mamba_prefix_caching
355
- and vllm_config.speculative_config is not None
356
- ):
357
- cache_config.mamba_cache_mode = "align"
358
- logger.warning(
359
- "Mamba cache mode is set to 'align' for %s by default "
360
- "when prefix caching and speculative decoding are enabled",
361
- model_config.architecture,
362
- )
363
- else:
364
- cache_config.mamba_cache_mode = (
365
- "all" if model_config.supports_mamba_prefix_caching else "align"
366
- )
367
- logger.warning(
368
- "Mamba cache mode is set to '%s' for %s by default "
369
- "when prefix caching is enabled",
370
- cache_config.mamba_cache_mode,
371
- model_config.architecture,
372
- )
353
+ cache_config.mamba_cache_mode = (
354
+ "all" if model_config.supports_mamba_prefix_caching else "align"
355
+ )
356
+ logger.warning(
357
+ "Mamba cache mode is set to '%s' for %s by default "
358
+ "when prefix caching is enabled",
359
+ cache_config.mamba_cache_mode,
360
+ model_config.architecture,
361
+ )
362
if (
363
cache_config.mamba_cache_mode == "all"
364
and not model_config.supports_mamba_prefix_caching
vllm/v1/attention/backends/mamba2_attn.py
+3
-1
@@ -137,7 +137,9 @@ class Mamba2AttentionMetadataBuilder(
137
**kwargs: Any,
138
) -> Mamba2AttentionMetadata:
139
common = self._compute_common_metadata(
140
- common_attn_metadata, num_accepted_tokens=kwargs.get("num_accepted_tokens")
140
+ common_attn_metadata,
141
+ num_accepted_tokens=kwargs.get("num_accepted_tokens"),
142
+ prev_last_scheduled_idx=kwargs.get("prev_last_scheduled_idx"),
143
)
144
145
seq_idx_p = None
vllm/v1/attention/backends/mamba_attn.py
+77
-9
@@ -56,6 +56,7 @@ class BaseMambaAttentionMetadata:
56
block_idx_last_scheduled_token: torch.Tensor | None
57
block_idx_first_scheduled_token_p: torch.Tensor | None
58
block_idx_last_computed_token: torch.Tensor | None
59
+ block_idx_last_scheduled_token_prev_step: torch.Tensor | None
60
61
# The following tensor is only used for prefix caching in align mode
62
seq_lens: torch.Tensor
@@ -108,12 +109,13 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
109
)
110
111
if self.vllm_config.cache_config.mamba_cache_mode == "all":
111
- max_num_blocks = cdiv(
112
- self.vllm_config.model_config.max_model_len,
113
- self.kv_cache_spec.block_size,
112
+ max_num_blocks = (
113
+ cdiv(
114
+ self.vllm_config.model_config.max_model_len,
115
+ kv_cache_spec.block_size,
116
+ )
117
+ + kv_cache_spec.num_speculative_blocks
118
)
115
- # Speculative decoding not supported with prefix caching,
116
- # so keep shape consistent with prefill buffer
119
# TODO: reduce this size as needed for decode-only cudagraph capture
120
self.state_indices_tensor_d: torch.Tensor = torch.empty(
121
(
@@ -133,6 +135,14 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
135
dtype=torch.int32,
136
device=device,
137
)
138
+ if self.use_spec_decode:
139
+ self.block_idx_last_scheduled_token_prev_step: torch.Tensor = (
140
+ torch.empty(
141
+ (self.decode_cudagraph_max_bs,),
142
+ dtype=torch.int32,
143
+ device=device,
144
+ )
145
+ )
146
else:
147
self.state_indices_tensor_d = torch.empty(
148
(self.decode_cudagraph_max_bs, 1 + self.num_spec_tokens),
@@ -176,7 +186,23 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
186
if self.num_spec_tokens > 0:
187
num_accepted_tokens = torch.diff(m.query_start_loc)
188
179
- return self.build(0, m, num_accepted_tokens=num_accepted_tokens)
189
+ prev_last_scheduled_idx = None
190
+ if (
191
+ self.use_spec_decode
192
+ and self.vllm_config.cache_config.mamba_cache_mode == "all"
193
+ ):
194
+ prev_last_scheduled_idx = torch.zeros(
195
+ (m.num_reqs,),
196
+ dtype=torch.int32,
197
+ device=m.query_start_loc.device,
198
+ )
199
+
200
+ return self.build(
201
+ 0,
202
+ m,
203
+ num_accepted_tokens=num_accepted_tokens,
204
+ prev_last_scheduled_idx=prev_last_scheduled_idx,
205
+ )
206
207
def build(
208
self,
@@ -185,6 +211,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
211
fast_build: bool = False,
212
*,
213
num_accepted_tokens: torch.Tensor | None = None,
214
+ prev_last_scheduled_idx: torch.Tensor | None = None,
215
**kwargs: Any,
216
) -> M:
217
"""
@@ -192,7 +219,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
219
Subclasses (e.g., Mamba2) can override to add additional metadata.
220
"""
221
return self._compute_common_metadata(
195
- common_attn_metadata, num_accepted_tokens=num_accepted_tokens
222
+ common_attn_metadata,
223
+ num_accepted_tokens=num_accepted_tokens,
224
+ prev_last_scheduled_idx=prev_last_scheduled_idx,
225
)
226
227
def _compute_chunk_metadata(
@@ -341,6 +370,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
370
common_attn_metadata: CommonAttentionMetadata,
371
*,
372
num_accepted_tokens: torch.Tensor | None = None,
373
+ prev_last_scheduled_idx: torch.Tensor | None = None,
374
) -> M:
375
"""
376
Compute metadata common to both Mamba1 and Mamba2.
@@ -375,6 +405,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
405
block_idx_first_scheduled_token_p = None
406
block_idx_last_computed_token = None
407
block_idx_last_scheduled_token = None
408
+ block_idx_last_scheduled_token_prev_step = None
409
410
# for causal_conv1d
411
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
@@ -393,6 +424,15 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
424
) = self._compute_prefix_caching_block_indices(
425
common_attn_metadata, mamba_block_size
426
)
427
+ if self.use_spec_decode and prev_last_scheduled_idx is not None:
428
+ fallback = torch.clamp(
429
+ (num_computed_tokens - 1) // mamba_block_size, min=0
430
+ )
431
+ block_idx_last_scheduled_token_prev_step = torch.where(
432
+ prev_last_scheduled_idx >= 0,
433
+ prev_last_scheduled_idx,
434
+ fallback,
435
+ )
436
else:
437
state_indices_tensor = mamba_get_block_table_tensor(
438
common_attn_metadata.block_table_tensor,
@@ -470,6 +510,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
510
block_idx_last_scheduled_token=block_idx_last_scheduled_token,
511
block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
512
block_idx_last_computed_token=block_idx_last_computed_token,
513
+ block_idx_last_scheduled_token_prev_step=(
514
+ block_idx_last_scheduled_token_prev_step
515
+ ),
516
num_computed_tokens_p=num_computed_tokens_p,
517
num_reqs=num_reqs,
518
seq_lens=common_attn_metadata.seq_lens,
@@ -493,6 +536,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
536
num_accepted_tokens = metadata.num_accepted_tokens
537
block_idx_last_scheduled_token = metadata.block_idx_last_scheduled_token
538
block_idx_last_computed_token = metadata.block_idx_last_computed_token
539
+ block_idx_last_scheduled_token_prev_step = (
540
+ metadata.block_idx_last_scheduled_token_prev_step
541
+ )
542
if (
543
metadata.num_prefills == 0
544
and metadata.num_decodes <= self.decode_cudagraph_max_bs
@@ -524,16 +570,35 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
570
non_blocking=True,
571
)
572
block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
527
- : metadata.num_decode_tokens
573
+ :padded_bs
574
]
575
+ block_idx_last_scheduled_token[metadata.num_decodes :] = 0
576
577
self.block_idx_last_computed_token[: metadata.num_decodes].copy_(
578
block_idx_last_computed_token[: metadata.num_decodes],
579
non_blocking=True,
580
)
581
block_idx_last_computed_token = self.block_idx_last_computed_token[
535
- : metadata.num_decode_tokens
582
+ :padded_bs
583
]
584
+ block_idx_last_computed_token[metadata.num_decodes :] = 0
585
+
586
+ if (
587
+ self.use_spec_decode
588
+ and block_idx_last_scheduled_token_prev_step is not None
589
+ ):
590
+ self.block_idx_last_scheduled_token_prev_step[
591
+ : metadata.num_decodes
592
+ ].copy_(
593
+ block_idx_last_scheduled_token_prev_step[
594
+ : metadata.num_decodes
595
+ ],
596
+ non_blocking=True,
597
+ )
598
+ block_idx_last_scheduled_token_prev_step = (
599
+ self.block_idx_last_scheduled_token_prev_step[:padded_bs]
600
+ )
601
+ block_idx_last_scheduled_token_prev_step[metadata.num_decodes :] = 0
602
603
return replace(
604
metadata,
@@ -542,6 +607,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
607
num_accepted_tokens=num_accepted_tokens,
608
block_idx_last_scheduled_token=block_idx_last_scheduled_token,
609
block_idx_last_computed_token=block_idx_last_computed_token,
610
+ block_idx_last_scheduled_token_prev_step=(
611
+ block_idx_last_scheduled_token_prev_step
612
+ ),
613
)
614
615
def update_block_table(
vllm/v1/attention/backends/utils.py
+4
-2
@@ -878,8 +878,10 @@ def mamba_get_block_table_tensor(
878
Get the block table tensor for mamba kernels from the input
879
common_attn_metadata.block_table_tensor given different mamba cache modes.
880
881
- - "all": input (#requests, cdiv(max_model_len, block_size));
882
- output (#requests, cdiv(max_model_len, block_size)).
881
+ - "all": input (#requests, cdiv(max_model_len, block_size)
882
+ + num_speculative_blocks);
883
+ output (#requests, cdiv(max_model_len, block_size)
884
+ + num_speculative_blocks).
885
886
- "none": input (#requests, 1 + num_speculative_blocks);
887
output (#requests, 1 + num_speculative_blocks).
vllm/v1/kv_cache_interface.py
+3
-1
@@ -582,7 +582,9 @@ class MambaSpec(KVCacheSpec):
582
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
583
if vllm_config.cache_config.mamba_cache_mode == "all":
584
max_model_len = vllm_config.model_config.max_model_len
585
- return cdiv(max_model_len, self.block_size) * self.page_size_bytes
585
+ return (
586
+ cdiv(max_model_len, self.block_size) + self.num_speculative_blocks
587
+ ) * self.page_size_bytes
588
elif vllm_config.cache_config.mamba_cache_mode == "align":
589
return self.page_size_bytes * (2 + self.num_speculative_blocks)
590
else:
vllm/v1/worker/gpu_model_runner.py
+41
-11
@@ -888,6 +888,11 @@ class GPUModelRunner(
888
self.kv_connector_output: KVConnectorOutput | None = None
889
self.mamba_state_idx: dict[str, int] = {}
890
self._mamba_copy_bufs: mamba_utils.MambaCopyBuffers | None = None
891
+ self.mamba_prev_last_scheduled_idx: CpuGpuBuffer | None = None
892
+ if self.cache_config.mamba_cache_mode == "all" and self.num_spec_tokens > 0:
893
+ self.mamba_prev_last_scheduled_idx = self._make_buffer(
894
+ self.max_num_reqs, dtype=torch.int32
895
+ )
896
self.layerwise_nvtx_hooks_registered = False
897
898
def update_max_model_len(self, max_model_len: int) -> None:
@@ -1493,21 +1498,12 @@ class GPUModelRunner(
1498
num_reqs = output_token_ids.size(0)
1499
self.num_accepted_tokens.gpu[:num_reqs] = (output_token_ids != -1).sum(dim=1)
1500
1496
- if self.cache_config.mamba_cache_mode == "align":
1501
+ is_align = self.cache_config.mamba_cache_mode == "align"
1502
+ if is_align:
1503
for i, num_tokens in enumerate(
1504
self.num_accepted_tokens.gpu[:num_reqs].cpu().numpy()
1505
):
1506
self.input_batch.num_accepted_tokens_cpu[i] = num_tokens
1501
- mamba_utils.postprocess_mamba(
1502
- scheduler_output,
1503
- self.kv_cache_config,
1504
- self.input_batch,
1505
- self.requests,
1506
- self.mamba_state_idx,
1507
- self.compilation_config.static_forward_context,
1508
- self.model.get_mamba_state_copy_func(),
1509
- self._get_mamba_copy_bufs(),
1510
- )
1507
else:
1508
self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_(
1509
self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True
@@ -1515,6 +1511,24 @@ class GPUModelRunner(
1511
assert self.num_accepted_tokens_event is not None
1512
self.num_accepted_tokens_event.record()
1513
1514
+ mamba_utils.postprocess_mamba(
1515
+ scheduler_output,
1516
+ self.kv_cache_config,
1517
+ self.cache_config,
1518
+ self.input_batch,
1519
+ self.requests,
1520
+ self.mamba_state_idx,
1521
+ self.num_spec_tokens,
1522
+ num_reqs,
1523
+ forward_context=(
1524
+ self.compilation_config.static_forward_context if is_align else None
1525
+ ),
1526
+ mamba_state_copy_funcs=(
1527
+ self.model.get_mamba_state_copy_func() if is_align else None
1528
+ ),
1529
+ copy_bufs=self._get_mamba_copy_bufs() if is_align else None,
1530
+ )
1531
+
1532
def _update_streaming_request(
1533
self, req_id: str, new_req_data: NewRequestData
1534
) -> CachedRequestState:
@@ -2012,6 +2026,15 @@ class GPUModelRunner(
2026
self.num_accepted_tokens.np.fill(1)
2027
self.num_accepted_tokens.gpu.fill_(1)
2028
2029
+ if self.mamba_prev_last_scheduled_idx is not None:
2030
+ mamba_utils.preprocess_mamba_all_specdec(
2031
+ scheduler_output,
2032
+ self.input_batch,
2033
+ self.mamba_state_idx,
2034
+ num_reqs,
2035
+ self.mamba_prev_last_scheduled_idx,
2036
+ )
2037
+
2038
# Update num_computed_tokens on GPU. In async spec decode,
2039
# CPU values are optimistic (all drafts accepted). The kernel
2040
# corrects on GPU using the previous step's
@@ -2319,6 +2342,13 @@ class GPUModelRunner(
2342
:num_reqs_padded
2343
],
2344
)
2345
+ if (
2346
+ isinstance(builder, Mamba2AttentionMetadataBuilder)
2347
+ and self.mamba_prev_last_scheduled_idx is not None
2348
+ ):
2349
+ extra_attn_metadata_args["prev_last_scheduled_idx"] = (
2350
+ self.mamba_prev_last_scheduled_idx.gpu[:num_reqs_padded]
2351
+ )
2352
2353
if for_cudagraph_capture:
2354
attn_metadata_i = builder.build_for_cudagraph_capture(
vllm/v1/worker/mamba_utils.py
+108
-51
@@ -144,6 +144,24 @@ def do_mamba_copy_block(copy_bufs: MambaCopyBuffers):
144
)
145
146
147
+def cleanup_mamba_state_idx(
148
+ scheduler_output: SchedulerOutput,
149
+ mamba_state_idx: dict[str, int],
150
+) -> None:
151
+ """Pop stale `mamba_state_idx` entries for finished/preempted/resumed reqs.
152
+
153
+ Force-preempted requests (e.g., during reset_prefix_cache / KV cache
154
+ flush) appear in resumed_req_ids without a corresponding entry in
155
+ preempted_req_ids, leaving stale entries that can point to block
156
+ indices beyond the new (smaller) block allocation.
157
+ """
158
+ finished_req_ids = scheduler_output.finished_req_ids
159
+ preempted_req_ids = scheduler_output.preempted_req_ids or set()
160
+ resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
161
+ for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids):
162
+ mamba_state_idx.pop(req_id, None)
163
+
164
+
165
def preprocess_mamba(
166
scheduler_output: SchedulerOutput,
167
kv_cache_config: KVCacheConfig,
@@ -165,16 +183,7 @@ def preprocess_mamba(
183
# TODO(Chen): we need to optimize this function a lot
184
assert cache_config.enable_prefix_caching
185
block_size = mamba_spec.block_size
168
- finished_req_ids = scheduler_output.finished_req_ids
169
- preempted_req_ids = scheduler_output.preempted_req_ids or set()
170
- # We need to clear mamba_state_idx for resumed requests. When requests are
171
- # force-preempted (e.g., during reset_prefix_cache / KV cache flush),
172
- # they appear in resumed_req_ids without a corresponding entry in
173
- # preempted_req_ids, leaving stale mamba_state_idx entries that can
174
- # point to block indices beyond the new (smaller) block allocation.
175
- resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
176
- for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids):
177
- mamba_state_idx.pop(req_id, None)
186
+ cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)
187
188
copy_bufs.offset = 0
189
for i, req_id in enumerate(input_batch.req_ids):
@@ -222,52 +231,100 @@ def preprocess_mamba(
231
def postprocess_mamba(
232
scheduler_output: SchedulerOutput,
233
kv_cache_config: KVCacheConfig,
234
+ cache_config: CacheConfig,
235
input_batch: GPUInputBatch,
236
requests: dict[str, CachedRequestState],
237
mamba_state_idx: dict[str, int],
228
- forward_context: dict[str, Any],
229
- mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
230
- copy_bufs: MambaCopyBuffers,
238
+ num_spec_tokens: int,
239
+ num_reqs: int,
240
+ *,
241
+ forward_context: dict[str, Any] | None = None,
242
+ mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...] | None = None,
243
+ copy_bufs: MambaCopyBuffers | None = None,
244
):
245
"""
233
- If a blocks is converted from partial block to full block in this step, copy the
234
- state from the block for running state to the new full block.
246
+ Post-model-execute mamba prefix-caching bookkeeping. Dispatched by
247
+ cache_config.mamba_cache_mode:
248
+ - "align": if a block is converted from partial to full this step,
249
+ copy the running state into the new full block.
250
+ - "all" + num_spec_tokens > 0: record per-request the block index of
251
+ the last token scheduled this step, so the next step can anchor
252
+ its in-place writes when accepted drafts leave the sequence at a
253
+ non-block-aligned position.
254
"""
236
- num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
237
- scheduled_spec_decode_tokens_dict = scheduler_output.scheduled_spec_decode_tokens
238
- num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu
239
- mamba_group_ids = copy_bufs.mamba_group_ids
240
- mamba_spec = copy_bufs.mamba_spec
241
- copy_bufs.offset = 0
242
- for i, req_id in enumerate(input_batch.req_ids):
243
- req_state = requests[req_id]
244
- num_computed_tokens = req_state.num_computed_tokens
245
- num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, []))
246
- num_scheduled_tokens = num_scheduled_tokens_dict[req_id]
247
- num_accepted_tokens = num_accepted_tokens_cpu[i]
248
- num_tokens_running_state = (
249
- num_computed_tokens + num_scheduled_tokens - num_draft_tokens
255
+ if cache_config.mamba_cache_mode == "align":
256
+ assert forward_context is not None
257
+ assert mamba_state_copy_funcs is not None
258
+ assert copy_bufs is not None
259
+ num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
260
+ scheduled_spec_decode_tokens_dict = (
261
+ scheduler_output.scheduled_spec_decode_tokens
262
)
251
- new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1
252
- aligned_new_computed_tokens = (
253
- new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size
254
- )
255
- # TODO: how to ensure all blocks that cache_blocks called are cached here?
256
- if aligned_new_computed_tokens >= num_tokens_running_state:
257
- accept_token_bias = aligned_new_computed_tokens - num_tokens_running_state
258
- src_block_idx = mamba_state_idx[req_id]
259
- dest_block_idx = aligned_new_computed_tokens // mamba_spec.block_size - 1
260
- collect_mamba_copy_meta(
261
- copy_bufs,
262
- kv_cache_config,
263
- mamba_state_copy_funcs,
264
- mamba_group_ids,
265
- src_block_idx,
266
- dest_block_idx,
267
- accept_token_bias,
268
- req_state,
269
- forward_context,
263
+ num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu
264
+ mamba_group_ids = copy_bufs.mamba_group_ids
265
+ mamba_spec = copy_bufs.mamba_spec
266
+ copy_bufs.offset = 0
267
+ for i, req_id in enumerate(input_batch.req_ids):
268
+ req_state = requests[req_id]
269
+ num_computed_tokens = req_state.num_computed_tokens
270
+ num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, []))
271
+ num_scheduled_tokens = num_scheduled_tokens_dict[req_id]
272
+ num_accepted_tokens = num_accepted_tokens_cpu[i]
273
+ num_tokens_running_state = (
274
+ num_computed_tokens + num_scheduled_tokens - num_draft_tokens
275
)
271
- if src_block_idx == dest_block_idx:
272
- num_accepted_tokens_cpu[i] = 1
273
- do_mamba_copy_block(copy_bufs)
276
+ new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1
277
+ aligned_new_computed_tokens = (
278
+ new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size
279
+ )
280
+ # TODO: how to ensure all blocks that cache_blocks called are cached here?
281
+ if aligned_new_computed_tokens >= num_tokens_running_state:
282
+ accept_token_bias = (
283
+ aligned_new_computed_tokens - num_tokens_running_state
284
+ )
285
+ src_block_idx = mamba_state_idx[req_id]
286
+ dest_block_idx = (
287
+ aligned_new_computed_tokens // mamba_spec.block_size - 1
288
+ )
289
+ collect_mamba_copy_meta(
290
+ copy_bufs,
291
+ kv_cache_config,
292
+ mamba_state_copy_funcs,
293
+ mamba_group_ids,
294
+ src_block_idx,
295
+ dest_block_idx,
296
+ accept_token_bias,
297
+ req_state,
298
+ forward_context,
299
+ )
300
+ if src_block_idx == dest_block_idx:
301
+ num_accepted_tokens_cpu[i] = 1
302
+ do_mamba_copy_block(copy_bufs)
303
+ elif cache_config.mamba_cache_mode == "all" and num_spec_tokens > 0:
304
+ _, mamba_spec = get_mamba_groups(kv_cache_config)
305
+ block_size = mamba_spec.block_size
306
+ full_decode_len = 1 + num_spec_tokens
307
+ scheduled = scheduler_output.num_scheduled_tokens
308
+ for req_id in input_batch.req_ids[:num_reqs]:
309
+ num_query = scheduled.get(req_id, 0)
310
+ if num_query == full_decode_len:
311
+ req = requests[req_id]
312
+ seq_len = req.num_computed_tokens + num_query
313
+ mamba_state_idx[req_id] = max(0, (seq_len - 1) // block_size)
314
+ else:
315
+ mamba_state_idx.pop(req_id, None)
316
+
317
+
318
+def preprocess_mamba_all_specdec(
319
+ scheduler_output: SchedulerOutput,
320
+ input_batch: GPUInputBatch,
321
+ mamba_state_idx: dict[str, int],
322
+ num_reqs: int,
323
+ prev_last_scheduled_idx_buf: CpuGpuBuffer,
324
+) -> None:
325
+ cleanup_mamba_state_idx(scheduler_output, mamba_state_idx)
326
+ np_view = prev_last_scheduled_idx_buf.np
327
+ for i, req_id in enumerate(input_batch.req_ids[:num_reqs]):
328
+ np_view[i] = mamba_state_idx.get(req_id, -1)
329
+ np_view[num_reqs:].fill(-1)
330
+ prev_last_scheduled_idx_buf.copy_to_gpu()