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

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