[Diffusion][HunyuanImage3] Add request paged prefix caching (#7719)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>

tests/diffusion/diffusion_kv/test_config.py

@@ -21,6 +21,86 @@ def test_dense_legacy_is_default() -> None:
     assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
 
 
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+@pytest.mark.parametrize("mode", [None, "dense_legacy", DiffusionKVCacheMode.DENSE_LEGACY])
+def test_prefix_caching_rejects_non_paged_mode(config_cls, mode) -> None:
+    kwargs = {} if mode is None else {"diffusion_kv_mode": mode}
+    with pytest.raises(
+        ValueError,
+        match="set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching",
+    ):
+        config_cls.from_kwargs(enable_prefix_caching=True, **kwargs)
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_prefix_caching_accepts_paged_scheduler(config_cls) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=True,
+    )
+
+    assert config.enable_prefix_caching is True
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("build_config", [OmniDiffusionConfig, OmniDiffusionConfig.from_kwargs])
+def test_prefix_caching_rejects_sleep_mode(build_config) -> None:
+    with pytest.raises(ValueError, match="disable enable_prefix_caching or enable_sleep_mode"):
+        build_config(
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+            enable_prefix_caching=True,
+            enable_sleep_mode=True,
+        )
+
+
+@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
+def test_sleep_mode_accepts_disabled_prefix_caching(mode) -> None:
+    config = OmniDiffusionConfig.from_kwargs(
+        diffusion_kv_mode=mode,
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=False,
+        enable_sleep_mode=True,
+    )
+
+    assert config.enable_sleep_mode is True
+    assert config.enable_prefix_caching is False
+
+
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+def test_native_kv_transfer_requires_prefix_caching_disabled(enable_prefix_caching) -> None:
+    kwargs = dict(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=enable_prefix_caching,
+        kv_transfer_config={
+            "kv_connector": "MooncakeConnector",
+            "kv_role": "kv_consumer",
+            "engine_id": "test-diffusion",
+        },
+    )
+    if enable_prefix_caching:
+        with pytest.raises(ValueError, match="disable enable_prefix_caching for AR KV import"):
+            OmniDiffusionConfig.from_kwargs(**kwargs)
+    else:
+        config = OmniDiffusionConfig.from_kwargs(**kwargs)
+        assert config.kv_transfer_config.kv_role == "kv_consumer"
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
+def test_disabled_prefix_caching_accepts_both_modes(config_cls, mode) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode=mode,
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=False,
+    )
+
+    assert config.enable_prefix_caching is False
+    assert config.diffusion_kv_mode.value == mode
+
+
 def test_paged_worker_local_is_rejected_until_implemented() -> None:
     with pytest.raises(ValueError, match="reserved but not implemented"):
         OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local")

