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

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)