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

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

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