tests/diffusion/diffusion_kv/test_initialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -225,6 +225,22 @@ def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -
     assert vllm_config.cache_config.gpu_memory_utilization == 0.42
 
 
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+@pytest.mark.parametrize("requests,rows", [(1, 2), (2, 3)])
+def test_native_config_counts_cfg_rows_once(enable_prefix_caching, requests, rows) -> None:
+    od_config = _od_config(
+        max_num_seqs=requests,
+        diffusion_kv_max_rows_per_request=rows,
+        enable_prefix_caching=enable_prefix_caching,
+    )
+    config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    # Engine and Worker must see the same row capacity, including on reconfiguration.
+    diffusion_vllm_config.configure_diffusion_vllm_config(config, od_config)
+    assert config.scheduler_config.max_num_seqs == requests * rows
+    assert od_config.max_num_seqs == requests
+    assert config.cache_config.enable_prefix_caching is enable_prefix_caching
+
+
 def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
     quantization_config = build_quant_config(
         {

tests/diffusion/diffusion_kv/test_kv_cache_utils.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native multimodal block hashing and diffusion-specific identity helpers."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
+
+import vllm_omni.diffusion.diffusion_engine as engine_module
+import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
+import vllm_omni.diffusion.diffusion_kv.request as request_module
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+def _request(*, num_rows=1, seed=7, adapter_id=None, scale=1.0):
+    request = OmniDiffusionRequest(
+        request_id="fake", prompt="fake model", sampling_params=OmniDiffusionSamplingParams(seed=seed)
+    )
+    request.sampling_params.lora_request = None if adapter_id is None else SimpleNamespace(lora_int_id=adapter_id)
+    request.sampling_params.lora_scale = scale
+    request.diffusion_kv_requests = tuple(
+        DiffusionKVRequest(f"fake/{row}", sequence_id=row, prefix_len=8, target_len=4, seq_len=12)
+        for row in range(num_rows)
+    )
+    return request
+
+
+def _feature(*, value="image", start=4, end=8):
+    return MultiModalFeatureSpec(
+        data=None,
+        modality="image",
+        identifier=hash_prefix_cache_value((value, start, end)).hex(),
+        mm_position=PlaceholderRange(offset=start, length=end - start),
+    )
+
+
+def _prepare(request, feature=None):
+    namespace = get_cache_namespace("fake-model-v1", request.sampling_params)
+    for row in request.diffusion_kv_requests:
+        row.cache_token_ids = tuple(range(row.prefix_len))
+        row.mm_features = [feature or _feature()]
+        row.cache_namespace = namespace
+
+
+def _hashes(request, feature=None):
+    _prepare(request, feature)
+    row = request.diffusion_kv_requests[0]
+    row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+    return row.block_hashes
+
+
+def test_partial_hits_and_randomness_are_model_opt_in():
+    base = _hashes(_request())
+    assert base == _hashes(_request(seed=999))
+    changed = _hashes(_request(), _feature(value="other image"))
+    assert base[0] == changed[0]
+    assert base[1] != changed[1]
+    changed_geometry = _hashes(_request(), _feature(end=9))
+    assert base[0] == changed_geometry[0]
+    assert base[1] != changed_geometry[1]
+
+
+@pytest.mark.parametrize(("adapter_id", "scale"), [(None, 1.0), (2, 0.5), (1, 1.0)])
+def test_lora_is_common_request_wide_identity(adapter_id, scale):
+    base = _hashes(_request(adapter_id=1, scale=0.5))
+    changed = _hashes(_request(adapter_id=adapter_id, scale=scale))
+    assert all(a != b for a, b in zip(base, changed, strict=True))
+    assert _hashes(_request()) == _hashes(_request(scale=0.5))
+
+
+def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
+    row = DiffusionKVRequest(
+        "native",
+        sequence_id=0,
+        prefix_len=14,
+        target_len=2,
+        seq_len=16,
+        cache_token_ids=range(16),
+        mm_features=[
+            _feature(start=2, end=7),
+            _feature(start=7, end=10),
+        ],
+    )
+    upstream = Mock(wraps=generate_block_hash_extra_keys)
+    monkeypatch.setattr(request_module, "generate_block_hash_extra_keys", upstream)
+    hash_fn = get_hash_fn_by_name("sha256")
+    row.build_block_hashes(4, hash_fn)
+    assert [(call.args[1], call.args[2]) for call in upstream.call_args_list] == [(0, 4), (4, 8), (8, 12)]
+    # Independent expected native extras: a crossing item has a negative
+    # relative offset; a block can include the end of one item and the next.
+    first, second = (feature.identifier for feature in row.mm_features)
+    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
+    parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
+    expected = []
+    for start, keys in zip(range(0, 12, 4), extras, strict=True):
+        parent = hash_block_tokens(hash_fn, parent, list(range(start, start + 4)), keys)
+        expected.append(parent)
+    assert row.block_hashes == expected
+    # Never hash the partial prefix tail or mutable target [14, 16), even
+    # though cache_token_ids deliberately contains the whole sequence.
+    assert len(row.block_hashes) == 3
+
+
+def test_same_identifier_at_different_offsets_changes_block_hash():
+    first = _feature(start=4, end=6)
+    shifted = _feature(start=5, end=7)
+    shifted.identifier = first.identifier
+    a = _hashes(_request(), first)
+    b = _hashes(_request(), shifted)
+    assert a[0] == b[0]
+    assert a[1] != b[1]
+
+
+def test_multimodal_identity_beyond_prefix_does_not_affect_hashes():
+    assert _hashes(_request(), _feature(value="a", start=8, end=12)) == _hashes(
+        _request(), _feature(value="b", start=8, end=12)
+    )
+
+
+@pytest.mark.parametrize(
+    "positions",
+    [
+        [(-1, 2)],
+        [(4, 4)],
+        [(4, 13)],
+        [(5, 7), (4, 5)],
+        [(3, 7), (4, 6)],
+    ],
+)
+def test_invalid_or_overlapping_positions_fail_before_hashing(positions):
+    row = _request().diffusion_kv_requests[0]
+    row.cache_token_ids = tuple(range(8))
+    row.mm_features = [_feature(start=start, end=end) for start, end in positions]
+    fail = Mock(side_effect=AssertionError("invalid positions must fail before hashing"))
+    with pytest.raises(ValueError, match="sorted, non-overlapping"):
+        row.build_block_hashes(4, fail)
+    fail.assert_not_called()
+    assert row.block_hashes == []
+
+
+def test_empty_multimodal_identifier_is_rejected():
+    row = _request().diffusion_kv_requests[0]
+    row.cache_token_ids = tuple(range(8))
+    feature = _feature()
+    feature.identifier = ""
+    row.mm_features = [feature]
+    with pytest.raises(ValueError, match="content identifiers"):
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.int64, torch.bool])
+def test_tensor_identity_includes_dtype_shape_and_values(dtype):
+    tensor = torch.tensor([[0, 1], [1, 0]], dtype=dtype)
+    assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.clone())
+    assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.t().contiguous().t())
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.reshape(4))
+    changed = tensor.clone()
+    changed[0, 0] = 1
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(changed)
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.to(torch.float64))
+
+
+def test_numpy_scalar_shape_and_nested_dict_order():
+    assert hash_prefix_cache_value({"b": {"y": 2, "x": 1}, "a": None}) == hash_prefix_cache_value(
+        {"a": None, "b": {"x": 1, "y": 2}}
+    )
+    assert hash_prefix_cache_value(np.array(1)) != hash_prefix_cache_value(np.array([1]))
+    assert hash_prefix_cache_value(np.int64(1)) == hash_prefix_cache_value(1)
+
+
+@pytest.mark.parametrize(
+    ("first", "second"), [(1, True), (1, 1.0), ("1", b"1"), ([], {}), ([None], []), (["ab", "c"], ["a", "bc"])]
+)
+def test_identity_has_unambiguous_type_and_container_boundaries(first, second):
+    assert hash_prefix_cache_value(first) != hash_prefix_cache_value(second)
+
+
+@pytest.mark.parametrize("value", [object(), {1: "value"}, np.array([object()]), np.longdouble(1)])
+def test_identity_rejects_unsupported_values_without_pickle(value):
+    with pytest.raises(TypeError):
+        hash_prefix_cache_value(value)
+
+
+@pytest.mark.parametrize(
+    ("mode", "enabled", "expect_hook"),
+    [
+        (DiffusionKVCacheMode.PAGED_SCHEDULER, False, False),
+        (DiffusionKVCacheMode.PAGED_SCHEDULER, True, True),
+        (DiffusionKVCacheMode.DENSE_LEGACY, True, False),
+    ],
+)
+def test_engine_gates_hook_loading_and_admission(monkeypatch, mode, enabled, expect_hook):
+    events = []
+    request = _request()
+
+    def preprocess(req):
+        events.append("layout")
+        req.prepared_layout = "fake layout"
+        return req
+
+    def prepare(req):
+        assert req.prepared_layout == "fake layout"
+        assert req.diffusion_kv_requests[0].cache_token_ids == ()
+        events.append("cache_inputs")
+        _prepare(req)
+
+    hook = Mock(side_effect=prepare if expect_hook else AssertionError("disabled hook"))
+    load_hook = Mock(return_value=hook)
+    if not expect_hook:
+        load_hook.side_effect = AssertionError("disabled hook factory")
+        monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", Mock(side_effect=AssertionError("disabled hash")))
+    monkeypatch.setattr(engine_module, "get_diffusion_pre_process_func", lambda _config: preprocess)
+    monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
+    monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", load_hook)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=mode, enable_prefix_caching=enabled)
+    engine._init_process_hooks(engine.od_config)
+    engine.prefix_cache_func = hook
+    assert engine._prepare_request_for_admission(request) is request
+    row = request.diffusion_kv_requests[0]
+    if expect_hook:
+        assert events == ["layout", "cache_inputs"]
+        assert row.cache_token_ids == tuple(range(8))
+        load_hook.assert_called_once_with(engine.od_config)
+        hook.assert_called_once_with(request)
+    else:
+        assert events == ["layout"]
+        assert row.cache_token_ids == ()
+        assert row.mm_features == []
+        load_hook.assert_not_called()
+        hook.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "model_class_name,load_format",
+    [(None, "default"), ("UnregisteredPipeline", "default"), ("HunyuanImage3ForCausalMM", "diffusers")],
+)
+def test_engine_rejects_enabled_prefix_caching_without_hook(monkeypatch, model_class_name, load_format):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.from_kwargs(
+        model_class_name=model_class_name,
+        diffusion_load_format=load_format,
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=True,
+    )
+    with pytest.raises(ValueError, match="requires a registered prefix-cache hook"):
+        engine._init_process_hooks(engine.od_config)
+
+
+def test_engine_rejects_out_of_profile_request_before_hashing():
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, enable_prefix_caching=True
+    )
+    engine._diffusion_kv_profile_limits = (1, 8, 4)
+    engine.prefix_cache_func = Mock(side_effect=AssertionError("invalid request must fail before identity work"))
+    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+        engine._prepare_request_for_admission(_request())
+    engine.prefix_cache_func.assert_not_called()

