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 7, 19:00 UTC
[Linear Attention] fix bug for linear attention + prefix caching + reset_prefix_cache (#35157)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Feb 24, 2026 at 22:00 UTC
8fae54faff485e446dc8d1a700417f07659ef89e
2 files changed
+74
-1
tests/v1/worker/test_mamba_utils.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from unittest.mock import MagicMock, patch
4
+
5
+from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
6
+from vllm.v1.worker.mamba_utils import preprocess_mamba
7
+
8
+
9
+def _make_scheduler_output(
10
+ finished_req_ids: set[str],
11
+ preempted_req_ids: set[str] | None,
12
+ resumed_req_ids: set[str],
13
+) -> SchedulerOutput:
14
+ cached = CachedRequestData.make_empty()
15
+ cached.resumed_req_ids = resumed_req_ids
16
+ return SchedulerOutput(
17
+ scheduled_new_reqs=[],
18
+ scheduled_cached_reqs=cached,
19
+ num_scheduled_tokens={},
20
+ total_num_scheduled_tokens=0,
21
+ scheduled_spec_decode_tokens={},
22
+ scheduled_encoder_inputs={},
23
+ num_common_prefix_blocks=[],
24
+ finished_req_ids=finished_req_ids,
25
+ free_encoder_mm_hashes=[],
26
+ preempted_req_ids=preempted_req_ids,
27
+ )
28
+
29
+
30
+def test_resumed_req_ids_cleared_from_mamba_state_idx():
31
+ """When a request is force-preempted (e.g. reset_prefix_cache),
32
+ it appears in resumed_req_ids but NOT in preempted_req_ids.
33
+ preprocess_mamba must still clear its mamba_state_idx entry,
34
+ otherwise stale indices can point beyond the new block allocation.
35
+ """
36
+ spec = MagicMock(block_size=64, num_speculative_blocks=0)
37
+ cache_config = MagicMock(enable_prefix_caching=True)
38
+ input_batch = MagicMock(req_ids=[])
39
+
40
+ mamba_state_idx = {
41
+ "finished": 1,
42
+ "preempted": 2,
43
+ "resumed": 3, # only in resumed_req_ids, NOT in preempted
44
+ "keep": 99,
45
+ }
46
+ sched = _make_scheduler_output(
47
+ finished_req_ids={"finished"},
48
+ preempted_req_ids={"preempted"},
49
+ resumed_req_ids={"resumed"},
50
+ )
51
+
52
+ with patch(
53
+ "vllm.v1.worker.mamba_utils.get_mamba_groups",
54
+ return_value=([0], spec),
55
+ ):
56
+ preprocess_mamba(
57
+ sched,
58
+ MagicMock(),
59
+ cache_config,
60
+ mamba_state_idx,
61
+ input_batch,
62
+ {},
63
+ {},
64
+ (),
65
+ )
66
+
67
+ assert mamba_state_idx == {"keep": 99}
vllm/v1/worker/mamba_utils.py
+7
-1
@@ -129,7 +129,13 @@ def preprocess_mamba(
129
block_size = mamba_spec.block_size
130
finished_req_ids = scheduler_output.finished_req_ids
131
preempted_req_ids = scheduler_output.preempted_req_ids or set()
132
- for req_id in itertools.chain(finished_req_ids, preempted_req_ids):
132
+ # We need to clear mamba_state_idx for resumed requests. When requests are
133
+ # force-preempted (e.g., during reset_prefix_cache / KV cache flush),
134
+ # they appear in resumed_req_ids without a corresponding entry in
135
+ # preempted_req_ids, leaving stale mamba_state_idx entries that can
136
+ # point to block indices beyond the new (smaller) block allocation.
137
+ resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
138
+ for req_id in itertools.chain(finished_req_ids, preempted_req_ids, resumed_req_ids):
139
mamba_state_idx.pop(req_id, None)
140
141
src_state_list: list[int] = []