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