tests/diffusion/diffusion_kv/test_manager.py

@@ -30,26 +30,159 @@ def _config(num_blocks: int) -> KVCacheConfig:
     )
 
 
-def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+def _request(
+    public_id: str,
+    sequence_id: int,
+    *,
+    prefix_len: int = 4,
+    target_len: int = 4,
+    seq_len: int = 8,
+    cache_token_ids=(),
+    mm_features=(),
+    kv_contexts=(),
+) -> DiffusionKVRequest:
     return DiffusionKVRequest(
         f"{public_id}/diffusion-kv/{sequence_id}",
         sequence_id=sequence_id,
-        prefix_len=4,
-        target_len=4,
+        prefix_len=prefix_len,
+        target_len=target_len,
         seq_len=seq_len,
+        cache_token_ids=cache_token_ids,
+        mm_features=mm_features,
         kv_contexts=kv_contexts,
     )
 
 
-def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+def _manager(
+    num_blocks: int,
+    *,
+    max_model_len: int = 64,
+    enable_prefix_caching: bool = False,
+) -> DiffusionKVCacheManager:
     return DiffusionKVCacheManager(
         _config(num_blocks),
         max_model_len=max_model_len,
         scheduler_block_size=BLOCK_SIZE,
         hash_block_size=BLOCK_SIZE,
+        enable_prefix_caching=enable_prefix_caching,
     )
 
 
