[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

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py

@@ -2,13 +2,24 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from transformers.utils.generic import ModelOutput
+from vllm.lora.request import LoRARequest
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+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.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
+import vllm_omni.diffusion.models.hunyuan_image3.request_layout as layout_module
+from tests.helpers.kv_layout import build_kv_cache_tensor
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
@@ -27,6 +38,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
     hunyuan_num_special_tokens,
     normalize_hunyuan_cot_text,
     prepare_hunyuan_layout,
+    prepare_hunyuan_prefix_cache,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -68,18 +80,26 @@ class _FakeTokenizerWrapper:
         seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens]
         padded_seq_len = max(seq_lens)
         positions = torch.arange(padded_seq_len)
-        joint_image_slices = [[slice(1, 3), slice(3, 5)] if has_cond_image else [] for _ in range(rows)]
+        # Like the real tokenizer, joint spans enclose VAE + separator + ViT,
+        # while all_image_slices contains the individual RoPE spans. Leave one
+        # full text block before the image so partial reuse can be tested.
+        joint_image_slices = [[slice(5, 10)] if has_cond_image else [] for _ in range(rows)]
+        cond_vae_image_slices = [[slice(5, 7)] if has_cond_image else [] for _ in range(rows)]
+        cond_vit_image_slices = [[slice(8, 10)] if has_cond_image else [] for _ in range(rows)]
         gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens]
-        all_image_slices = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)]
+        all_image_slices = [
+            vae + vit + generated
+            for vae, vit, generated in zip(cond_vae_image_slices, cond_vit_image_slices, gen_image_slices)
+        ]
         cond_vae_image_mask = (
-            torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None
+            torch.stack([(positions >= 5) & (positions < 7) for _ in range(rows)]) if has_cond_image else None
         )
         cond_vit_image_mask = (
-            torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None
+            torch.stack([(positions >= 8) & (positions < 10) for _ in range(rows)]) if has_cond_image else None
         )
         sections = []
         for _ in range(rows):
-            row_sections = []
+            row_sections: list[dict[str, object]] = []
             if has_cond_image:
                 row_sections.append(
                     {
@@ -101,6 +121,8 @@ class _FakeTokenizerWrapper:
                 ),
                 cond_vae_image_mask=cond_vae_image_mask,
                 cond_vit_image_mask=cond_vit_image_mask,
+                cond_vae_image_slices=cond_vae_image_slices,
+                cond_vit_image_slices=cond_vit_image_slices,
                 joint_image_slices=joint_image_slices,
                 gen_image_slices=gen_image_slices,
                 real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
@@ -155,6 +177,7 @@ def _request(
     guidance_scale: float,
     prompt="draw a cat",
     request_id: str = "req",
+    seed: int | None = None,
 ) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompt=prompt,
@@ -163,6 +186,7 @@ def _request(
             width=1024,
             guidance_scale=guidance_scale,
             num_inference_steps=4,
+            seed=seed,
         ),
         request_id=request_id,
     )
@@ -205,6 +229,8 @@ def test_builds_kv_request_lengths_without_model_execution() -> None:
     assert tokenizer.calls[0]["sequence_template"] == "instruct"
     assert tokenizer.calls[0]["cfg_factor"] == 1
     assert kv_requests[0].block_hashes == []
+    assert kv_requests[0].cache_token_ids == ()
+    assert kv_requests[0].mm_features == []
     assert kv_requests[0].kv_contexts == ()
     assert kv_requests[0].skip_reading_prefix_cache is True
 
@@ -226,12 +252,14 @@ def _reference_image() -> JointImageInfo:
     return JointImageInfo(
         vae_image_info=ImageInfo(
             image_type="vae",
+            image_tensor=torch.arange(12, dtype=torch.float32).reshape(1, 3, 2, 2),
             token_width=8,
             token_height=8,
             image_token_length=64,
         ),
         vision_image_info=ImageInfo(
             image_type="siglip2",
+            image_tensor=torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2),
             token_width=4,
             token_height=4,
             image_token_length=16,
@@ -243,6 +271,342 @@ def _reference_image() -> JointImageInfo:
     )
 
 
