[Diffusion] Prepare HunyuanImage3 for Scheduler-managed paged KV cache (#5541)

Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/diffusion_kv/test_config.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+@pytest.fixture(autouse=True)
+def _fixed_master_port(monkeypatch) -> None:
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+
+
+def test_dense_legacy_is_default() -> None:
+    config = OmniDiffusionConfig.from_kwargs()
+
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
+
+
+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")
+
+
+def test_unknown_cache_mode_is_rejected() -> None:
+    with pytest.raises(ValueError, match="Unsupported Diffusion KV diffusion_kv_mode"):
+        OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="unknown")
+
+
+def test_non_mapping_omni_kv_config_is_rejected() -> None:
+    with pytest.raises(TypeError, match="omni_kv_config must be a mapping"):
+        OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
+
+
+def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
+    with pytest.raises(ValueError, match="does not support imported AR KV"):
+        OmniDiffusionConfig.from_kwargs(
+            diffusion_kv_mode="paged_scheduler",
+            omni_kv_config={"need_recv_cache": True},
+        )
+
+
+def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
+    config = OmniDiffusionConfig.from_kwargs(
+        model_class_name="FutureDiffusionModel",
+        diffusion_kv_mode="paged_scheduler",
+    )
+
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER

tests/diffusion/diffusion_kv/test_request.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.kv_cache_interface import (
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    KVCacheTensor,
+)
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+BLOCK_SIZE = 4
+
+
+def _request(**overrides) -> DiffusionKVRequest:
+    values = dict(
+        request_id="req/diffusion-kv/0",
+        sequence_id=0,
+        prefix_len=4,
+        target_len=4,
+        seq_len=8,
+    )
+    values.update(overrides)
+    return DiffusionKVRequest(**values)
+
+
+def test_request_exposes_native_request_and_diffusion_semantics() -> None:
+    request = _request()
+
+    assert request.sequence_id == 0
+    assert request.prefix_len == 4
+    assert request.target_len == 4
+    assert request.num_tokens == request.seq_len == 8
+    assert request.num_prompt_tokens == 4
+    assert request.num_computed_tokens == 0
+    assert request.block_hashes == []
+    assert request.kv_contexts == ()
+    assert request.skip_reading_prefix_cache is True
+    assert request.shared_prefix_boundary == 0
+    assert request.status is RequestStatus.WAITING
+
+
+@pytest.mark.parametrize(
+    ("overrides", "message"),
+    [
+        ({"sequence_id": -1}, "sequence_id"),
+        ({"prefix_len": -1}, "prefix_len"),
+        ({"target_len": 0}, "target_len"),
+        ({"seq_len": 0}, "seq_len"),
+        ({"prefix_len": 5, "target_len": 4, "seq_len": 8}, "must not exceed"),
+    ],
+)
+def test_request_rejects_invalid_lengths(overrides, message: str) -> None:
+    with pytest.raises(ValueError, match=message):
+        _request(**overrides)
+
+
+@pytest.mark.parametrize(
+    ("overrides", "message"),
+    [
+        ({"context_id": ""}, "context_id"),
+        ({"cache_role": ""}, "cache_role"),
+        ({"num_tokens": 0}, "num_tokens"),
+    ],
+)
+def test_context_rejects_invalid_identity_or_length(overrides, message: str) -> None:
+    values = dict(context_id="text", cache_role="cross.text", num_tokens=8)
+    values.update(overrides)
+
+    with pytest.raises(ValueError, match=message):
+        DiffusionKVContext(**values)
+
+
+def test_request_owns_independent_kv_contexts() -> None:
+    text = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+    image = DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=4)
+
+    request = _request(kv_contexts=(text, image))
+
+    assert request.kv_contexts == (text, image)
+
+
+def test_request_rejects_duplicate_context_ids() -> None:
+    contexts = (
+        DiffusionKVContext(context_id="condition", cache_role="cross.text", num_tokens=8),
+        DiffusionKVContext(context_id="condition", cache_role="cross.image", num_tokens=4),
+    )
+
+    with pytest.raises(ValueError, match="unique context_id"):
+        _request(kv_contexts=contexts)
+
+
+def test_request_rejects_non_context_values() -> None:
+    with pytest.raises(TypeError, match="DiffusionKVContext"):
+        _request(kv_contexts=(object(),))
+
+
+def _manager(*, enable_caching: bool) -> KVCacheManager:
+    # vLLM 0.26 lazily registers its built-in KV specs. Another test may have
+    # already registered an out-of-tree spec, making the registry non-empty
+    # before FullAttentionSpec is registered. Initialize the native mappings
+    # explicitly; vLLM 0.25 has no registration hook and needs no setup.
+    register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None)
+    if register_all is not None:
+        register_all(None)
+
+    spec = FullAttentionSpec(
+        block_size=BLOCK_SIZE,
+        num_kv_heads=2,
+        head_size=8,
+        dtype=torch.bfloat16,
+    )
+    group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)
+    config = KVCacheConfig(
+        num_blocks=8,
+        kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])],
+        kv_cache_groups=[group],
+    )
+    return KVCacheManager(
+        config,
+        max_model_len=64,
+        scheduler_block_size=BLOCK_SIZE,
+        hash_block_size=BLOCK_SIZE,
+        enable_caching=enable_caching,
+    )
+
+
+def test_empty_hash_request_conforms_to_native_request_local_allocation() -> None:
+    manager = _manager(enable_caching=True)
+    free_before = manager.block_pool.get_num_free_blocks()
+    request = _request()
+
+    # vLLM 0.26 adds a shared-prefix-boundary return value.
+    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
+    assert num_computed_tokens == 0
+    blocks = manager.allocate_slots(
+        request,
+        num_new_tokens=request.num_tokens,
+        new_computed_blocks=computed_blocks,
+        delay_cache_blocks=True,
+        full_sequence_must_fit=True,
+    )
+
+    assert blocks is not None
+    assert len(manager.get_block_ids(request.request_id)[0]) == 2
+    # Empty hashes intentionally mean request-local residency only. Calling
+    # native cache_blocks() would require a canonical hash for the full prefix
+    # block and is therefore outside this request's contract.
+    manager.free(request)
+    assert manager.block_pool.get_num_free_blocks() == free_before
+
+
+def test_hashed_request_conforms_to_native_delayed_prefix_commit() -> None:
+    manager = _manager(enable_caching=True)
+    prefix_hash = BlockHash(b"prefix".ljust(32, b"\0"))
+    target_hash = BlockHash(b"target".ljust(32, b"\0"))
+    request = _request(block_hashes=[prefix_hash, target_hash])
+
+    computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request)
+    assert num_computed_tokens == 0
+    blocks = manager.allocate_slots(
+        request,
+        num_new_tokens=request.num_tokens,
+        new_computed_blocks=computed_blocks,
+        delay_cache_blocks=True,
+        full_sequence_must_fit=True,
+    )
+
+    assert blocks is not None
+    manager.cache_blocks(request, request.prefix_len)
+    manager.free(request)
+
+    cached_request = _request(
+        request_id="req/diffusion-kv/cached",
+        block_hashes=[prefix_hash, target_hash],
+    )
+    cached_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(cached_request)
+
+    # Only the stable prefix was published. The dynamic target remains a miss.
+    assert num_computed_tokens == cached_request.prefix_len
+    assert len(cached_blocks.blocks[0]) == 1
+    cached_allocation = manager.allocate_slots(
+        cached_request,
+        num_new_tokens=cached_request.num_tokens - num_computed_tokens,
+        num_new_computed_tokens=num_computed_tokens,
+        new_computed_blocks=cached_blocks,
+        delay_cache_blocks=True,
+        full_sequence_must_fit=True,
+    )
+    assert cached_allocation is not None
+    manager.free(cached_request)

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py