+def test_successful_request_publishes_prefix_for_a_warm_hit() -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    cold = _request("cold", 0, cache_token_ids=range(4))
+
+    cold_metadata = manager.reserve_request("cold", (cold,))
+    assert cold_metadata is not None
+    assert cold_metadata.sequences[0].cached_prefix_len == 0
+    cold_prefix_block = cold_metadata.sequences[0].block_ids[0][0]
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request("warm", 0, cache_token_ids=range(4))
+    warm_metadata = manager.reserve_request("warm", (warm,))
+
+    assert warm_metadata is not None
+    assert warm_metadata.sequences[0].cached_prefix_len == 4
+    assert warm_metadata.sequences[0].block_ids[0][0] == cold_prefix_block
+    # Metadata contains both the shared hit and the newly allocated target.
+    assert len(warm_metadata.sequences[0].block_ids[0]) == 2
+
+    manager.free_request("warm")
+    assert warm.num_computed_tokens == 0
+    assert warm.shared_prefix_boundary == 0
+
+
+def test_prefix_hash_root_is_stable_across_manager_instances() -> None:
+    first_manager = _manager(8, enable_prefix_caching=True)
+    second_manager = _manager(8, enable_prefix_caching=True)
+    first = _request("first", 0, cache_token_ids=range(4))
+    second = _request("second", 0, cache_token_ids=range(4))
+
+    first_manager._prepare_block_hashes((first,))
+    second_manager._prepare_block_hashes((second,))
+
+    assert first.block_hashes == second.block_hashes
+
+
+def test_warm_hit_fits_at_the_same_capacity_as_the_cold_request() -> None:
+    manager = _manager(3, enable_prefix_caching=True)
+    cold = _request("cold", 0, cache_token_ids=range(4))
+    assert manager.reserve_request("cold", (cold,)) is not None
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request("warm", 0, cache_token_ids=range(4))
+    metadata = manager.reserve_request("warm", (warm,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 4
+
+
+def test_unpublished_request_does_not_become_a_prefix_hit() -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    failed = _request("failed", 0, cache_token_ids=range(4))
+    assert manager.reserve_request("failed", (failed,)) is not None
+    manager.free_request("failed")
+
+    retry = _request("retry", 0, cache_token_ids=range(4))
+    metadata = manager.reserve_request("retry", (retry,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 0
+
+
+def test_partial_prefix_hit_recomputes_from_the_first_changed_block() -> None:
+    manager = _manager(10, enable_prefix_caching=True)
+    cold = _request(
+        "cold",
+        0,
+        prefix_len=8,
+        seq_len=12,
+        cache_token_ids=range(8),
+    )
+    assert manager.reserve_request("cold", (cold,)) is not None
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request(
+        "warm",
+        0,
+        prefix_len=8,
+        seq_len=12,
+        cache_token_ids=(*range(4), 100, 101, 102, 103),
+    )
+    metadata = manager.reserve_request("warm", (warm,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 4
+    assert len(metadata.sequences[0].block_ids[0]) == 3
+
+
+def test_cfg_rows_use_the_minimum_warm_hit_as_one_execution_boundary() -> None:
+    manager = _manager(16, enable_prefix_caching=True)
+    for public_id, tokens in (("seed-a", range(8)), ("seed-b", (*range(4), 20, 21, 22, 23))):
+        request = _request(public_id, 0, prefix_len=8, seq_len=12, cache_token_ids=tokens)
+        assert manager.reserve_request(public_id, (request,)) is not None
+        manager.publish_request(public_id)
+        manager.free_request(public_id)
+
+    requests = (
+        _request("warm", 0, prefix_len=8, seq_len=12, cache_token_ids=range(8)),
+        _request(
+            "warm",
+            1,
+            prefix_len=8,
+            seq_len=12,
+            cache_token_ids=(*range(4), 30, 31, 32, 33),
+        ),
+    )
+    metadata = manager.reserve_request("warm", requests)
+
+    assert metadata is not None
+    assert [sequence.cached_prefix_len for sequence in metadata.sequences] == [4, 4]
+
+
 def test_reserve_and_free_multi_cfg_request() -> None:
     manager = _manager(8)
     free_before = manager.native_manager.block_pool.get_num_free_blocks()

tests/diffusion/diffusion_kv/test_native_prefetch.py

@@ -256,7 +256,11 @@ def _scheduler(monkeypatch, num_blocks=24):
         kv_cache_config=config,
         scheduler_block_size=4,
         hash_block_size=4,
-        kv_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64), max_in_flight_tokens=64),
+        kv_vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(max_model_len=64),
+            max_in_flight_tokens=64,
+            cache_config=SimpleNamespace(enable_prefix_caching=False),
+        ),
     )
     return scheduler, connector
 

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -113,7 +113,7 @@ class _FakeLayer:
         self.head_size_v = 4
         self.spec = _FakeSpec(non_causal=non_causal)
         self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
-        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = []
+        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, _FakeNativeMetadata]] = []
         self.native_events: list[str] = []
         self.layer_name = "layer-0"
         self.kv_cache = object()