+def _prepare_cached_requests(request, tokenizer, image_processor):
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    prepare_hunyuan_prefix_cache(request)
+    return request.diffusion_kv_requests
+
+
+def test_reference_prefix_identity_includes_image_content_and_vae_seed() -> None:
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id: str, image: JointImageInfo, seed: int):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            seed=seed,
+            prompt={
+                "prompt": "edit this image",
+                "additional_information": {"batch_cond_image_info": [image]},
+            },
+        )
+        return _prepare_cached_requests(request, tokenizer, image_processor)[0]
+
+    first = build("first", _reference_image(), 7)
+    same = build("same", _reference_image(), 7)
+    changed_seed = build("changed-seed", _reference_image(), 8)
+    changed_image_value = _reference_image()
+    changed_image_value.vae_image_info.image_tensor[0, 0, 0, 0] = 99
+    changed_image = build("changed-image", changed_image_value, 7)
+
+    assert first.mm_features == same.mm_features
+    assert [(feature.mm_position.offset, feature.mm_position.length) for feature in first.mm_features] == [
+        (5, 2),
+        (8, 2),
+    ]
+    for changed in (changed_seed, changed_image):
+        # Both subspans must carry the identity, not just the first one.
+        assert all(a.identifier != b.identifier for a, b in zip(first.mm_features, changed.mm_features, strict=True))
+
+
+def test_reference_prefix_rejects_missing_inputs_with_nested_spans() -> None:
+    tokenizer, image_processor = _components([20])
+    request = _request(
+        guidance_scale=1.0,
+        prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    request.prompt["additional_information"].pop("batch_cond_image_info")
+
+    with pytest.raises(ValueError, match="missing its canonical inputs"):
+        prepare_hunyuan_prefix_cache(request)
+
+
+def test_reference_content_is_hashed_once_across_cfg_rows(monkeypatch):
+    tokenizer, image_processor = _components([20, 20])
+    image = _reference_image()
+    hasher = Mock(wraps=layout_module.hash_prefix_cache_value)
+    monkeypatch.setattr(layout_module, "hash_prefix_cache_value", hasher)
+
+    def build(request_id):
+        request = _request(
+            guidance_scale=5.0,
+            request_id=request_id,
+            seed=7,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [image]}},
+        )
+        return _prepare_cached_requests(request, tokenizer, image_processor)
+
+    first = build("first")
+    assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 1
+    assert len(first[0].mm_features) == len(first[1].mm_features) == 2
+    assert first[0].mm_features == first[1].mm_features
+    image.vae_image_info.image_tensor[0, 0, 0, 0] += 1
+    changed = build("changed")
+    assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 2
+    assert all(a.identifier != b.identifier for a, b in zip(first[0].mm_features, changed[0].mm_features, strict=True))
+
+
+def test_failed_cfg_cache_preparation_does_not_install_partial_inputs():
+    tokenizer, image_processor = _components([20, 20])
+    request = _request(
+        guidance_scale=5.0,
+        prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    # First row prepares successfully; the second exposes inconsistent inputs.
+    layout.rope_image_info[0] = []
+    request.prompt["additional_information"].pop("batch_cond_image_info")
+    with pytest.raises(ValueError, match="missing its canonical inputs"):
+        prepare_hunyuan_prefix_cache(request)
+    assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests)
+
+
+@pytest.mark.parametrize("generator_list", [False, True])
+def test_reference_identity_uses_explicit_generator_state_before_seed(generator_list):
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id, seed, advance):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            seed=seed,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+        )
+        generator = torch.Generator(device="cpu").manual_seed(123)
+        if advance:
+            torch.rand(1, generator=generator)
+        request.sampling_params.generator = [generator] if generator_list else generator
+        before = generator.get_state().clone()
+        row = _prepare_cached_requests(request, tokenizer, image_processor)[0]
+        torch.testing.assert_close(before, generator.get_state())
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return row.block_hashes
+
+    first = build("first", 7, False)
+    assert first == build("same-generator-different-unused-seed", 8, False)
+    changed = build("advanced-generator-same-seed", 7, True)
+    assert first[:1] == changed[:1]
+    assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True))
+
+
+def test_unseeded_reference_identity_is_request_local():
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+        )
+        # Real requests normally receive an automatic seed in __post_init__.
+        request.sampling_params.seed = None
+        row = _prepare_cached_requests(request, tokenizer, image_processor)[0]
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return row.block_hashes
+
+    first = build("first")
+    assert first == build("first")
+    changed = build("second")
+    assert first[:1] == changed[:1]
+    assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True))
+
+
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+def test_text_prefix_hashes_include_lora_identity_and_scale(guidance_scale: float) -> None:
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+
+    def build(request_id: str, adapter_id: int | None, scale: float):
+        request = _request(guidance_scale=guidance_scale, request_id=request_id)
+        if adapter_id is not None:
+            request.sampling_params.lora_request = LoRARequest(
+                lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}"
+            )
+        request.sampling_params.lora_scale = scale
+        kv_requests = _prepare_cached_requests(request, tokenizer, image_processor)
+        for kv_request in kv_requests:
+            kv_request.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return kv_requests
+
+    base = build("base", None, 1.0)
+    same_base = build("same-base", None, 0.5)  # Scale has no effect without an adapter.
+    first = build("first", 1, 0.5)
+    same = build("same", 1, 0.5)
+    changed_id = build("changed-id", 2, 0.5)
+    changed_scale = build("changed-scale", 1, 1.0)
+
+    for row in range(len(first)):
+        assert base[row].block_hashes == same_base[row].block_hashes
+        assert first[row].block_hashes == same[row].block_hashes
+        assert len(first[row].block_hashes) == 5
+        for changed in (base, changed_id, changed_scale):
+            assert first[row].cache_token_ids == changed[row].cache_token_ids
+            # LoRA changes even the text prefix, so every block must differ.
+            assert all(a != b for a, b in zip(first[row].block_hashes, changed[row].block_hashes, strict=True))
+
+
+@pytest.fixture
+def prefix_cache_manager(request):
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
+    num_blocks = 64
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+    manager = DiffusionKVCacheManager(
+        config,
+        max_model_len=64,
+        scheduler_block_size=4,
+        hash_block_size=4,
+        enable_prefix_caching=getattr(request, "param", True),
+    )
+    yield manager
+    manager.close()
+
+
+@pytest.mark.parametrize("prefix_cache_manager", [False], indirect=True)
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+def test_paged_cache_off_prepares_and_allocates_without_identity_work(
+    monkeypatch, prefix_cache_manager, guidance_scale
+):
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+    config = SimpleNamespace(
+        model="fake",
+        model_class_name="HunyuanImage3ForCausalMM",
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        enable_prefix_caching=False,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_config",
+        lambda *_args, **_kwargs: SimpleNamespace(
+            vae_downsample_factor=(8, 8),
+            patch_size=2,
+            image_base_size=1024,
+        ),
+    )
+    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
+    monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer)
+    monkeypatch.setattr(pipeline_module.GenerationConfig, "from_pretrained", lambda _model: SimpleNamespace())
+    monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
+    fail = Mock(side_effect=AssertionError("disabled prefix cache must not prepare identity"))
+    monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", fail)
+    monkeypatch.setattr(layout_module, "prepare_hunyuan_prefix_cache", fail)
+    monkeypatch.setattr(layout_module, "hash_prefix_cache_value", fail)
+    monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", fail)
+    monkeypatch.setattr(kv_utils.MultiModalHasher, "hash_kwargs", fail)
+    monkeypatch.setattr(torch.Tensor, "tolist", fail)
+    monkeypatch.setattr(prefix_cache_manager, "_hash_function", fail)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = config
+    engine._init_process_hooks(config)
+    # Guard admission too, even if someone injects a hook after initialization.
+    engine.prefix_cache_func = fail
+    request = _request(
+        guidance_scale=guidance_scale,
+        prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    request.sampling_params.seed = None
+    request.sampling_params.generator = SimpleNamespace(get_state=fail)
+    engine._prepare_request_for_admission(request)
+    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
+    assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests)
+    metadata = prefix_cache_manager.reserve_request(request.request_id, request.diffusion_kv_requests)
+    assert metadata is not None
+    assert all(row.cached_prefix_len == 0 for row in metadata.sequences)
+    prefix_cache_manager.publish_request(request.request_id)
+    prefix_cache_manager.free_request(request.request_id)
+    fail.assert_not_called()
+
+
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+@pytest.mark.parametrize(
+    ("change", "expected_hit"),
+    [
+        ("same", 20),
+        ("prompt_suffix", 12),
+        ("target", 20),
+        ("vae", 4),
+        ("vit", 4),
+        ("vision_kwargs", 4),
+        ("rope", 4),
+        ("seed", 4),
+        ("lora_id", 0),
+        ("lora_scale", 0),
+        ("base", 0),
+    ],
+)
+def test_prefix_lookup_respects_reference_and_lora_identity(
+    prefix_cache_manager: DiffusionKVCacheManager, guidance_scale: float, change: str, expected_hit: int
+) -> None:
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+
+    def build(request_id: str, change: str):
+        image = _reference_image()
+        if change == "vae":
+            image.vae_image_info.image_tensor[0, 0, 0, 0] = 99
+        elif change == "vit":
+            image.vision_image_info.image_tensor[0, 0, 0, 0] = 99
+        elif change == "vision_kwargs":
+            image.vision_encoder_kwargs["pixel_attention_mask"][0] = False
+        request = _request(
+            guidance_scale=guidance_scale,
+            request_id=request_id,
+            seed=8 if change == "seed" else 7,
+            prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [image]}},
+        )
+        if change != "base":
+            adapter_id = 2 if change == "lora_id" else 1
+            request.sampling_params.lora_request = LoRARequest(
+                lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}"
+            )
+        request.sampling_params.lora_scale = 1.0 if change == "lora_scale" else 0.5
+        layout = _prepare(request, tokenizer, image_processor)
+        if change == "rope":
+            for spans in layout.rope_image_info:
+                spans[0] = (slice(5, 7), (2, 1))
+        elif change == "prompt_suffix":
+            layout.tokenizer_output.tokens[:, 12] += 100
+        elif change == "target":
+            layout.tokenizer_output.tokens[:, 20:] += 100
+        request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+        prepare_hunyuan_prefix_cache(request)
+        return request.diffusion_kv_requests
+
+    manager = prefix_cache_manager
+    cold = build("cold", "same")
+    cold_metadata = manager.reserve_request("cold", cold)
+    assert cold_metadata is not None
+    assert all(sequence.cached_prefix_len == 0 for sequence in cold_metadata.sequences)
+    # CPU-only: simulate successful materialization to test native cache lookup.
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = build("warm", change)
+    warm_metadata = manager.reserve_request("warm", warm)
+    assert warm_metadata is not None
+    for cold_row, warm_row, cold_sequence, warm_sequence in zip(
+        cold, warm, cold_metadata.sequences, warm_metadata.sequences, strict=True
+    ):
+        if change == "prompt_suffix":
+            assert cold_row.cache_token_ids[:12] == warm_row.cache_token_ids[:12]
+            assert cold_row.cache_token_ids[12] != warm_row.cache_token_ids[12]
+        else:
+            assert cold_row.cache_token_ids == warm_row.cache_token_ids
+        assert warm_sequence.cached_prefix_len == expected_hit
+        hit_blocks = expected_hit // 4
+        assert cold_sequence.block_ids[0][:hit_blocks] == warm_sequence.block_ids[0][:hit_blocks]
+        assert cold_row.block_hashes[:hit_blocks] == warm_row.block_hashes[:hit_blocks]
+        assert all(
+            a != b for a, b in zip(cold_row.block_hashes[hit_blocks:], warm_row.block_hashes[hit_blocks:], strict=True)
+        )
+
+
 def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
     tokenizer, image_processor = _components([20])
     joint_image = _reference_image()
