[Feat] add extra info into request hash for multimodal/lora/salt support

test/test_ucm_blend_hash.py

@@ -0,0 +1,87 @@
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("vllm")
+pytest.importorskip("triton")
+
+from ucm.integration.vllm.blend_connector import UCMBlendConnector
+from ucm.integration.vllm.request_hasher import RequestHasher, RequestHashError
+
+
+def _config():
+    return SimpleNamespace(
+        model_config=SimpleNamespace(model="org/model", dtype="bfloat16"),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        speculative_config=None,
+        additional_config={},
+    )
+
+
+def _plain_text_request():
+    return SimpleNamespace(
+        all_token_ids=list(range(8)),
+        mm_features=[],
+        lora_request=None,
+        cache_salt=None,
+        prompt_embeds=None,
+    )
+
+
+def _connector_for_hash_test():
+    connector = UCMBlendConnector.__new__(UCMBlendConnector)
+    connector.block_size = 4
+    connector.request_hasher = RequestHasher(_config(), 0)
+    connector._seed = connector.request_hasher.seed
+    connector.request_block_hasher = connector.request_hasher.make_request_block_hasher(
+        connector.block_size, connector._seed
+    )
+    return connector
+
+
+def test_chunk_build_and_lookup_use_the_same_plain_text_hash_schema():
+    connector = _connector_for_hash_test()
+    request = _plain_text_request()
+
+    prefix_hashes = connector.request_block_hasher(request)
+    chunk_hashes = connector._generate_chunk_hashes(request.all_token_ids)
+
+    assert chunk_hashes == prefix_hashes
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [
+        ("mm_features", [SimpleNamespace(identifier="image")]),
+        ("lora_request", SimpleNamespace(lora_name="adapter")),
+        ("cache_salt", "tenant"),
+        ("prompt_embeds", object()),
+    ],
+)
+def test_cacheblend_rejects_request_semantics_it_cannot_hash(field, value):
+    request = _plain_text_request()
+    setattr(request, field, value)
+
+    with pytest.raises(RequestHashError):
+        UCMBlendConnector._validate_supported_request(request)
+
+
+def test_get_num_new_matched_tokens_returns_zero_on_unsupported_request():
+    """Connector-level fail-closed (design §10): an unhashable request returns
+    zero external hit tokens instead of raising through the scheduler."""
+    connector = _connector_for_hash_test()
+    request = SimpleNamespace(
+        all_token_ids=list(range(8)),
+        mm_features=[
+            SimpleNamespace(
+                identifier="image",
+                mm_position=SimpleNamespace(offset=0, length=4),
+            )
+        ],
+        lora_request=None,
+        cache_salt=None,
+        prompt_embeds=None,
+        request_id="req-mm-1",
+    )
+
+    assert connector.get_num_new_matched_tokens(request, 0) == (0, False)

test/test_ucm_hla_hash.py

@@ -0,0 +1,116 @@
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("vllm")
+
+from ucm.integration.vllm.hla_connector import GroupInfo, KVCacheGroupManager
+from ucm.integration.vllm.request_hasher import RequestHasher, RequestHashError
+
+
+def _config():
+    return SimpleNamespace(
+        model_config=SimpleNamespace(model="org/model", dtype="bfloat16"),
+        parallel_config=SimpleNamespace(tensor_parallel_size=1),
+        speculative_config=None,
+        additional_config={},
+    )
+
+
+def _request(identifier):
+    return SimpleNamespace(
+        all_token_ids=list(range(16)),
+        mm_features=[
+            SimpleNamespace(
+                identifier=identifier,
+                mm_position=SimpleNamespace(offset=4, length=5),
+            )
+        ],
+        lora_request=None,
+        cache_salt=None,
+        prompt_embeds=None,
+    )
+
+
+def _manager_with_groups(full_attention_block_sizes):
+    hasher = RequestHasher(_config(), 0)
+    manager = KVCacheGroupManager.__new__(KVCacheGroupManager)
+    manager.request_hasher = hasher
+    manager.groups_by_id = []
+    manager.full_attn_groups = []
+    manager.state_groups = []
+
+    for group_id, block_size in enumerate(full_attention_block_sizes):
+        seed = hasher((b"UCM_GROUP_SEED", hasher.seed, group_id))
+        group = GroupInfo(
+            group_id=group_id,
+            block_size=block_size,
+            layer_names=(f"fa.{group_id}",),
+            seed=seed,
+            block_hasher=hasher.make_request_block_hasher(block_size, seed),
+        )
+        manager.groups_by_id.append(group)
+        manager.full_attn_groups.append(group)
+
+    state_id = len(manager.groups_by_id)
+    state_seed = hasher((b"UCM_GROUP_SEED", hasher.seed, state_id))
+    state_group = GroupInfo(
+        group_id=state_id,
+        block_size=8,
+        layer_names=("mamba.0",),
+        seed=state_seed,
+        is_mamba_align=True,
+    )
+    manager.groups_by_id.append(state_group)
+    manager.state_groups.append(state_group)
+    manager.lcm_block_size = 8
+    return manager, state_group
+
+
+def test_full_attention_groups_hash_request_at_their_own_boundaries():
+    manager, _ = _manager_with_groups([4, 8])
+
+    small = manager.compute_all_group_block_ids(_request("small-image"))
+    large = manager.compute_all_group_block_ids(_request("large-image"))
+
+    assert len(small[0]) == 4
+    assert len(small[1]) == 2
+    assert small[0][0] == large[0][0]
+    assert small[0][1:] != large[0][1:]
+    assert small[1] != large[1]
+
+
+def test_mamba_state_key_inherits_primary_attention_request_semantics():
+    manager, state_group = _manager_with_groups([4])
+    small_blocks = manager.compute_all_group_block_ids(_request("small-image"))
+    large_blocks = manager.compute_all_group_block_ids(_request("large-image"))
+
+    small_state = manager.compute_mamba_align_state_hash(state_group, 8, small_blocks)
+    large_state = manager.compute_mamba_align_state_hash(state_group, 8, large_blocks)
+
+    assert small_state is not None
+    assert large_state is not None
+    assert small_state != large_state
+
+
+def test_hash_failure_in_any_group_disables_whole_lookup_without_partial_keys():
+    """Design §10: an HLA hash failure in any group must fail the whole
+    external lookup; a partial set of group keys is never returned."""
+    manager, _ = _manager_with_groups([4, 8])
+    request = SimpleNamespace(all_token_ids=list(range(16)))
+
+    attempted = {"n": 0}
+
+    def good_hasher(req):
+        attempted["n"] += 1
+        return [b"\x00"] * (len(req.all_token_ids) // 4)
+
+    def failing_hasher(_req):
+        raise RequestHashError("simulated group hash failure")
+
+    manager.groups_by_id[0].block_hasher = good_hasher
+    manager.groups_by_id[1].block_hasher = failing_hasher
+
+    with pytest.raises(RequestHashError):
+        manager.compute_all_group_block_ids(request)
+    assert attempted["n"] == 1