@@ -0,0 +1,447 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from transformers.utils.generic import ModelOutput
+
+import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+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 (
+    ImageInfo,
+    JointImageInfo,
+)
+from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
+    HunyuanImage3Pipeline,
+    get_hunyuan_image_3_pre_process_func,
+)
+from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
+    HunyuanPreparedLayout,
+    build_hunyuan_diffusion_kv_requests,
+    extract_hunyuan_prompt_inputs,
+    hunyuan_num_image_tokens,
+    normalize_hunyuan_cot_text,
+    prepare_hunyuan_layout,
+)
+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 test_dense_legacy_preprocess_does_not_initialize_layout_tokenizer(monkeypatch) -> None:
+    hf_config = SimpleNamespace(vae_downsample_factor=(8, 8), patch_size=2)
+    image_processor = SimpleNamespace(vision_encoder_processor=SimpleNamespace(patch_size=1))
+    monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
+    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
+    monkeypatch.setattr(
+        pipeline_module,
+        "TokenizerWrapper",
+        lambda _model: pytest.fail("dense_legacy must not initialize the Diffusion KV tokenizer"),
+    )
+
+    get_hunyuan_image_3_pre_process_func(
+        SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY)
+    )
+
+
+class _FakeTokenizerWrapper:
+    def __init__(self, *, prefix_lens: list[int]) -> None:
+        self.prefix_lens = prefix_lens
+        self.calls: list[dict] = []
+        self.pad_token_id = 0
+        self.eos_token_id = 2
+        self.boi_token_id = 3
+        self.end_recaption_token_id = 4
+        self.end_answer_token_id = 5
+        self.special_token_map = {f"<img_ratio_{index}>": 100 + index for index in range(33)}
+
+    def apply_chat_template(self, **kwargs):
+        self.calls.append(kwargs)
+        rows = len(self.prefix_lens)
+        has_cond_image = kwargs.get("batch_cond_image_info") is not None
+        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)]
+        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)]
+        cond_vae_image_mask = (
+            torch.stack([(positions >= 1) & (positions < 3) 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
+        )
+        sections = []
+        for _ in range(rows):
+            row_sections = []
+            if has_cond_image:
+                row_sections.append(
+                    {
+                        "type": "joint_image",
+                        "token_height": [1, 1],
+                        "token_width": [2, 2],
+                    }
+                )
+            row_sections.append({"type": "gen_image", "token_height": 4, "token_width": 4})
+            sections.append(row_sections)
+        return {
+            "output": TokenizerEncodeOutput(
+                tokens=torch.arange(rows * padded_seq_len, dtype=torch.long).reshape(rows, padded_seq_len),
+                gen_timestep_scatter_index=torch.tensor(self.prefix_lens, dtype=torch.long).reshape(rows, 1),
+                cond_timestep_scatter_index=(torch.zeros(rows, 1, dtype=torch.long) if has_cond_image else None),
+                all_image_slices=all_image_slices,
+                gen_image_mask=torch.stack(
+                    [(positions >= length + 1) & (positions < length + 17) for length in self.prefix_lens]
+                ),
+                cond_vae_image_mask=cond_vae_image_mask,
+                cond_vit_image_mask=cond_vit_image_mask,
+                joint_image_slices=joint_image_slices,
+                gen_image_slices=gen_image_slices,
+                real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
+            ),
+            "sections": sections,
+        }
+
+
+class _FakeImageProcessor:
+    def __init__(self, image_token_length: int = 16) -> None:
+        self.image_token_length = image_token_length
+        self.image_sizes: list[tuple[int, int]] = []
+        self.vision_encoder_processor = SimpleNamespace(patch_size=1)
+
+    def build_image_info(self, image_size):
+        self.image_sizes.append(image_size)
+        return ImageInfo(
+            image_type="gen_image",
+            image_width=image_size[1],
+            image_height=image_size[0],
+            token_width=4,
+            token_height=4,
+            image_token_length=self.image_token_length,
+            base_size=1024,
+            ratio_index=0,
+        )
+
+
+def _components(prefix_lens: list[int]):
+    return _FakeTokenizerWrapper(prefix_lens=prefix_lens), _FakeImageProcessor()
+
+
+def _prepare(
+    request: OmniDiffusionRequest,
+    tokenizer: _FakeTokenizerWrapper,
+    image_processor: _FakeImageProcessor,
+) -> HunyuanPreparedLayout:
+    prepared_layout = prepare_hunyuan_layout(
+        request,
+        tokenizer_wrapper=tokenizer,
+        image_processor=image_processor,
+        generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
+        image_base_size=1024,
+    )
+    request.prepared_layout = prepared_layout
+    return prepared_layout
+
+
+def _request(
+    *,
+    guidance_scale: float,
+    prompt="draw a cat",
+    request_id: str = "req",
+) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompt=prompt,
+        sampling_params=OmniDiffusionSamplingParams(
+            height=768,
+            width=1024,
+            guidance_scale=guidance_scale,
+            num_inference_steps=4,
+        ),
+        request_id=request_id,
+    )
+
+
+def test_builds_kv_request_lengths_without_model_execution() -> None:
+    tokenizer, image_processor = _components([12])
+    request = _request(guidance_scale=1.0)
+    prepared_layout = _prepare(request, tokenizer, image_processor)
+
+    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
+
+    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
+    assert all(isinstance(item, DiffusionKVRequest) for item in kv_requests)
+    assert [(item.prefix_len, item.target_len, item.seq_len) for item in kv_requests] == [(12, 17, 32)]
+    assert image_processor.image_sizes == [(768, 1024)]
+    assert tokenizer.calls[0]["sequence_template"] == "instruct"
+    assert tokenizer.calls[0]["cfg_factor"] == 1
+    assert kv_requests[0].block_hashes == []
+    assert kv_requests[0].kv_contexts == ()
+    assert kv_requests[0].skip_reading_prefix_cache is True
+
+
+def test_builds_one_kv_request_per_cfg_row() -> None:
+    tokenizer, image_processor = _components([12, 14])
+    request = _request(guidance_scale=5.0)
+    prepared_layout = _prepare(request, tokenizer, image_processor)
+
+    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
+
+    assert [item.sequence_id for item in kv_requests] == [0, 1]
+    assert [item.prefix_len for item in kv_requests] == [12, 14]
+    assert [item.seq_len for item in kv_requests] == [32, 34]
+    assert tokenizer.calls[0]["cfg_factor"] == 2
+
+
+def _reference_image() -> JointImageInfo:
+    return JointImageInfo(
+        vae_image_info=ImageInfo(
+            image_type="vae",
+            token_width=8,
+            token_height=8,
+            image_token_length=64,
+        ),
+        vision_image_info=ImageInfo(
+            image_type="siglip2",
+            token_width=4,
+            token_height=4,
+            image_token_length=16,
+        ),
+        vision_encoder_kwargs={
+            "spatial_shapes": torch.tensor([4, 4]),
+            "pixel_attention_mask": torch.ones(4, dtype=torch.bool),
+        },
+    )
+
+
+def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
+    tokenizer, image_processor = _components([20])
+    joint_image = _reference_image()
+    prompt = {
+        "prompt": "edit this image",
+        "additional_information": {"batch_cond_image_info": [joint_image]},
+    }
+    request = _request(guidance_scale=1.0, prompt=prompt)
+    prepared_layout = _prepare(request, tokenizer, image_processor)
+
+    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
+
+    assert tokenizer.calls[0]["batch_cond_image_info"] == [[joint_image]]
+    assert kv_requests[0].prefix_len == 20
+    assert kv_requests[0].target_len == 17
+    assert kv_requests[0].seq_len == 40
+    assert kv_requests[0].kv_contexts == ()
+
+
+def _assert_nested_equal(actual, expected) -> None:
+    if isinstance(expected, torch.Tensor):
+        torch.testing.assert_close(actual, expected)
+    elif isinstance(expected, dict):
+        assert actual.keys() == expected.keys()
+        for key in expected:
+            _assert_nested_equal(actual[key], expected[key])
+    elif isinstance(expected, list | tuple):
+        assert isinstance(actual, type(expected))
+        assert len(actual) == len(expected)
+        for actual_item, expected_item in zip(actual, expected):
+            _assert_nested_equal(actual_item, expected_item)
+    else:
+        assert actual == expected
+
+
+@pytest.mark.parametrize(
+    ("prefix_lens", "guidance_scale", "with_reference_image"),
+    [
+        pytest.param([12], 1.0, False, id="cfg-1"),
+        pytest.param([12, 14], 5.0, False, id="cfg-2"),
+        pytest.param([20], 1.0, True, id="cfg-1-reference-image"),
+        pytest.param([20, 22], 5.0, True, id="cfg-2-reference-image"),
+    ],
+)
+def test_prepared_model_inputs_match_local_tokenization(
+    monkeypatch,
+    prefix_lens: list[int],
+    guidance_scale: float,
+    with_reference_image: bool,
+) -> None:
+    tokenizer, image_processor = _components(prefix_lens)
+    prompt: str | dict = "draw a cat"
+    if with_reference_image:
+        prompt = {
+            "prompt": "edit this image",
+            "additional_information": {"batch_cond_image_info": [_reference_image()]},
+        }
+    request = _request(guidance_scale=guidance_scale, prompt=prompt)
+    prepared_layout = _prepare(request, tokenizer, image_processor)
+
+    prompts, cot_texts, system_prompt, batch_cond_image_info, bot_task = extract_hunyuan_prompt_inputs(
+        [request.prompt],
+        request.sampling_params.extra_args or {},
+        request_id=request.request_id,
+        allow_cond_image=True,
+    )
+    cot_text = (
+        [normalize_hunyuan_cot_text(text) for text in cot_texts]
+        if any(text is not None for text in cot_texts)
+        else None
+    )
+
+    pipeline = object.__new__(HunyuanImage3Pipeline)
+    pipeline._tkwrapper = tokenizer
+    pipeline.image_processor = image_processor
+    pipeline.generation_config = SimpleNamespace(sequence_template="instruct", drop_think=False)
+    pipeline.config = SimpleNamespace(
+        image_base_size=1024,
+        attention_head_dim=2,
+        rope_theta=10000.0,
+    )
+    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda _self: torch.device("cpu")))
+
+    def fake_encode_cond_image(batch_cond_image_info_list, cfg_factor=1, generator=None):
+        del generator
+        rows = len(batch_cond_image_info_list) * cfg_factor
+        return (
+            torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 1, 2, 2),
+            torch.arange(rows, dtype=torch.float32),
+            torch.arange(rows * 3, dtype=torch.float32).reshape(rows, 3),
+        )
+
+    pipeline._encode_cond_image = fake_encode_cond_image
+
+    def fake_build_batch_2d_rope(*, image_infos, seq_len, **_kwargs):
+        geometry_sum = sum(
+            int(image_slice.start or 0) + int(image_slice.stop or 0) + height + width
+            for row in image_infos
+            for image_slice, (height, width) in row
+        )
+        values = torch.tensor([len(image_infos), seq_len, geometry_sum], dtype=torch.float32)
+        return values, values + 1
+
+    monkeypatch.setattr(pipeline_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
+    common_kwargs = dict(
+        prompt=prompts,
+        cot_text=cot_text,
+        system_prompt=system_prompt,
+        mode="gen_image",
+        guidance_scale=guidance_scale,
+        image_size=(request.sampling_params.height, request.sampling_params.width),
+        generator=[torch.Generator().manual_seed(0)],
+        batch_cond_image_info=batch_cond_image_info,
+        bot_task=bot_task,
+    )
+
+    local_inputs = pipeline.prepare_model_inputs(**common_kwargs)
+    prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout)
+
+    comparable_fields = (
+        "input_ids",
+        "position_ids",
+        "custom_pos_emb",
+        "image_mask",
+        "gen_timestep_scatter_index",
+        "cond_vae_images",
+        "cond_timestep",
+        "cond_vae_image_mask",
+        "cond_vit_images",
+        "cond_vit_image_mask",
+        "vit_kwargs",
+        "cond_timestep_scatter_index",
+    )
+    for field in comparable_fields:
+        _assert_nested_equal(prepared_inputs[field], local_inputs[field])
+
+    tokenizer_fields = (
+        "tokens",
+        "gen_image_mask",
+        "cond_vae_image_mask",
+        "cond_vit_image_mask",
+        "gen_timestep_scatter_index",
+        "cond_timestep_scatter_index",
+        "real_pos",
+        "all_image_slices",
+        "joint_image_slices",
+        "gen_image_slices",
+    )
+    for field in tokenizer_fields:
+        _assert_nested_equal(
+            getattr(prepared_inputs["tokenizer_output"], field),
+            getattr(local_inputs["tokenizer_output"], field),
+        )
+
+    local_mask = pipeline._prepare_attention_mask_for_generation(
+        local_inputs["input_ids"],
+        pipeline.generation_config,
+        local_inputs,
+    )
+    prepared_mask = pipeline._prepare_attention_mask_for_generation(
+        prepared_inputs["input_ids"],
+        pipeline.generation_config,
+        prepared_inputs,
+    )
+    torch.testing.assert_close(prepared_mask, local_mask)
+    assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"]
+
+    num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
+    local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens)
+    prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens)
+    local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
+    prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)
+
+    for field in ("position_ids", "attention_mask", "gen_timestep_scatter_index", "full_attn_spans"):
+        _assert_nested_equal(prepared_step_inputs[field], local_step_inputs[field])
+
+
+def test_rejects_tokenizer_cfg_row_mismatch() -> None:
+    tokenizer, image_processor = _components([12])
+    request = _request(guidance_scale=5.0)
+    prepared_layout = _prepare(request, tokenizer, image_processor)
+
+    with pytest.raises(ValueError, match="sequence count does not match"):
+        build_hunyuan_diffusion_kv_requests(request, prepared_layout)
+
+
+def test_sender_endpoint_does_not_change_local_kv_request() -> None:
+    tokenizer, image_processor = _components([12])
+    request = _request(guidance_scale=1.0)
+    request.kv_sender_info = {"host": "127.0.0.1", "zmq_port": 5000}
+    prepared_layout = _prepare(request, tokenizer, image_processor)
+
+    kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout)
+
+    assert kv_requests[0].seq_len == 32
+    assert kv_requests[0].num_computed_tokens == 0
+
+
+def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch) -> None:
+    tokenizer, image_processor = _components([12])
+    hf_config = SimpleNamespace(
+        vae_downsample_factor=(8, 8),
+        patch_size=2,
+        image_base_size=1024,
+    )
+    monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config)
+    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(sequence_template="instruct", drop_think=False),
+    )
+    preprocess = get_hunyuan_image_3_pre_process_func(
+        SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
+    )
+    request = _request(guidance_scale=1.0)
+
+    prepared_request = preprocess(request)
+
+    assert prepared_request is request
+    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
+    assert request.diffusion_kv_requests is not None
+    assert [(item.prefix_len, item.target_len, item.seq_len) for item in request.diffusion_kv_requests] == [
+        (12, 17, 32)
+    ]
+    assert len(tokenizer.calls) == 1

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -7,7 +7,10 @@ import pytest
 import torch
 
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module
+import vllm_omni.diffusion.models.hunyuan_image3.request_layout as hy3_layout_module
 from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec
+from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
+from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageInfo
 from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
     _STEP_AR_KV,
     _STEP_CFG_FACTOR,
@@ -18,6 +21,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
     _STEP_PROMPT_KV,
     HunyuanImage3Pipeline,
 )
+from vllm_omni.diffusion.models.hunyuan_image3.request_layout import HunyuanPreparedLayout
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.diffusion.worker.utils import StepRequestState
@@ -76,6 +80,72 @@ def _sampling_params(**extra_args):
     )
 
 
+def _prepared_layout() -> HunyuanPreparedLayout:
+    return HunyuanPreparedLayout(
+        tokenizer_output=TokenizerEncodeOutput(
+            tokens=torch.arange(21).reshape(1, 21),
+            gen_image_mask=torch.ones(1, 21, dtype=torch.bool),
+            gen_timestep_scatter_index=torch.tensor([[4]]),
+            all_image_slices=[[slice(5, 21)]],
+            joint_image_slices=[[]],
+            gen_image_slices=[[slice(5, 21)]],
+            real_pos=torch.tensor([[21]]),
+        ),
+        rope_image_info=[[(slice(5, 21), (4, 4))]],
+        generated_image_info=ImageInfo(
+            image_type="gen_image",
+            image_width=512,
+            image_height=512,
+            token_width=4,
+            token_height=4,
+            image_token_length=16,
+        ),
+    )
+
+
+def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
+    pipeline = _pipeline()
+    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
+    rope_kwargs = {}
+
+    def fake_build_batch_2d_rope(**kwargs):
+        rope_kwargs.update(kwargs)
+        return torch.zeros(1), torch.zeros(1)
+
+    monkeypatch.setattr(hy3_module, "build_batch_2d_rope", fake_build_batch_2d_rope)
+
+    def fail_apply_chat_template(**_kwargs):
+        pytest.fail("prepared Hunyuan execution must not apply the chat template again")
+
+    pipeline._tkwrapper = SimpleNamespace(
+        apply_chat_template=fail_apply_chat_template,
+        eos_token_id=2,
+        boi_token_id=3,
+        end_recaption_token_id=4,
+        end_answer_token_id=5,
+    )
+    pipeline.config = SimpleNamespace(
+        attention_head_dim=2,
+        rope_theta=10000.0,
+    )
+    prepared = _prepared_layout()
+
+    model_inputs = pipeline.prepare_model_inputs(
+        prompt="prompt",
+        mode="gen_image",
+        guidance_scale=1.0,
+        image_size=(512, 512),
+        device=torch.device("cpu"),
+        generator=[torch.Generator().manual_seed(0)],
+        prepared_layout=prepared,
+    )
+
+    assert model_inputs["tokenizer_output"] is prepared.tokenizer_output
+    assert model_inputs["batch_gen_image_info"] == [prepared.generated_image_info]
+    assert rope_kwargs["image_infos"] is prepared.rope_image_info
+    torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
+
+
 def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
     pipeline = _pipeline()
     states = [_state("req-0", 1), _state("req-1", 3)]