@@ -366,13 +730,17 @@ def test_prepared_model_inputs_match_local_tokenization(
         # Preparation runs without a forward context. Coverage must come from
         # the scheduler metadata copied onto this request, including every CFG row.
         reused = pipeline.prepare_model_inputs(
-            **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(5,) * len(prefix_lens)
+            # The image span is [5, 10); native KV must cover its end, not
+            # merely reach the first image token.
+            **common_kwargs,
+            prepared_layout=prepared_layout,
+            kv_computed_tokens=(10,) * len(prefix_lens),
         )
         assert reused["cond_vae_images"] is None and reused["cond_vit_images"] is None
         partial = pipeline.prepare_model_inputs(
             **common_kwargs,
             prepared_layout=prepared_layout,
-            kv_computed_tokens=(5,) * (len(prefix_lens) - 1) + (4,),
+            kv_computed_tokens=(10,) * (len(prefix_lens) - 1) + (9,),
         )
         assert partial["cond_vae_images"] is not None and partial["cond_vit_images"] is not None
 
@@ -496,6 +864,9 @@ def test_cfg_rows_preserve_reusable_token_ids_and_native_transfer_boundary(trans
 
     tokenizer, image_processor = _components([12, 14])
     request = _request(guidance_scale=5.0)
+    # Prompt token IDs are only materialized for the native AR -> DiT
+    # transfer path; local prefix caching does not need them.
+    request.kv_transfer_params = {"num_transfer_tokens": 0}
     layout = _prepare(request, tokenizer, image_processor)
     layout.tokenizer_output.think_recaption_end_pos = [[7], [5]]
     layout.tokenizer_output.uncond_cfg_start_pos = [[None], [5]]

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -147,6 +147,80 @@ def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
     torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
 
 
+def test_slice_cached_prefix_inputs_keeps_full_kv_axis() -> None:
+    inputs_embeds = torch.arange(12, dtype=torch.float32).reshape(1, 6, 2)
+    attention_mask = torch.arange(36).reshape(1, 1, 6, 6)
+    position_ids = torch.arange(6).reshape(1, 6)
+    custom_pos_emb = (torch.arange(6).reshape(1, 6), torch.arange(10, 16).reshape(1, 6))
+    image_mask = torch.tensor([[False, False, True, True, True, True]])
+    gen_timestep_scatter_index = torch.tensor([[2]])
+
+    sliced = HunyuanImage3Pipeline._slice_cached_prefix_inputs(
+        inputs_embeds,
+        attention_mask,
+        position_ids,
+        custom_pos_emb,
+        image_mask,
+        gen_timestep_scatter_index,
+        [6],
+        [6],
+        2,
+    )
+
+    (
+        sliced_embeds,
+        sliced_mask,
+        sliced_positions,
+        sliced_rope,
+        sliced_image_mask,
+        sliced_scatter,
+        sliced_query_lens,
+        sliced_seq_len,
+    ) = sliced
+    torch.testing.assert_close(sliced_embeds, inputs_embeds[:, 2:])
+    torch.testing.assert_close(sliced_mask, attention_mask[:, :, 2:, :])
+    torch.testing.assert_close(sliced_positions, position_ids[:, 2:])
+    torch.testing.assert_close(sliced_rope[0], custom_pos_emb[0][:, 2:])
+    torch.testing.assert_close(sliced_rope[1], custom_pos_emb[1][:, 2:])
+    torch.testing.assert_close(sliced_image_mask, image_mask[:, 2:])
+    torch.testing.assert_close(sliced_scatter, torch.tensor([[0]]))
+    assert sliced_query_lens == [4]
+    assert sliced_seq_len == 4
+
+
+@pytest.mark.parametrize("local_prefix_hit", [True, False])
+def test_ar_reuse_does_not_consume_local_prefix_hits(local_prefix_hit):
+    from vllm_omni.diffusion.forward_context import set_forward_context
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+        HunyuanImage3Text2ImagePipeline,
+    )
+
+    pipe = object.__new__(HunyuanImage3Text2ImagePipeline)
+    input_ids = torch.arange(12).reshape(1, 12)
+    cond_image = torch.ones(1, 1)
+    kwargs = dict(
+        query_lens=[12],
+        attention_mask=torch.ones(1, 1, 12, 12),
+        position_ids=input_ids.clone(),
+        image_mask=torch.zeros_like(input_ids, dtype=torch.bool),
+        gen_timestep_scatter_index=torch.tensor([[8]]),
+        cond_vae_images=cond_image,
+    )
+    runtime = SimpleNamespace(metadata=SimpleNamespace(prefill_rows=[SimpleNamespace(kv_start_pos=4)]))
+    with set_forward_context(paged_kv_runtime=runtime, paged_kv_cached_prefix_len=4 if local_prefix_hit else 0):
+        output, ar_reuse_len = pipe._maybe_handle_ar_kv_reuse(input_ids, kwargs, 1, False, None, torch.device("cpu"))
+    if local_prefix_hit:
+        assert ar_reuse_len == 0
+        assert output is input_ids
+        assert kwargs["query_lens"] == [12]
+        assert kwargs["cond_vae_images"] is cond_image
+    else:
+        assert ar_reuse_len == 4
+        torch.testing.assert_close(output, input_ids[:, 4:])
+        assert kwargs["query_lens"] == [8]
+        assert "cond_vae_images" not in kwargs
+
+
 def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
     pipeline = _pipeline()
     states = [_state("req-0", 1), _state("req-1", 3)]
@@ -416,7 +490,7 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
             }
         ]
 