@@ -122,6 +122,8 @@ class _FakeLayer:
 
 
 class _FakeNativeImpl:
+    vllm_flash_attn_version: int
+
     def __init__(self, layer: _FakeLayer) -> None:
         self.layer = layer
 
@@ -143,7 +145,7 @@ class _FakeNativeImpl:
         key: torch.Tensor,
         value: torch.Tensor,
         _kv_cache,
-        metadata: object,
+        metadata: _FakeNativeMetadata,
         output: torch.Tensor,
     ) -> torch.Tensor:
         self.layer.native_events.append("forward")
@@ -202,7 +204,7 @@ def _make_adapter(
 
     config = SimpleNamespace(
         name="vllm-config",
-        model_config=SimpleNamespace(dtype=torch.float32),
+        model_config=SimpleNamespace(dtype=torch.float32, rswa_window=None),
     )
     adapter = DiffusionPagedAttentionAdapter(
         vllm_config=config,
@@ -1222,9 +1224,15 @@ def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: st
         DiffusionPagedAttentionRow(**values)
 
 
+@pytest.mark.parametrize("pcp_size", [1, 2])
+@pytest.mark.parametrize("selection_fails", [False, True])
 def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads(
     monkeypatch: pytest.MonkeyPatch,
+    pcp_size: int,
+    selection_fails: bool,
 ) -> None:
+    from vllm.v1.attention import selector
+
     selected_backends = []
     specialized_backends = []
     impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
@@ -1237,28 +1245,39 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
 
     original_backend_per_kind = {"full": object()}
     config = SimpleNamespace(
-        attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
-        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
-        model_config=SimpleNamespace(dtype=torch.float16),
-        cache_config=SimpleNamespace(cache_dtype="auto"),
+        attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind, use_non_causal=True),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size, decode_context_parallel_size=1),
+        model_config=SimpleNamespace(dtype=torch.float16, rswa_window=None),
+        cache_config=SimpleNamespace(cache_dtype="auto", user_specified_block_size=False),
+        kv_transfer_config=None,
+        speculative_config=None,
     )
 