@@ -131,12 +201,13 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
         captured.update(kwargs)
         raise RuntimeError("stop after prepare_model_inputs")
 
-    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
+    monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
     pipeline.prepare_model_inputs = fake_prepare_model_inputs
     state = StepRequestState(
         request_id="req-bot-task",
         sampling=sampling,
         prompt=prompt_item,
+        prepared_layout=_prepared_layout(),
     )
 
     with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
@@ -144,6 +215,7 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics(
 
     assert captured["bot_task"] == expected_model_bot_task
     assert captured["system_prompt_bot_task"] == expected_system_bot_task
+    assert captured["prepared_layout"] is state.prepared_layout
 
 
 def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
@@ -159,23 +231,22 @@ def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch):
         captured.update(kwargs)
         raise RuntimeError("stop after prepare_model_inputs")
 
-    monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt)
+    monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt)
     pipeline.prepare_model_inputs = fake_prepare_model_inputs
-    req = DiffusionRequestBatch(
-        requests=[
-            SimpleNamespace(
-                request_id="req-forward-bot-task",
-                sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
-                prompt={"prompt": "prompt", "bot_task": "vanilla"},
-            )
-        ]
+    request = SimpleNamespace(
+        request_id="req-forward-bot-task",
+        sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"),
+        prompt={"prompt": "prompt", "bot_task": "vanilla"},
+        prepared_layout=_prepared_layout(),
     )
+    req = DiffusionRequestBatch(requests=[request])
 
     with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"):
         pipeline.forward(req)
 
     assert captured["bot_task"] == "think"
     assert captured["system_prompt_bot_task"] == "think"
+    assert captured["prepared_layout"] is request.prepared_layout
 
 
 def test_grouped_denoise_rejects_non_sdpa_attention_backend():

tests/diffusion/test_diffusion_config_propagation.py

@@ -15,6 +15,7 @@ from vllm_omni.diffusion.data import (
     DiffusionParallelConfig,
     OmniDiffusionConfig,
 )
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.model_metadata import (
     HUNYUAN_IMAGE3_MAX_INPUT_IMAGES,
     QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES,
@@ -124,6 +125,11 @@ class TestCreateDefaultDiffusion:
         assert ea["enforce_eager"] is True
         assert ea["lora_path"] == "/tmp/lora"
 
+    def test_diffusion_kv_mode_roundtrip(self):
+        od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
+
+        assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
 
 def test_qwen_image_edit_plus_sets_generic_multimodal_limit():
     od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline")

tests/diffusion/test_diffusion_engine.py

@@ -108,6 +108,17 @@ class MockScheduler:
         pass
 
 
+def _make_admission_engine(pre_process_func) -> DiffusionEngine:
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = MockScheduler()
+    engine.pre_process_func = pre_process_func
+    engine._out_streams = {}
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._closed = False
+    return engine
+
+
 class _BatchCapablePipeline:
     supports_request_batch = True
 
@@ -683,6 +694,34 @@ async def _consume_final_output(generator):
     return final_output
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"])
+def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None:
+    raw_request = OmniDiffusionRequest(
+        prompt="raw",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="admission",
+    )
+    prepared_request = OmniDiffusionRequest(
+        prompt="prepared",
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        request_id="admission",
+        prepared_layout=object(),
+    )
+    preprocess_calls = []
+
+    def preprocess(request):
+        preprocess_calls.append(request)
+        return prepared_request
+
+    engine = _make_admission_engine(preprocess)
+
+    getattr(engine, entrypoint)(raw_request)
+
+    assert preprocess_calls == [raw_request]
+    assert engine.scheduler._waiting_queue == [prepared_request]
+
+
 @pytest.mark.cpu
 @pytest.mark.asyncio
 async def test_async_add_req_and_stream_response():

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -127,6 +127,7 @@ def _make_request(tag: str):
         request_id=tag,
         prompt=f"prompt_{tag}",
         sampling_params=SimpleNamespace(**sampling_params),
+        diffusion_kv_requests=None,
     )
 
 

tests/diffusion/test_diffusion_model_runner.py

@@ -218,6 +218,25 @@ def _make_compile_runner(
     return runner
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_update_states_carries_prepared_layout() -> None:
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    request = _make_request()
+    prepared_layout = object()
+    request.prepared_layout = prepared_layout
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request)],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+
+    states, new_request_ids = DiffusionModelRunner._update_states(runner, scheduler_output)
+
+    assert new_request_ids == [request.request_id]
+    assert states[0].prepared_layout is prepared_layout
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 @pytest.mark.parametrize("use_hsdp", [False, True])

