[UMBP] Enable partial hash hits by default for Mamba hybrid models

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. |
 

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(

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