[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/config/test_omni_config.py

@@ -39,6 +39,7 @@ from vllm_omni.config.stage_config import (
     load_deploy_config,
     merge_pipeline_deploy,
 )
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -966,20 +967,24 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
 
     cfg = omni_config_module._DiffusionConfigProjection.from_kwargs(
         diffusion_attention_backend="flash_attn",
+        fa_deterministic=True,
         kv_cache_dtype="fp8",
         kv_cache_skip_steps="0-1",
         kv_cache_skip_layers=[2],
         static_lora_scale=0.25,
+        diffusion_kv_mode="paged_scheduler",
         diffusers_load_kwargs=None,
         diffusers_call_kwargs=None,
     )
 
     assert cfg.diffusion_attention_config.default.backend == "flash_attn"
+    assert cfg.fa_deterministic is True
     assert cfg.diffusion_kv_cache_dtype == "fp8"
     assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1}
     assert cfg.diffusion_kv_cache_skip_layer_indices == {2}
     assert cfg.lora_scale == 0.25
     assert cfg.cache_backend == "tea_cache"
+    assert cfg.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
     assert cfg.diffusers_load_kwargs == {}
     assert cfg.diffusers_call_kwargs == {}
 
@@ -994,6 +999,8 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
                 "stages:",
                 "  - stage_id: 0",
                 "    diffusion_attention_backend: flash_attn",
+                "    fa_deterministic: true",
+                "    diffusion_kv_mode: paged_scheduler",
             ]
         )
     )
@@ -1005,6 +1012,29 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar
 
     assert isinstance(stage, VllmOmniDiffusionStageConfig)
     assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn"
+    assert stage.diffusion_config.fa_deterministic is True
+    assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path):
+    deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml"
+    deploy_path.write_text(
+        "\n".join(
+            [
+                "pipeline: dreamzero",
+                "async_chunk: false",
+                "stages:",
+                "  - stage_id: 0",
+                "    diffusion_kv_mode: paged_worker_local",
+            ]
+        )
+    )
+
+    with pytest.raises(ValueError, match="reserved but not implemented"):
+        _from_pipeline_key(
+            "dreamzero",
+            deploy_config_path=str(deploy_path),
+        )
 
 
 def test_diffusion_config_field_classification_covers_current_fields():
@@ -1025,6 +1055,7 @@ def test_diffusion_config_field_classification_covers_current_fields():
         "enable_session_state_manager",
         "prompt_embed_cache_size",
         "diffusion_kv_cache_dtype",
+        "diffusion_kv_mode",
     } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS
     assert {
         "revision",

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,
     )
 
 

vllm_omni/config/omni_config.py

@@ -37,6 +37,7 @@ from vllm_omni.config.stage_config import (
     build_stage_runtime_overrides,
     load_deploy_config,
 )
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 
 _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = {
     StageExecutionType.LLM_AR: (StageType.LLM, "ar"),
@@ -549,6 +550,7 @@ class _DiffusionConfigProjection:
     cache_backend: str = "none"
     cache_config: Any = field(default_factory=dict)
     enable_cache_dit_summary: bool = False
+    diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
     enable_prompt_embed_cache: bool = False
     prompt_embed_cache_size: int = Field(default=32, ge=1)
     enable_session_state_manager: bool = False
@@ -623,6 +625,7 @@ class _DiffusionConfigProjection:
             build_attention_config,
             parse_kv_cache_skip_selector,
         )
+        from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
 
         if self.tf_model_config is None:
@@ -684,6 +687,7 @@ class _DiffusionConfigProjection:
                 f"got {type(self.diffusion_attention_config)!r}"
             )
 
+        self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
         self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps)
         self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers)
 

vllm_omni/diffusion/data.py

@@ -23,6 +23,10 @@ from vllm.model_executor.layers.quantization.base_config import (
 )
 from vllm.transformers_utils.repo_utils import get_model_path
 
+from vllm_omni.diffusion.diffusion_kv.config import (
+    DiffusionKVCacheMode,
+    parse_diffusion_kv_cache_mode,
+)
 from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 from vllm_omni.diffusion.utils.network_utils import is_port_available
 from vllm_omni.errors import client_error_metadata
@@ -658,6 +662,9 @@ class OmniDiffusionConfig:
     # the string form only).
     engine_backend: str | type = "default"
 
+    # Local Diffusion KV ownership and cache-layout mode.
+    diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+
     # Optional override for the diffusion model runner class (import path).
     # Precedence in the worker: this override > the runner declared by the
     # selected engine class (``default_diffusion_model_runner_cls``) > the
@@ -923,6 +930,22 @@ class OmniDiffusionConfig:
             )
         if not isinstance(self.diffusion_compile_dynamic, bool):
             raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}")
+        self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+
+        if self.omni_kv_config is None:
+            self.omni_kv_config = {}
+        elif isinstance(self.omni_kv_config, Mapping):
+            self.omni_kv_config = dict(self.omni_kv_config)
+        else:
+            raise TypeError("omni_kv_config must be a mapping")
+        if self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and self.omni_kv_config.get(
+            "need_recv_cache", False
+        ):
+            raise ValueError(
+                "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; "
+                "disable need_recv_cache until connector-aware admission is implemented"
+            )
+
         self.master_port = self._resolve_master_port()
         self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0)
         if self.request_batch_max_wait_ms < 0:

vllm_omni/diffusion/diffusion_engine.py

@@ -310,14 +310,16 @@ class DiffusionEngine:
         diffusion_engine_start_time = time.perf_counter()
 
         preprocess_time = 0.0
-        if self.pre_process_func is not None:
-            preprocess_start_time = time.perf_counter()
-            request = self.pre_process_func(request)
+        has_preprocessor = getattr(self, "pre_process_func", None) is not None
+        preprocess_start_time = time.perf_counter() if has_preprocessor else None
+        request = self._prepare_request_for_admission(request)
+        if preprocess_start_time is not None:
             preprocess_time = time.perf_counter() - preprocess_start_time
             logger.debug("Pre-processing completed in %.4f seconds", preprocess_time)
 
         exec_start_time = time.perf_counter()
-        generator = self.async_add_req_and_stream_response(request)
+        request_id = self._add_prepared_request(request)
+        generator = self.get_output_stream(request_id)
         async for output in generator:
             exec_total_time = time.perf_counter() - exec_start_time
             # Async mode: wait for background D2H/SHM to complete.
@@ -689,7 +691,17 @@ class DiffusionEngine:
         engine.run_startup_warmup()
         return engine
 
-    def add_request(self, request: OmniDiffusionRequest) -> str:
+    def _prepare_request_for_admission(self, request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        """Run model-owned preprocessing once, before entering Engine locks."""
+
+        pre_process_func = getattr(self, "pre_process_func", None)
+        if pre_process_func is None:
+            return request
+        return pre_process_func(request)
+
+    def _add_prepared_request(self, request: OmniDiffusionRequest) -> str:
+        """Admit a request whose model-owned preprocessing is complete."""
+
         with self._cv:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
@@ -700,6 +712,10 @@ class DiffusionEngine:
 
         return request_id
 
+    def add_request(self, request: OmniDiffusionRequest) -> str:
+        request = self._prepare_request_for_admission(request)
+        return self._add_prepared_request(request)
+
     async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
         with self._cv:
             queue = self._out_streams.get(request_id)
@@ -742,6 +758,7 @@ class DiffusionEngine:
         return final_output
 
     def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
+        request = self._prepare_request_for_admission(request)
         with self._rpc_lock:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
@@ -886,8 +903,7 @@ class DiffusionEngine:
             ),
         )
         logger.info("dummy run to warm up the model")
