[Bugfix] Reject a prefix_match_unit that a single KV cache group cannot honor (#58021)
Signed-off-by: QHarshil <harshil_c@hotmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Harshil committed
Sep 29, 2026 at 16:12 UTC
be255076d0498c163441d4aa3fc9d2da142f30fd
2 files changed
+21
-2
tests/v1/kv_connector/unit/test_mooncake_store_worker.py
+3
-1
@@ -310,7 +310,9 @@ def _make_vllm_config(
310
kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC,
311
disable_hybrid_kv_cache_manager: bool = True,
312
) -> SimpleNamespace:
313
- cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10)
313
+ cache_config = SimpleNamespace(
314
+ block_size=16, num_gpu_blocks=10, prefix_match_unit=None
315
+ )
316
cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout
317
return SimpleNamespace(
318
model_config=_FakeModelConfig(),
vllm/v1/core/kv_cache_utils.py
+18
-1
@@ -714,7 +714,9 @@ def resolve_kv_cache_block_sizes(
714
group's effective block size. Attention groups are scaled by DCP;
715
Mamba groups keep their full per-rank state and are not scaled.
716
- ``hash_block_size`` is the granularity at which ``Request.block_hashes``
717
- is computed. Single group: equals scheduler block size. Multiple groups:
717
+ is computed. Single group: equals scheduler block size, and any other
718
+ ``cache_config.prefix_match_unit`` is rejected while block hashing is
719
+ active. Multiple groups:
720
``cache_config.prefix_match_unit`` override if set, else the GCD of
721
group block sizes; every group's block size must be divisible by it.
722
Returns the scheduler block size (i.e. disables finer hashing) if block
@@ -728,6 +730,21 @@ def resolve_kv_cache_block_sizes(
730
if groups and not groups[0].kv_cache_spec.dcp_sharded:
731
dcp = 1
732
bs = cache_config.block_size * dcp
733
+ # The Mamba prefill checkpoint builder reads prefix_match_unit directly,
734
+ # so a value dropped here puts its checkpoints off the scheduler's grid.
735
+ if (
736
+ cache_config.prefix_match_unit not in (None, bs)
737
+ and len(groups) == 1
738
+ and (
739
+ cache_config.enable_prefix_caching
740
+ or vllm_config.kv_transfer_config is not None
741
+ )
742
+ ):
743
+ raise ValueError(
744
+ f"Invalid prefix_match_unit={cache_config.prefix_match_unit}; "
745
+ "with a single KV cache group, prefix-cache hits land on the "
746
+ f"block size ({bs}). Unset it or set it to {bs}."
747
+ )
748
return bs, bs
749
750
group_block_sizes = [