-    captured = {}
+    captured: dict[str, object] = {}
 
     def fake_restore_prompt_kv_cache(states_arg, row_state_indexes, row_branches):
         del states_arg
@@ -443,7 +517,9 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
     assert captured["row_state_indexes"] == [0, 1, 0, 1]
     assert captured["row_branches"] == [0, 0, 1, 1]
     assert captured["input_ids"] is None
+    assert isinstance(captured["images"], torch.Tensor)
     assert tuple(captured["images"].shape) == (4, 1)
+    assert isinstance(captured["timestep"], torch.Tensor)
     assert captured["timestep"].tolist() == [0.5, 0.0, 0.5, 0.0]
     assert captured["merged_attention_mask_shape"] == (4, 1, 2, 6)
     assert captured["merged_full_attn_spans"] == [[(4, 6)], [(4, 6)], [(4, 6)], [(4, 6)]]

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -8,7 +8,7 @@ Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP,
 from __future__ import annotations
 
 import math
-from contextlib import ExitStack, contextmanager
+from contextlib import AbstractContextManager, ExitStack, contextmanager
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -59,7 +59,7 @@ class MockAttention(nn.Module):
 @contextmanager
 def patched_mgr_env(sp_size=1):
     target = _TRANSFORMER_MODULE
-    patches = [
+    patches: list[AbstractContextManager] = [
         patch(f"{target}.get_sequence_parallel_world_size", return_value=sp_size),
         patch(f"{target}.get_allgather_parallel_world_size", return_value=sp_size),
         patch(f"{target}.get_ulysses_parallel_world_size", return_value=1, create=True),
@@ -242,6 +242,39 @@ def test_legacy_tensor_prefix_still_uses_dense_cache_and_reuse() -> None:
     assert mgr.attn.calls[-1][1].shape == (1, prefix_len + first_len, NUM_HEADS, HEAD_DIM)
 
 
+@pytest.mark.parametrize("prefix_len", [0, 19])
+@pytest.mark.parametrize("prompt_query_len", [0, 3])
+def test_paged_sp_splits_uncached_queries(prefix_len, prompt_query_len) -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr(sp_size=2)
+    mgr.attn.paged_kv_active = True
+    shard_image_size = IMAGE_TOKEN_LEN // 2
+    query_len = prompt_query_len + shard_image_size
+    query = torch.randn(2 * query_len, NUM_HEADS, HEAD_DIM)
+    key, value = _make_known_kv(2 * query_len)
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        output = mgr(
+            query,
+            key,
+            value,
+            first_step=True,
+            query_lens=[query_len] * 2,
+            seq_lens=[prefix_len + query_len] * 2,
+            shard_image_size=shard_image_size,
+            full_attn_spans=[[(prefix_len + prompt_query_len, prefix_len + query_len)]] * 2,
+        )
+
+    torch.testing.assert_close(output, query)
+    metadata = mgr.attn.paged_metadata[0]
+    torch.testing.assert_close(
+        metadata.joint_key, key.reshape(2, query_len, NUM_KV_HEADS, HEAD_DIM)[:, :prompt_query_len]
+    )
+    assert metadata.joint_query.shape[1] == prompt_query_len
+    assert mgr.attn.paged_calls[0][0].shape[1] == shard_image_size
+    assert mgr.image_kv_cache_map is None
+
+
 def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
     from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
 
@@ -391,6 +424,7 @@ def test_no_ar_kv(bs):
         gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
     )
 
+    assert mgr.image_kv_cache_map is not None
     cached_key, cached_value = mgr.image_kv_cache_map
     # 3 prompt tokens cached per batch
     assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
@@ -478,12 +512,14 @@ def test_ar_kv_no_cfg(sp_size):
         gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
     )
 