-        request = self.pre_process_func(req) if self.pre_process_func is not None else req
-        output = self.add_req_and_wait_for_response(request)
+        output = self.add_req_and_wait_for_response(req)
         if output.error:
             raise RuntimeError(f"Dummy run failed: {output.error}")
 

vllm_omni/diffusion/diffusion_kv/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/diffusion/diffusion_kv/config.py

@@ -0,0 +1,39 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from enum import Enum
+
+
+class DiffusionKVCacheMode(str, Enum):
+    """Migration mode for diffusion KV ownership."""
+
+    DENSE_LEGACY = "dense_legacy"
+    # Reserved for the experimental Worker-owned paging design. Production
+    # configuration rejects it until that engine is migrated to the common
+    # Scheduler-owned path.
+    PAGED_WORKER_LOCAL = "paged_worker_local"
+    PAGED_SCHEDULER = "paged_scheduler"
+
+
+def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode:
+    """Parse a selectable Diffusion KV ownership mode.
+
+    ``PAGED_WORKER_LOCAL`` remains in the enum to document the migration
+    topology, but it is not a valid production selection until that runtime is
+    wired. Keep this validation shared by structured and runtime configs so an
+    unsupported mode fails before managed-process startup.
+    """
+
+    try:
+        mode = DiffusionKVCacheMode(value)
+    except (TypeError, ValueError) as exc:
+        supported = ", ".join(mode.value for mode in DiffusionKVCacheMode)
+        raise ValueError(f"Unsupported Diffusion KV diffusion_kv_mode {value!r}; expected one of: {supported}") from exc
+    if mode is DiffusionKVCacheMode.PAGED_WORKER_LOCAL:
+        raise ValueError(
+            "Diffusion KV diffusion_kv_mode 'paged_worker_local' is reserved but not implemented; "
+            "use 'dense_legacy' or 'paged_scheduler'"
+        )
+    return mode

vllm_omni/diffusion/diffusion_kv/request.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.request import RequestStatus
+
+
+@dataclass(frozen=True)
+class DiffusionKVContext:
+    """An independently managed K/V context outside the primary sequence.
+
+    ``context_id`` identifies the logical context within one execution
+    sequence. ``cache_role`` binds it to a logical attention cache role exposed
+    by the Worker. Physical cache geometry remains native ``KVCacheSpec`` /
+    ``KVCacheConfig`` state and is deliberately absent here.
+    """
+
+    context_id: str
+    cache_role: str
+    num_tokens: int
+    block_hashes: tuple[BlockHash, ...] = ()
+
+    def __post_init__(self) -> None:
+        if not self.context_id:
+            raise ValueError("context_id must be non-empty")
+        if not self.cache_role:
+            raise ValueError("cache_role must be non-empty")
+        if self.num_tokens <= 0:
+            raise ValueError(f"num_tokens must be positive, got {self.num_tokens}")
+        object.__setattr__(self, "block_hashes", tuple(self.block_hashes))
+
+
+class DiffusionKVRequest:
+    """Scheduler-owned KV state for one diffusion execution sequence.
+
+    The primary sequence follows an ordered ``[prefix | target]`` policy, such
+    as one Hunyuan CFG row. ``prefix_len`` is the contiguous reusable prefix;
+    ``target_len`` is overwritten by every denoise step; ``num_tokens`` is the
+    complete first-step allocation boundary.
+
+    Independent cross/joint-attention K/V does not belong to that token axis
+    and is described by ``kv_contexts``. This object also exposes the minimal
+    mutable Request surface consumed by native ``KVCacheManager``.
+
+    An empty ``block_hashes`` sequence means the prefix has no canonical cache
+    identity yet. Such a request may use native request-local page allocation,
+    but consumers must not publish it through ``KVCacheManager.cache_blocks``.
+    Prefix publication becomes valid only after model preprocessing supplies
+    one canonical hash for every cacheable full block.
+    """
+
+    def __init__(
+        self,
+        request_id: str,
+        *,
+        sequence_id: int,
+        prefix_len: int,
+        target_len: int,
+        seq_len: int,
+        block_hashes: Sequence[BlockHash] = (),
+        kv_contexts: Sequence[DiffusionKVContext] = (),
+    ) -> None:
+        if not request_id:
+            raise ValueError("request_id must be non-empty")
+        if sequence_id < 0:
+            raise ValueError(f"sequence_id must be non-negative, got {sequence_id}")
+        if prefix_len < 0:
+            raise ValueError(f"prefix_len must be non-negative, got {prefix_len}")
+        if target_len <= 0:
+            raise ValueError(f"target_len must be positive, got {target_len}")
+        if seq_len <= 0:
+            raise ValueError(f"seq_len must be positive, got {seq_len}")
+        if prefix_len + target_len > seq_len:
+            raise ValueError(
+                "prefix_len + target_len must not exceed seq_len: "
+                f"prefix_len={prefix_len}, target_len={target_len}, seq_len={seq_len}"
+            )
+
+        contexts = tuple(kv_contexts)
+        if any(not isinstance(context, DiffusionKVContext) for context in contexts):
+            raise TypeError("kv_contexts must contain only DiffusionKVContext values")
+        context_ids = [context.context_id for context in contexts]
+        if len(context_ids) != len(set(context_ids)):
+            raise ValueError(f"kv_contexts must have unique context_id values, got {context_ids}")
+
+        # Diffusion semantics consumed by the Scheduler-side facade.
+        self.sequence_id = sequence_id
+        self.prefix_len = prefix_len
+        self.target_len = target_len
+        self.kv_contexts = contexts
+
+        # Native vLLM Request surface. Keep this list intentionally small and
+        # cover it with real-KVCacheManager conformance tests.
+        self.request_id = request_id
+        self.num_tokens = seq_len
+        self.num_prompt_tokens = prefix_len
+        self.num_computed_tokens = 0
+        self.block_hashes = list(block_hashes)
+        self.skip_reading_prefix_cache = not self.block_hashes
+        self.status = RequestStatus.WAITING
+        self.num_preemptions = 0
+        self.num_in_flight_tokens = 0
+        # vLLM 0.26 uses this optional boundary when publishing cached blocks;
+        # zero means that no sparse-retention boundary is active.
+        self.shared_prefix_boundary = 0
+
+    @property
+    def seq_len(self) -> int:
+        """Complete first-step sequence length and allocation boundary."""
+
+        return self.num_tokens

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -21,6 +21,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 from vllm.transformers_utils.config import get_config
 
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -33,6 +34,7 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 from vllm_omni.inputs.data import OmniTextPrompt
 from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
 