-    def select_backend(**_kwargs):
-        selected_backends.append(
-            (
-                config.attention_config.backend,
-                config.attention_config.backend_per_kind,
-                config.parallel_config.prefill_context_parallel_size,
-            )
-        )
+    def select_backend(*, backend, attn_selector_config, num_heads):
+        # Exercise the real 0.29 selector, replacing only platform resolution.
+        # Ulysses has already gathered tokens and sharded heads before the
+        # native kernel; the MoE PCP mapping must not request PCP attention.
+        assert not attn_selector_config.use_pcp
+        assert not attn_selector_config.use_dcp
+        assert attn_selector_config.use_non_causal
+        assert num_heads == 4
+        selected_backends.append((backend, config.attention_config.backend_per_kind))
+        if selection_fails:
+            raise ValueError("test backend unavailable")
         return native_backend
 
-    monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
+    monkeypatch.setattr("vllm.config.get_current_vllm_config", lambda: config)
+    monkeypatch.setattr(selector, "_cached_get_attn_backend", select_backend)
     monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+    def specialize_backend(backend, *, ulysses_degree):
+        specialized_backends.append((backend, ulysses_degree))
+        return backend
+
     monkeypatch.setattr(
         adapter_module.current_omni_platform,
         "get_diffusion_paged_kv_attn_backend",
-        lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
+        specialize_backend,
     )
     layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
     spec = FullAttentionSpec(
@@ -1269,7 +1288,7 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         non_causal=True,
     )
 