tests/diffusion/test_diffusion_scheduler.py

@@ -12,6 +12,8 @@ from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -259,6 +261,90 @@ class TestRequestScheduler:
         assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
         assert self.scheduler.has_requests() is False
 
+    def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None:
+        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        request = _make_request("diffusion-kv")
+        prepared_layout = object()
+        request.prepared_layout = prepared_layout
+        text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8)
+        kv_request = DiffusionKVRequest(
+            "diffusion-kv/diffusion-kv/0",
+            sequence_id=0,
+            prefix_len=4,
+            target_len=8,
+            seq_len=16,
+            kv_contexts=(text_context,),
+        )
+        request.diffusion_kv_requests = (kv_request,)
+
+        self.scheduler.add_request(request)
+        state = self.scheduler.get_request_state(request.request_id)
+        scheduler_output = self.scheduler.schedule()
+
+        assert state is not None
+        assert state.diffusion_kv_requests == (kv_request,)
+        assert state.diffusion_kv_requests[0].kv_contexts == (text_context,)
+        assert request.diffusion_kv_requests is None
+        assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout
+
+    def test_paged_scheduler_rejects_missing_kv_request(self) -> None:
+        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        request = _make_request("diffusion-kv")
+
+        with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"):
+            self.scheduler.add_request(request)
+
+    @pytest.mark.parametrize(
+        ("owner_name", "field_name"),
+        [
+            ("request", "past_key_values"),
+            ("sampling_params", "past_key_values"),
+            ("sampling_params", "cfg_text_past_key_values"),
+            ("sampling_params", "cfg_img_past_key_values"),
+            ("sampling_params", "cfg_branch_past_key_values"),
+        ],
+    )
+    def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None:
+        self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER))
+        request = _make_request("diffusion-kv")
+        owner = request if owner_name == "request" else request.sampling_params
+        setattr(owner, field_name, object())
+        request.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "diffusion-kv/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=4,
+                target_len=8,
+                seq_len=16,
+            ),
+        )
+
+        with pytest.raises(ValueError, match=rf"{owner_name}\.{field_name}"):
+            self.scheduler.add_request(request)
+
+    def test_dense_scheduler_rejects_scheduler_only_kv_request(self) -> None:
+        request = _make_request("dense-kv-request")
+        request.diffusion_kv_requests = (
+            DiffusionKVRequest(
+                "dense-kv-request/diffusion-kv/0",
+                sequence_id=0,
+                prefix_len=4,
+                target_len=8,
+                seq_len=16,
+            ),
+        )
+
+        with pytest.raises(ValueError, match="dense_legacy request unexpectedly contains"):
+            self.scheduler.add_request(request)
+
+    def test_dense_scheduler_accepts_legacy_dense_kv(self) -> None:
+        request = _make_request("dense-kv")
+        request.sampling_params.past_key_values = object()
+
+        request_id = self.scheduler.add_request(request)
+
+        assert request_id == request.request_id
+
     def test_error_output_marks_finished_error(self) -> None:
         req_id = self.scheduler.add_request(_make_request("err"))
 