+from . import request_layout as request_layout_utils
 from .hunyuan_image3_tokenizer import TokenizerEncodeOutput, TokenizerWrapper
 from .hunyuan_image3_transformer import (
     CausalMMOutputWithPast,
@@ -50,7 +52,7 @@ from .hunyuan_image3_transformer import (
     real_batched_index_select,
     retrieve_timesteps,
 )
-from .system_prompt import get_system_prompt
+from .request_layout import HunyuanPreparedLayout
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.worker.input_batch import InputBatch
@@ -146,12 +148,6 @@ def _resize_and_crop_center(image: PILImage.Image, target_width: int, target_hei
     return resized.crop((crop_left, crop_top, crop_left + tw, crop_top + th))
 
 
-def _to_python_scalar(value: Any) -> Any:
-    if isinstance(value, np.generic):
-        return value.item()
-    return value
-
-
 def _shift_full_attn_spans(
     spans: list[tuple[int, int]],
     prefix_len: int | None,
@@ -165,80 +161,12 @@ def _shift_full_attn_spans(
     return [(start + shift, end + shift) if start >= prefix_len else (start, end) for start, end in spans]
 
 
-def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]:
-    return {
-        "image_type": image_info.image_type,
-        "image_tensor": image_info.image_tensor,
-        "image_width": _to_python_scalar(image_info.image_width),
-        "image_height": _to_python_scalar(image_info.image_height),
-        "token_width": _to_python_scalar(image_info.token_width),
-        "token_height": _to_python_scalar(image_info.token_height),
-        "image_token_length": _to_python_scalar(image_info.image_token_length),
-        "base_size": _to_python_scalar(image_info.base_size),
-        "ratio_index": _to_python_scalar(image_info.ratio_index),
-        "add_timestep_token": image_info.add_timestep_token,
-        "add_guidance_token": image_info.add_guidance_token,
-        "use_front_boi_token": image_info.use_front_boi_token,
-        "add_image_shape_token": image_info.add_image_shape_token,
-    }
-
-
-def _to_tensor_if_needed(value: Any) -> Any:
-    if isinstance(value, np.generic):
-        return value.item()
-    if isinstance(value, list):
-        return torch.tensor(value)
-    return value
-
-
-def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo:
-    return ImageInfo(
-        image_type=payload.get("image_type"),
-        image_tensor=_to_tensor_if_needed(payload.get("image_tensor")),
-        image_width=payload.get("image_width"),
-        image_height=payload.get("image_height"),
-        token_width=payload.get("token_width"),
-        token_height=payload.get("token_height"),
-        image_token_length=payload.get("image_token_length"),
-        base_size=payload.get("base_size"),
-        ratio_index=payload.get("ratio_index"),
-        add_timestep_token=payload.get("add_timestep_token", True),
-        add_guidance_token=payload.get("add_guidance_token", False),
-        use_front_boi_token=payload.get("use_front_boi_token", True),
-        add_image_shape_token=payload.get("add_image_shape_token", True),
-    )
-
-
-def _joint_image_info_to_payload(joint_image_info: JointImageInfo) -> dict[str, Any]:
-    return {
-        "type": "joint_image_info",
-        "vae_image_info": _image_info_to_payload(joint_image_info.vae_image_info),
-        "vision_image_info": _image_info_to_payload(joint_image_info.vision_image_info),
-        "vision_encoder_kwargs": joint_image_info.vision_encoder_kwargs,
-    }
-
-
-def _joint_image_info_from_payload(payload: Any) -> JointImageInfo:
-    if isinstance(payload, JointImageInfo):
-        return payload
-    if not isinstance(payload, dict):
-        raise TypeError(f"Expected dict or JointImageInfo for conditional image payload, got {type(payload)}.")
-
-    vae_image_info = _image_info_from_payload(payload["vae_image_info"])
-    vision_image_info = _image_info_from_payload(payload["vision_image_info"])
-    vision_encoder_kwargs = payload.get("vision_encoder_kwargs") or {}
-    if isinstance(vision_encoder_kwargs, dict):
-        vision_encoder_kwargs = {key: _to_tensor_if_needed(value) for key, value in vision_encoder_kwargs.items()}
-    return JointImageInfo(
-        vae_image_info=vae_image_info,
-        vision_image_info=vision_image_info,
-        vision_encoder_kwargs=vision_encoder_kwargs,
-    )
-
-
 def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
     hf_config = get_config(od_config.model, trust_remote_code=True)
     image_processor = HunyuanImage3ImageProcessor(hf_config)
+    prepare_diffusion_kv_layout = od_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+    tokenizer_wrapper = TokenizerWrapper(od_config.model) if prepare_diffusion_kv_layout else None
+    generation_config = GenerationConfig.from_pretrained(od_config.model) if prepare_diffusion_kv_layout else None
     vae_h_factor = hf_config.vae_downsample_factor[0] * hf_config.patch_size
     vae_w_factor = hf_config.vae_downsample_factor[1] * hf_config.patch_size
     vit_patch_size = getattr(image_processor.vision_encoder_processor, "patch_size", 1)
@@ -291,7 +219,7 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
             image_token_length=int(vit_tensor.shape[1]),
         )
 
-        return _joint_image_info_to_payload(
+        return request_layout_utils.joint_image_info_to_payload(
             JointImageInfo(
                 vae_image_info=vae_info,
                 vision_image_info=vit_info,
@@ -329,6 +257,20 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig):
                 request.sampling_params.height = int(bridge_h or first_image_h)
 
         request.prompt = prompt
+        if prepare_diffusion_kv_layout:
+            assert tokenizer_wrapper is not None and generation_config is not None
+            prepared_layout = request_layout_utils.prepare_hunyuan_layout(
+                request,
+                tokenizer_wrapper=tokenizer_wrapper,
+                image_processor=image_processor,
+                generation_config=generation_config,
+                image_base_size=hf_config.image_base_size,
+            )
+            request.prepared_layout = prepared_layout
+            request.diffusion_kv_requests = request_layout_utils.build_hunyuan_diffusion_kv_requests(
+                request,
+                prepared_layout,
+            )
 
         return request
 
@@ -543,94 +485,6 @@ class HunyuanImage3Pipeline(
         if getattr(self.od_config, "diffusion_kv_cache_skip_step_indices", None):
             raise ValueError("HunyuanImage3 step execution does not support diffusion KV cache step skips yet.")
 
-    @staticmethod
-    def _normalize_single_stage_bot_task(bot_task: Any) -> str:
-        if isinstance(bot_task, str) and bot_task.lower() == "none":
-            bot_task = None
-        tokenizer_bot_task = bot_task
-        if tokenizer_bot_task == "think_recaption":
-            tokenizer_bot_task = "think"
-        elif tokenizer_bot_task == "vanilla":
-            tokenizer_bot_task = "image"
-        supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"}
-        if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks:
-            raise ValueError(
-                f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. "
-                f"Supported values are: {sorted(supported_bot_tasks)}."
-            )
-        return tokenizer_bot_task or "auto"
-
-    def _extract_prompt_inputs(
-        self,
-        prompts: list[Any],
-        extra_args: dict[str, Any],
-        *,
-        request_id: str,
-        allow_cond_image: bool,
-    ) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]:
-        is_dummy_warmup = OmniDiffusionRequest.is_dummy_run_request_id(request_id)
-        bot_task = extra_args.get("bot_task")
-        use_system_prompt = extra_args.get("use_system_prompt")
-        system_prompt = extra_args.get("system_prompt")
-
-        first_prompt = prompts[0] if prompts else None
-        if isinstance(first_prompt, dict):
-            if bot_task is None:
-                bot_task = first_prompt.get("bot_task")
-            if use_system_prompt is None:
-                use_system_prompt = first_prompt.get("use_system_prompt")
-            if system_prompt is None:
-                system_prompt = first_prompt.get("system_prompt")
-        tokenizer_bot_task = self._normalize_single_stage_bot_task(bot_task)
-        if use_system_prompt is not None:
-            system_prompt_bot_task = "image" if tokenizer_bot_task == "auto" else tokenizer_bot_task
-            system_prompt = get_system_prompt(use_system_prompt, system_prompt_bot_task, system_prompt)
-            system_prompt = system_prompt.strip() if system_prompt is not None else ""
-
-        prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts]
-        cot_text_list = [
-            (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in prompts
-        ]
-
-        batch_cond_image_info: list[list[JointImageInfo]] | None = None
-        if any(not isinstance(p, str) for p in prompts):
-            batch_cond_image_info = []
-            for prompt_item in prompts:
-                if isinstance(prompt_item, str):
-                    batch_cond_image_info.append([])
-                    continue
-                additional_info = prompt_item.get("additional_information") or {}
-                cond_infos = additional_info.get("batch_cond_image_info", [])
-                if isinstance(cond_infos, JointImageInfo | dict):
-                    cond_infos = [cond_infos]
-                if cond_infos is None:
-                    cond_infos = []
-                batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in cond_infos])
-
-            has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info]
-            if any(has_cond_image) and (not allow_cond_image) and not is_dummy_warmup:
-                raise ValueError("HunyuanImage3 step execution does not support image editing requests yet.")
-            if allow_cond_image and any(has_cond_image) and not all(has_cond_image):
-                raise ValueError(
-                    "When batching Hunyuan image editing requests, every prompt must include input image(s)."
-                )
-            if not allow_cond_image or not any(has_cond_image):
-                batch_cond_image_info = None
-
-        return prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task
-
-    def _extract_step_prompt_inputs(
-        self,
-        state: "StepRequestState",
-    ) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]:
-        sampling = state.sampling
-        return self._extract_prompt_inputs(
-            [state.prompt] if state.prompt is not None else [],
-            getattr(sampling, "extra_args", {}) or {},
-            request_id=state.request_id,
-            allow_cond_image=False,
-        )
-
     def _snapshot_injected_ar_kv(self) -> list[list[tuple[torch.Tensor, torch.Tensor]] | None] | None:
         snapshot: list[list[tuple[torch.Tensor, torch.Tensor]] | None] = []
         found = False
