[Test][UMBP] Restore a Mamba2 hybrid end to end through UMBP

tests/v1/kv_connector/umbp_integration/test_mori_token_e2e.py

@@ -3,7 +3,8 @@
 """End-to-end restore through the MORI-backed UMBP runtime.
 
 Batch-invariant kernels make a restore bitwise comparable with recomputing the
-prompt, so restored KV must reproduce the cold run's tokens and logprobs.
+prompt, so restored KV must reproduce the cold run's tokens and logprobs. Mamba
+hybrids have no batch-invariant kernels, so they compare with a GPU prefix hit.
 """
 
 import time
@@ -126,3 +127,52 @@ def test_restore_extends_a_local_prefix_hit(umbp_llm):
     assert mixed[0].num_cached_tokens == 8 * BLOCK_SIZE
     assert _external_tokens(umbp_llm) - before == 6 * BLOCK_SIZE
     _assert_same_outputs(computed, mixed)
+
+
+HYBRID_MODEL = "ibm-granite/granite-4.0-h-350m"
+HYBRID_REVISION = "3b17b717b8f2f5d305b0a92c1491e239aeda19c8"
+
+
+@pytest.fixture
+def umbp_hybrid_llm(vllm_runner, socket_dir):
+    with vllm_runner(
+        HYBRID_MODEL,
+        revision=HYBRID_REVISION,
+        enforce_eager=True,
+        max_model_len=2048,
+        enable_chunked_prefill=True,
+        enable_prefix_caching=True,
+        disable_log_stats=False,
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="UMBPStoreConnector",
+            kv_role="kv_both",
+            kv_connector_extra_config={
+                "mode": "embedded",
+                "capacity_bytes": 1 << 30,
+                "lookup_dir": socket_dir,
+            },
+        ),
+    ) as runner:
+        yield runner.llm
+
+
+def test_hybrid_restore_matches_local_prefix_hit(umbp_hybrid_llm):
+    """Mamba boundary states restored through UMBP resume like a GPU hit.
+
+    A prompt ending exactly on a block boundary still recomputes its last block.
+    """
+    llm = umbp_hybrid_llm
+    block_size = llm.llm_engine.vllm_config.cache_config.block_size
+    for seed, tail in enumerate((7, block_size)):
+        prompt = _prompt(seed, 2 * block_size + tail)
+        llm.generate(prompt, PARAMS, use_tqdm=False)
+        local_hit = llm.generate(prompt, PARAMS, use_tqdm=False)
+        assert local_hit[0].num_cached_tokens == 2 * block_size
+        _reset_gpu_prefix_cache(llm)
+
+        before = _external_tokens(llm)
+        restored = llm.generate(prompt, PARAMS, use_tqdm=False)
+
+        assert restored[0].num_cached_tokens == 2 * block_size
+        assert _external_tokens(llm) - before == 2 * block_size
+        _assert_same_outputs(local_hit, restored)