diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index f08213444e..1ef77e285a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py @@ -310,7 +310,9 @@ def _make_vllm_config( kv_cache_layout: KVCacheLayout = KVCacheLayout.LBHNC, disable_hybrid_kv_cache_manager: bool = True, ) -> SimpleNamespace: - cache_config = SimpleNamespace(block_size=16, num_gpu_blocks=10) + cache_config = SimpleNamespace( + block_size=16, num_gpu_blocks=10, prefix_match_unit=None + ) cache_config.get_resolved_kv_cache_layout = lambda: kv_cache_layout return SimpleNamespace( model_config=_FakeModelConfig(), diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 3f4f4e7bb1..ddeda3ed19 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -714,7 +714,9 @@ def resolve_kv_cache_block_sizes( group's effective block size. Attention groups are scaled by DCP; Mamba groups keep their full per-rank state and are not scaled. - ``hash_block_size`` is the granularity at which ``Request.block_hashes`` - is computed. Single group: equals scheduler block size. Multiple groups: + is computed. Single group: equals scheduler block size, and any other + ``cache_config.prefix_match_unit`` is rejected while block hashing is + active. Multiple groups: ``cache_config.prefix_match_unit`` override if set, else the GCD of group block sizes; every group's block size must be divisible by it. Returns the scheduler block size (i.e. disables finer hashing) if block @@ -728,6 +730,21 @@ def resolve_kv_cache_block_sizes( if groups and not groups[0].kv_cache_spec.dcp_sharded: dcp = 1 bs = cache_config.block_size * dcp + # The Mamba prefill checkpoint builder reads prefix_match_unit directly, + # so a value dropped here puts its checkpoints off the scheduler's grid. + if ( + cache_config.prefix_match_unit not in (None, bs) + and len(groups) == 1 + and ( + cache_config.enable_prefix_caching + or vllm_config.kv_transfer_config is not None + ) + ): + raise ValueError( + f"Invalid prefix_match_unit={cache_config.prefix_match_unit}; " + "with a single KV cache group, prefix-cache hits land on the " + f"block size ({bs}). Unset it or set it to {bs}." + ) return bs, bs group_block_sizes = [