+    assert mgr.image_kv_cache_map is not None
     cached_key, cached_value = mgr.image_kv_cache_map
     if sp_size == 1:
         # cached = ar(5) + prompt(3) = 8
         expected_cached_len = 8
     else:
         # cached = seq_len - shard_image_size = 17 - 4 = 13
+        assert shard_image_size is not None
         expected_cached_len = seq_len - shard_image_size
 
     assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)

tests/diffusion/test_diffusion_config_propagation.py

@@ -5,6 +5,7 @@
 Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
 """
 
+import json
 from collections.abc import Mapping
 
 import pytest
@@ -266,6 +267,21 @@ def test_architecture_name_resolves_via_pipeline_class_fallback():
     assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
 
 
+def test_architecture_only_checkpoint_propagates_multimodal_limit(tmp_path):
+    """A config.json-only checkpoint must expose its image-input capability."""
+    (tmp_path / "config.json").write_text(
+        json.dumps({"architectures": ["HunyuanImage3ForCausalMM"]}),
+        encoding="utf-8",
+    )
+
+    od_config = OmniDiffusionConfig(model=str(tmp_path))
+    od_config.enrich_config()
+
+    assert od_config.model_class_name == "HunyuanImage3ForCausalMM"
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
+
+
 def test_additional_config_roundtrip():
     additional_config = {"torchair_graph_config": {"enabled": True}}
     od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -149,6 +149,24 @@ def test_dense_mode_does_not_build_kv_profile_request() -> None:
     engine._make_dummy_request.assert_not_called()
 
 
+def test_explicit_profile_frames_preserve_warmup_skip(monkeypatch):
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.__new__(OmniDiffusionConfig)
+    engine.od_config.model_class_name = "HunyuanImage3ForCausalMM"
+    monkeypatch.setattr("vllm_omni.diffusion.diffusion_engine.supports_multimodal_input", lambda _: (True, False))
+    monkeypatch.setattr("vllm_omni.diffusion.diffusion_engine.image_color_format", lambda _: "RGB")
+    monkeypatch.setattr("vllm_omni.diffusion.diffusion_engine.get_dummy_run_num_frames", lambda *_: 0)
+    kwargs = dict(height=1024, width=1024, guidance_scale=5.0, num_image_inputs=3)
+
+    assert engine._make_dummy_request(**kwargs) is None
+    request = engine._make_dummy_request(**kwargs, num_frames=1)
+
+    assert request is not None
+    assert request.sampling_params.num_frames == 1
+    assert request.sampling_params.num_inference_steps == 1
+    assert len(request.prompt["multi_modal_data"]["image"]) == 3
+
+
 @pytest.mark.parametrize(
     ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
     [

tests/diffusion/test_diffusion_model_runner.py

@@ -223,7 +223,7 @@ def test_release_captured_graphs_tolerates_a_pipeline_without_captures():
     assert not hasattr(runner, "graph_runners")
 
 
-def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
+def _make_runner(cache_backend, cache_backend_name: str | None, enable_cache_dit_summary: bool = True):
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()
     runner.device = torch.device("cpu")
@@ -253,6 +253,55 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     return runner
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("imported_tokens", [None, (4, 4), (4, 2)])
+def test_request_forward_separates_imported_kv_and_local_hits(monkeypatch, imported_tokens):
+    from contextlib import nullcontext
+
+    from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata, DiffusionKVSequenceMetadata
+
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.od_config.kv_transfer_config = object() if imported_tokens else None
+    request = _make_request()
+    metadata = DiffusionKVMetadata(
+        request_id=request.request_id,
+        allocation_generation=1,
+        sequences=tuple(
+            DiffusionKVSequenceMetadata(
+                sequence_id=i,
+                prefix_len=8,
+                target_len=4,
+                seq_len=12,
+                block_ids=([1, 2, 3],),
+                cached_prefix_len=0 if imported_tokens else 4,
+                num_computed_tokens=imported_tokens[i] if imported_tokens else 4,
+            )
+            for i in range(2)
+        ),
+    )
+    captured = {}
+
+    @contextmanager
+    def capture_context(**kwargs):
+        captured.update(kwargs)
+        yield
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", capture_context)
+    runner.diffusion_kv_backend.activate_paged_attention_metadata = Mock(return_value=(object(), nullcontext()))
+    runner._execute_request_list(
+        [request],
+        od_config=runner.od_config,
+        allow_single_output=True,
+        require_request_batch_support=False,
+        record_name="test",
+        record_output_peak_memory=False,
+        diffusion_kv_metadata=[metadata],
+    )
+    assert captured["paged_kv_cached_prefix_len"] == (0 if imported_tokens else 4)
+    assert getattr(request, "kv_computed_tokens", ()) == (imported_tokens or ())
+
+
 def _make_compile_runner(
     model=None,
     *,
@@ -310,7 +359,8 @@ def test_refresh_cache_prefers_request_steps_then_schedule_then_pipeline_default
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_refresh_cache_without_request_or_pipeline_default_warns(caplog):
+def test_refresh_cache_without_request_or_pipeline_default_warns(caplog, monkeypatch):
+    monkeypatch.setattr(model_runner_module.logger, "handlers", [*model_runner_module.logger.handlers, caplog.handler])
     cache_backend = _EnabledCacheBackend()
     runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit")
     req = _make_request()
@@ -828,7 +878,7 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
         observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
         return original_denoise_step(input_batch, states)
 
-    runner.pipeline.denoise_step = denoise_step
+    monkeypatch.setattr(runner.pipeline, "denoise_step", denoise_step)
     runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
     forward_context_calls = []
 
@@ -1313,7 +1363,7 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
         device_config=DeviceConfig(device="cpu"),
         compilation_config=CompilationConfig(),
     )
-    calls = []
+    calls: list[tuple[str, object]] = []
 
     @contextmanager
     def _set_priority():

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -20,8 +20,10 @@ from vllm_omni.diffusion.registry import (
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
     _DIFFUSION_PRE_PROCESS_FUNCS,
+    _DIFFUSION_PREFIX_CACHE_FUNCS,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
+    get_diffusion_prefix_cache_func,
     register_diffusion_model,
 )
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
@@ -70,6 +72,7 @@ class TestRegisterDiffusionModel:
         original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
         original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
         original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
+        original_prefix_cache = _DIFFUSION_PREFIX_CACHE_FUNCS.copy()
         yield
         _DIFFUSION_MODELS.clear()
         _DIFFUSION_MODELS.update(original_models)
@@ -79,6 +82,8 @@ class TestRegisterDiffusionModel:
         _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
+        _DIFFUSION_PREFIX_CACHE_FUNCS.clear()
+        _DIFFUSION_PREFIX_CACHE_FUNCS.update(original_prefix_cache)
 
     def test_register_new_model(self):
         """Test registering a new diffusion model with pre/post process functions."""
@@ -89,6 +94,7 @@ class TestRegisterDiffusionModel:
             pre_process_func_name="test_pre_process",
             post_process_func_name="test_post_process",
             ir_op_priority_func_name="test_ir_op_priority",
+            prefix_cache_func_name="test_prefix_cache",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
         assert _DIFFUSION_MODELS["TestPipeline"] == (
@@ -99,6 +105,29 @@ class TestRegisterDiffusionModel:
         assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
         assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
         assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
+        assert _DIFFUSION_PREFIX_CACHE_FUNCS["TestPipeline"] == "test_prefix_cache"
+
+    def test_registered_prefix_hook_is_loaded_from_plugin(self):
+        hook = Mock()
+        factory = Mock(return_value=hook)
+        register_diffusion_model(
+            "TestPrefixPipeline", "test_plugin.pipeline", "TestPrefixPipeline", prefix_cache_func_name="prefix_factory"
+        )
+        config = SimpleNamespace(model_class_name="TestPrefixPipeline")
+        with patch(
+            "vllm_omni.diffusion.registry.importlib.import_module", return_value=SimpleNamespace(prefix_factory=factory)
+        ) as load:
+            assert get_diffusion_prefix_cache_func(config) is hook
+        load.assert_called_once_with("test_plugin.pipeline")
+        factory.assert_called_once_with(config)
+
+    def test_builtin_hunyuan_prefix_hook(self):
+        from vllm_omni.diffusion.models.hunyuan_image3.request_layout import prepare_hunyuan_prefix_cache
+
+        assert (
+            get_diffusion_prefix_cache_func(SimpleNamespace(model_class_name="HunyuanImage3ForCausalMM"))
+            is prepare_hunyuan_prefix_cache
+        )
 
     def test_register_model_accepts_deprecated_action_postprocess_keyword(self):
         """Deprecated action hook keyword is accepted but not registered."""

tests/diffusion/test_diffusion_scheduler.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.sched import (
     Scheduler,
     StepScheduler,
 )
-from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
+from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData, SchedulerRequestState
 from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -96,6 +96,7 @@ def _initialize_paged_scheduler(
     *,
     num_blocks: int = 64,
     max_num_seqs: int = 1,
+    enable_prefix_caching: bool = False,
 ) -> None:
     native_kv_managers.register_all_kvcache_specs(None)
     spec = FullAttentionSpec(
@@ -121,22 +122,42 @@ def _initialize_paged_scheduler(
         kv_vllm_config=SimpleNamespace(
             model_config=SimpleNamespace(max_model_len=64),
             max_in_flight_tokens=64,
+            cache_config=SimpleNamespace(
+                enable_prefix_caching=enable_prefix_caching,
+                prefix_caching_hash_algo="sha256",
+            ),
         ),
     )
 
 
-def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+def _attach_diffusion_kv(
+    request: OmniDiffusionRequest,
+    *,
+    seq_len: int = 8,
+    prefix_len: int = 4,
+    cache_token_ids=(),
+) -> None:
     request.diffusion_kv_requests = (
         DiffusionKVRequest(
             f"{request.request_id}/diffusion-kv/0",
             sequence_id=0,
-            prefix_len=4,
+            prefix_len=prefix_len,
             target_len=4,
             seq_len=seq_len,
+            cache_token_ids=cache_token_ids,
         ),
     )
 
 
+def _make_aborted_request_output(req_id: str) -> RunnerOutput:
+    return RunnerOutput(
+        request_id=req_id,
+        step_index=None,
+        finished=True,
+        result=DiffusionOutput(output=None, aborted=True),
+    )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("scheduler_cls", [RequestScheduler, StepScheduler])
 @pytest.mark.parametrize("failure", ["timeout", "registration"])
@@ -168,7 +189,7 @@ async def test_single_native_kv_failure_reaches_output_stream(mocker, scheduler_
     engine.abort_queue = queue.Queue()
     engine._cv = threading.Condition()
     engine.main_loop = asyncio.get_running_loop()
-    stream = asyncio.Queue()
+    stream: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
     engine._out_streams = {"failed": stream}
     engine.executor = mocker.Mock()
     engine.executor.prepare_kv_for_forward.return_value = KVConnectorOutput()
@@ -192,7 +213,7 @@ class _StubScheduler:
         self._output = output
         self.initialized_with = None
         self._request_id = request.request_id
-        self._state = None
+        self._state: SchedulerRequestState | None = None
         self._scheduled = False
         self.max_num_running_reqs = 1
 
@@ -201,7 +222,7 @@ class _StubScheduler:
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         assert request is self._request
-        self._state = SimpleNamespace(request_id=self._request_id, req=request)
+        self._state = SchedulerRequestState(request_id=self._request_id, req=request)
         return self._request_id
 
     def schedule(self):
@@ -223,6 +244,7 @@ class _StubScheduler:
     def update_from_output(self, sched_output, output) -> set[str]:
         del sched_output
         assert output is self._output
+        assert self._state is not None
         self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
         return {self._request_id}
 
@@ -553,6 +575,197 @@ class TestRequestScheduler:
         assert metadata.request_id == request.request_id
         assert len(metadata.sequences[0].block_ids[0]) == 4
 
+    def test_diffusion_kv_publishes_only_after_successful_completion(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, enable_prefix_caching=True)
+        first = _make_request("publish-success")
+        _attach_diffusion_kv(first, cache_token_ids=range(4))
+        self.scheduler.add_request(first)
+        first_schedule = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first_schedule, _make_request_output(first.request_id)) == {
+            first.request_id
+        }
+
+        warm = _make_request("publish-warm")
+        _attach_diffusion_kv(warm, cache_token_ids=range(4))
+        self.scheduler.add_request(warm)
+        warm_schedule = self.scheduler.schedule()
+        metadata = warm_schedule.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.sequences[0].cached_prefix_len == 4
+
+    @pytest.mark.parametrize("boundaries", [(8, 4), (8, 0), (0, 8), (4, 8), (8, 8), (4, 4), (0, 0)])
+    @pytest.mark.parametrize("num_branches", [1, 2])
+    def test_diffusion_kv_batches_only_matching_prefix_boundaries(self, boundaries, num_branches) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        pool = manager.native_manager.block_pool
+        empty_free_blocks = pool.get_num_free_blocks()
+
+        warmup = _make_request("warmup")
+        _attach_diffusion_kv(warmup, seq_len=12, prefix_len=8, cache_token_ids=range(8))
+        self.scheduler.add_request(warmup)
+        self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        for request_id, boundary in zip(("first", "second"), boundaries):
+            request = _make_request(request_id)
+            # Each request's CFG branches share a boundary, but requests may not.
+            tokens = tuple(range(boundary)) + tuple(range(100, 108 - boundary))
+            request.diffusion_kv_requests = tuple(
+                DiffusionKVRequest(
+                    f"{request_id}/diffusion-kv/{branch}",
+                    sequence_id=branch,
+                    prefix_len=8,
+                    target_len=4,
+                    seq_len=12,
+                    cache_token_ids=tokens,
+                )
+                for branch in range(num_branches)
+            )
+            self.scheduler.add_request(request)
+
+        scheduled = self.scheduler.schedule()
+        matching = boundaries[0] == boundaries[1]
+        expected_ids = ["first", "second"] if matching else ["first"]
+        assert _new_ids(scheduled) == expected_ids
+        assert {
+            sequence.cached_prefix_len
+            for req in scheduled.scheduled_new_reqs
+            for sequence in req.diffusion_kv_metadata.sequences
+        } == {boundaries[0]}
+
+        if not matching:
+            second_state = self.scheduler.get_request_state("second")
+            assert second_state.status is DiffusionRequestStatus.WAITING
+            assert not manager.has_request("second")
+            # Retrying admission cannot leak reservations or publish its suffix.
+            free_blocks = pool.get_num_free_blocks()
+            for _ in range(3):
+                retry = self.scheduler.schedule()
+                assert _new_ids(retry) == [] and _cached_ids(retry) == ["first"]
+                assert pool.get_num_free_blocks() == free_blocks
+                assert not manager.has_request("second")
+                for row in second_state.diffusion_kv_requests:
+                    assert row.num_computed_tokens == 0
+                    assert manager.native_manager.get_computed_blocks(row)[1] == boundaries[1]
+
+        self.scheduler.update_from_output(
+            scheduled,
+            BatchRunnerOutput.from_list([_make_request_output(req_id) for req_id in expected_ids]),
+        )
+        if not matching:
+            next_wave = self.scheduler.schedule()
+            assert _new_ids(next_wave) == ["second"]
+            assert next_wave.scheduled_new_reqs[0].diffusion_kv_metadata.sequences[0].cached_prefix_len == boundaries[1]
+            self.scheduler.update_from_output(next_wave, _make_request_output("second"))
+        assert not self.scheduler.has_requests()
+        assert pool.get_num_free_blocks() == empty_free_blocks
+
+    def test_deferred_allocation_is_not_registered_for_transfer(self, mocker) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2)
+        manager = self.scheduler._diffusion_kv_manager
+        connector = mocker.Mock()
+        connector.get_num_new_matched_tokens.return_value = (4, True)
+        self.scheduler._kv_connector = connector
+        # Model a compatibility decision that needs the completed lookup.
+        can_schedule = self.scheduler._can_schedule_waiting
+        mocker.patch.object(
+            self.scheduler,
+            "_can_schedule_waiting",
+            side_effect=lambda state: can_schedule(state)
+            and not (state.request_id == "deferred" and manager.has_request("deferred")),
+        )
+        for request_id in ("admitted", "deferred"):
+            request = _make_request(request_id)
+            _attach_diffusion_kv(request)
+            request.diffusion_kv_requests[0].prompt_token_ids = [1, 2, 3, 4]
+            request.kv_transfer_params = {"num_transfer_tokens": 4}
+            self.scheduler.add_request(request)
+
+        assert _new_ids(self.scheduler.schedule()) == ["admitted"]
+        assert connector.update_state_after_alloc.call_count == 1
+        assert not manager.has_request("deferred")
+        assert "deferred" not in self.scheduler._kv_request_generations
+        assert "deferred" not in self.scheduler._kv_loading_request_ids
+        assert "deferred/diffusion-kv/0" not in self.scheduler._kv_transfer_request_ids
+        assert self.scheduler.get_request_state("deferred").status is DiffusionRequestStatus.WAITING
+
+    def test_diffusion_kv_deferred_request_can_be_cancelled(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        pool = manager.native_manager.block_pool
+        empty_free_blocks = pool.get_num_free_blocks()
+        for request_id, tokens in (("warmup", range(4)), ("hit", range(4)), ("miss", range(10, 14))):
+            request = _make_request(request_id)
+            _attach_diffusion_kv(request, cache_token_ids=tokens)
+            self.scheduler.add_request(request)
+            if request_id == "warmup":
+                self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        scheduled = self.scheduler.schedule()
+        assert _new_ids(scheduled) == ["hit"]
+        self.scheduler.finish_requests("miss", DiffusionRequestStatus.FINISHED_ABORTED)
+        self.scheduler.update_from_output(scheduled, _make_request_output("hit", error="worker failed"))
+        assert not manager.has_request("miss")
+        assert not self.scheduler.has_requests()
+        assert pool.get_num_free_blocks() == empty_free_blocks
+
+    def test_diffusion_kv_preempted_request_keeps_reservation_when_deferred(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        warmup = _make_request("warmup")
+        _attach_diffusion_kv(warmup, cache_token_ids=range(4))
+        self.scheduler.add_request(warmup)
+        self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        hit = _make_request("hit")
+        _attach_diffusion_kv(hit, cache_token_ids=range(4))
+        self.scheduler.add_request(hit)
+        self.scheduler.schedule()
+        miss = _make_request("miss")
+        _attach_diffusion_kv(miss, cache_token_ids=range(10, 14))
+        self.scheduler.add_request(miss)
+        # Install retained state as if this request were resuming an earlier wave.
+        miss_state = self.scheduler.get_request_state("miss")
+        miss_state.status = DiffusionRequestStatus.PREEMPTED
+        miss_metadata = manager.reserve_request("miss", miss_state.diffusion_kv_requests)
+        assert _cached_ids(self.scheduler.schedule()) == ["hit"]
+        assert self.scheduler.get_request_state("miss").status is DiffusionRequestStatus.PREEMPTED
+        assert manager.get_metadata("miss") is miss_metadata
+        self.scheduler.finish_requests("hit", DiffusionRequestStatus.FINISHED_ABORTED)
+        assert _cached_ids(self.scheduler.schedule()) == ["miss"]
+        assert manager.get_metadata("miss") is miss_metadata
+        self.scheduler.close()
+
+    @pytest.mark.parametrize(
+        "terminal_output",
+        [
+            pytest.param(
+                _make_request_output("publish-error", error="worker failed"),
+                id="error",
+            ),
+            pytest.param(_make_aborted_request_output("publish-abort"), id="abort"),
+        ],
+    )
+    def test_diffusion_kv_does_not_publish_failed_or_aborted_request(
+        self,
+        terminal_output: RunnerOutput,
+    ) -> None:
+        _initialize_paged_scheduler(self.scheduler, enable_prefix_caching=True)
+        first = _make_request(terminal_output.request_id)
+        _attach_diffusion_kv(first, cache_token_ids=range(4))
+        self.scheduler.add_request(first)
+        first_schedule = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first_schedule, terminal_output) == {first.request_id}
+
+        warm = _make_request(f"{first.request_id}-warm")
+        _attach_diffusion_kv(warm, cache_token_ids=range(4))
+        self.scheduler.add_request(warm)
+        warm_schedule = self.scheduler.schedule()
+        metadata = warm_schedule.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.sequences[0].cached_prefix_len == 0
+
     def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
         _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
         first_request = _make_request("first")