diff --git a/docs/features/umbp_store_connector_usage.md b/docs/features/umbp_store_connector_usage.md index ff91c2bab315..3c34428a5c90 100644 --- a/docs/features/umbp_store_connector_usage.md +++ b/docs/features/umbp_store_connector_usage.md @@ -126,7 +126,7 @@ Set these options in `kv_connector_extra_config`: | `kv_event_max_objects` | `65536` | Maximum objects tracked for KV events. When full, new stored-object notifications are suppressed until tracked objects leave the pool. Storage and lookup continue normally. | | `enable_lookup` | `true` | Restore from the pool. With `false` the connector only stores. | | `save_decode_cache` | `false` | Also store blocks produced by decoding. | -| `enable_partial_hash_hits` | `false` | Allow hits that end inside a block, at core's prefix-match unit (`cache_config.prefix_match_unit`). | +| `enable_partial_hash_hits` | `false`; `true` for Mamba "align" hybrid models whose prefix-match unit is finer than the scheduler block | Allow hits that end inside a block, at core's prefix-match unit (`cache_config.prefix_match_unit`). A hybrid model's state checkpoints sit on prompt-tail hash boundaries, so block-granular lookups almost never hit it. | | `num_workers`, `timeout_ms` | `4`, `30000` | Transfer threads per worker, and the transfer timeout. A synchronous wait that times out fails the step rather than reusing its blocks. | | `key_namespace` | derived | Overrides the namespace derived from the model, revision, KV layout, cache groups, parallel topology and draft model. | diff --git a/tests/v1/kv_connector/unit/test_umbp_hybrid.py b/tests/v1/kv_connector/unit/test_umbp_hybrid.py index 6d211f2a68a8..0f7161b49588 100644 --- a/tests/v1/kv_connector/unit/test_umbp_hybrid.py +++ b/tests/v1/kv_connector/unit/test_umbp_hybrid.py @@ -838,6 +838,33 @@ def test_hybrid_fine_hash_lookup_is_not_limited_by_gpu_page_count(): scheduler.close() +@pytest.mark.parametrize( + ("hybrid", "prefix_match_unit", "configured", "expected"), + [ + (True, 4, None, True), + (True, 4, False, False), + (True, None, None, False), + (False, None, None, False), + ], +) +def test_partial_hash_hits_default_on_for_mamba_align_hybrids( + hybrid, prefix_match_unit, configured, expected +): + extra = {"mode": "embedded"} + if configured is not None: + extra["enable_partial_hash_hits"] = configured + scheduler = UMBPStoreConnectorScheduler( + _vllm_config(extra, prefix_match_unit=prefix_match_unit), + _hybrid_kv_cache_config() if hybrid else _kv_cache_config(), + _SchedulerHandle([]), + BlockIdentityCodec(UMBPNamespace("partial-default")), + ) + try: + assert scheduler.enable_partial_hash_hits is expected + finally: + scheduler.close() + + def test_finished_hybrid_tail_pins_full_pages_until_all_workers_finish(): scheduler = UMBPStoreConnectorScheduler( _vllm_config( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py index a7e88e3a4308..2ada0b1110f0 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/umbp/scheduler.py @@ -241,11 +241,17 @@ def __init__( # A kv_consumer only restores from the pool, as with other store # connectors; it never offloads its own KV. self.store_enabled = getattr(transfer_config, "kv_role", None) != "kv_consumer" - self.enable_partial_hash_hits = bool( - extra.get("enable_partial_hash_hits", False) - ) # Request.block_hashes are computed at this size; keys index into them. self.hash_block_size = resolved_hash_block_size + # Mamba "align" checkpoints sit on prompt-tail hash boundaries, which + # block-granular lookups almost never reach. + self.enable_partial_hash_hits = bool( + extra.get( + "enable_partial_hash_hits", + bool(self._boundary_handoff_groups) + and self.hash_block_size < self.block_size, + ) + ) lookup_groups = [ kv_cache_config.kv_cache_groups[group_id] for group_id in kv_cache_config.prefix_cacheable_group_ids