@@ -1081,27 +935,6 @@ class HunyuanImage3Pipeline(
         pred = self.final_layer(image_output, timestep_emb, token_h, token_w)
         return pred
 
-    @staticmethod
-    def build_batch_rope_image_info(output, sections):
-        rope_image_info = []
-        for image_slices, sections_i in zip(output.all_image_slices, sections):
-            image_shapes = []
-            for section in sections_i:
-                if "image" in section["type"]:
-                    if isinstance(section["token_height"], list):
-                        assert len(section["token_height"]) == len(section["token_width"]), (
-                            f"token_height and token_width should have the same length, "
-                            f"but got {len(section['token_height'])} and {len(section['token_width'])}"
-                        )
-                        image_shapes.extend(list(zip(section["token_height"], section["token_width"])))
-                    else:
-                        image_shapes.append((section["token_height"], section["token_width"]))
-            assert len(image_slices) == len(image_shapes), (
-                f"Size miss matching: Image slices({len(image_slices)}) != image shapes({len(image_shapes)})"
-            )
-            rope_image_info.append(list(zip(image_slices, image_shapes)))
-        return rope_image_info
-
     def vae_encode(self, image, cfg_factor=1, generator=None):
         config = self.vae.config
 
@@ -1210,26 +1043,6 @@ class HunyuanImage3Pipeline(
                     f"Each message should be a list of dicts or a dict, but got {type(message)}."
                 )
 
-    @staticmethod
-    def _normalize_cot_text(cot: str | None) -> str | None:
-        """
-        Ensure cot_text starts with the appropriate opening tag.
-
-        AR-generated text may omit the leading <think> or <recaption> tag
-        (since it was used as a generation trigger). This normalizes the text
-        so downstream parsing in get_cot_sections works correctly.
-        """
-        if not cot:
-            return cot
-
-        if "</think>" in cot and not cot.startswith("<think>"):
-            cot = "<think>" + cot
-            return cot
-        if "</recaption>" in cot and not cot.startswith("<recaption>"):
-            cot = "<recaption>" + cot
-            return cot
-        return cot
-
     def prepare_model_inputs(
         self,
         prompt=None,
@@ -1242,6 +1055,7 @@ class HunyuanImage3Pipeline(
         message_list=None,
         device=None,
         max_new_tokens=None,
+        prepared_layout: HunyuanPreparedLayout | None = None,
         **kwargs,
     ):
         # 1. Sanity check
@@ -1297,7 +1111,16 @@ class HunyuanImage3Pipeline(
                     )
 
             if mode == "gen_image":
-                batch_gen_image_info = [self.image_processor.build_image_info(image_size) for _ in range(batch_size)]
+                if prepared_layout is not None:
+                    # TODO: Support batch_size > 1 once prepared layouts can
+                    # preserve the public-request-to-CFG-branch mapping.
+                    if batch_size != 1:
+                        raise ValueError("Hunyuan prepared layout currently supports one public request at a time.")
+                    batch_gen_image_info = [prepared_layout.generated_image_info]
+                else:
+                    batch_gen_image_info = [
+                        self.image_processor.build_image_info(image_size) for _ in range(batch_size)
+                    ]
 
             if batch_cond_image_info is not None:
                 assert isinstance(batch_cond_image_info, list) and len(batch_cond_image_info) == batch_size, (
@@ -1314,30 +1137,41 @@ class HunyuanImage3Pipeline(
         # 3. apply chat template
         cfg_factor = {"gen_text": 1, "gen_image": 1 + int(guidance_scale > 1.0)}
         bot_task = kwargs.pop("bot_task", "auto")
-        # If `drop_think` enabled, always drop <think> parts in the context.
-        drop_think = kwargs.get("drop_think", self.generation_config.drop_think)
-        # Pull sequence_template from the model's generation_config so the DiT
-        # text prefix matches how the model was trained (Instruct for the
-        # HunyuanImage-3.0-Instruct checkpoint).  Falling back to "pretrain"
-        # only if the config does not specify it.
-        sequence_template = getattr(self.generation_config, "sequence_template", "pretrain")
-        # Apply batched prompt or batched message_list to build input sequence with associated info.
-        out = self._tkwrapper.apply_chat_template(
-            batch_prompt=batch_prompt,
-            batch_message_list=batch_message_list,
-            mode=mode,
-            batch_gen_image_info=batch_gen_image_info,
-            batch_cond_image_info=batch_cond_image_info,
-            batch_system_prompt=batch_system_prompt,
-            batch_cot_text=batch_cot_text,
-            max_length=kwargs.get("max_length"),
-            bot_task=bot_task,
-            image_base_size=self.config.image_base_size,
-            sequence_template=sequence_template,
-            cfg_factor=cfg_factor[mode],
-            drop_think=drop_think,
-        )
-        output, sections = out["output"], out["sections"]
+        if prepared_layout is None:
+            # Pull template options from generation_config only when the
+            # Worker must tokenize locally.
+            drop_think = kwargs.get("drop_think", self.generation_config.drop_think)
+            sequence_template = getattr(self.generation_config, "sequence_template", "pretrain")
+            # Apply batched prompt or batched message_list to build input
+            # sequence with associated info.
+            out = self._tkwrapper.apply_chat_template(
+                batch_prompt=batch_prompt,
+                batch_message_list=batch_message_list,
+                mode=mode,
+                batch_gen_image_info=batch_gen_image_info,
+                batch_cond_image_info=batch_cond_image_info,
+                batch_system_prompt=batch_system_prompt,
+                batch_cot_text=batch_cot_text,
+                max_length=kwargs.get("max_length"),
+                bot_task=bot_task,
+                image_base_size=self.config.image_base_size,
+                sequence_template=sequence_template,
+                cfg_factor=cfg_factor[mode],
+                drop_think=drop_think,
+            )
+            output = out["output"]
+            rope_image_info = request_layout_utils.build_hunyuan_batch_rope_image_info(output, out["sections"])
+        else:
+            if mode != "gen_image" or batch_message_list is not None:
+                raise ValueError("Hunyuan prepared layout only supports prompt-based gen_image execution.")
+            output = prepared_layout.tokenizer_output
+            rope_image_info = prepared_layout.rope_image_info
+            expected_rows = batch_size * cfg_factor[mode]
+            if prepared_layout.num_branches != expected_rows:
+                raise ValueError(
+                    "Hunyuan prepared branches do not match execution CFG layout: "
+                    f"expected_rows={expected_rows}, prepared_branches={prepared_layout.num_branches}"
+                )
 
         # 4. Encode conditional images
         # Skip encoding if AR KV reuse is enabled
@@ -1362,7 +1196,6 @@ class HunyuanImage3Pipeline(
             vit_kwargs = None
 
         # 5. Build position embeddings
-        rope_image_info = self.build_batch_rope_image_info(output, sections)
         if mode == "gen_text":
             seq_len = self.generation_config.max_length
         else:
@@ -1616,12 +1449,7 @@ class HunyuanImage3Pipeline(
                 raise ValueError("`batch_gen_image_info` should be provided when `mode` is `gen_image`.")
 
             image_info: ImageInfo = batch_gen_image_info[0]
-            num_image_tokens = (
-                image_info.image_token_length
-                + (1 if image_info.add_timestep_token else 0)
-                + (1 if image_info.add_guidance_token else 0)
-            )
-            kwargs["num_image_tokens"] = num_image_tokens
+            kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info)
             # 50 and 5.0 hard code
             results = self.pipeline(
                 batch_size=len(batch_gen_image_info),
@@ -1888,9 +1716,14 @@ class HunyuanImage3Pipeline(
             system_prompt,
             batch_cond_image_info,
             tokenizer_bot_task,
-        ) = self._extract_step_prompt_inputs(state)
+        ) = request_layout_utils.extract_hunyuan_prompt_inputs(
+            [state.prompt] if state.prompt is not None else [],
+            getattr(sampling, "extra_args", {}) or {},
+            request_id=state.request_id,
+            allow_cond_image=False,
+        )
         cot_text = (
-            [self._normalize_cot_text(text) for text in cot_text_list]
+            [request_layout_utils.normalize_hunyuan_cot_text(text) for text in cot_text_list]
             if any(text is not None for text in cot_text_list)
             else None
         )
@@ -1899,7 +1732,7 @@ class HunyuanImage3Pipeline(
         width = sampling.width or 1024
         image_size = (height, width)
         num_inference_steps = sampling.num_inference_steps or 50
-        guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 5.0
+        guidance_scale = request_layout_utils.resolve_hunyuan_guidance_scale(sampling)
         if guidance_scale <= 1.0:
             logger.info("HunyuanImage3.0 step execution runs without classifier-free guidance.")
         pipe._guidance_scale = guidance_scale
@@ -1916,6 +1749,7 @@ class HunyuanImage3Pipeline(
             guidance_scale=guidance_scale,
             batch_cond_image_info=batch_cond_image_info,
             bot_task=tokenizer_bot_task,
+            prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(state),
         )
         model_kwargs.update(self._extract_ar_kv_from_sampling(sampling))
         model_kwargs["use_cache"] = False
@@ -1923,14 +1757,9 @@ class HunyuanImage3Pipeline(
         input_ids = model_kwargs.pop("input_ids")
         batch_gen_image_info: list[ImageInfo] = model_kwargs["batch_gen_image_info"]
         image_info = batch_gen_image_info[0]
-        target_height = int(_to_python_scalar(image_info.image_height))
-        target_width = int(_to_python_scalar(image_info.image_width))
-        num_image_tokens = (
-            image_info.image_token_length
-            + (1 if image_info.add_timestep_token else 0)
-            + (1 if image_info.add_guidance_token else 0)
-        )
-        model_kwargs["num_image_tokens"] = num_image_tokens
+        target_height = int(image_info.image_height)
+        target_width = int(image_info.image_width)
+        model_kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info)
 
         timesteps, _ = retrieve_timesteps(
             self.scheduler,
@@ -2294,7 +2123,7 @@ class HunyuanImage3Pipeline(
             system_prompt,
             batch_cond_image_info,
             tokenizer_bot_task,
-        ) = self._extract_prompt_inputs(
+        ) = request_layout_utils.extract_hunyuan_prompt_inputs(
             req.prompts,
             extra_args,
             request_id=req.request_id,
@@ -2302,15 +2131,16 @@ class HunyuanImage3Pipeline(
         )
         prompt = prompt_from_req or prompt
         cot_text = (
-            [self._normalize_cot_text(t) for t in cot_text_list] if any(t is not None for t in cot_text_list) else None
+            [request_layout_utils.normalize_hunyuan_cot_text(t) for t in cot_text_list]
+            if any(t is not None for t in cot_text_list)
+            else None
         )
 
         generator = req.sampling_params.generator or generator
         height = req.sampling_params.height or height
         width = req.sampling_params.width or width
         num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
-        if req.sampling_params.guidance_scale_provided:
-            guidance_scale = req.sampling_params.guidance_scale
+        guidance_scale = request_layout_utils.resolve_hunyuan_guidance_scale(req.sampling_params, guidance_scale)
         if guidance_scale <= 1.0:
             logger.info("HunyuanImage3.0 runs without classifier-free guidance when guidance_scale <= 1.0.")
         image_size = (height, width)
@@ -2329,6 +2159,7 @@ class HunyuanImage3Pipeline(
             guidance_scale=guidance_scale,
             batch_cond_image_info=batch_cond_image_info,
             bot_task=tokenizer_bot_task,
+            prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(req.requests[0]),
             **ar_kv_kwargs,
         )
 

vllm_omni/diffusion/models/hunyuan_image3/request_layout.py

@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+from .hunyuan_image3_tokenizer import TokenizerEncodeOutput
+from .hunyuan_image3_transformer import ImageInfo, JointImageInfo
+from .system_prompt import get_system_prompt
+
+if TYPE_CHECKING:
+    from transformers.generation.configuration_utils import GenerationConfig
+
+    from .hunyuan_image3_tokenizer import TokenizerWrapper
+    from .hunyuan_image3_transformer import HunyuanImage3ImageProcessor
+
+
+def _to_python_scalar(value: Any) -> Any:
+    if isinstance(value, np.generic):
+        return value.item()
+    return value
+
+
+def _to_tensor_if_needed(value: Any) -> Any:
+    if isinstance(value, np.generic):
+        return value.item()
+    if isinstance(value, list):
+        return torch.tensor(value)
+    return value
+
+
+def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]:
+    return {
+        "image_type": image_info.image_type,
+        "image_tensor": image_info.image_tensor,
+        "image_width": _to_python_scalar(image_info.image_width),
+        "image_height": _to_python_scalar(image_info.image_height),
+        "token_width": _to_python_scalar(image_info.token_width),
+        "token_height": _to_python_scalar(image_info.token_height),
+        "image_token_length": _to_python_scalar(image_info.image_token_length),
+        "base_size": _to_python_scalar(image_info.base_size),
+        "ratio_index": _to_python_scalar(image_info.ratio_index),
+        "add_timestep_token": image_info.add_timestep_token,
+        "add_guidance_token": image_info.add_guidance_token,
+        "use_front_boi_token": image_info.use_front_boi_token,
+        "add_image_shape_token": image_info.add_image_shape_token,
+    }
+
+
+def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo:
+    return ImageInfo(
+        image_type=payload.get("image_type"),
+        image_tensor=_to_tensor_if_needed(payload.get("image_tensor")),
+        image_width=payload.get("image_width"),
+        image_height=payload.get("image_height"),
+        token_width=payload.get("token_width"),
+        token_height=payload.get("token_height"),
+        image_token_length=payload.get("image_token_length"),
+        base_size=payload.get("base_size"),
+        ratio_index=payload.get("ratio_index"),
+        add_timestep_token=payload.get("add_timestep_token", True),
+        add_guidance_token=payload.get("add_guidance_token", False),
+        use_front_boi_token=payload.get("use_front_boi_token", True),
+        add_image_shape_token=payload.get("add_image_shape_token", True),
+    )
+
+
+def joint_image_info_to_payload(joint_image_info: JointImageInfo) -> dict[str, Any]:
+    return {
+        "type": "joint_image_info",
+        "vae_image_info": _image_info_to_payload(joint_image_info.vae_image_info),
+        "vision_image_info": _image_info_to_payload(joint_image_info.vision_image_info),
+        "vision_encoder_kwargs": joint_image_info.vision_encoder_kwargs,
+    }
+
+
+def _joint_image_info_from_payload(payload: Any) -> JointImageInfo:
+    if isinstance(payload, JointImageInfo):
+        return payload
+    if not isinstance(payload, dict):
+        raise TypeError(f"Expected dict or JointImageInfo for conditional image payload, got {type(payload)}.")
+
+    vae_image_info = _image_info_from_payload(payload["vae_image_info"])
+    vision_image_info = _image_info_from_payload(payload["vision_image_info"])
+    vision_encoder_kwargs = payload.get("vision_encoder_kwargs") or {}
+    if isinstance(vision_encoder_kwargs, dict):
+        vision_encoder_kwargs = {key: _to_tensor_if_needed(value) for key, value in vision_encoder_kwargs.items()}
+    return JointImageInfo(
+        vae_image_info=vae_image_info,
+        vision_image_info=vision_image_info,
+        vision_encoder_kwargs=vision_encoder_kwargs,
+    )
+
+
+def normalize_hunyuan_single_stage_bot_task(bot_task: Any) -> str:
+    if isinstance(bot_task, str) and bot_task.lower() == "none":
+        bot_task = None
+    tokenizer_bot_task = bot_task
+    if tokenizer_bot_task == "think_recaption":
+        tokenizer_bot_task = "think"
+    elif tokenizer_bot_task == "vanilla":
+        tokenizer_bot_task = "image"
+    supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"}
+    if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks:
+        raise ValueError(
+            f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. "
+            f"Supported values are: {sorted(supported_bot_tasks)}."
+        )
+    return tokenizer_bot_task or "auto"
+
+
+def normalize_hunyuan_cot_text(cot: str | None) -> str | None:
+    """Restore an AR generation trigger tag omitted from generated text."""
+
+    if not cot:
+        return cot
+    if "</think>" in cot and not cot.startswith("<think>"):
+        return "<think>" + cot
+    if "</recaption>" in cot and not cot.startswith("<recaption>"):
+        return "<recaption>" + cot
+    return cot
+
+
+def extract_hunyuan_prompt_inputs(
+    prompts: list[Any],
+    extra_args: dict[str, Any],
+    *,
+    request_id: str,
+    allow_cond_image: bool,
+) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]:
+    """Normalize request prompt fields shared by planning and execution."""
+
+    is_dummy_warmup = OmniDiffusionRequest.is_dummy_run_request_id(request_id)
+    bot_task = extra_args.get("bot_task")
+    use_system_prompt = extra_args.get("use_system_prompt")
+    system_prompt = extra_args.get("system_prompt")
+
+    first_prompt = prompts[0] if prompts else None
+    if isinstance(first_prompt, dict):
+        if bot_task is None:
+            bot_task = first_prompt.get("bot_task")
+        if use_system_prompt is None:
+            use_system_prompt = first_prompt.get("use_system_prompt")
+        if system_prompt is None:
+            system_prompt = first_prompt.get("system_prompt")
+    tokenizer_bot_task = normalize_hunyuan_single_stage_bot_task(bot_task)
+    if use_system_prompt is not None:
+        system_prompt_bot_task = "image" if tokenizer_bot_task == "auto" else tokenizer_bot_task
+        system_prompt = get_system_prompt(use_system_prompt, system_prompt_bot_task, system_prompt)
+        system_prompt = system_prompt.strip() if system_prompt is not None else ""
+
+    prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts]
+    cot_text_list = [
+        (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in prompts
+    ]
+
+    batch_cond_image_info: list[list[JointImageInfo]] | None = None
+    if any(not isinstance(p, str) for p in prompts):
+        batch_cond_image_info = []
+        for prompt_item in prompts:
+            if isinstance(prompt_item, str):
+                batch_cond_image_info.append([])
+                continue
+            additional_info = prompt_item.get("additional_information") or {}
+            cond_infos = additional_info.get("batch_cond_image_info", [])
+            if isinstance(cond_infos, JointImageInfo | dict):
+                cond_infos = [cond_infos]
+            if cond_infos is None:
+                cond_infos = []
+            batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in cond_infos])
+
+        has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info]
+        if any(has_cond_image) and (not allow_cond_image) and not is_dummy_warmup:
+            raise ValueError("HunyuanImage3 step execution does not support image editing requests yet.")
+        if allow_cond_image and any(has_cond_image) and not all(has_cond_image):
+            raise ValueError("When batching Hunyuan image editing requests, every prompt must include input image(s).")
+        if not allow_cond_image or not any(has_cond_image):
+            batch_cond_image_info = None
+
+    return prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task
+
+
+def resolve_hunyuan_guidance_scale(sampling: Any, default_scale: float = 5.0) -> float:
+    if getattr(sampling, "guidance_scale_provided", False):
+        return sampling.guidance_scale
+    return default_scale
+
+
+def hunyuan_num_image_tokens(image_info: ImageInfo) -> int:
+    """Return the generated-image span overwritten on every denoise step."""
+
+    return int(image_info.image_token_length) + int(image_info.add_timestep_token) + int(image_info.add_guidance_token)
+
+
+def build_hunyuan_batch_rope_image_info(
+    output: TokenizerEncodeOutput,
+    sections: list[list[dict[str, Any]]],
+) -> list[list[tuple[slice, tuple[int, int]]]]:
+    if output.all_image_slices is None:
+        raise ValueError("Hunyuan tokenizer output is missing all_image_slices.")
+    if len(output.all_image_slices) != len(sections):
+        raise ValueError(
+            "Hunyuan image-slice rows do not match template section rows: "
+            f"image_slices={len(output.all_image_slices)}, sections={len(sections)}"
+        )
+    rope_image_info: list[list[tuple[slice, tuple[int, int]]]] = []
+    for image_slices, sections_i in zip(output.all_image_slices, sections):
+        image_shapes: list[tuple[int, int]] = []
+        for section in sections_i:
+            if "image" not in section["type"]:
+                continue
+            if isinstance(section["token_height"], list):
+                if len(section["token_height"]) != len(section["token_width"]):
+                    raise ValueError(
+                        "token_height and token_width should have the same length, "
+                        f"but got {len(section['token_height'])} and {len(section['token_width'])}"
+                    )
+                image_shapes.extend(
+                    (int(token_height), int(token_width))
+                    for token_height, token_width in zip(section["token_height"], section["token_width"])
+                )
+            else:
+                image_shapes.append((int(section["token_height"]), int(section["token_width"])))
+        if len(image_slices) != len(image_shapes):
+            raise ValueError(f"Image slices({len(image_slices)}) do not match image shapes({len(image_shapes)}).")
+        rope_image_info.append(list(zip(image_slices, image_shapes)))
+    return rope_image_info
+
+
+def prepare_hunyuan_layout(
+    request: OmniDiffusionRequest,
+    *,
+    tokenizer_wrapper: TokenizerWrapper,
+    image_processor: HunyuanImage3ImageProcessor,
+    generation_config: GenerationConfig,
+    image_base_size: int,
+) -> HunyuanPreparedLayout:
+    """Build the CPU token/image layout reused by Scheduler and Worker."""
+
+    sampling = request.sampling_params
+    extra_args = getattr(sampling, "extra_args", {}) or {}
+    prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task = extract_hunyuan_prompt_inputs(
+        [request.prompt],
+        extra_args,
+        request_id=request.request_id,
+        allow_cond_image=True,
+    )
+    cot_text = (
+        [normalize_hunyuan_cot_text(text) for text in cot_text_list]
+        if any(text is not None for text in cot_text_list)
+        else None
+    )
+    height = sampling.height or 1024
+    width = sampling.width or 1024
+    guidance_scale = resolve_hunyuan_guidance_scale(sampling)
+    generated_image_info = image_processor.build_image_info((height, width))
+    result = tokenizer_wrapper.apply_chat_template(
+        batch_prompt=prompt,
+        mode="gen_image",
+        batch_gen_image_info=[generated_image_info],
+        batch_cond_image_info=batch_cond_image_info,
+        batch_system_prompt=[system_prompt],
+        batch_cot_text=cot_text,
+        max_length=None,
+        bot_task=tokenizer_bot_task,
+        image_base_size=image_base_size,
+        sequence_template=getattr(generation_config, "sequence_template", "pretrain"),
+        cfg_factor=1 + int(guidance_scale > 1.0),
+        drop_think=getattr(generation_config, "drop_think", False),
+    )
+    tokenizer_output = result["output"]
+    return HunyuanPreparedLayout(
+        tokenizer_output=tokenizer_output,
+        rope_image_info=build_hunyuan_batch_rope_image_info(tokenizer_output, result["sections"]),
+        generated_image_info=generated_image_info,
+    )
+
+
+@dataclass(frozen=True)
+class HunyuanPreparedLayout:
+    """CPU execution layout prepared once before Scheduler admission."""
+
+    tokenizer_output: TokenizerEncodeOutput
+    rope_image_info: list[list[tuple[slice, tuple[int, int]]]]
+    generated_image_info: ImageInfo
+
+    def __post_init__(self) -> None:
+        tokens = self.tokenizer_output.tokens
+        if tokens is None or tokens.ndim != 2:
+            shape = None if tokens is None else tuple(tokens.shape)
+            raise ValueError(f"Hunyuan prepared tokenizer tokens must be 2-D, got {shape}")
+        if len(self.rope_image_info) != int(tokens.shape[0]):
+            raise ValueError(
+                "Hunyuan prepared RoPE rows do not match tokenizer rows: "
+                f"rope_rows={len(self.rope_image_info)}, tokens={int(tokens.shape[0])}"
+            )
+        prefix_positions = self.tokenizer_output.gen_timestep_scatter_index
+        if prefix_positions is None or prefix_positions.ndim != 2 or prefix_positions.shape[1] == 0:
+            raise ValueError("Hunyuan prepared layout requires a non-empty 2-D gen_timestep_scatter_index")
+        if prefix_positions.shape[0] != int(tokens.shape[0]):
+            raise ValueError(
+                "Hunyuan prepared prefix-position rows do not match tokenizer rows: "
+                f"prefix_positions={prefix_positions.shape[0]}, tokens={int(tokens.shape[0])}"
+            )
+        real_pos = self.tokenizer_output.real_pos
+        if real_pos is None or real_pos.ndim != 2 or real_pos.shape[1] == 0:
+            raise ValueError("Hunyuan prepared layout requires a non-empty 2-D real_pos")
+        if real_pos.shape[0] != int(tokens.shape[0]):
+            raise ValueError(
+                "Hunyuan prepared valid-length rows do not match tokenizer rows: "
+                f"real_pos={real_pos.shape[0]}, tokens={int(tokens.shape[0])}"
+            )
+        seq_lens = [int(row[-1].item()) for row in real_pos]
+        if any(seq_len <= 0 or seq_len > int(tokens.shape[1]) for seq_len in seq_lens):
+            raise ValueError(
+                "Hunyuan prepared valid sequence lengths must fit the padded tokenizer width: "
+                f"seq_lens={seq_lens}, padded_width={int(tokens.shape[1])}"
+            )
+
+    @property
+    def num_branches(self) -> int:
+        return int(self.tokenizer_output.tokens.shape[0])
+
+
+def build_hunyuan_diffusion_kv_requests(
+    request: OmniDiffusionRequest,
+    prepared_layout: HunyuanPreparedLayout,
+) -> tuple[DiffusionKVRequest, ...]:
+    """Build one persistent Scheduler KV request per Hunyuan execution row."""
+
+    tokenizer_output = prepared_layout.tokenizer_output
+    cfg_factor = 1 + int(resolve_hunyuan_guidance_scale(request.sampling_params) > 1.0)
+    if prepared_layout.num_branches != cfg_factor:
+        raise ValueError(
+            "Hunyuan tokenizer sequence count does not match CFG execution: "
+            f"rows={prepared_layout.num_branches}, cfg_factor={cfg_factor}"
+        )
+    prefix_positions = tokenizer_output.gen_timestep_scatter_index
+    real_pos = tokenizer_output.real_pos
+    assert prefix_positions is not None and real_pos is not None
+
+    target_len = hunyuan_num_image_tokens(prepared_layout.generated_image_info)
+    return tuple(
+        DiffusionKVRequest(
+            f"{request.request_id}/diffusion-kv/{sequence_id}",
+            sequence_id=sequence_id,
+            # The generated-image timestep position terminates the reusable
+            # prompt/reference-image prefix for this execution row.
+            prefix_len=int(prefix_row[-1].item()),
+            target_len=target_len,
+            seq_len=int(valid_row[-1].item()),
+            # Prompt and reference-image tokens are already embedded in this
+            # row's primary self-attention sequence. Hunyuan therefore has no
+            # independently projected cross/joint-attention KV context.
+            kv_contexts=(),
+        )
+        for sequence_id, (prefix_row, valid_row) in enumerate(zip(prefix_positions, real_pos))
+    )
+
+
+def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None:
+    prepared_layout = getattr(source, "prepared_layout", None)
+    if prepared_layout is None:
+        return None
+    if not isinstance(prepared_layout, HunyuanPreparedLayout):
+        raise TypeError(
+            f"HunyuanImage3 expected prepared_layout to be HunyuanPreparedLayout, got {type(prepared_layout).__name__}"
+        )
+    return prepared_layout

vllm_omni/diffusion/request.py

@@ -2,12 +2,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from __future__ import annotations
+
 import random
 from dataclasses import dataclass
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+
 DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id"
 
 
@@ -30,6 +35,14 @@ class OmniDiffusionRequest:
     sampling_params: OmniDiffusionSamplingParams
     request_id: str
     kv_sender_info: dict[str, Any] | None = None
+    # Optional opaque, model-owned input prepared before Scheduler admission.
+    # Model code validates its concrete type when consuming it on the Worker.
+    prepared_layout: Any | None = None
+    # Scheduler-only native-compatible KV requests produced by model
+    # preprocessing.
+    # BaseScheduler takes ownership and clears this field before the request is
+    # sent to a Worker.
+    diffusion_kv_requests: tuple[DiffusionKVRequest, ...] | None = None
 
     def __post_init__(self):
         """Initialize dependent fields after dataclass initialization."""

vllm_omni/diffusion/sched/base_scheduler.py

@@ -10,6 +10,7 @@ from dataclasses import fields
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
@@ -47,6 +48,7 @@ class BaseScheduler(ABC):
         self._finished_req_ids: set[str] = set()
         self.max_num_running_reqs: int = 1
         self._prefetch_enabled: bool = False
+        self._diffusion_kv_enabled = False
 
     def initialize(self, od_config: OmniDiffusionConfig) -> None:
         self.od_config = od_config
@@ -63,6 +65,10 @@ class BaseScheduler(ABC):
             self.max_num_running_reqs = 1
         omni_kv = getattr(od_config, "omni_kv_config", None) or {}
         self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False))
+        self._diffusion_kv_enabled = (
+            getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY)
+            is DiffusionKVCacheMode.PAGED_SCHEDULER
+        )
         self._reset_scheduler_state()
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
@@ -251,12 +257,43 @@ class BaseScheduler(ABC):
         """Remove subclass-owned per-request state before popping request state."""
 
     def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> SchedulerRequestState:
+        kv_requests = request.diffusion_kv_requests or ()
+        if self._diffusion_kv_enabled:
+            self._reject_legacy_dense_kv(request)
+            if not kv_requests:
+                raise ValueError("paged_scheduler request preprocessing did not produce DiffusionKVRequest state")
+        elif kv_requests:
+            raise ValueError("dense_legacy request unexpectedly contains Scheduler Diffusion KV requests")
+
+        # DiffusionKVRequest objects are mutable Scheduler/native-KVCacheManager
+        # state and must never ride the normal request payload to a Worker.
+        request.diffusion_kv_requests = None
         return SchedulerRequestState(
             request_id=request_id,
             req=request,
             sampling_params_key=self._build_sampling_params_key(request),
+            diffusion_kv_requests=kv_requests,
         )
 
+    @staticmethod
+    def _reject_legacy_dense_kv(request: OmniDiffusionRequest) -> None:
+        """Keep dense injected KV out of the Scheduler-owned paged path."""
+
+        populated_fields: list[str] = []
+        for owner_name, owner in (
+            ("request", request),
+            ("sampling_params", request.sampling_params),
+        ):
+            for field_name, value in vars(owner).items():
+                if value is not None and (field_name == "past_key_values" or field_name.endswith("_past_key_values")):
+                    populated_fields.append(f"{owner_name}.{field_name}")
+        if populated_fields:
+            fields_text = ", ".join(sorted(populated_fields))
+            raise ValueError(
+                "paged_scheduler Diffusion KV does not accept legacy dense KV payloads; "
+                f"clear these fields before admission: {fields_text}"
+            )
+
     def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool:
         if not self._running:
             return True

