@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 263c4ff95f

[Bugfix][Spec Decode] Cache the Mamba state at the block-grid position of EAGLE resume (#53945)

Signed-off-by: wzhao18 <wzhao18.sz@gmail.com> Signed-off-by: Adam Shaver <ashaver@nvidia.com> Signed-off-by: akshaver <168006157+akshaver@users.noreply.github.com> Co-authored-by: wzhao18 <wzhao18.sz@gmail.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com> Co-authored-by: roikoren755 <26850796+roikoren755@users.noreply.github.com>

akshaver committed Sep 8, 2026 at 08:59 UTC 263c4ff95fadb62d80f315171b17d4f623494e56
12 files changed +710 -21
docs/features/automatic_prefix_caching.md
+22
@@ -13,6 +13,28 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl
13
14 [examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py](../../examples/features/automatic_prefix_caching/automatic_prefix_caching_offline.py)
15
16 +## Hybrid Mamba models
17 +
18 +Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-fine-grained-prefix-cache` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it.
19 +
20 +This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold:
21 +
22 +- `--mamba-cache-mode align`
23 +- EAGLE/MTP speculative decoding on the Mamba group
24 +- `--prefix-match-unit` smaller than the Mamba block size
25 +- the model does not use multi-module MTP
26 +
27 +```bash
28 +vllm serve <hybrid-model> \
29 + --mamba-cache-mode align \
30 + --prefix-match-unit 64 \
31 + --enable-mamba-fine-grained-prefix-cache
32 +```
33 +
34 +`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect.
35 +
36 +Choose a value that divides the block size of every prefix-cacheable KV cache group, and that is a multiple of the per-state compression ratio for models that use one, such as sparse MLA. vLLM validates both at startup and names the offending sizes in the error. Read the served block size from the startup log. 64 is a reasonable starting point.
37 +
38 ## Example workloads
39
40 We describe two example workloads, where APC can provide huge performance benefit:
tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py new
+359
@@ -0,0 +1,359 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Mamba "align" check-points at the position an EAGLE sibling resumes from.
4 +
5 +Full attention hits where it holds a key; EAGLE prunes one hash unit off that
6 +candidate and drops it. Mamba materializes state only on its own block grid, so
7 +the pruned position holds nothing and the hit floors back a whole Mamba block.
8 +
9 +Gated by ``CacheConfig.enable_mamba_fine_grained_prefix_cache``. Every test here
10 +drives the real ``Scheduler._mamba_block_aligned_split`` and the real
11 +``KVCacheManager``; no chunk boundary is hard-coded.
12 +"""
13 +
14 +from dataclasses import replace
15 +from types import SimpleNamespace
16 +
17 +import pytest
18 +
19 +from tests.v1.core.test_prefix_caching import (
20 + _make_hybrid_kv_cache_config,
21 + make_kv_cache_manager,
22 + make_request,
23 +)
24 +from vllm.utils.hashing import sha256
25 +from vllm.v1.core.kv_cache_utils import init_none_hash
26 +from vllm.v1.core.sched.scheduler import Scheduler
27 +
28 +PREFIX = list(range(1, 20_001))
29 +
30 +
31 +def _manager(
32 + block_size,
33 + hash_block_size,
34 + *,
35 + fine_grained=True,
36 + num_blocks=8192,
37 + eagle_group=None,
38 + num_prefill_lookahead=0,
39 +):
40 + init_none_hash(sha256)
41 + config = _make_hybrid_kv_cache_config(
42 + block_size, num_blocks, ["full", "mamba_align"]
43 + )
44 + if eagle_group is not None:
45 + groups = list(config.kv_cache_groups)
46 + groups[eagle_group] = replace(groups[eagle_group], is_eagle_group=True)
47 + config = replace(config, kv_cache_groups=groups)
48 + return make_kv_cache_manager(
49 + kv_cache_config=config,
50 + max_model_len=1 << 20,
51 + enable_caching=True,
52 + hash_block_size=hash_block_size,
53 + use_eagle=True,
54 + num_prefill_lookahead=num_prefill_lookahead,
55 + enable_mamba_fine_grained_prefix_cache=fine_grained,
56 + )
57 +
58 +
59 +def _stub(manager, block_size, hash_block_size, *, block_drop=True):
60 + """``self`` for the real ``Scheduler._mamba_block_aligned_split``.
61 +
62 + Derives both gates exactly as ``Scheduler.__init__`` does, so a manager that
63 + cannot honour the junction stop also stops the split from making one.
64 + """
65 + partial_hit = (
66 + hash_block_size < block_size and manager.coordinator.enable_partial_hash_hits
67 + )
68 + return SimpleNamespace(
69 + cache_config=SimpleNamespace(block_size=block_size),
70 + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
71 + max_num_scheduled_tokens=1 << 20,
72 + use_eagle=True,
73 + # The EAGLE adjustments key on the block-drop bit, not plain use_eagle:
74 + # they exist only to compensate for the drop.
75 + use_eagle_block_drop=block_drop,
76 + hash_block_size=hash_block_size,
77 + mamba_has_prefill_checkpoint_blocks=False, # forced False under eagle
78 + mamba_partial_cache_hit=partial_hit,
79 + mamba_fine_grained_prefix_cache=(
80 + partial_hit and manager.mamba_fine_grained_prefix_cache
81 + ),
82 + )
83 +
84 +
85 +def _prefill(manager, stub, request, *, external=0) -> list[int]:
86 + """Schedule ``request`` to completion the way ``Scheduler.schedule()`` does.
87 +
88 + Returns the chunk ends. Every boundary comes from the real splitter, so no
89 + test in this file hard-codes one.
90 + """
91 + blocks, local, junction = manager.get_computed_blocks(request)
92 + request.shared_prefix_boundary = junction
93 + ends = []
94 + first = True
95 + while request.num_computed_tokens < request.num_tokens:
96 + new_local, ext = (local, external) if first else (0, 0)
97 + start = request.num_computed_tokens + new_local + ext
98 + if start >= request.num_tokens:
99 + break
100 + num_new = Scheduler._mamba_block_aligned_split(
101 + stub, request, request.num_tokens - start, new_local, ext
102 + )
103 + if num_new <= 0:
104 + break
105 + assert (
106 + manager.allocate_slots(
107 + request,
108 + num_new,
109 + num_new_computed_tokens=new_local,
110 + new_computed_blocks=blocks if first else None,
111 + num_external_computed_tokens=ext,
112 + has_scheduled_reqs=False,
113 + )
114 + is not None
115 + )
116 + request.num_computed_tokens = start + num_new
117 + ends.append(request.num_computed_tokens)
118 + _, retained = manager.take_kv_cache_block_copies()
119 + if retained:
120 + manager.block_pool.free_blocks(retained)
121 + manager.new_step_starts()
122 + first = False
123 + return ends
124 +
125 +
126 +def _orphaned_full_attention_tail(manager, stub, prompt_len):
127 + """Produce the state a KV connector leaves behind, which creates a junction.
128 +
129 + NIXL reports ``num_prompt_tokens - 1`` external computed tokens for Mamba
130 + models (``nixl/base_scheduler.py::_get_remote_prefill_token_count``), so the
131 + producer starts past its own tail stop: full attention still registers a
132 + partial tail the Mamba group has no state for.
133 + """
134 + producer = make_request(
135 + "producer", PREFIX[:prompt_len], stub.hash_block_size, sha256
136 + )
137 + _prefill(manager, stub, producer, external=prompt_len - 1)
138 +
139 +
140 +def _sibling_hit(manager, shared, suffix, hash_block_size):
141 + request = make_request("sibling", PREFIX[:shared] + suffix, hash_block_size, sha256)
142 + return manager.get_computed_blocks(request)[1]
143 +
144 +
145 +# --------------------------------------------------------------------------
146 +# The two positions a sibling resumes at
147 +
148 +
149 +def test_sibling_resumes_from_the_observed_junction():
150 + """The scheduler splits at the junction; the manager must cache there."""
151 + block_size, hash_block_size = 512, 32
152 + manager = _manager(block_size, hash_block_size)
153 + stub = _stub(manager, block_size, hash_block_size)
154 + assert stub.mamba_fine_grained_prefix_cache, "the feature must be armed"
155 + _orphaned_full_attention_tail(manager, stub, 2020)
156 +
157 + consumer = make_request(
158 + "consumer", PREFIX[:2016] + [-1] * 584, hash_block_size, sha256
159 + )
160 + junction = manager.get_computed_blocks(consumer)[2]
161 + assert junction, "expected a junction against the orphaned full-attention tail"
162 + _prefill(manager, stub, consumer)
163 +
164 + hit = _sibling_hit(manager, 2016, [-2] * 584, hash_block_size)
165 + assert hit == junction, (
166 + f"the chunk stopped at {junction} but nothing was cached there: "
167 + f"sibling resumes at {hit}"
168 + )
169 +
170 +
171 +def test_sibling_resumes_below_the_block_grid_when_the_prefix_ends_early():
172 + """A system prompt followed by a per-request suffix -- the deployed shape.
173 +
174 + The owner's prompt tail sits over its own suffix, which no sibling shares, so
175 + a check-point there is unreachable. The next request's full-attention match
176 + stops at the last shared block boundary and EAGLE drops one hash unit below
177 + it, so state has to exist just under the block grid too.
178 + """
179 + block_size, hash_block_size = 16, 4
180 + manager = _manager(block_size, hash_block_size)
181 + stub = _stub(manager, block_size, hash_block_size)
182 +
183 + shared = 24
184 + owner = make_request("owner", PREFIX[:shared] + [-1] * 16, hash_block_size, sha256)
185 + _prefill(manager, stub, owner)
186 +
187 + # The first follower observes the junction (16 shared, dropped to 12) and
188 + # registers state there; the second one gets to resume from it.
189 + follower = make_request(
190 + "follower", PREFIX[:shared] + [-2] * 16, hash_block_size, sha256
191 + )
192 + _prefill(manager, stub, follower)
193 +
194 + resume = shared // block_size * block_size - hash_block_size
195 + hit = _sibling_hit(manager, shared, [-3] * 16, hash_block_size)
196 + assert hit == resume, f"expected the resume point at {resume}, got {hit}"
197 +
198 +
199 +# --------------------------------------------------------------------------
200 +# Invariants the two clauses have to keep
201 +
202 +
203 +def _armed_configs():
204 + """``(block_size, hash_block_size)`` pairs the junction path can arm on.
205 +
206 + The prefix match unit must divide every group's block size and be strictly
207 + smaller, or ``enable_partial_hash_hits`` is False and the path is inert.
208 + """
209 + for block_size in (16, 128, 512):
210 + for hash_block_size in (4, 32):
211 + if hash_block_size < block_size and block_size % hash_block_size == 0:
212 + yield block_size, hash_block_size
213 +
214 +
215 +def test_scheduler_never_stops_where_the_manager_refuses():
216 + """A junction stop must always be a position the manager caches.
217 +
218 + "The scheduler schedules a number of tokens it thinks can be cached, but the
219 + cache manager won't agree" is a property of the pair, not of one
220 + configuration, so sweep the armed parameter space.
221 + """
222 + evaluated, refused, skipped = 0, [], []
223 + for block_size, hash_block_size in _armed_configs():
224 + for prompt_len in (
225 + 2 * block_size - hash_block_size,
226 + 3 * block_size + hash_block_size + 1,
227 + 4 * block_size - 1,
228 + ):
229 + manager = _manager(block_size, hash_block_size)
230 + stub = _stub(manager, block_size, hash_block_size)
231 + _orphaned_full_attention_tail(manager, stub, prompt_len)
232 +
233 + shared = prompt_len // hash_block_size * hash_block_size
234 + suffix = [-1] * (block_size + 3)
235 + consumer = make_request(
236 + "consumer", PREFIX[:shared] + suffix, hash_block_size, sha256
237 + )
238 + junction = manager.get_computed_blocks(consumer)[2]
239 + if not junction:
240 + continue
241 + case = (block_size, hash_block_size, prompt_len, junction)
242 + # The stop and the check-point are both keyed on the junction, so it
243 + # has to be on the hash grid -- nothing else keeps them together.
244 + assert junction % hash_block_size == 0, case
245 + if junction not in _prefill(manager, stub, consumer):
246 + skipped.append(case)
247 + continue
248 + evaluated += 1
249 + hit = _sibling_hit(
250 + manager, shared, [-2] * (block_size + 3), hash_block_size
251 + )
252 + if hit < junction:
253 + refused.append(case + (hit,))
254 +
255 + assert not skipped, (
256 + f"the scheduler declined to stop at {len(skipped)}: {skipped[:3]}"
257 + )
258 + assert evaluated >= 10, f"sweep did not exercise the junction: {evaluated} cases"
259 + assert not refused, (
260 + f"{len(refused)} of {evaluated} configs split at a junction the manager then "
261 + f"refused (block, hash, prompt_len, junction, sibling_hit): {refused[:3]}"
262 + )
263 +
264 +
265 +def test_a_junction_past_the_prompt_falls_back_and_registers_nothing():
266 + """A resumed request's junction can land in its output tokens.
267 +
268 + The manager writes nothing there -- during decode the target is the running
269 + state block, mutated in place, which equals what its key promises only after
270 + that step's forward. So the scheduler must not stop AT it, but it must still
271 + take the block-floored stop stock vLLM would have made: dropping the stop
272 + entirely loses a check-point a sibling could have resumed from.
273 + """
274 + block_size, hash_block_size = 512, 32
275 + manager = _manager(block_size, hash_block_size)
276 + stub = _stub(manager, block_size, hash_block_size)
277 +
278 + request = make_request("r", PREFIX[:2000], hash_block_size, sha256)
279 + for _ in range(2000):
280 + request.append_output_token_ids(1)
281 + assert request.num_prompt_tokens == 2000 and request.num_tokens == 4000
282 +
283 + # 2208 is past the prompt: not used verbatim, but the block-floored stop
284 + # (2048) still applies -- not the last cacheable boundary (3072), which is
285 + # what dropping the stop would give.
286 + request.shared_prefix_boundary = 2208
287 + assert Scheduler._mamba_block_aligned_split(stub, request, 8192) == 2048
288 + # Inside the prompt it still stops, at the junction itself.
289 + request.shared_prefix_boundary = 992
290 + assert Scheduler._mamba_block_aligned_split(stub, request, 8192) == 992
291 +
292 + # And the manager refuses on its own at the position that matters: the
293 + # RUNNING state block mid-decode. Interior blocks are nulled in align mode,
294 + # so every other decode-side position is already refused for that reason --
295 + # here the block is live and holds exactly what a key would claim, and only
296 + # the prompt bound stops it being published a forward before it changes.
297 + _prefill(manager, stub, request)
298 + mamba = manager.coordinator.single_type_managers[1]
299 + running = request.num_computed_tokens
300 + assert not mamba.req_to_blocks["r"][running // block_size].is_null
301 + request.shared_prefix_boundary = running
302 + assert mamba._cache_partial_tail_block(request, running) is None
303 +
304 +
305 +@pytest.mark.parametrize("eagle_group", [None, 0])
306 +@pytest.mark.parametrize("num_prefill_lookahead", [0, 2, 33])
307 +def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead):
308 + """Turning the flag on may add check-points; it must never move one away.
309 +
310 + Two axes decide whether the manager can honour the junction stop the
311 + scheduler makes for it, and both fail silently -- as lost reuse, not as an
312 + error. ``is_eagle_group`` annotation: the manager takes the EAGLE drop per
313 + group while the scheduler flag is model-wide, so on an annotated model the
314 + Mamba group can be outside ``eagle_group_ids``. Multi-module MTP: the
315 + coordinator hands the manager ``num_computed - num_reprefillable``, not the
316 + chunk end. In either case the junction stop displaces the block-boundary
317 + stop and the check-point it replaced is never written.
318 + """
319 + block_size, hash_block_size = 512, 32
320 +
321 + def sibling_hit(fine_grained):
322 + manager = _manager(
323 + block_size,
324 + hash_block_size,
325 + fine_grained=fine_grained,
326 + eagle_group=eagle_group,
327 + num_prefill_lookahead=num_prefill_lookahead,
328 + )
329 + stub = _stub(manager, block_size, hash_block_size)
330 + _orphaned_full_attention_tail(manager, stub, 2020)
331 + consumer = make_request(
332 + "consumer", PREFIX[:2016] + [-1] * 584, hash_block_size, sha256
333 + )
334 + _prefill(manager, stub, consumer)
335 + return _sibling_hit(manager, 2016, [-2] * 584, hash_block_size)
336 +
337 + on, off = sibling_hit(True), sibling_hit(False)
338 + assert on >= off, f"enabling the flag cut the sibling's hit from {off} to {on}"
339 +
340 +
341 +def test_disabled_by_default():
342 + """Off, the junction is block-floored and no resume point is registered."""
343 + block_size, hash_block_size = 512, 32
344 + manager = _manager(block_size, hash_block_size, fine_grained=False)
345 + stub = _stub(manager, block_size, hash_block_size)
346 + _orphaned_full_attention_tail(manager, stub, 2020)
347 +
348 + consumer = make_request(
349 + "consumer", PREFIX[:2016] + [-1] * 584, hash_block_size, sha256
350 + )
351 + junction = manager.get_computed_blocks(consumer)[2]
352 + assert junction, "expected a junction"
353 + _prefill(manager, stub, consumer)
354 +
355 + assert _sibling_hit(manager, 2016, [-2] * 584, hash_block_size) < junction
356 +
357 + request = make_request("r", PREFIX[:10000], hash_block_size, sha256)
358 + request.shared_prefix_boundary = 2000
359 + assert Scheduler._mamba_block_aligned_split(stub, request, 8192) == 1536
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+175
@@ -184,6 +184,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
184 dcp_world_size=dcp_world_size,
185 scheduler_block_size=scheduler_block_size,
186 mamba_partial_cache_hit=True,
187 + mamba_fine_grained_prefix_cache=False,
188 mamba_has_prefill_checkpoint_blocks=False,
189 )
190 split = Scheduler._mamba_block_aligned_split
@@ -231,6 +232,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
232 use_eagle_block_drop=False,
233 hash_block_size=32,
234 mamba_partial_cache_hit=False,
235 + mamba_fine_grained_prefix_cache=False,
236 mamba_has_prefill_checkpoint_blocks=False,
237 )
238 req = make_request("0", [0] * prompt_length, 32, sha256)
@@ -271,6 +273,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
273 use_eagle_block_drop=False,
274 hash_block_size=32,
275 mamba_partial_cache_hit=False,
276 + mamba_fine_grained_prefix_cache=False,
277 mamba_has_prefill_checkpoint_blocks=False,
278 )
279 req = make_request("0", [0] * prompt_length, 32, sha256)
@@ -1798,6 +1801,178 @@ def test_hybrid_partial_hit_with_eagle_stays_within_group_blocks():
1801 assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None
1802
1803
1804 +def test_mamba_align_split_stops_below_eagle_proof_boundary():
1805 + """The chunk that registers the partial tail must end one hash unit below
1806 + the prompt's last hash boundary under Eagle, which drops one unit past the
1807 + candidate.
1808 + """
1809 + block_size = 1536
1810 + hash_block_size = 128
1811 + mock = SimpleNamespace(
1812 + cache_config=SimpleNamespace(block_size=block_size),
1813 + max_num_scheduled_tokens=8192,
1814 + scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
1815 + use_eagle=True,
1816 + use_eagle_block_drop=True,
1817 + hash_block_size=hash_block_size,
1818 + mamba_partial_cache_hit=True,
1819 + mamba_fine_grained_prefix_cache=False,
1820 + mamba_has_prefill_checkpoint_blocks=False,
1821 + )
1822 + split = Scheduler._mamba_block_aligned_split
1823 +
1824 + for prompt_len in (24576, 24577):
1825 + req = make_request("0", [0] * prompt_len, hash_block_size, sha256)
1826 + req.num_computed_tokens = 23040
1827 + assert (
1828 + split(self=mock, request=req, num_new_tokens=prompt_len - 23040)
1829 + == 24448 - 23040
1830 + )
1831 +
1832 + # Without the eagle block drop the tail stays at the prompt's own last hash
1833 + # boundary -- the shift exists only to compensate for the drop, so it keys
1834 + # on use_eagle_block_drop rather than plain use_eagle.
1835 + mock.use_eagle_block_drop = False
1836 + req = make_request("1", [0] * 24577, hash_block_size, sha256)
1837 + req.num_computed_tokens = 23040
1838 + assert split(self=mock, request=req, num_new_tokens=1537) == 24576 - 23040
1839 +
1840 +
1841 +def test_mamba_partial_tail_hits_below_eagle_proof_boundary():
1842 + """Regression: a prompt whose last hash boundary coincides with its last
1843 + mamba block boundary registered its only fine-grained checkpoint at a
1844 + position Eagle can never resume from, collapsing the next turn to the
1845 + coarse block boundary (or, under sparse retention, to a full miss).
1846 + """
1847 + hash_block_size = 2
1848 + mamba_block_size = 2 * hash_block_size
1849 + kv_cache_config = KVCacheConfig(
1850 + num_blocks=32,
1851 + kv_cache_tensors=[],
1852 + kv_cache_groups=[
1853 + KVCacheGroupSpec(
1854 + ["full"],
1855 + FullAttentionSpec(
1856 + block_size=hash_block_size,
1857 + num_kv_heads=1,
1858 + head_size=1,
1859 + dtype=torch.float32,
1860 + ),
1861 + ),
1862 + KVCacheGroupSpec(
1863 + ["mamba"],
1864 + MambaSpec(
1865 + block_size=mamba_block_size,
1866 + shapes=(1, 1),
1867 + dtypes=(torch.float32,),
1868 + mamba_cache_mode="align",
1869 + ),
1870 + ),
1871 + ],
1872 + )
1873 + manager = make_kv_cache_manager(
1874 + kv_cache_config=kv_cache_config,
1875 + max_model_len=8192,
1876 + enable_caching=True,
1877 + hash_block_size=hash_block_size,
1878 + use_eagle=True,
1879 + )
1880 +
1881 + # 9 tokens: the last hash boundary (8) is also a mamba block boundary, so
1882 + # the prompt's two checkpoints collapse. The reachable tail is 6.
1883 + req0 = make_request("0", [7] * 9, hash_block_size, sha256)
1884 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1885 + assert num_computed == 0
1886 + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None
1887 + req0.num_computed_tokens = 6
1888 + manager.new_step_starts()
1889 + assert manager.allocate_slots(req0, 3) is not None
1890 + req0.num_computed_tokens = 9
1891 + manager.new_step_starts()
1892 +
1893 + partial_hash = req0.block_hashes[6 // hash_block_size - 1]
1894 + partial_block = manager.block_pool.get_cached_block(
1895 + partial_hash, kv_cache_group_ids=[1]
1896 + )
1897 + assert partial_block is not None
1898 + assert partial_block[0].block_hash_num_tokens == 6
1899 +
1900 + # Full attention matches to 8 and drops one hash unit to 6, where the
1901 + # mamba tail now exists -- instead of falling back to the block boundary 4.
1902 + req1 = make_request("1", [7] * 9 + [9] * 4, hash_block_size, sha256)
1903 + _, num_computed, _ = manager.get_computed_blocks(req1)
1904 + assert num_computed == 6
1905 +
1906 +
1907 +def test_mamba_coarse_checkpoint_retained_below_eagle_proof_boundary():
1908 + """Without a prefix-match unit there is no partial tail, so the only
1909 + checkpoint comes from the retention mask. A prompt ending just past its
1910 + last block boundary pinned that checkpoint where eagle can never resume,
1911 + so under ``retention_interval=0`` the reachable block was masked out and
1912 + the next turn missed entirely.
1913 + """
1914 + block_size = 4
1915 + kv_cache_config = KVCacheConfig(
1916 + num_blocks=32,
1917 + kv_cache_tensors=[],
1918 + kv_cache_groups=[
1919 + KVCacheGroupSpec(
1920 + ["full"],
1921 + FullAttentionSpec(
1922 + block_size=block_size,
1923 + num_kv_heads=1,
1924 + head_size=1,
1925 + dtype=torch.float32,
1926 + ),
1927 + ),
1928 + KVCacheGroupSpec(
1929 + ["mamba"],
1930 + MambaSpec(
1931 + block_size=block_size,
1932 + shapes=(1, 1),
1933 + dtypes=(torch.float32,),
1934 + mamba_cache_mode="align",
1935 + ),
1936 + ),
1937 + ],
1938 + prefix_cache_retention_interval=0,
1939 + )
1940 + manager = make_kv_cache_manager(
1941 + kv_cache_config=kv_cache_config,
1942 + max_model_len=8192,
1943 + enable_caching=True,
1944 + hash_block_size=block_size,
1945 + use_eagle=True,
1946 + )
1947 +
1948 + # 9 tokens: the last aligned boundary is 8, which eagle cannot resume
1949 + # from, so block 0 (ending at 4) must stay retained. Prefill in
1950 + # scheduler-split style -- Mamba keeps one rolling state, so a boundary is
1951 + # only materialized by a chunk that ends exactly on it.
1952 + req0 = make_request("0", [7] * 9, block_size, sha256)
1953 + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0)
1954 + assert num_computed == 0
1955 + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None
1956 + req0.num_computed_tokens = 4
1957 + manager.new_step_starts()
1958 + assert manager.allocate_slots(req0, 5) is not None
1959 + req0.num_computed_tokens = 9
1960 + manager.free(req0)
1961 + manager.new_step_starts()
1962 +
1963 + # The reachable state (4) is retained; the unreachable one (8) is not.
1964 + assert (
1965 + manager.block_pool.get_cached_block(
1966 + req0.block_hashes[0], kv_cache_group_ids=[1]
1967 + )
1968 + is not None
1969 + )
1970 +
1971 + req1 = make_request("1", [7] * 9 + [9] * 4, block_size, sha256)
1972 + _, num_computed, _ = manager.get_computed_blocks(req1)
1973 + assert num_computed == 4
1974 +
1975 +
1976 def test_hybrid_sliding_window_group_disables_partial_hash_hits():
1977 hash_block_size = 2
1978 sliding_window_block_size = 2 * hash_block_size
tests/v1/core/test_mamba_align_chunk_split.py
+9 -3
@@ -144,6 +144,7 @@ def _split(
144 scheduler_config=SimpleNamespace(long_prefill_token_threshold=0),
145 # `prefix_match_unit` finer than the block size (#46384).
146 mamba_partial_cache_hit=partial_hit,
147 + mamba_fine_grained_prefix_cache=False,
148 hash_block_size=ATTN_BLOCK_SIZE,
149 mamba_has_prefill_checkpoint_blocks=(num_prefill_checkpoint_blocks > 0),
150 mamba_prefill_checkpoint_alignment=(
@@ -357,7 +358,12 @@ def test_unaligned_resume_never_runs_past_its_block(
358 """
359 prompt_len = 3602
360 (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE)
360 - tail_boundary = prompt_len // ATTN_BLOCK_SIZE * ATTN_BLOCK_SIZE
361 + # Under eagle the partial-tail stop sits one hash unit below the prompt's
362 + # last hash boundary: eagle matches a unit past its candidate and drops it,
363 + # so nothing proves that last boundary.
364 + tail_stop = prompt_len // ATTN_BLOCK_SIZE * ATTN_BLOCK_SIZE
365 + if use_eagle:
366 + tail_stop -= ATTN_BLOCK_SIZE
367
368 pos, ends = resume_at, []
369 while pos < prompt_len:
@@ -382,7 +388,7 @@ def test_unaligned_resume_never_runs_past_its_block(
388
389 for end in ends[:-1]:
390 aligned = end % MAMBA_BLOCK_SIZE == 0
385 - assert aligned or (partial_hit and end == tail_boundary), (
391 + assert aligned or (partial_hit and end == tail_stop), (
392 f"intermediate chunk end {end} is neither block-aligned nor the "
387 - f"partial-tail boundary"
393 + f"partial-tail stop ({tail_stop})"
394 )
tests/v1/core/test_prefix_caching.py
+8 -6
@@ -1169,6 +1169,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
1169 use_eagle_block_drop=False,
1170 hash_block_size=block_size,
1171 mamba_partial_cache_hit=False,
1172 + mamba_fine_grained_prefix_cache=False,
1173 mamba_has_prefill_checkpoint_blocks=False,
1174 )
1175 req_2.shared_prefix_boundary = shared_prefix_boundary
@@ -3607,9 +3608,10 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():
3608 use_eagle=True,
3609 )
3610
3610 - # 127 tokens: latest replay boundary is floor((127 - 1) / 32) * 32 = 96.
3611 - # The EAGLE/MTP SWA lookup group must cache the local tail ending at
3612 - # 104 tokens, and that tail is two 8-token blocks wide: hashes 11 and 12.
3611 + # 127 tokens: eagle proves a boundary only if an aligned unit exists past
3612 + # it, so the replay boundary rewinds one unit to
3613 + # floor(127 / 32) * 32 - 32 = 64. The SWA tail plus its proof block then
3614 + # ends at 72, two 8-token blocks wide: hashes 7 and 8.
3615 token_ids = [i for i in range(15) for _ in range(block_size)] + [15] * 7
3616 req0 = make_request("0", token_ids, block_size, sha256)
3617 computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0)
@@ -3623,7 +3625,7 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():
3625 assert blocks is not None
3626
3627 pool = manager.block_pool
3626 - expected_swa_cached = {11, 12}
3628 + expected_swa_cached = {7, 8}
3629 for i in range(15):
3630 cached = pool.get_cached_block(req0.block_hashes[i], kv_cache_group_ids=[1])
3631 if i in expected_swa_cached:
@@ -3635,8 +3637,8 @@ def test_hybrid_local_kv_retention_mtp_reuses_latest_boundary():
3637
3638 req1 = make_request("1", token_ids, block_size, sha256)
3639 computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1)
3638 - assert num_computed_tokens == 12 * block_size
3639 - assert [len(blocks) for blocks in computed_blocks.blocks] == [3, 12]
3640 + assert num_computed_tokens == 8 * block_size
3641 + assert [len(blocks) for blocks in computed_blocks.blocks] == [2, 8]
3642
3643
3644 def test_block_lookup_cache_single_block_per_key():
tests/v1/core/test_single_type_kv_cache_manager.py
+1 -1
@@ -192,7 +192,7 @@ def test_circular_buffer_allocates_one_block_for_the_request_lifetime():
192 )
193 assert manager.allocate_new_blocks(request_id, num_tokens, num_tokens) == []
194
195 - manager.cache_blocks(request_id, 1024)
195 + manager.cache_blocks(request_id, 1024, replay_boundary=0)
196 manager.remove_skipped_blocks(request_id, 1024)
197 assert manager.get_num_common_prefix_blocks(request_id) == 0
198 assert manager.get_num_skipped_tokens(1024) == 0
vllm/config/cache.py
+7
@@ -195,6 +195,12 @@ class CacheConfig:
195 when the token is at position i * block_size. This is the default when prefix
196 caching is enabled.
197 """
198 + enable_mamba_fine_grained_prefix_cache: bool = False
199 + """Also register a Mamba "align" checkpoint at the shared-prefix junction --
200 + where an EAGLE/MTP sibling was observed to resume -- instead of only at the
201 + prompt tail. Off by default; only takes effect with `mamba_cache_mode`
202 + "align", EAGLE on the Mamba group, and a prefix match unit smaller than the
203 + Mamba block size."""
204 replayssm_buffer_len: int = Field(default=16, gt=0)
205 """ReplaySSM logical history length B for Mamba2. Triton uses B physical
206 rows and FlashInfer uses B+1. Kimi-K3 speculative decode does not use B.
@@ -274,6 +280,7 @@ class CacheConfig:
280 "prefix_cache_retention_interval",
281 # Prefix-caching implementation detail (doesn't affect compiled graph).
282 "prefix_match_unit",
283 + "enable_mamba_fine_grained_prefix_cache",
284 "mamba_page_size_padded",
285 "skip_page_size_padded",
286 "user_specified_block_size",
vllm/engine/arg_utils.py
+10
@@ -726,6 +726,9 @@ class EngineArgs:
726 mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
727 prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit")
728 mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode
729 + enable_mamba_fine_grained_prefix_cache: bool = (
730 + CacheConfig.enable_mamba_fine_grained_prefix_cache
731 + )
732 replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len
733 use_replayssm: bool = CacheConfig.use_replayssm
734
@@ -1287,6 +1290,10 @@ class EngineArgs:
1290 cache_group.add_argument(
1291 "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]
1292 )
1293 + cache_group.add_argument(
1294 + "--enable-mamba-fine-grained-prefix-cache",
1295 + **cache_kwargs["enable_mamba_fine_grained_prefix_cache"],
1296 + )
1297 cache_group.add_argument(
1298 "--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"]
1299 )
@@ -2072,6 +2079,9 @@ class EngineArgs:
2079 mamba_block_size=self.mamba_block_size,
2080 prefix_match_unit=self.prefix_match_unit,
2081 mamba_cache_mode=self.mamba_cache_mode,
2082 + enable_mamba_fine_grained_prefix_cache=(
2083 + self.enable_mamba_fine_grained_prefix_cache
2084 + ),
2085 replayssm_buffer_len=self.replayssm_buffer_len,
2086 use_replayssm=self.use_replayssm,
2087 kv_offloading_size=self.kv_offloading_size,
vllm/v1/core/kv_cache_coordinator.py
+26
@@ -309,6 +309,28 @@ class KVCacheCoordinator(ABC):
309 for manager in self.single_type_managers
310 )
311
312 + def get_replay_boundary(self, request: Request) -> int:
313 + """Return the position a later request replaying this prompt resumes at.
314 +
315 + A cache hit is the shortest hit across all groups, so this is a
316 + model-level position: every group has to retain state here, whether or
317 + not it is the group that drops. Groups differ only in how much they
318 + keep around it -- EAGLE groups also keep the block above, which they
319 + match and drop back from (see ``reachable_block_mask``).
320 +
321 + Under EAGLE that block must exist, so the boundary sits one alignment
322 + unit below the prompt's last aligned position; every group's block size
323 + divides the alignment, so the block above always fits in the prompt.
324 + """
325 + if not self.eagle_group_ids:
326 + return request.num_prompt_tokens - 1
327 + aligned = (
328 + request.num_prompt_tokens
329 + // self.scheduler_block_size
330 + * self.scheduler_block_size
331 + )
332 + return max(aligned - self.scheduler_block_size, 0)
333 +
334 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
335 """
336 Cache the blocks for the request.
@@ -319,6 +341,7 @@ class KVCacheCoordinator(ABC):
341 that need to be cached
342 (including tokens that are already cached).
343 """
344 + replay_boundary = self.get_replay_boundary(request)
345 for manager in self.single_type_managers:
346 # Only cache tokens with finalized KV. The last num_reprefillable_tokens
347 # tokens can be re-prefilled during multi-module MTP.
@@ -329,6 +352,7 @@ class KVCacheCoordinator(ABC):
352 request,
353 num_tokens_to_cache,
354 retention_interval=self.retention_interval,
355 + replay_boundary=replay_boundary,
356 )
357
358 def free(self, request_id: str) -> None:
@@ -753,6 +777,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
777
778 def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
779 cached_num_computed_tokens = self._align_cacheable(num_computed_tokens)
780 + replay_boundary = self.get_replay_boundary(request)
781 for manager in self.single_type_managers:
782 num_tokens_to_cache = cached_num_computed_tokens
783 # EAGLE groups match one block past each aligned boundary and drop
@@ -779,6 +804,7 @@ class HybridKVCacheCoordinator(KVCacheCoordinator):
804 request,
805 num_tokens_to_cache,
806 retention_interval=self.retention_interval,
807 + replay_boundary=replay_boundary,
808 )
809
810 def find_longest_cache_hit(
vllm/v1/core/kv_cache_manager.py
+17
@@ -15,6 +15,7 @@ from vllm.v1.core.kv_cache_coordinator import (
15 )
16 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
17 from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy
18 +from vllm.v1.core.single_type_kv_cache_manager import MambaManager
19 from vllm.v1.kv_cache_interface import (
20 AttentionSpec,
21 CrossAttentionSpec,
@@ -132,6 +133,7 @@ class KVCacheManager:
133 pcp_world_size: int = 1,
134 metrics_collector: KVCacheMetricsCollector | None = None,
135 watermark: float = 0.0,
136 + enable_mamba_fine_grained_prefix_cache: bool = False,
137 ) -> None:
138 self.max_model_len = max_model_len
139 # When unset, fall back to `max_model_len` so the recycling-aware cap
@@ -164,6 +166,21 @@ class KVCacheManager:
166 metrics_collector=self.metrics_collector,
167 num_prefill_lookahead=num_prefill_lookahead,
168 )
169 + # One predicate, read by both sides of the feature, so the scheduler
170 + # cannot end a chunk at a junction the manager would refuse -- a refused
171 + # junction costs a forward pass and displaces the block-boundary stop.
172 + # Multi-module MTP is excluded because ``cache_blocks`` then hands the
173 + # manager ``num_computed - num_reprefillable`` rather than the chunk end.
174 + self.mamba_fine_grained_prefix_cache = (
175 + enable_mamba_fine_grained_prefix_cache
176 + and bool(self.coordinator.eagle_group_ids)
177 + and self.coordinator.enable_partial_hash_hits
178 + and self.coordinator.num_reprefillable_tokens == 0
179 + )
180 + if self.mamba_fine_grained_prefix_cache:
181 + for manager in self.coordinator.single_type_managers:
182 + if isinstance(manager, MambaManager):
183 + manager.fine_grained_prefix_cache = True
184 self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
185 self.block_pool = self.coordinator.block_pool
186 self.kv_cache_config = kv_cache_config
vllm/v1/core/sched/scheduler.py
+34 -6
@@ -312,6 +312,9 @@ class Scheduler(SchedulerInterface):
312 hash_block_size=hash_block_size,
313 metrics_collector=self.kv_metrics_collector,
314 watermark=self.scheduler_config.watermark,
315 + enable_mamba_fine_grained_prefix_cache=(
316 + self.cache_config.enable_mamba_fine_grained_prefix_cache
317 + ),
318 )
319 # Bind GPU block pool to the KV connector. This must happen after
320 # kv_cache_manager is constructed so block_pool is available.
@@ -362,6 +365,14 @@ class Scheduler(SchedulerInterface):
365 and self.hash_block_size < self.block_size
366 and self.kv_cache_manager.coordinator.enable_partial_hash_hits
367 )
368 + # Opt-in: also stop at the junction, where an eagle sibling resumes. The
369 + # manager decides whether it can check-point there (per-group eagle bit,
370 + # no MTP re-prefill tail); splitting for a stop it would refuse costs a
371 + # forward pass and displaces the block-boundary stop.
372 + self.mamba_fine_grained_prefix_cache = (
373 + self.mamba_partial_cache_hit
374 + and self.kv_cache_manager.mamba_fine_grained_prefix_cache
375 + )
376
377 # Counts of non-empty steps scheduled / processed. update_from_output
378 # is called once per scheduled step in FIFO order, so these stay in sync.
@@ -475,6 +486,26 @@ class Scheduler(SchedulerInterface):
486 if self.mamba_partial_cache_hit and not use_internal_checkpoint
487 else 0
488 )
489 + if tail_boundary and self.use_eagle_block_drop:
490 + # Eagle matches one hash unit past the candidate and drops it, so
491 + # nothing proves the prompt's own last hash boundary. Materialize
492 + # the state one unit lower, where the hit can actually land. Keyed on
493 + # the block-drop bit, not plain use_eagle: this shift exists only to
494 + # compensate for the drop, and the Mamba manager's matching gate
495 + # reads the same bit (the coordinator is handed use_eagle_block_drop).
496 + tail_boundary = max(tail_boundary - self.hash_block_size, 0)
497 + junction = request.shared_prefix_boundary
498 + # Block-floored: a sub-block junction's state is not separately cacheable.
499 + block_floored = start + (junction - start) // block_size * block_size
500 + # Past the prompt the manager writes nothing, so fall back to the
501 + # block-floored stop rather than dropping it: a resumed request replaying
502 + # output tokens can still observe a junction there.
503 + junction_stop = (
504 + junction
505 + if self.mamba_fine_grained_prefix_cache
506 + and junction <= request.num_prompt_tokens
507 + else block_floored
508 + )
509 stops = (
510 # Same invariant: a chunk starting mid-block stops at the boundary
511 # rather than running past it.
@@ -488,12 +519,9 @@ class Scheduler(SchedulerInterface):
519 tail_boundary
520 if last_cache_position < tail_boundary < request.num_prompt_tokens
521 else 0,
491 - # Marconi shared-prefix junction, block-floored (a sub-block
492 - # junction's state is not separately cacheable): cache its state
493 - # so sibling requests sharing the prefix can reuse it.
494 - start + (request.shared_prefix_boundary - start) // block_size * block_size
495 - if start < request.shared_prefix_boundary < end
496 - else 0,
522 + # Marconi shared-prefix junction: cache its state so sibling
523 + # requests sharing the prefix can reuse it.
524 + junction_stop if start < junction < end else 0,
525 )
526 # Stop at the earliest mandatory position strictly inside the chunk.
527 end = min((s for s in stops if start < s < end), default=end)
vllm/v1/core/single_type_kv_cache_manager.py
+42 -5
@@ -119,7 +119,9 @@ class SingleTypeKVCacheManager(ABC):
119 # aligned segment (SWA). Initialized lazily by the coordinator after
120 # determining the attention groups.
121 self.use_eagle = False
122 -
122 + # ``CacheConfig.enable_mamba_fine_grained_prefix_cache``, narrowed and set
123 + # by ``KVCacheManager``; only an EAGLE Mamba "align" group ever gets it.
124 + self.fine_grained_prefix_cache = False
125 # Partial-hit copy-on-write bookkeeping. Populated only by fine-grained
126 # managers (full attention, mamba "align"); harmlessly empty elsewhere.
127 self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {}
@@ -449,6 +451,8 @@ class SingleTypeKVCacheManager(ABC):
451 request: Request,
452 num_tokens: int,
453 retention_interval: int | None = None,
454 + *,
455 + replay_boundary: int,
456 ) -> None:
457 """
458 Cache the blocks for the request.
@@ -471,7 +475,7 @@ class SingleTypeKVCacheManager(ABC):
475 # Token boundaries whose reachable tail must be retained under sparse
476 # retention: the replay boundary (``num_prompt - 1``, capped by
477 # ``get_computed_blocks``) and any detected shared-prefix junction.
474 - reachable_boundaries = [request.num_prompt_tokens - 1]
478 + reachable_boundaries = [replay_boundary]
479 if request.shared_prefix_boundary:
480 reachable_boundaries.append(request.shared_prefix_boundary)
481
@@ -813,8 +817,15 @@ class FullAttentionManager(SingleTypeKVCacheManager):
817 request: Request,
818 num_tokens: int,
819 retention_interval: int | None = None,
820 + *,
821 + replay_boundary: int,
822 ) -> None:
817 - super().cache_blocks(request, num_tokens, retention_interval=retention_interval)
823 + super().cache_blocks(
824 + request,
825 + num_tokens,
826 + retention_interval=retention_interval,
827 + replay_boundary=replay_boundary,
828 + )
829 hash_block_size = self.block_pool.hash_block_size
830 if self.block_size == hash_block_size:
831 return
@@ -1216,6 +1227,8 @@ class CircularBufferManager(FullAttentionManager):
1227 request: Request,
1228 num_tokens: int,
1229 retention_interval: int | None = None,
1230 + *,
1231 + replay_boundary: int,
1232 ) -> None:
1233 return
1234
@@ -1926,9 +1939,16 @@ class MambaManager(SingleTypeKVCacheManager):
1939 request: Request,
1940 num_tokens: int,
1941 retention_interval: int | None = None,
1942 + *,
1943 + replay_boundary: int,
1944 ) -> None:
1945 num_cached_blocks_before = self.num_cached_block.get(request.request_id, 0)
1931 - super().cache_blocks(request, num_tokens, retention_interval=retention_interval)
1946 + super().cache_blocks(
1947 + request,
1948 + num_tokens,
1949 + retention_interval=retention_interval,
1950 + replay_boundary=replay_boundary,
1951 + )
1952 num_cached_blocks_after = self.num_cached_block.get(request.request_id, 0)
1953 if self.mamba_cache_mode == "align":
1954 partial_hash = self._cache_partial_tail_block(request, num_tokens)
@@ -1994,7 +2014,22 @@ class MambaManager(SingleTypeKVCacheManager):
2014 latest_prompt_hash_boundary = (
2015 request.num_prompt_tokens // hash_block_size
2016 ) * hash_block_size
1997 - if num_tokens != latest_prompt_hash_boundary:
2017 + if self.use_eagle:
2018 + # Eagle groups match one hash unit past the candidate and drop it,
2019 + # so register the tail one unit lower.
2020 + latest_prompt_hash_boundary = max(
2021 + latest_prompt_hash_boundary - hash_block_size, 0
2022 + )
2023 + # The junction is the other position a sibling resumes at: where one was
2024 + # observed to stop, and where the scheduler already ends a chunk. Bounded
2025 + # to the prompt chunk being computed -- during decode the target is the
2026 + # running state block, mutated in place, which equals what its key
2027 + # promises only after that step's forward.
2028 + if num_tokens != latest_prompt_hash_boundary and not (
2029 + self.fine_grained_prefix_cache
2030 + and num_tokens == request.shared_prefix_boundary
2031 + and request.num_computed_tokens < num_tokens <= request.num_prompt_tokens
2032 + ):
2033 return None
2034
2035 block_idx = num_tokens // self.block_size
@@ -2054,6 +2089,8 @@ class CrossAttentionManager(SingleTypeKVCacheManager):
2089 request: Request,
2090 num_tokens: int,
2091 retention_interval: int | None = None,
2092 + *,
2093 + replay_boundary: int,
2094 ) -> None:
2095 # We do not cache blocks for cross-attention to be shared between
2096 # requests, so this method is not relevant.