@@ -562,12 +648,20 @@ class TestDiffusionEngine:
         engine.abort_queue = queue.Queue()
 
         request = _make_request("engine")
+        prepared_layout = object()
+
+        def preprocess(req):
+            req.prepared_layout = prepared_layout
+            return req
+
+        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
         runner_output = _make_request_output("engine")
         engine.execute_fn = mocker.Mock(return_value=runner_output)
 
         output = engine.add_req_and_wait_for_response(request)
 
         assert output is runner_output.result
+        engine.pre_process_func.assert_called_once_with(request)
         engine.execute_fn.assert_called_once()
 
     def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None:
@@ -631,17 +725,29 @@ class TestDiffusionEngine:
     async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._check_and_start_background_loop = mocker.AsyncMock()
-        engine.pre_process_func = None
+        request = _make_request("req-abort")
+        prepared_layout = object()
+
+        def preprocess(req):
+            req.prepared_layout = prepared_layout
+            return req
+
+        engine.pre_process_func = mocker.Mock(side_effect=preprocess)
 
         async def _stream(_request):
             yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
 
-        engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None))
+        engine._add_prepared_request = mocker.Mock(return_value=request.request_id)
+        engine.get_output_stream = mocker.Mock(return_value=_stream(None))
 
         with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            async for _ in engine.step_streaming(_make_request("req-abort")):
+            async for _ in engine.step_streaming(request):
                 pass
 
+        engine.pre_process_func.assert_called_once_with(request)
+        engine._add_prepared_request.assert_called_once_with(request)
+        engine.get_output_stream.assert_called_once_with(request.request_id)
+
     def test_abort_queue_marks_request_finished_aborted(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._rpc_lock = threading.RLock()
@@ -774,6 +880,18 @@ class TestDiffusionEngine:
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
 
+    def test_dummy_run_delegates_preprocessing_to_sync_admission(self, mocker: MockerFixture) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.od_config = SimpleNamespace(model_class_name="mock_model", diffusion_load_format="default")
+        engine.pre_process_func = mocker.Mock()
+        engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(output=None))
+
+        engine._dummy_run()
+
+        engine.pre_process_func.assert_not_called()
+        admitted_request = engine.add_req_and_wait_for_response.call_args.args[0]
+        assert admitted_request.request_id == "dummy_req_id"
+
 
 class TestStepScheduler:
     def setup_method(self) -> None:

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -48,6 +48,7 @@ def _mock_request(tag: str):
         request_id=tag,
         prompt=f"prompt_{tag}",
         sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+        diffusion_kv_requests=None,
     )