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 8, 06:15 UTC
[Hybrid] Enable mamba prefix cache "align" mode with async scheduling (#33997)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Feb 14, 2026 at 22:15 UTC
d5fe3f702c2f4392515cdfde6fd0442271c74dda
4 files changed
+77
-32
tests/v1/e2e/test_mamba_prefix_cache.py
+58
-19
@@ -76,14 +76,11 @@ def get_fake_sample_fn() -> SamplerOutput:
76
),
77
logprobs_tensors=None,
78
)
79
- num_sampled_tokens = spec_decode_metadata.cu_num_sampled_tokens[0].item() + 1
79
accpeted_tokens = prompt_token_ids[
80
first_token_id_index : first_token_id_index
81
+ min(num_accepted_tokens, logits.shape[0])
82
]
84
- sampled_token_ids = accpeted_tokens + [-1] * (
85
- num_sampled_tokens - len(accpeted_tokens)
86
- )
83
+ sampled_token_ids = accpeted_tokens
84
return SamplerOutput(
85
sampled_token_ids=torch.tensor(
86
[sampled_token_ids], device="cuda", dtype=torch.int32
@@ -124,7 +121,24 @@ def get_fake_propose_draft_token_ids_fn():
121
first_token_id_index : first_token_id_index + num_speculative_tokens
122
]
123
]
127
- return proposed_draft_token_ids
124
+
125
+ next_token_ids = torch.tensor(
126
+ prompt_token_ids[
127
+ first_token_id_index - 1 : first_token_id_index
128
+ - 1
129
+ + num_accepted_tokens
130
+ ],
131
+ device="cuda",
132
+ dtype=torch.int32,
133
+ )
134
+
135
+ valid_sampled_tokens_count = torch.tensor(
136
+ [num_accepted_tokens], device="cuda", dtype=torch.int32
137
+ )
138
+
139
+ self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
140
+
141
+ return torch.tensor(proposed_draft_token_ids, device="cuda", dtype=torch.int32)
142
143
return fake_propose_draft_token_ids_fn
144
@@ -184,6 +198,7 @@ mamba_kv_cache_dict = {}
198
199
def get_fake_execute_model_fn(original_execute_model_fn: Callable):
200
last_num_computed_tokens = 0
201
+ num_prompt_tokens = None
202
203
def fake_execute_model_fn(
204
self: GPUModelRunner,
@@ -201,10 +216,30 @@ def get_fake_execute_model_fn(original_execute_model_fn: Callable):
216
mamba_group_id
217
].layer_names[0]
218
nonlocal last_num_computed_tokens
219
+ nonlocal num_prompt_tokens
220
+
221
+ if (
222
+ len(scheduler_output.scheduled_new_reqs) > 0
223
+ and scheduler_output.scheduled_new_reqs[0].prompt_token_ids is not None
224
+ ):
225
+ # record number of prompt tokens
226
+ num_prompt_tokens = len(
227
+ scheduler_output.scheduled_new_reqs[0].prompt_token_ids
228
+ )
229
+
230
if len(scheduler_output.scheduled_cached_reqs.req_ids) > 0:
231
num_computed_tokens = (
232
scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
233
)
234
+ if (
235
+ self.num_spec_tokens
236
+ and num_prompt_tokens is not None
237
+ and num_computed_tokens > num_prompt_tokens
238
+ ):
239
+ # NOTE (tdoublep) with async scheduling, the scheduler does not have an
240
+ # accurate measure of the number of computed tokens; we need to subtract
241
+ # the number of reject tokens from the previous timestep.
242
+ num_computed_tokens -= num_speculative_tokens + 1 - num_accepted_tokens
243
if (
244
num_computed_tokens // BLOCK_SIZE
245
> last_num_computed_tokens // BLOCK_SIZE
@@ -493,9 +528,9 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
528
step_actions=[
529
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
530
StepAction(554, 4, [], (-1, -1), (-1, -1)),
496
- StepAction(555, 4, [], (-1, -1), (-1, -1)),
531
+ StepAction(555, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
532
StepAction(556, 4, [], (-1, -1), (-1, -1)),
498
- StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
533
+ StepAction(557, 4, [], (0, 1), (-1, -1)),
534
StepAction(558, 4, [], (-1, -1), (-1, -1)),
535
StepAction(559, 4, [], (-1, -1), (1, 0)),
536
StepAction(560, 4, [], (-1, -1), (-1, -1)),
@@ -510,8 +545,8 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
545
step_actions=[
546
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
547
StepAction(554, 4, [], (-1, -1), (-1, -1)),
513
- StepAction(556, 4, [], (-1, -1), (-1, -1)),
514
- StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
548
+ StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
549
+ StepAction(558, 4, [], (1, 1), (2, 0)),
550
StepAction(560, 4, [], (-1, -1), (-1, -1)),
551
StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
552
],
@@ -526,7 +561,8 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
561
StepAction(555, 4, [], (-1, -1), (-1, -1)),
562
StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
563
StepAction(559, 4, [], (-1, -1), (1, 0)),
529
- StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
564
+ StepAction(561, 4, [], (-1, -1), (-1, -1)),
565
+ StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
566
],
567
),
568
"accept_3_1": TestConfig(
@@ -536,9 +572,10 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
572
step_actions=[
573
StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
574
StepAction(553, 4, [], (-1, -1), (-1, -1)),
539
- StepAction(556, 4, [], (-1, -1), (-1, -1)),
540
- StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
541
- StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
575
+ StepAction(556, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
576
+ StepAction(559, 4, [], (2, 1), (1, 0)),
577
+ StepAction(562, 4, [], (-1, -1), (-1, -1)),
578
+ StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
579
],
580
),
581
"accept_3_2": TestConfig(
@@ -561,7 +598,8 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
598
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
599
StepAction(555, 4, [], (-1, -1), (-1, -1)),
600
StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
564
- StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
601
+ StepAction(561, 4, [], (-1, -1), (-1, -1)),
602
+ StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
603
],
604
),
605
"accept_4_1": TestConfig(
@@ -572,8 +610,8 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
610
StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
611
StepAction(553, 4, [], (-1, -1), (-1, -1)),
612
StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
575
- StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
576
- StepAction(565, 4, [], (-1, -1), (-1, -1)),
613
+ StepAction(561, 4, [], (-1, -1), (-1, -1)),
614
+ StepAction(565, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
615
],
616
),
617
"accept_4_2": TestConfig(
@@ -584,8 +622,8 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
622
StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
623
StepAction(554, 4, [], (-1, -1), (-1, -1)),
624
StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
587
- StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
588
- StepAction(566, 4, [], (-1, -1), (-1, -1)),
625
+ StepAction(562, 4, [], (-1, -1), (-1, -1)),
626
+ StepAction(566, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
627
],
628
),
629
"accept_4_3": TestConfig(
@@ -596,7 +634,8 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
634
StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
635
StepAction(555, 4, [], (-1, -1), (-1, -1)),
636
StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
599
- StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
637
+ StepAction(563, 4, [], (-1, -1), (-1, -1)),
638
+ StepAction(567, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
639
],
640
),
641
"accept_4_4": TestConfig(
vllm/config/vllm.py
-12
@@ -648,11 +648,6 @@ class VllmConfig:
648
"`external_launcher` distributed executor backend, but you chose "
649
f"`{executor_backend}`."
650
)
651
- if self.cache_config.mamba_cache_mode != "none":
652
- raise ValueError(
653
- "Currently, async scheduling is not compatible with "
654
- "prefix caching for Mamba models."
655
- )
651
elif self.scheduler_config.async_scheduling is None:
652
# Enable async scheduling unless there is an incompatible option.
653
if (
@@ -685,13 +680,6 @@ class VllmConfig:
680
scope="local",
681
)
682
self.scheduler_config.async_scheduling = False
688
- elif self.cache_config.mamba_cache_mode != "none":
689
- logger.warning_once(
690
- "Async scheduling is not compatible with "
691
- "prefix caching for Mamba models and will be disabled.",
692
- scope="local",
693
- )
694
- self.scheduler_config.async_scheduling = False
683
else:
684
self.scheduler_config.async_scheduling = True
685
vllm/v1/core/single_type_kv_cache_manager.py
+13
@@ -814,6 +814,14 @@ class MambaManager(SingleTypeKVCacheManager):
814
815
def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
816
assert isinstance(self.kv_cache_spec, MambaSpec)
817
+
818
+ # NOTE (tdoublep) with async scheduling, the num_computed_tokens can contain
819
+ # draft tokens from the previous step that may or may not be rejected later.
820
+ # This can make us think we are further ahead in the sequence than we actually
821
+ # are, so let's assume that all tokens are rejected so we don't free blocks
822
+ # that we might actually need.
823
+ num_computed_tokens = max(0, num_computed_tokens - self.num_speculative_blocks)
824
+
825
super().remove_skipped_blocks(request_id, num_computed_tokens)
826
if self.mamba_cache_mode == "align":
827
# `last_state_block_idx` refers to the block index allocated two steps ago.
@@ -879,6 +887,9 @@ class MambaManager(SingleTypeKVCacheManager):
887
# We can ignore lookahead tokens because current draft models don't have
888
# mamba layers.
889
num_tokens = num_tokens_main_model
890
+
891
+ # NOTE(tdouble): this is an over-estimate of how many blocks we need because
892
+ # num_tokens can include draft tokens that will later be rejected.
893
num_required_blocks = (
894
cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
895
)
@@ -922,6 +933,8 @@ class MambaManager(SingleTypeKVCacheManager):
933
# mamba layers.
934
num_tokens = num_tokens_main_model
935
req_blocks: list[KVCacheBlock] = self.req_to_blocks[request_id]
936
+ # NOTE(tdouble): this is an over-estimate of how many blocks we need because
937
+ # num_tokens can include draft tokens that will later be rejected.
938
num_required_blocks = (
939
cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
940
)
vllm/v1/worker/mamba_utils.py
+6
-1
@@ -10,6 +10,7 @@ from vllm.model_executor.layers.mamba.mamba_utils import (
10
MambaStateCopyFunc,
11
)
12
from vllm.triton_utils import tl, triton
13
+from vllm.utils.math_utils import cdiv
14
from vllm.v1.core.sched.output import SchedulerOutput
15
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
16
from vllm.v1.worker.gpu_input_batch import CachedRequestState
@@ -142,7 +143,11 @@ def preprocess_mamba(
143
# if num_computed_tokens is 0, prev_state_idx will be -1
144
prev_state_idx = (req_state.num_computed_tokens - 1) // block_size
145
145
- num_blocks = len(req_state.block_ids[mamba_group_ids[0]])
146
+ num_scheduled_tokens = scheduler_output.num_scheduled_tokens[req_id]
147
+ num_blocks: int = (
148
+ cdiv(req_state.num_computed_tokens + num_scheduled_tokens, block_size)
149
+ + num_speculative_blocks
150
+ )
151
152
# We always save the current running state at the last
153
# (1 + num_speculative_blocks) block.