-    native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(
+    kwargs = dict(
         layer_name="layer-0",
         layer=layer,
         spec=spec,
@@ -1277,12 +1296,21 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         device=torch.device("cpu"),
         ulysses_degree=2,
     )
+    if selection_fails:
+        with pytest.raises(ValueError, match="test backend unavailable"):
+            adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
+    else:
+        native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 1)]
-    assert specialized_backends == [(native_backend, 2)]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind
-    assert config.parallel_config.prefill_context_parallel_size == 2
+    assert config.parallel_config.prefill_context_parallel_size == pcp_size
+    if selection_fails:
+        assert not specialized_backends
+        impl_cls.assert_not_called()
+        return
+    assert specialized_backends == [(native_backend, 2)]
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
@@ -1432,7 +1460,12 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
     layer._no_parallel_strategy = object()
     layer._get_active_parallel_strategy = lambda: Strategy()
     layer._with_kv_cache_dtype = lambda metadata: metadata
-    layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query
+
+    def run_local_attention(query, _key, _value, _metadata):
+        events.append("dense")
+        return query
+
+    layer._run_local_attention = run_local_attention
     qkv = torch.zeros(1, 2, 2, 4)
 
     assert not layer.is_paged_kv_active()

tests/diffusion/diffusion_kv/test_request.py

@@ -4,8 +4,10 @@
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.hashing import get_hash_fn_by_name
 from vllm.v1.core.kv_cache_manager import KVCacheManager
-from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.core.kv_cache_utils import BlockHash, init_none_hash
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
@@ -49,6 +51,51 @@ def test_request_exposes_native_request_and_diffusion_semantics() -> None:
     assert request.status is RequestStatus.WAITING
 
 
+def test_request_builds_only_full_prefix_block_hashes() -> None:
+    hash_function = get_hash_fn_by_name("sha256")
+    init_none_hash(hash_function)
+    request = _request(
+        prefix_len=6,
+        target_len=2,
+        cache_token_ids=range(6),
+        mm_features=[
+            MultiModalFeatureSpec(
+                data=None,
+                modality="image",
+                identifier="tail",
+                mm_position=PlaceholderRange(offset=4, length=2),
+            )
+        ],
+    )
+
+    request.build_block_hashes(BLOCK_SIZE, hash_function)
+
+    assert len(request.block_hashes) == 1
+    assert request.skip_reading_prefix_cache is False
+
+
+def test_changed_multimodal_identity_changes_the_block_hash() -> None:
+    hash_function = get_hash_fn_by_name("sha256")
+    init_none_hash(hash_function)
+    first = _request(cache_token_ids=range(4))
+    second = _request(
+        cache_token_ids=range(4),
+        mm_features=[
+            MultiModalFeatureSpec(
+                data=None,
+                modality="image",
+                identifier="image-a",
+                mm_position=PlaceholderRange(offset=2, length=2),
+            )
+        ],
+    )
+
+    first.build_block_hashes(BLOCK_SIZE, hash_function)
+    second.build_block_hashes(BLOCK_SIZE, hash_function)
+
+    assert first.block_hashes != second.block_hashes
+
+
 @pytest.mark.parametrize(
     ("overrides", "message"),
     [

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -111,6 +111,29 @@ def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> Non
     ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
 
 
+def test_runner_builds_suffix_prefill_row_for_a_prefix_hit() -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(
+                sequence_id=0,
+                prefix_len=8,
+                target_len=4,
+                seq_len=12,
+                block_ids=([1, 2, 3],),
+                cached_prefix_len=4,
+            ),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    row = attn_metadata.prefill_rows[0]
+    assert (row.kv_start_pos, row.query_len, row.seq_len) == (4, 8, 12)
+
+
 def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner.od_config.parallel_config.cfg_parallel_size = 2