[Misc] Rename --enable-mamba-fine-grained-prefix-cache (#53945 follow-up) (#57382)
Signed-off-by: Roi Koren <roik@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
roikoren755 committed
Sep 20, 2026 at 19:50 UTC
49ee12d742b4a5524d2a3a3ee6c3fb77e4913073
10 files changed
+38
-38
docs/features/automatic_prefix_caching.md
+2
-2
@@ -15,7 +15,7 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl
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.
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-shared-prefix-checkpoint` 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
@@ -28,7 +28,7 @@ This helps when many requests share a long system prompt and then diverge. It is
28
vllm serve <hybrid-model> \
29
--mamba-cache-mode align \
30
--prefix-match-unit 64 \
31
- --enable-mamba-fine-grained-prefix-cache
31
+ --enable-mamba-shared-prefix-checkpoint
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.
tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py
+9
-9
@@ -6,7 +6,7 @@ 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
9
+Gated by ``CacheConfig.enable_mamba_shared_prefix_checkpoint``. Every test here
10
drives the real ``Scheduler._mamba_block_aligned_split`` and the real
11
``KVCacheManager``; no chunk boundary is hard-coded.
12
"""
@@ -32,7 +32,7 @@ def _manager(
32
block_size,
33
hash_block_size,
34
*,
35
- fine_grained=True,
35
+ shared_prefix_checkpoint=True,
36
num_blocks=8192,
37
eagle_group=None,
38
num_prefill_lookahead=0,
@@ -52,7 +52,7 @@ def _manager(
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,
55
+ enable_mamba_shared_prefix_checkpoint=shared_prefix_checkpoint,
56
)
57
58
@@ -76,8 +76,8 @@ def _stub(manager, block_size, hash_block_size, *, block_drop=True):
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
79
+ mamba_shared_prefix_checkpoint=(
80
+ partial_hit and manager.mamba_shared_prefix_checkpoint
81
),
82
)
83
@@ -151,7 +151,7 @@ def test_sibling_resumes_from_the_observed_junction():
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"
154
+ assert stub.mamba_shared_prefix_checkpoint, "the feature must be armed"
155
_orphaned_full_attention_tail(manager, stub, 2020)
156
157
consumer = make_request(
@@ -325,11 +325,11 @@ def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead):
325
"""
326
block_size, hash_block_size = 512, 32
327
328
- def sibling_hit(fine_grained):
328
+ def sibling_hit(shared_prefix_checkpoint):
329
manager = _manager(
330
block_size,
331
hash_block_size,
332
- fine_grained=fine_grained,
332
+ shared_prefix_checkpoint=shared_prefix_checkpoint,
333
eagle_group=eagle_group,
334
num_prefill_lookahead=num_prefill_lookahead,
335
)
@@ -348,7 +348,7 @@ def test_enabling_never_reduces_reuse(eagle_group, num_prefill_lookahead):
348
def test_disabled_by_default():
349
"""Off, the junction is block-floored and no resume point is registered."""
350
block_size, hash_block_size = 512, 32
351
- manager = _manager(block_size, hash_block_size, fine_grained=False)
351
+ manager = _manager(block_size, hash_block_size, shared_prefix_checkpoint=False)
352
stub = _stub(manager, block_size, hash_block_size)
353
_orphaned_full_attention_tail(manager, stub, 2020)
354
tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py
+4
-4
@@ -188,7 +188,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int):
188
dcp_world_size=dcp_world_size,
189
scheduler_block_size=scheduler_block_size,
190
mamba_partial_cache_hit=True,
191
- mamba_fine_grained_prefix_cache=False,
191
+ mamba_shared_prefix_checkpoint=False,
192
mamba_has_prefill_checkpoint_blocks=False,
193
)
194
split = Scheduler._mamba_block_aligned_split
@@ -236,7 +236,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget():
236
use_eagle_block_drop=False,
237
hash_block_size=32,
238
mamba_partial_cache_hit=False,
239
- mamba_fine_grained_prefix_cache=False,
239
+ mamba_shared_prefix_checkpoint=False,
240
mamba_has_prefill_checkpoint_blocks=False,
241
)
242
req = make_request("0", [0] * prompt_length, 32, sha256)
@@ -277,7 +277,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold():
277
use_eagle_block_drop=False,
278
hash_block_size=32,
279
mamba_partial_cache_hit=False,
280
- mamba_fine_grained_prefix_cache=False,
280
+ mamba_shared_prefix_checkpoint=False,
281
mamba_has_prefill_checkpoint_blocks=False,
282
)
283
req = make_request("0", [0] * prompt_length, 32, sha256)
@@ -1826,7 +1826,7 @@ def test_mamba_align_split_stops_below_eagle_proof_boundary():
1826
use_eagle_block_drop=True,
1827
hash_block_size=hash_block_size,
1828
mamba_partial_cache_hit=True,
1829
- mamba_fine_grained_prefix_cache=False,
1829
+ mamba_shared_prefix_checkpoint=False,
1830
mamba_has_prefill_checkpoint_blocks=False,
1831
)
1832
split = Scheduler._mamba_block_aligned_split
tests/v1/core/test_mamba_align_chunk_split.py
+1
-1
@@ -144,7 +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,
147
+ mamba_shared_prefix_checkpoint=False,
148
hash_block_size=ATTN_BLOCK_SIZE,
149
mamba_has_prefill_checkpoint_blocks=(num_prefill_checkpoint_blocks > 0),
150
mamba_prefill_checkpoint_alignment=(
tests/v1/core/test_prefix_caching.py
+1
-1
@@ -2269,7 +2269,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection():
2269
use_eagle_block_drop=False,
2270
hash_block_size=block_size,
2271
mamba_partial_cache_hit=False,
2272
- mamba_fine_grained_prefix_cache=False,
2272
+ mamba_shared_prefix_checkpoint=False,
2273
mamba_has_prefill_checkpoint_blocks=False,
2274
)
2275
req_2.shared_prefix_boundary = shared_prefix_boundary
vllm/config/cache.py
+2
-2
@@ -184,7 +184,7 @@ class CacheConfig:
184
when the token is at position i * block_size. This is the default when prefix
185
caching is enabled.
186
"""
187
- enable_mamba_fine_grained_prefix_cache: bool = False
187
+ enable_mamba_shared_prefix_checkpoint: bool = False
188
"""Also register a Mamba "align" checkpoint at the shared-prefix junction --
189
where an EAGLE/MTP sibling was observed to resume -- instead of only at the
190
prompt tail. Off by default; only takes effect with `mamba_cache_mode`
@@ -274,7 +274,7 @@ class CacheConfig:
274
"prefix_cache_retention_interval",
275
# Prefix-caching implementation detail (doesn't affect compiled graph).
276
"prefix_match_unit",
277
- "enable_mamba_fine_grained_prefix_cache",
277
+ "enable_mamba_shared_prefix_checkpoint",
278
"mamba_page_size_padded",
279
"skip_page_size_padded",
280
"user_specified_block_size",
vllm/engine/arg_utils.py
+6
-6
@@ -753,8 +753,8 @@ class EngineArgs:
753
mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
754
prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit")
755
mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode
756
- enable_mamba_fine_grained_prefix_cache: bool = (
757
- CacheConfig.enable_mamba_fine_grained_prefix_cache
756
+ enable_mamba_shared_prefix_checkpoint: bool = (
757
+ CacheConfig.enable_mamba_shared_prefix_checkpoint
758
)
759
replayssm_buffer_len: int = CacheConfig.replayssm_buffer_len
760
use_replayssm: bool = CacheConfig.use_replayssm
@@ -1330,8 +1330,8 @@ class EngineArgs:
1330
"--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]
1331
)
1332
cache_group.add_argument(
1333
- "--enable-mamba-fine-grained-prefix-cache",
1334
- **cache_kwargs["enable_mamba_fine_grained_prefix_cache"],
1333
+ "--enable-mamba-shared-prefix-checkpoint",
1334
+ **cache_kwargs["enable_mamba_shared_prefix_checkpoint"],
1335
)
1336
cache_group.add_argument(
1337
"--replayssm-buffer-len", **cache_kwargs["replayssm_buffer_len"]
@@ -2128,8 +2128,8 @@ class EngineArgs:
2128
mamba_block_size=self.mamba_block_size,
2129
prefix_match_unit=self.prefix_match_unit,
2130
mamba_cache_mode=self.mamba_cache_mode,
2131
- enable_mamba_fine_grained_prefix_cache=(
2132
- self.enable_mamba_fine_grained_prefix_cache
2131
+ enable_mamba_shared_prefix_checkpoint=(
2132
+ self.enable_mamba_shared_prefix_checkpoint
2133
),
2134
replayssm_buffer_len=self.replayssm_buffer_len,
2135
use_replayssm=self.use_replayssm,
vllm/v1/core/kv_cache_manager.py
+5
-5
@@ -145,7 +145,7 @@ class KVCacheManager:
145
pcp_world_size: int = 1,
146
metrics_collector: KVCacheMetricsCollector | None = None,
147
watermark: float = 0.0,
148
- enable_mamba_fine_grained_prefix_cache: bool = False,
148
+ enable_mamba_shared_prefix_checkpoint: bool = False,
149
) -> None:
150
self.max_model_len = max_model_len
151
# When unset, fall back to `max_model_len` so the recycling-aware cap
@@ -183,16 +183,16 @@ class KVCacheManager:
183
# junction costs a forward pass and displaces the block-boundary stop.
184
# Multi-module MTP is excluded because ``cache_blocks`` then hands the
185
# manager ``num_computed - num_reprefillable`` rather than the chunk end.
186
- self.mamba_fine_grained_prefix_cache = (
187
- enable_mamba_fine_grained_prefix_cache
186
+ self.mamba_shared_prefix_checkpoint = (
187
+ enable_mamba_shared_prefix_checkpoint
188
and bool(self.coordinator.eagle_group_ids)
189
and self.coordinator.enable_partial_hash_hits
190
and self.coordinator.num_reprefillable_tokens == 0
191
)
192
- if self.mamba_fine_grained_prefix_cache:
192
+ if self.mamba_shared_prefix_checkpoint:
193
for manager in self.coordinator.single_type_managers:
194
if isinstance(manager, MambaManager):
195
- manager.fine_grained_prefix_cache = True
195
+ manager.shared_prefix_checkpoint = True
196
self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
197
self.block_pool = self.coordinator.block_pool
198
self.retained_hit_group_ids = tuple(
vllm/v1/core/sched/scheduler.py
+5
-5
@@ -320,8 +320,8 @@ class Scheduler(SchedulerInterface):
320
hash_block_size=hash_block_size,
321
metrics_collector=self.kv_metrics_collector,
322
watermark=self.scheduler_config.watermark,
323
- enable_mamba_fine_grained_prefix_cache=(
324
- self.cache_config.enable_mamba_fine_grained_prefix_cache
323
+ enable_mamba_shared_prefix_checkpoint=(
324
+ self.cache_config.enable_mamba_shared_prefix_checkpoint
325
),
326
)
327
# Bind after construction so connectors can access the cache manager.
@@ -376,9 +376,9 @@ class Scheduler(SchedulerInterface):
376
# manager decides whether it can check-point there (per-group eagle bit,
377
# no MTP re-prefill tail); splitting for a stop it would refuse costs a
378
# forward pass and displaces the block-boundary stop.
379
- self.mamba_fine_grained_prefix_cache = (
379
+ self.mamba_shared_prefix_checkpoint = (
380
self.mamba_partial_cache_hit
381
- and self.kv_cache_manager.mamba_fine_grained_prefix_cache
381
+ and self.kv_cache_manager.mamba_shared_prefix_checkpoint
382
)
383
384
# Counts of non-empty steps scheduled / processed. update_from_output
@@ -509,7 +509,7 @@ class Scheduler(SchedulerInterface):
509
# output tokens can still observe a junction there.
510
junction_stop = (
511
junction
512
- if self.mamba_fine_grained_prefix_cache
512
+ if self.mamba_shared_prefix_checkpoint
513
and junction <= request.num_prompt_tokens
514
else block_floored
515
)
vllm/v1/core/single_type_kv_cache_manager.py
+3
-3
@@ -143,9 +143,9 @@ class SingleTypeKVCacheManager(ABC):
143
# aligned segment (SWA). Initialized lazily by the coordinator after
144
# determining the attention groups.
145
self.use_eagle = False
146
- # ``CacheConfig.enable_mamba_fine_grained_prefix_cache``, narrowed and set
146
+ # ``CacheConfig.enable_mamba_shared_prefix_checkpoint``, narrowed and set
147
# by ``KVCacheManager``; only an EAGLE Mamba "align" group ever gets it.
148
- self.fine_grained_prefix_cache = False
148
+ self.shared_prefix_checkpoint = False
149
# Partial-hit copy-on-write bookkeeping. Populated only by fine-grained
150
# managers (full attention, mamba "align"); harmlessly empty elsewhere.
151
self._partial_hit_reqs: dict[str, tuple[int, KVCacheBlock]] = {}
@@ -2090,7 +2090,7 @@ class MambaManager(SingleTypeKVCacheManager):
2090
# running state block, mutated in place, which equals what its key
2091
# promises only after that step's forward.
2092
if num_tokens != latest_prompt_hash_boundary and not (
2093
- self.fine_grained_prefix_cache
2093
+ self.shared_prefix_checkpoint
2094
and num_tokens == request.shared_prefix_boundary
2095
and request.num_computed_tokens < num_tokens <= request.num_prompt_tokens
2096
):