vllm_omni/diffusion/sched/interface.py

@@ -6,10 +6,13 @@ from __future__ import annotations
 import enum
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Any, TypedDict
+from typing import TYPE_CHECKING, Any, TypedDict
 
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+
 
 class DiffusionRequestStatus(enum.IntEnum):
     """Request status tracked by diffusion scheduler."""
@@ -129,6 +132,7 @@ class SchedulerRequestState:
     request_id: str
     req: OmniDiffusionRequest
     sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None
+    diffusion_kv_requests: tuple[DiffusionKVRequest, ...] = ()
     status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
     error: str | None = None
 

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -641,6 +641,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     sampling=copy.deepcopy(sched_new_req.req.sampling_params),
                     prompt=sched_new_req.req.prompt,
                     kv_sender_info=sched_new_req.req.kv_sender_info,
+                    prepared_layout=getattr(sched_new_req.req, "prepared_layout", None),
                 )
                 state_req = copy.copy(sched_new_req.req)
                 state_req.sampling_params = new_state.sampling

vllm_omni/diffusion/worker/utils.py

@@ -77,6 +77,9 @@ class StepRequestState:
     sampling: OmniDiffusionSamplingParams
     prompt: OmniPromptType | None = None
     kv_sender_info: dict | None = None
+    # Opaque model-owned preprocessing result. Scheduler-owned Diffusion KV
+    # request state is removed before this Worker state is built.
+    prepared_layout: Any | None = None
 
     # ── Encoded prompts (set once by prepare_encode) ──
     prompt_embeds: torch.Tensor | None = None