[Bugfix] Keep mamba align prefill chunks block-aligned past last_cache_position (#51113)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Hernan <kodek@eat1337.com> Co-authored-by: yanghui1-arch <3053034939@qq.com> Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Yifan Qiao committed
Aug 6, 2026 at 10:21 UTC
c56f169d9ae46ca420617e2cf5f0c9135da0f651
2 files changed
+258
-11
tests/v1/core/test_mamba_align_chunk_split.py
new
+247
@@ -0,0 +1,247 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Mamba "align" prefill chunk splitting (`_mamba_block_aligned_split`).
4
+
5
+Invariant: slot `p` holds the state after exactly `(p + 1) * block_size` tokens.
6
+State is written at chunk ends, so a chunk ending mid-block leaves its slot at
7
+the wrong offset, and a later chunk crossing that boundary publishes it anyway.
8
+Requests resuming from it then restore a truncated state (#43559).
9
+"""
10
+
11
+from types import SimpleNamespace
12
+
13
+import pytest
14
+import torch
15
+
16
+from vllm.utils.math_utils import cdiv
17
+from vllm.v1.core.kv_cache_manager import KVCacheManager
18
+from vllm.v1.core.sched.scheduler import Scheduler
19
+from vllm.v1.kv_cache_interface import (
20
+ FullAttentionSpec,
21
+ KVCacheConfig,
22
+ KVCacheGroupSpec,
23
+ MambaSpec,
24
+)
25
+from vllm.v1.request import Request
26
+
27
+from .utils import create_requests
28
+
29
+pytestmark = pytest.mark.cpu_test
30
+
31
+# Mirrors the deployment where the poisoning was observed (Qwen3.6-27B): mamba
32
+# block 1600, MTP with 3 draft tokens, prompts shorter than 2 mamba blocks.
33
+ATTN_BLOCK_SIZE = 16
34
+MAMBA_BLOCK_SIZE = 1600
35
+NUM_SPEC = 3
36
+PROMPT_LEN = 2002
37
+MAMBA_GROUP_ID = 1
38
+
39
+
40
+def _make_hybrid_kv_cache_manager() -> KVCacheManager:
41
+ config = KVCacheConfig(
42
+ num_blocks=10000,
43
+ kv_cache_tensors=[],
44
+ kv_cache_groups=[
45
+ KVCacheGroupSpec(
46
+ ["full_layer"],
47
+ FullAttentionSpec(
48
+ block_size=ATTN_BLOCK_SIZE,
49
+ num_kv_heads=1,
50
+ head_size=1,
51
+ dtype=torch.float32,
52
+ ),
53
+ ),
54
+ KVCacheGroupSpec(
55
+ ["mamba_layer"],
56
+ MambaSpec(
57
+ block_size=MAMBA_BLOCK_SIZE,
58
+ shapes=((1, 1),),
59
+ dtypes=(torch.float32,),
60
+ mamba_cache_mode="align",
61
+ num_speculative_blocks=NUM_SPEC,
62
+ ),
63
+ ),
64
+ ],
65
+ )
66
+ return KVCacheManager(
67
+ config,
68
+ max_model_len=262144,
69
+ scheduler_block_size=MAMBA_BLOCK_SIZE,
70
+ hash_block_size=ATTN_BLOCK_SIZE,
71
+ enable_caching=True,
72
+ use_eagle=True,
73
+ )
74
+
75
+
76
+def _split(
77
+ request: Request,
78
+ num_new_tokens: int,
79
+ use_eagle: bool = True,
80
+ partial_hit: bool = False,
81
+) -> int:
82
+ """Call the real `Scheduler._mamba_block_aligned_split` on a stub self."""
83
+ stub = SimpleNamespace(
84
+ cache_config=SimpleNamespace(block_size=MAMBA_BLOCK_SIZE),
85
+ use_eagle=use_eagle,
86
+ max_num_scheduled_tokens=16384,
87
+ scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
88
+ # `prefix_match_unit` finer than the block size (#46384).
89
+ mamba_partial_cache_hit=partial_hit,
90
+ hash_block_size=ATTN_BLOCK_SIZE,
91
+ )
92
+ return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens)
93
+
94
+
95
+def _run_chunked_prefill(
96
+ manager: KVCacheManager, request: Request, budgets: list[int]
97
+) -> dict[int, int]:
98
+ """Prefill `request`, one step per entry in `budgets`.
99
+
100
+ A zero-token split means the budget cannot fund an aligned chunk; the
101
+ scheduler defers the request to a later step, so this does too.
102
+
103
+ Returns physical mamba block id -> the token offset of the state it holds,
104
+ mirroring the GDN kernel: the running slot ends up at the chunk end.
105
+ """
106
+ mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID]
107
+ state_at: dict[int, int] = {}
108
+ # `budgets` fragments the first steps; afterwards the request is alone and
109
+ # gets as much as it can use, so the prefill always finishes.
110
+ for step in range(len(budgets) + 64):
111
+ computed = request.num_computed_tokens
112
+ if computed >= request.num_tokens:
113
+ break
114
+ budget = budgets[step] if step < len(budgets) else request.num_tokens
115
+ num_new = _split(request, min(request.num_tokens - computed, budget))
116
+ if num_new == 0:
117
+ continue
118
+ assert (
119
+ manager.allocate_slots(request, num_new, num_lookahead_tokens=NUM_SPEC)
120
+ is not None
121
+ )
122
+ request.num_computed_tokens = computed + num_new
123
+ blocks = mamba_manager.req_to_blocks[request.request_id]
124
+ running = cdiv(request.num_computed_tokens, MAMBA_BLOCK_SIZE) - 1
125
+ state_at[blocks[running].block_id] = request.num_computed_tokens
126
+ return state_at
127
+
128
+
129
+def _count_cached_boundary_states(
130
+ manager: KVCacheManager, request: Request, state_at: dict[int, int]
131
+) -> int:
132
+ """Assert every hash-cached mamba slot holds the state its hash claims.
133
+
134
+ Covers both full-block snapshots (`(p + 1) * block_size`) and the
135
+ partial-tail entries align mode registers at an exact token count.
136
+
137
+ Returns the number of cached slots checked.
138
+ """
139
+ mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID]
140
+ checked = 0
141
+ for pos, block in enumerate(mamba_manager.req_to_blocks[request.request_id]):
142
+ if block.is_null or block.block_hash is None:
143
+ continue
144
+ claimed = block.block_hash_num_tokens
145
+ assert state_at.get(block.block_id) == claimed, (
146
+ f"mamba slot {pos} is hashed as state@{claimed} but holds "
147
+ f"state@{state_at.get(block.block_id)}"
148
+ )
149
+ checked += 1
150
+ return checked
151
+
152
+
153
+def _prefill(prompt_len: int, budgets: list[int]) -> int:
154
+ manager = _make_hybrid_kv_cache_manager()
155
+ (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
156
+ state_at = _run_chunked_prefill(manager, request, budgets)
157
+ assert request.num_computed_tokens == prompt_len, "prefill did not complete"
158
+ return _count_cached_boundary_states(manager, request, state_at)
159
+
160
+
161
+def test_fragmented_first_chunk_does_not_poison_mamba_prefix_cache() -> None:
162
+ """EAGLE zeroes `last_cache_position` for prompts under two blocks.
163
+
164
+ Past that point any chunk end used to be accepted, so a short first chunk
165
+ (concurrent prefills sharing the budget) left slot 0 at state@364 while the
166
+ next chunk crossed 1600 and published slot 0 as state@1600.
167
+ """
168
+ _prefill(PROMPT_LEN, budgets=[364, PROMPT_LEN])
169
+
170
+
171
+def test_fragmented_tail_chunk_does_not_poison_mamba_prefix_cache() -> None:
172
+ """Same poisoning one block in, where a hit is still cacheable.
173
+
174
+ `last_cache_position` is 1600, so the chunk ending there is cached. The next
175
+ chunk used to be free to stop mid-block (slot 1 at state@2600) and the one
176
+ after it crossed 3200, publishing slot 1 as state@3200.
177
+ """
178
+ assert _prefill(3602, budgets=[1600, 1000, 3602]) > 0
179
+
180
+
181
+@pytest.mark.parametrize("first_chunk", [800, 900, 1599, 1601, 2000])
182
+def test_intermediate_chunk_ends_stay_block_aligned(first_chunk: int) -> None:
183
+ """Every non-final prefill chunk must end on a mamba block boundary."""
184
+ _prefill(PROMPT_LEN, budgets=[first_chunk, PROMPT_LEN, PROMPT_LEN])
185
+
186
+
187
+@pytest.mark.parametrize(
188
+ ("block_size", "prompt_len", "budgets"),
189
+ [
190
+ # Kimi-K3-scale mamba blocks: TP8 shards the recurrent state 8 ways
191
+ # (~1.5k block), DEP16 keeps it whole (~12k block). The first budget
192
+ # walks the request up to `last_cache_position`; the second is the
193
+ # sub-block fragment that lands in the unguarded tail region.
194
+ (1536, 30000, [27648, 1024, 30000]),
195
+ (12288, 30000, [12288, 4000, 30000]),
196
+ (12288, 41000, [24576, 8000, 41000]),
197
+ ],
198
+)
199
+def test_poisoning_is_block_size_independent(
200
+ monkeypatch: pytest.MonkeyPatch,
201
+ block_size: int,
202
+ prompt_len: int,
203
+ budgets: list[int],
204
+) -> None:
205
+ """The invariant is per-block, so large mamba blocks are not safer."""
206
+ import sys
207
+
208
+ monkeypatch.setattr(sys.modules[__name__], "MAMBA_BLOCK_SIZE", block_size)
209
+ assert _prefill(prompt_len, budgets=budgets) > 0
210
+
211
+
212
+@pytest.mark.parametrize("partial_hit", [False, True])
213
+@pytest.mark.parametrize("resume_at", [331, 1599, 1601, 2531, 3011])
214
+def test_unaligned_resume_never_runs_past_its_block(
215
+ partial_hit: bool, resume_at: int
216
+) -> None:
217
+ """A prefill resuming mid-block must re-align before crossing a boundary.
218
+
219
+ Reachable with a finer `prefix_match_unit` (its partial-tail stop ends a
220
+ chunk off-grid by design) and with unaligned external tokens from a KV
221
+ connector.
222
+ """
223
+ prompt_len = 3602
224
+ (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
225
+ tail_boundary = prompt_len // ATTN_BLOCK_SIZE * ATTN_BLOCK_SIZE
226
+
227
+ pos, ends = resume_at, []
228
+ while pos < prompt_len:
229
+ request.num_computed_tokens = pos
230
+ num_new = _split(request, prompt_len - pos, partial_hit=partial_hit)
231
+ assert num_new > 0, f"no progress at {pos}"
232
+ if pos % MAMBA_BLOCK_SIZE != 0:
233
+ block_end = (pos // MAMBA_BLOCK_SIZE + 1) * MAMBA_BLOCK_SIZE
234
+ assert pos + num_new <= block_end, (
235
+ f"chunk [{pos}, {pos + num_new}) starts mid-block and runs past "
236
+ f"{block_end}; the slot holding state@{pos} gets hashed as "
237
+ f"state@{block_end}"
238
+ )
239
+ pos += num_new
240
+ ends.append(pos)
241
+
242
+ for end in ends[:-1]:
243
+ aligned = end % MAMBA_BLOCK_SIZE == 0
244
+ assert aligned or (partial_hit and end == tail_boundary), (
245
+ f"intermediate chunk end {end} is neither block-aligned nor the "
246
+ f"partial-tail boundary"
247
+ )
vllm/v1/core/sched/scheduler.py
+11
-11
@@ -382,7 +382,8 @@ class Scheduler(SchedulerInterface):
382
)
383
# Split only during prefill: `request.num_tokens - 1` extends this to
384
# resumed requests replaying their output tokens.
385
- if start >= max(request.num_prompt_tokens, request.num_tokens - 1):
385
+ prefill_end = max(request.num_prompt_tokens, request.num_tokens - 1)
386
+ if start >= prefill_end:
387
return num_new_tokens
388
389
block_size = self.cache_config.block_size
@@ -394,10 +395,12 @@ class Scheduler(SchedulerInterface):
395
last_cache_position = max(last_cache_position - block_size, 0)
396
397
end = start + num_new_tokens
397
- # Until `last_cache_position`, prefer chunks ending on block
398
- # boundaries. When a block cannot fit in any configured prefill chunk,
399
- # allow sub-block progress and re-align at the next reachable boundary.
400
- if end < last_cache_position:
398
+ # Invariant: slot p holds the state after exactly (p + 1) * block_size
399
+ # tokens. State is written at chunk ends, so chunk ends must be block
400
+ # aligned. Exempt: the prompt's last chunk, whose slot decode advances
401
+ # to the boundary. A block too wide for one chunk advances sub-block
402
+ # and re-aligns at the next boundary.
403
+ if end < prefill_end:
404
max_prefill_tokens = self.max_num_scheduled_tokens
405
long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold
406
if long_prefill_threshold > 0:
@@ -413,12 +416,9 @@ class Scheduler(SchedulerInterface):
416
else 0
417
)
418
stops = (
416
- # Resumed mid-block (fine-grained partial hash hit): re-align to
417
- # the block grid before running on, so the crossed boundary's
418
- # state is materialized (unless it is past the cacheable range).
419
- next_block_boundary
420
- if start % block_size != 0 and next_block_boundary <= last_cache_position
421
- else 0,
419
+ # Same invariant: a chunk starting mid-block stops at the boundary
420
+ # rather than running past it.
421
+ next_block_boundary if start % block_size != 0 else 0,
422
# Never run past the last cacheable block boundary mid-chunk.
423
last_cache_position,
424
# Fine-grained hits: the prompt's partial-tail entry can only be