[Diffusion][HunyuanImage3] Add request paged prefix caching (#7719)

Signed-off-by: zwhzzz0821 <2831474076@qq.com>

.gitignore

@@ -248,6 +248,8 @@ datasets/
 # Omni-DuplexEval CI config JSON (e.g. accuracy/omni_duplex_eval)
 !tests/e2e/accuracy/omni_duplex_eval/*.json
 *.jsonl
+# Shared Hunyuan IT2I input for accuracy tests and vllm bench.
+!tests/assets/hunyuan_image3/it2i.jsonl
 *.parquet
 
 # Output files

benchmarks/accuracy/common.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 import base64
@@ -200,7 +203,16 @@ class VllmOmniImageClient:
             )
             last_response = response
             if response.status_code != 404:
-                response.raise_for_status()
+                try:
+                    response.raise_for_status()
+                except requests.HTTPError as error:
+                    # Keep the server's validation detail in the exception;
+                    # otherwise multipart image tests only report a bare
+                    # ``400 Bad Request`` and hide the actionable cause.
+                    raise requests.HTTPError(
+                        f"{error}: {response.text}",
+                        response=response,
+                    ) from error
                 return decode_base64_image(response.json()["data"][0]["b64_json"])
 
         assert last_response is not None

benchmarks/diffusion/README.md

@@ -164,3 +164,22 @@ batch may still pay compile or CUDA-graph capture cost.
 
 For a Qwen-Image continuous-batching replay example, see
 [`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).
+
+## HunyuanImage3 reference-prefix reuse
+
+The prefix-cache benchmark uses the unified `vllm bench serve --omni` runner,
+not this legacy diffusion benchmark. It reuses the two-image IT2I input on a
+single DiT stage, comparing dense, paged without caching, and paged with caching.
+See [Shared-reference benchmark](../../docs/design/feature/prefix_caching.md#shared-reference-benchmark)
+for the configuration, commands, and separate partial-hit accuracy coverage.
+
+The accuracy regression compares paged-no-cache, partial-hit and exact-repeat
+outputs against the checked-in official-repository reference-image goldens at
+**50 steps, CFG=2.5, seeds 43/45** (both CFG branches remain on CFGP1).
+It uses the AR-to-DiT image criteria: CLIP ≥ 90, SSIM ≥ 0.26, PSNR ≥ 12.5 dB.
+`HUNYUAN_IMAGE3_INCLUDE_DENSE=1` adds dense as another comparison to the goldens.
+It also requires actual reference-image hits and model-side query slicing,
+and checks bitwise repeatability for repeated requests at the same hit boundary.
+Images, logs, deployment YAMLs, hit traces and quality metrics are saved under
+pytest's temporary output directory. `HUNYUAN_IMAGE3_MODEL` selects a local model;
+DFX uses the repository's normal model/cache resolution.

docs/design/feature/prefix_caching.md

@@ -8,6 +8,7 @@
 - [High-Level Approach](#high-level-approach)
 - [Example](#example)
 - [What About Multimodal Inputs?](#what-about-multimodal-inputs)
+- [Diffusion KV Prefix Caching](#diffusion-kv-prefix-caching)
 - [Implementation](#implementation)
 - [Related Files](#related-files)
 
@@ -25,7 +26,7 @@ vLLM implements automatic prefix caching for managing its kv-cache, which is bes
 - Model / stage specific multimodal data
 
 !!! note "Note 1"
-    This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice.
+    The following sections describe caching tensor outputs passed between AR stages. This is distinct from multimodal encoder caching and from [diffusion KV prefix caching](#diffusion-kv-prefix-caching), which reuses GPU KV pages inside a diffusion stage.
 
 ### High-Level Approach
 
@@ -45,7 +46,7 @@ Host footprint: each cached key costs `num_blocks × block_size × D × dtype_by
 ### Example
 
 !!! note "Note 3"
-    Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config.
+    The stage-output tensor cache described here supports AutoRegressive stages with one kv-cache group. Configure it with the pipeline-wide `enable_prefix_caching` field in the deploy config. Diffusion KV caching has the separate requirements below.
 
 The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following:
 
@@ -167,6 +168,65 @@ When we pass our multimodal tensors to the language model component in the same
 
 Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled.
 
+### Diffusion KV Prefix Caching
+
+HunyuanImage3's standalone DiT pipeline can reuse stable text/reference-image KV
+across requests. Enable it on the diffusion stage in the deploy config:
+
+```yaml
+pipeline: hunyuan_image3_dit
+stages:
+  - stage_id: 0
+    diffusion_kv_mode: paged_scheduler
+    enable_prefix_caching: true
+```
+
+The Scheduler uses vLLM's native `KVCacheManager` to look up and retain complete
+prefix blocks. The Worker owns the GPU pages and block tables. On a hit, the first
+denoise forward computes only the uncached query suffix while attending to both
+cached and new KV. Dynamic target-image KV is never published as a reusable prefix.
+
+Cache identity includes token IDs, reference-image content and VAE random state,
+plus model/layout and LoRA context. The same image with different prompts can
+reuse the common leading blocks; it does not imply that every image span or CFG
+branch is interchangeable. Disabling prefix caching skips cache-identity hashing;
+`dense_legacy` remains the default. Enabling prefix caching with a mode other than
+`paged_scheduler`, without a registered model hook, or together with native
+`kv_transfer_config` raises a configuration error.
+Combining it with `enable_sleep_mode: true` is also rejected: sleep discards KV
+pages without invalidating the Scheduler's prefix-cache index. Disable either
+prefix caching or sleep mode.
+
+The current scope is local DiT reuse, not AR-imported KV or missing-page-only
+cross-stage transfer. Prefix-hit accuracy has been exercised with TP4/SP1 and
+TP2/SP2 (Ulysses), both with EP and CFGP1. Other combinations still require validation. See the
+[diffusion compatibility notes](../../user_guide/diffusion_features.md#diffusion-kv-prefix-caching).
+
+#### Shared-reference benchmark
+
+From the repository root, use the unified `vllm bench serve --omni` test runner:
+
+```bash
+pytest tests/dfx/perf/scripts/run_benchmark.py \
+  --test-config-file tests/dfx/perf/tests/test_hunyuan_image3_prefix_caching.json
+```
+
+Run GPU tests through your environment's GPU scheduler. The suite reuses the
+two-image IT2I accuracy input on a single DiT stage, comparing dense, paged without
+prefix caching, and paged with prefix caching. Each mode runs two ordinary upstream
+warmups followed by eight identical requests at CFG 2.5, 8 denoise steps and seed 42.
+Warmups are excluded from latency and request throughput. This measures full-prefix
+reuse, not the earlier distinct-prompt workload; its results must be reported separately.
+
+`tests/e2e/accuracy/test_hunyuan_image3_prefix_cache_accuracy.py` uses the same input
+at CFG 2.5, 50 steps and seed 42, comparing all modes against the existing official
+Instruct output image and IT2I thresholds (CLIP ≥90, SSIM ≥0.26, PSNR ≥12.5).
+A different prompt warms the cache; the original prompt then exercises partial
+and repeated hits, with reference-image reuse and query slicing verified.
+Only outputs for the original prompt are scored. The golden comes from AR-to-DiT;
+this single-DiT quality check does not establish identical conditioning between
+the pipelines. The existing AR-to-DiT test remains unchanged.
+
 ### Implementation
 
 The block/slot model is `vllm_omni/core/prefix_cache/`.

docs/user_guide/diffusion_features.md

@@ -32,6 +32,30 @@ Cache methods trade minimal quality for significant speedup. Quality loss is typ
 | **[TeaCache](diffusion/cache_acceleration/teacache.md)** | Adaptive caching using modulated inputs | Quick setup, balanced quality/speed on single GPU |
 | **[Cache-DiT](diffusion/cache_acceleration/cache_dit.md)** | Multiple caching techniques: DBCache, TaylorSeer, SCM | Fine-grained control, tunable quality-speed tradeoff |
 
+#### Diffusion KV Prefix Caching
+
+[KV prefix caching](../design/feature/prefix_caching.md#diffusion-kv-prefix-caching)
+reuses stable context KV across requests, rather than approximating denoise
+steps. Enable `diffusion_kv_mode: paged_scheduler` and `enable_prefix_caching: true`
+on the HunyuanImage3 standalone DiT stage. This is separate from the AR stage-output
+cache, TeaCache and Cache-DiT.
+
+| Model / combination | Scope |
+|---------------------|-------|
+| HunyuanImage3 standalone DiT | Stable text/reference-image prefix reuse; dynamic image tokens are excluded |
+| TP4/SP1 + EP, CFGP1 | Validated reference-image prefix hits |
+| TP2/SP2 (Ulysses) + EP, CFGP1 | Validated reference-image prefix hits; CFG guidance does not require CFG parallelism |
+| Other SP/TP combinations, CFGP>1 | Prefix-hit E2E accuracy not yet validated |
+| AR-imported KV / cross-stage missing-page transfer | Not covered by local prefix caching |
+| Sleep mode | Rejected with prefix caching; discarded KV pages would leave stale cache hits |
+| TeaCache / Cache-DiT, offload, quantization | Not validated with prefix hits |
+| Step execution | Not supported with HunyuanImage3 paged KV; use request-level execution |
+| Other diffusion models | No prefix-cache model adapter yet |
+
+This table describes prefix-cache scope; the general parallelism tables below do
+not establish prefix-hit compatibility. Floating-point kernel differences still
+require accuracy validation. `dense_legacy` remains the default.
+
 #### Lossless Acceleration
 
 Parallelism methods distribute computation across GPUs without quality loss (mathematically equivalent to single-GPU).

tests/assets/hunyuan_image3/it2i.jsonl

@@ -0,0 +1 @@
+{"prompt": "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴", "image_files": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png", "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png"]}

tests/benchmarks/test_accuracy_bench_utils.py

@@ -9,6 +9,7 @@ import sys
 from pathlib import Path
 
 import pytest
+import requests
 from PIL import Image
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -311,7 +312,8 @@ def test_trajectory_judge_payload_collapses_six_frames_into_single_storyboard():
     assert judge_images[0].size == (24, 12)
 
 
-def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch):
+@pytest.mark.parametrize("num_images", [1, 2])
+def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch, num_images):
     captured = {}
 
     class FakeResponse:
@@ -344,16 +346,46 @@ def test_image_edit_client_uses_openai_image_edit_endpoint(monkeypatch):
     output = client.generate_image_edit(
         model="Qwen/Qwen-Image-Edit",
         prompt="edit this image",
-        images=image,
+        images=image if num_images == 1 else [image] * num_images,
         width=512,
         height=512,
+        guidance_scale=2.5,
+        seed=42,
+        bot_task="think_recaption",
+        sys_type="en_unified",
     )
 
     assert output.size == (1, 1)
     assert captured["url"] == "http://127.0.0.1:8093/v1/images/edits"
     assert captured["data"]["prompt"] == "edit this image"
     assert captured["data"]["size"] == "512x512"
-    assert captured["files"][0][0] == "image"
+    assert len(captured["files"]) == num_images
+    assert all(field == ("image" if num_images == 1 else "image[]") for field, _ in captured["files"])
+    assert captured["data"]["guidance_scale"] == "2.5"
+    assert captured["data"]["seed"] == "42"
+    assert captured["data"]["bot_task"] == "think_recaption"
+    assert captured["data"]["sys_type"] == "en_unified"
+
+
+@pytest.mark.parametrize("status_code", [400, 500])
+def test_image_edit_client_preserves_server_error_detail(mocker, status_code):
+    response = requests.Response()
+    response.status_code = status_code
+    response._content = b'{"error": "image input limit exceeded"}'
+    post = mocker.patch("benchmarks.accuracy.common.requests.post", return_value=response)
+    client = VllmOmniImageClient(base_url="http://127.0.0.1:8093")
+
+    with pytest.raises(requests.HTTPError, match="image input limit exceeded") as exc:
+        client.generate_image_edit(
+            model="test-model",
+            prompt="edit this image",
+            images=Image.new("RGB", (2, 2)),
+            width=512,
+            height=512,
+        )
+
+    assert exc.value.response is response
+    post.assert_called_once()
 
 
 def test_text_to_image_client_forwards_output_compression(monkeypatch):

tests/benchmarks/test_serve_cli.py

@@ -6,6 +6,7 @@ import json
 import os
 import subprocess
 import textwrap
+import time
 from pathlib import Path
 
 import pytest
@@ -21,6 +22,76 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
 
 
+def test_it2i_dataset_uses_upstream_warmups(monkeypatch, tmp_path):
+    """Exercise real dataset/benchmark code; replace only the HTTP backend."""
+    from vllm_omni.benchmarks import serve
+    from vllm_omni.benchmarks.patch import patch
+    from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
+
+    dataset_path = Path(__file__).parents[1] / "assets/hunyuan_image3/it2i.jsonl"
+    sample = json.loads(dataset_path.read_text(encoding="utf-8"))
+    parser = argparse.ArgumentParser()
+    OmniBenchmarkServingSubcommand.add_cli_args(parser)
+    args = parser.parse_args(
+        [
+            "--model",
+            "test-model",
+            "--endpoint",
+            "/v1/images/edits",
+            "--dataset-name",
+            "custom_image",
+            "--dataset-path",
+            str(dataset_path),
+            "--skip-tokenizer-init",
+            "--disable-shuffle",
+            "--disable-tqdm",
+            "--num-prompts",
+            "8",
+            "--num-warmups",
+            "2",
+            "--max-concurrency",
+            "1",
+            "--ready-check-timeout-sec",
+            "0",
+            "--percentile-metrics",
+            "e2el",
+            "--extra-body",
+            '{"seed":42,"guidance_scale":2.5,"bot_task":"think_recaption"}',
+            "--save-result",
+            "--result-dir",
+            str(tmp_path),
+            "--result-filename",
+            "result.json",
+        ]
+    )
+    requests = []
+
+    async def request_func(request_func_input, session, pbar=None):
+        requests.append(request_func_input)
+        if pbar is not None:
+            pbar.update(1)
+        return patch.MixRequestFuncOutput(
+            success=True,
+            prompt_len=1,
+            output_tokens=1,
+            image_count=1,
+            start_time=time.perf_counter(),
+            latency=99.0 if len(requests) <= 2 else 0.1,
+        )
+
+    monkeypatch.setitem(patch.ASYNC_REQUEST_FUNCS, "/v1/images/edits", request_func)
+    result = serve.main(args)
+    assert len(requests) == 10
+    for request in requests:
+        assert request.prompt == sample["prompt"]
+        assert [image["image_url"]["url"] for image in request.multi_modal_content] == sample["image_files"]
+        assert request.extra_body == args.extra_body
+    assert result["completed"] == 8
+    saved = json.loads((tmp_path / "result.json").read_text())
+    assert saved["completed"] == 8
+    assert saved["mean_e2el_ms"] == pytest.approx(100.0)  # Excludes both 99-second warmups.
+
+
 @pytest.mark.parametrize(
     "dataset_name",
     [

tests/config/test_omni_config.py

@@ -1578,6 +1578,68 @@ def test_diffusion_config_preserves_existing_coercion_hooks():
     assert cfg.max_cpu_loras == 1
 
 
+@pytest.mark.parametrize("enabled", [False, True])
+def test_diffusion_projection_retains_prefix_caching(enabled):
+    projection = omni_config_module._DiffusionConfigProjection
+    assert projection.from_kwargs().enable_prefix_caching is False
+    config = projection.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=enabled,
+    )
+    assert config.enable_prefix_caching is enabled
+
+
+@pytest.mark.parametrize(
+    "pipeline_value,stage_value,cli_value,expected",
+    [
+        (None, None, None, False),
+        (True, None, None, True),
+        (False, True, None, True),
+        (True, False, None, False),
+        (None, None, True, True),
+        (True, True, False, False),
+        (False, False, True, True),
+    ],
+)
+def test_diffusion_prefix_caching_precedence_and_transport(
+    tmp_path,
+    pipeline_value,
+    stage_value,
+    cli_value,
+    expected,
+):
+    import yaml
+
+    from vllm_omni.engine.stage_init_utils import _project_omni_stage_engine_args
+
+    stage_deploy = {
+        "stage_id": 0,
+        "diffusion_kv_mode": "paged_scheduler",
+        "diffusion_kv_max_rows_per_request": 2,
+    }
+    deploy = {"pipeline": "hunyuan_image3_dit", "async_chunk": False, "stages": [stage_deploy]}
+    if pipeline_value is not None:
+        deploy["enable_prefix_caching"] = pipeline_value
+    if stage_value is not None:
+        stage_deploy["enable_prefix_caching"] = stage_value
+    deploy_path = tmp_path / "prefix.yaml"
+    deploy_path.write_text(yaml.safe_dump(deploy))
+    cli = {} if cli_value is None else {"stage_0_enable_prefix_caching": cli_value}
+    stage = _from_pipeline_key(
+        "hunyuan_image3_dit",
+        deploy_config_path=str(deploy_path),
+        cli_overrides=cli,
+    ).stage_by_id(0)
+
+    assert bool(stage.cache_config.enable_prefix_caching) is expected
+    assert stage.diffusion_config.enable_prefix_caching is expected
+    assert _project_omni_stage_engine_args(stage)["enable_prefix_caching"] is expected
+    serialized = _serialize_stage_config(stage)
+    restored = omni_config_module._DiffusionConfigProjection.from_kwargs(**serialized["diffusion_config"])
+    assert restored.enable_prefix_caching is expected
+
+
 def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
     from vllm_omni.platforms import current_omni_platform
 

tests/dfx/perf/tests/test_hunyuan_image3_prefix_caching.json

@@ -0,0 +1,131 @@
+[
+  {
+    "test_name": "test_hunyuan_image3_prefix_dense",
+    "mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
+    "description": "Repeated two-image IT2I request on a single DiT stage: dense baseline (TP4, CFGP1)",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
+      },
+      "stage_overrides": {
+        "0": {
+          "max_model_len": 22800,
+          "max_num_batched_tokens": 32768,
+          "diffusion_kv_mode": "dense_legacy",
+          "diffusion_kv_max_rows_per_request": 2,
+          "enable_prefix_caching": false,
+          "vae_use_tiling": true,
+          "parallel_config": {"vae_patch_parallel_size": 4}
+        }
+      }
+    },
+    "benchmark_params": [{
+      "dataset_name": "custom_image",
+      "dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
+      "endpoint": "/v1/images/edits",
+      "skip_tokenizer_init": true,
+      "disable_shuffle": true,
+      "num_prompts": 8,
+      "max_concurrency": 1,
+      "num_warmups": 2,
+      "seed": 42,
+      "percentile_metrics": "e2el",
+      "extra_body": {
+        "guidance_scale": 2.5,
+        "sys_type": "en_unified",
+        "seed": 42,
+        "size": "1280x720",
+        "num_inference_steps": 8,
+        "bot_task": "think_recaption"
+      }
+    }]
+  },
+  {
+    "test_name": "test_hunyuan_image3_prefix_paged_no_cache",
+    "mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
+    "description": "Repeated two-image IT2I request on a single DiT stage: paged without prefix caching (TP4, CFGP1)",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
+      },
+      "stage_overrides": {
+        "0": {
+          "max_model_len": 22800,
+          "max_num_batched_tokens": 32768,
+          "diffusion_kv_mode": "paged_scheduler",
+          "diffusion_kv_max_rows_per_request": 2,
+          "enable_prefix_caching": false,
+          "vae_use_tiling": true,
+          "parallel_config": {"vae_patch_parallel_size": 4}
+        }
+      }
+    },
+    "benchmark_params": [{
+      "dataset_name": "custom_image",
+      "dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
+      "endpoint": "/v1/images/edits",
+      "skip_tokenizer_init": true,
+      "disable_shuffle": true,
+      "num_prompts": 8,
+      "max_concurrency": 1,
+      "num_warmups": 2,
+      "seed": 42,
+      "percentile_metrics": "e2el",
+      "extra_body": {
+        "guidance_scale": 2.5,
+        "sys_type": "en_unified",
+        "seed": 42,
+        "size": "1280x720",
+        "num_inference_steps": 8,
+        "bot_task": "think_recaption"
+      }
+    }]
+  },
+  {
+    "test_name": "test_hunyuan_image3_prefix_paged_prefix",
+    "mark": [{"hardware_marks": {"res": {"cuda": "H100"}, "num_cards": 4}}, "diffusion", "local_model"],
+    "description": "Repeated two-image IT2I request on a single DiT stage: paged with full prefix reuse (TP4, CFGP1)",
+    "server_params": {
+      "model": "tencent/HunyuanImage-3.0-Instruct",
+      "serve_args": {
+        "trust-remote-code": true,
+        "deploy-config": "vllm_omni/deploy/hunyuan_image3_dit.yaml"
+      },
+      "stage_overrides": {
+        "0": {
+          "max_model_len": 22800,
+          "max_num_batched_tokens": 32768,
+          "diffusion_kv_mode": "paged_scheduler",
+          "diffusion_kv_max_rows_per_request": 2,
+          "enable_prefix_caching": true,
+          "vae_use_tiling": true,
+          "parallel_config": {"vae_patch_parallel_size": 4}
+        }
+      }
+    },
+    "benchmark_params": [{
+      "dataset_name": "custom_image",
+      "dataset_path": "tests/assets/hunyuan_image3/it2i.jsonl",
+      "endpoint": "/v1/images/edits",
+      "skip_tokenizer_init": true,
+      "disable_shuffle": true,
+      "num_prompts": 8,
+      "max_concurrency": 1,
+      "num_warmups": 2,
+      "seed": 42,
+      "percentile_metrics": "e2el",
+      "extra_body": {
+        "guidance_scale": 2.5,
+        "sys_type": "en_unified",
+        "seed": 42,
+        "size": "1280x720",
+        "num_inference_steps": 8,
+        "bot_task": "think_recaption"
+      }
+    }]
+  }
+]

tests/dfx/perf/tests/test_runner_metadata.py

@@ -254,6 +254,38 @@ def test_merge_omni_default_server_args_respects_json():
     ]
 
 
+def test_prefix_benchmark_uses_unified_runner(monkeypatch):
+    from pathlib import Path
+
+    from tests.dfx.conftest import is_diffusion_perf_config, load_benchmark_configs
+    from tests.dfx.perf.scripts import run_benchmark
+
+    configs = load_benchmark_configs(str(Path(__file__).with_name("test_hunyuan_image3_prefix_caching.json")))
+    assert len(configs) == 3
+    assert all(not is_diffusion_perf_config(config) for config in configs)
+    assert all(config["benchmark_params"] == configs[0]["benchmark_params"] for config in configs)
+    monkeypatch.setattr(run_benchmark, "BENCHMARK_CONFIGS", configs)
+    monkeypatch.setattr(run_benchmark, "get_runtime_resource_label", lambda: "H100")
+    calls = []
+
+    def benchmark(**kwargs):
+        calls.append(kwargs)
+        assert "--name" not in kwargs["args"]
+        assert "--warmup-dataset-path" not in kwargs["args"]
+        assert "tests/assets/hunyuan_image3/it2i.jsonl" in kwargs["args"]
+        assert kwargs["num_warmups"] == 2
+        return {"completed": 8}
+
+    monkeypatch.setattr(run_benchmark, "run_benchmark", benchmark)
+    for config in configs:
+        for params in config["benchmark_params"]:
+            run_benchmark.test_performance_benchmark(
+                SimpleNamespace(host="localhost", port=8000, model="test-model"),
+                {"test_name": config["test_name"], "params": params},
+            )
+    assert len(calls) == 3
+
+
 def test_benchmark_param_id_suffix_from_task_eval_phase():
     from tests.dfx.conftest import _unique_benchmark_param_id_suffixes
 

tests/diffusion/diffusion_kv/test_config.py

@@ -21,6 +21,86 @@ def test_dense_legacy_is_default() -> None:
     assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY
 
 
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+@pytest.mark.parametrize("mode", [None, "dense_legacy", DiffusionKVCacheMode.DENSE_LEGACY])
+def test_prefix_caching_rejects_non_paged_mode(config_cls, mode) -> None:
+    kwargs = {} if mode is None else {"diffusion_kv_mode": mode}
+    with pytest.raises(
+        ValueError,
+        match="set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching",
+    ):
+        config_cls.from_kwargs(enable_prefix_caching=True, **kwargs)
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+def test_prefix_caching_accepts_paged_scheduler(config_cls) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=True,
+    )
+
+    assert config.enable_prefix_caching is True
+    assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
+
+
+@pytest.mark.parametrize("build_config", [OmniDiffusionConfig, OmniDiffusionConfig.from_kwargs])
+def test_prefix_caching_rejects_sleep_mode(build_config) -> None:
+    with pytest.raises(ValueError, match="disable enable_prefix_caching or enable_sleep_mode"):
+        build_config(
+            diffusion_kv_mode="paged_scheduler",
+            diffusion_kv_max_rows_per_request=2,
+            enable_prefix_caching=True,
+            enable_sleep_mode=True,
+        )
+
+
+@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
+def test_sleep_mode_accepts_disabled_prefix_caching(mode) -> None:
+    config = OmniDiffusionConfig.from_kwargs(
+        diffusion_kv_mode=mode,
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=False,
+        enable_sleep_mode=True,
+    )
+
+    assert config.enable_sleep_mode is True
+    assert config.enable_prefix_caching is False
+
+
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+def test_native_kv_transfer_requires_prefix_caching_disabled(enable_prefix_caching) -> None:
+    kwargs = dict(
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=enable_prefix_caching,
+        kv_transfer_config={
+            "kv_connector": "MooncakeConnector",
+            "kv_role": "kv_consumer",
+            "engine_id": "test-diffusion",
+        },
+    )
+    if enable_prefix_caching:
+        with pytest.raises(ValueError, match="disable enable_prefix_caching for AR KV import"):
+            OmniDiffusionConfig.from_kwargs(**kwargs)
+    else:
+        config = OmniDiffusionConfig.from_kwargs(**kwargs)
+        assert config.kv_transfer_config.kv_role == "kv_consumer"
+
+
+@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
+@pytest.mark.parametrize("mode", ["dense_legacy", "paged_scheduler"])
+def test_disabled_prefix_caching_accepts_both_modes(config_cls, mode) -> None:
+    config = config_cls.from_kwargs(
+        diffusion_kv_mode=mode,
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=False,
+    )
+
+    assert config.enable_prefix_caching is False
+    assert config.diffusion_kv_mode.value == mode
+
+
 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")

tests/diffusion/diffusion_kv/test_initialization.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -225,6 +225,22 @@ def test_paged_config_forwards_gpu_memory_utilization_to_native_cache_config() -
     assert vllm_config.cache_config.gpu_memory_utilization == 0.42
 
 
+@pytest.mark.parametrize("enable_prefix_caching", [False, True])
+@pytest.mark.parametrize("requests,rows", [(1, 2), (2, 3)])
+def test_native_config_counts_cfg_rows_once(enable_prefix_caching, requests, rows) -> None:
+    od_config = _od_config(
+        max_num_seqs=requests,
+        diffusion_kv_max_rows_per_request=rows,
+        enable_prefix_caching=enable_prefix_caching,
+    )
+    config = diffusion_vllm_config.create_diffusion_vllm_config(torch.device("cpu"), od_config)
+    # Engine and Worker must see the same row capacity, including on reconfiguration.
+    diffusion_vllm_config.configure_diffusion_vllm_config(config, od_config)
+    assert config.scheduler_config.max_num_seqs == requests * rows
+    assert od_config.max_num_seqs == requests
+    assert config.cache_config.enable_prefix_caching is enable_prefix_caching
+
+
 def test_diffusion_vllm_model_config_supplies_dtype_for_quant_methods() -> None:
     quantization_config = build_quant_config(
         {

tests/diffusion/diffusion_kv/test_kv_cache_utils.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Native multimodal block hashing and diffusion-specific identity helpers."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import numpy as np
+import pytest
+import torch
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
+
+import vllm_omni.diffusion.diffusion_engine as engine_module
+import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
+import vllm_omni.diffusion.diffusion_kv.request as request_module
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
+from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
+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 _request(*, num_rows=1, seed=7, adapter_id=None, scale=1.0):
+    request = OmniDiffusionRequest(
+        request_id="fake", prompt="fake model", sampling_params=OmniDiffusionSamplingParams(seed=seed)
+    )
+    request.sampling_params.lora_request = None if adapter_id is None else SimpleNamespace(lora_int_id=adapter_id)
+    request.sampling_params.lora_scale = scale
+    request.diffusion_kv_requests = tuple(
+        DiffusionKVRequest(f"fake/{row}", sequence_id=row, prefix_len=8, target_len=4, seq_len=12)
+        for row in range(num_rows)
+    )
+    return request
+
+
+def _feature(*, value="image", start=4, end=8):
+    return MultiModalFeatureSpec(
+        data=None,
+        modality="image",
+        identifier=hash_prefix_cache_value((value, start, end)).hex(),
+        mm_position=PlaceholderRange(offset=start, length=end - start),
+    )
+
+
+def _prepare(request, feature=None):
+    namespace = get_cache_namespace("fake-model-v1", request.sampling_params)
+    for row in request.diffusion_kv_requests:
+        row.cache_token_ids = tuple(range(row.prefix_len))
+        row.mm_features = [feature or _feature()]
+        row.cache_namespace = namespace
+
+
+def _hashes(request, feature=None):
+    _prepare(request, feature)
+    row = request.diffusion_kv_requests[0]
+    row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+    return row.block_hashes
+
+
+def test_partial_hits_and_randomness_are_model_opt_in():
+    base = _hashes(_request())
+    assert base == _hashes(_request(seed=999))
+    changed = _hashes(_request(), _feature(value="other image"))
+    assert base[0] == changed[0]
+    assert base[1] != changed[1]
+    changed_geometry = _hashes(_request(), _feature(end=9))
+    assert base[0] == changed_geometry[0]
+    assert base[1] != changed_geometry[1]
+
+
+@pytest.mark.parametrize(("adapter_id", "scale"), [(None, 1.0), (2, 0.5), (1, 1.0)])
+def test_lora_is_common_request_wide_identity(adapter_id, scale):
+    base = _hashes(_request(adapter_id=1, scale=0.5))
+    changed = _hashes(_request(adapter_id=adapter_id, scale=scale))
+    assert all(a != b for a, b in zip(base, changed, strict=True))
+    assert _hashes(_request()) == _hashes(_request(scale=0.5))
+
+
+def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
+    row = DiffusionKVRequest(
+        "native",
+        sequence_id=0,
+        prefix_len=14,
+        target_len=2,
+        seq_len=16,
+        cache_token_ids=range(16),
+        mm_features=[
+            _feature(start=2, end=7),
+            _feature(start=7, end=10),
+        ],
+    )
+    upstream = Mock(wraps=generate_block_hash_extra_keys)
+    monkeypatch.setattr(request_module, "generate_block_hash_extra_keys", upstream)
+    hash_fn = get_hash_fn_by_name("sha256")
+    row.build_block_hashes(4, hash_fn)
+    assert [(call.args[1], call.args[2]) for call in upstream.call_args_list] == [(0, 4), (4, 8), (8, 12)]
+    # Independent expected native extras: a crossing item has a negative
+    # relative offset; a block can include the end of one item and the next.
+    first, second = (feature.identifier for feature in row.mm_features)
+    extras = [((first, 2),), ((first, -2), (second, 3)), ((second, -1),)]
+    parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
+    expected = []
+    for start, keys in zip(range(0, 12, 4), extras, strict=True):
+        parent = hash_block_tokens(hash_fn, parent, list(range(start, start + 4)), keys)
+        expected.append(parent)
+    assert row.block_hashes == expected
+    # Never hash the partial prefix tail or mutable target [14, 16), even
+    # though cache_token_ids deliberately contains the whole sequence.
+    assert len(row.block_hashes) == 3
+
+
+def test_same_identifier_at_different_offsets_changes_block_hash():
+    first = _feature(start=4, end=6)
+    shifted = _feature(start=5, end=7)
+    shifted.identifier = first.identifier
+    a = _hashes(_request(), first)
+    b = _hashes(_request(), shifted)
+    assert a[0] == b[0]
+    assert a[1] != b[1]
+
+
+def test_multimodal_identity_beyond_prefix_does_not_affect_hashes():
+    assert _hashes(_request(), _feature(value="a", start=8, end=12)) == _hashes(
+        _request(), _feature(value="b", start=8, end=12)
+    )
+
+
+@pytest.mark.parametrize(
+    "positions",
+    [
+        [(-1, 2)],
+        [(4, 4)],
+        [(4, 13)],
+        [(5, 7), (4, 5)],
+        [(3, 7), (4, 6)],
+    ],
+)
+def test_invalid_or_overlapping_positions_fail_before_hashing(positions):
+    row = _request().diffusion_kv_requests[0]
+    row.cache_token_ids = tuple(range(8))
+    row.mm_features = [_feature(start=start, end=end) for start, end in positions]
+    fail = Mock(side_effect=AssertionError("invalid positions must fail before hashing"))
+    with pytest.raises(ValueError, match="sorted, non-overlapping"):
+        row.build_block_hashes(4, fail)
+    fail.assert_not_called()
+    assert row.block_hashes == []
+
+
+def test_empty_multimodal_identifier_is_rejected():
+    row = _request().diffusion_kv_requests[0]
+    row.cache_token_ids = tuple(range(8))
+    feature = _feature()
+    feature.identifier = ""
+    row.mm_features = [feature]
+    with pytest.raises(ValueError, match="content identifiers"):
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.int64, torch.bool])
+def test_tensor_identity_includes_dtype_shape_and_values(dtype):
+    tensor = torch.tensor([[0, 1], [1, 0]], dtype=dtype)
+    assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.clone())
+    assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.t().contiguous().t())
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.reshape(4))
+    changed = tensor.clone()
+    changed[0, 0] = 1
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(changed)
+    assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.to(torch.float64))
+
+
+def test_numpy_scalar_shape_and_nested_dict_order():
+    assert hash_prefix_cache_value({"b": {"y": 2, "x": 1}, "a": None}) == hash_prefix_cache_value(
+        {"a": None, "b": {"x": 1, "y": 2}}
+    )
+    assert hash_prefix_cache_value(np.array(1)) != hash_prefix_cache_value(np.array([1]))
+    assert hash_prefix_cache_value(np.int64(1)) == hash_prefix_cache_value(1)
+
+
+@pytest.mark.parametrize(
+    ("first", "second"), [(1, True), (1, 1.0), ("1", b"1"), ([], {}), ([None], []), (["ab", "c"], ["a", "bc"])]
+)
+def test_identity_has_unambiguous_type_and_container_boundaries(first, second):
+    assert hash_prefix_cache_value(first) != hash_prefix_cache_value(second)
+
+
+@pytest.mark.parametrize("value", [object(), {1: "value"}, np.array([object()]), np.longdouble(1)])
+def test_identity_rejects_unsupported_values_without_pickle(value):
+    with pytest.raises(TypeError):
+        hash_prefix_cache_value(value)
+
+
+@pytest.mark.parametrize(
+    ("mode", "enabled", "expect_hook"),
+    [
+        (DiffusionKVCacheMode.PAGED_SCHEDULER, False, False),
+        (DiffusionKVCacheMode.PAGED_SCHEDULER, True, True),
+        (DiffusionKVCacheMode.DENSE_LEGACY, True, False),
+    ],
+)
+def test_engine_gates_hook_loading_and_admission(monkeypatch, mode, enabled, expect_hook):
+    events = []
+    request = _request()
+
+    def preprocess(req):
+        events.append("layout")
+        req.prepared_layout = "fake layout"
+        return req
+
+    def prepare(req):
+        assert req.prepared_layout == "fake layout"
+        assert req.diffusion_kv_requests[0].cache_token_ids == ()
+        events.append("cache_inputs")
+        _prepare(req)
+
+    hook = Mock(side_effect=prepare if expect_hook else AssertionError("disabled hook"))
+    load_hook = Mock(return_value=hook)
+    if not expect_hook:
+        load_hook.side_effect = AssertionError("disabled hook factory")
+        monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", Mock(side_effect=AssertionError("disabled hash")))
+    monkeypatch.setattr(engine_module, "get_diffusion_pre_process_func", lambda _config: preprocess)
+    monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
+    monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", load_hook)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(diffusion_kv_mode=mode, enable_prefix_caching=enabled)
+    engine._init_process_hooks(engine.od_config)
+    engine.prefix_cache_func = hook
+    assert engine._prepare_request_for_admission(request) is request
+    row = request.diffusion_kv_requests[0]
+    if expect_hook:
+        assert events == ["layout", "cache_inputs"]
+        assert row.cache_token_ids == tuple(range(8))
+        load_hook.assert_called_once_with(engine.od_config)
+        hook.assert_called_once_with(request)
+    else:
+        assert events == ["layout"]
+        assert row.cache_token_ids == ()
+        assert row.mm_features == []
+        load_hook.assert_not_called()
+        hook.assert_not_called()
+
+
+@pytest.mark.parametrize(
+    "model_class_name,load_format",
+    [(None, "default"), ("UnregisteredPipeline", "default"), ("HunyuanImage3ForCausalMM", "diffusers")],
+)
+def test_engine_rejects_enabled_prefix_caching_without_hook(monkeypatch, model_class_name, load_format):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+
+    monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.from_kwargs(
+        model_class_name=model_class_name,
+        diffusion_load_format=load_format,
+        diffusion_kv_mode="paged_scheduler",
+        diffusion_kv_max_rows_per_request=2,
+        enable_prefix_caching=True,
+    )
+    with pytest.raises(ValueError, match="requires a registered prefix-cache hook"):
+        engine._init_process_hooks(engine.od_config)
+
+
+def test_engine_rejects_out_of_profile_request_before_hashing():
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, enable_prefix_caching=True
+    )
+    engine._diffusion_kv_profile_limits = (1, 8, 4)
+    engine.prefix_cache_func = Mock(side_effect=AssertionError("invalid request must fail before identity work"))
+    with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
+        engine._prepare_request_for_admission(_request())
+    engine.prefix_cache_func.assert_not_called()

tests/diffusion/diffusion_kv/test_manager.py

@@ -30,26 +30,159 @@ def _config(num_blocks: int) -> KVCacheConfig:
     )
 
 
-def _request(public_id: str, sequence_id: int, *, seq_len: int = 8, kv_contexts=()) -> DiffusionKVRequest:
+def _request(
+    public_id: str,
+    sequence_id: int,
+    *,
+    prefix_len: int = 4,
+    target_len: int = 4,
+    seq_len: int = 8,
+    cache_token_ids=(),
+    mm_features=(),
+    kv_contexts=(),
+) -> DiffusionKVRequest:
     return DiffusionKVRequest(
         f"{public_id}/diffusion-kv/{sequence_id}",
         sequence_id=sequence_id,
-        prefix_len=4,
-        target_len=4,
+        prefix_len=prefix_len,
+        target_len=target_len,
         seq_len=seq_len,
+        cache_token_ids=cache_token_ids,
+        mm_features=mm_features,
         kv_contexts=kv_contexts,
     )
 
 
-def _manager(num_blocks: int, *, max_model_len: int = 64) -> DiffusionKVCacheManager:
+def _manager(
+    num_blocks: int,
+    *,
+    max_model_len: int = 64,
+    enable_prefix_caching: bool = False,
+) -> DiffusionKVCacheManager:
     return DiffusionKVCacheManager(
         _config(num_blocks),
         max_model_len=max_model_len,
         scheduler_block_size=BLOCK_SIZE,
         hash_block_size=BLOCK_SIZE,
+        enable_prefix_caching=enable_prefix_caching,
     )
 
 
+def test_successful_request_publishes_prefix_for_a_warm_hit() -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    cold = _request("cold", 0, cache_token_ids=range(4))
+
+    cold_metadata = manager.reserve_request("cold", (cold,))
+    assert cold_metadata is not None
+    assert cold_metadata.sequences[0].cached_prefix_len == 0
+    cold_prefix_block = cold_metadata.sequences[0].block_ids[0][0]
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request("warm", 0, cache_token_ids=range(4))
+    warm_metadata = manager.reserve_request("warm", (warm,))
+
+    assert warm_metadata is not None
+    assert warm_metadata.sequences[0].cached_prefix_len == 4
+    assert warm_metadata.sequences[0].block_ids[0][0] == cold_prefix_block
+    # Metadata contains both the shared hit and the newly allocated target.
+    assert len(warm_metadata.sequences[0].block_ids[0]) == 2
+
+    manager.free_request("warm")
+    assert warm.num_computed_tokens == 0
+    assert warm.shared_prefix_boundary == 0
+
+
+def test_prefix_hash_root_is_stable_across_manager_instances() -> None:
+    first_manager = _manager(8, enable_prefix_caching=True)
+    second_manager = _manager(8, enable_prefix_caching=True)
+    first = _request("first", 0, cache_token_ids=range(4))
+    second = _request("second", 0, cache_token_ids=range(4))
+
+    first_manager._prepare_block_hashes((first,))
+    second_manager._prepare_block_hashes((second,))
+
+    assert first.block_hashes == second.block_hashes
+
+
+def test_warm_hit_fits_at_the_same_capacity_as_the_cold_request() -> None:
+    manager = _manager(3, enable_prefix_caching=True)
+    cold = _request("cold", 0, cache_token_ids=range(4))
+    assert manager.reserve_request("cold", (cold,)) is not None
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request("warm", 0, cache_token_ids=range(4))
+    metadata = manager.reserve_request("warm", (warm,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 4
+
+
+def test_unpublished_request_does_not_become_a_prefix_hit() -> None:
+    manager = _manager(8, enable_prefix_caching=True)
+    failed = _request("failed", 0, cache_token_ids=range(4))
+    assert manager.reserve_request("failed", (failed,)) is not None
+    manager.free_request("failed")
+
+    retry = _request("retry", 0, cache_token_ids=range(4))
+    metadata = manager.reserve_request("retry", (retry,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 0
+
+
+def test_partial_prefix_hit_recomputes_from_the_first_changed_block() -> None:
+    manager = _manager(10, enable_prefix_caching=True)
+    cold = _request(
+        "cold",
+        0,
+        prefix_len=8,
+        seq_len=12,
+        cache_token_ids=range(8),
+    )
+    assert manager.reserve_request("cold", (cold,)) is not None
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = _request(
+        "warm",
+        0,
+        prefix_len=8,
+        seq_len=12,
+        cache_token_ids=(*range(4), 100, 101, 102, 103),
+    )
+    metadata = manager.reserve_request("warm", (warm,))
+
+    assert metadata is not None
+    assert metadata.sequences[0].cached_prefix_len == 4
+    assert len(metadata.sequences[0].block_ids[0]) == 3
+
+
+def test_cfg_rows_use_the_minimum_warm_hit_as_one_execution_boundary() -> None:
+    manager = _manager(16, enable_prefix_caching=True)
+    for public_id, tokens in (("seed-a", range(8)), ("seed-b", (*range(4), 20, 21, 22, 23))):
+        request = _request(public_id, 0, prefix_len=8, seq_len=12, cache_token_ids=tokens)
+        assert manager.reserve_request(public_id, (request,)) is not None
+        manager.publish_request(public_id)
+        manager.free_request(public_id)
+
+    requests = (
+        _request("warm", 0, prefix_len=8, seq_len=12, cache_token_ids=range(8)),
+        _request(
+            "warm",
+            1,
+            prefix_len=8,
+            seq_len=12,
+            cache_token_ids=(*range(4), 30, 31, 32, 33),
+        ),
+    )
+    metadata = manager.reserve_request("warm", requests)
+
+    assert metadata is not None
+    assert [sequence.cached_prefix_len for sequence in metadata.sequences] == [4, 4]
+
+
 def test_reserve_and_free_multi_cfg_request() -> None:
     manager = _manager(8)
     free_before = manager.native_manager.block_pool.get_num_free_blocks()

tests/diffusion/diffusion_kv/test_native_prefetch.py

@@ -256,7 +256,11 @@ def _scheduler(monkeypatch, num_blocks=24):
         kv_cache_config=config,
         scheduler_block_size=4,
         hash_block_size=4,
-        kv_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64), max_in_flight_tokens=64),
+        kv_vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(max_model_len=64),
+            max_in_flight_tokens=64,
+            cache_config=SimpleNamespace(enable_prefix_caching=False),
+        ),
     )
     return scheduler, connector
 

tests/diffusion/diffusion_kv/test_paged_attention_adapter.py

@@ -113,7 +113,7 @@ class _FakeLayer:
         self.head_size_v = 4
         self.spec = _FakeSpec(non_causal=non_causal)
         self.updates: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
-        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, object]] = []
+        self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor, _FakeNativeMetadata]] = []
         self.native_events: list[str] = []
         self.layer_name = "layer-0"
         self.kv_cache = object()
@@ -122,6 +122,8 @@ class _FakeLayer:
 
 
 class _FakeNativeImpl:
+    vllm_flash_attn_version: int
+
     def __init__(self, layer: _FakeLayer) -> None:
         self.layer = layer
 
@@ -143,7 +145,7 @@ class _FakeNativeImpl:
         key: torch.Tensor,
         value: torch.Tensor,
         _kv_cache,
-        metadata: object,
+        metadata: _FakeNativeMetadata,
         output: torch.Tensor,
     ) -> torch.Tensor:
         self.layer.native_events.append("forward")
@@ -202,7 +204,7 @@ def _make_adapter(
 
     config = SimpleNamespace(
         name="vllm-config",
-        model_config=SimpleNamespace(dtype=torch.float32),
+        model_config=SimpleNamespace(dtype=torch.float32, rswa_window=None),
     )
     adapter = DiffusionPagedAttentionAdapter(
         vllm_config=config,
@@ -1222,9 +1224,15 @@ def test_row_contract_rejects_invalid_identity_or_span(kwargs: dict, message: st
         DiffusionPagedAttentionRow(**values)
 
 
+@pytest.mark.parametrize("pcp_size", [1, 2])
+@pytest.mark.parametrize("selection_fails", [False, True])
 def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads(
     monkeypatch: pytest.MonkeyPatch,
+    pcp_size: int,
+    selection_fails: bool,
 ) -> None:
+    from vllm.v1.attention import selector
+
     selected_backends = []
     specialized_backends = []
     impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
@@ -1237,28 +1245,39 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
 
     original_backend_per_kind = {"full": object()}
     config = SimpleNamespace(
-        attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
-        parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
-        model_config=SimpleNamespace(dtype=torch.float16),
-        cache_config=SimpleNamespace(cache_dtype="auto"),
+        attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind, use_non_causal=True),
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=pcp_size, decode_context_parallel_size=1),
+        model_config=SimpleNamespace(dtype=torch.float16, rswa_window=None),
+        cache_config=SimpleNamespace(cache_dtype="auto", user_specified_block_size=False),
+        kv_transfer_config=None,
+        speculative_config=None,
     )
 
-    def select_backend(**_kwargs):
-        selected_backends.append(
-            (
-                config.attention_config.backend,
-                config.attention_config.backend_per_kind,
-                config.parallel_config.prefill_context_parallel_size,
-            )
-        )
+    def select_backend(*, backend, attn_selector_config, num_heads):
+        # Exercise the real 0.29 selector, replacing only platform resolution.
+        # Ulysses has already gathered tokens and sharded heads before the
+        # native kernel; the MoE PCP mapping must not request PCP attention.
+        assert not attn_selector_config.use_pcp
+        assert not attn_selector_config.use_dcp
+        assert attn_selector_config.use_non_causal
+        assert num_heads == 4
+        selected_backends.append((backend, config.attention_config.backend_per_kind))
+        if selection_fails:
+            raise ValueError("test backend unavailable")
         return native_backend
 
-    monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
+    monkeypatch.setattr("vllm.config.get_current_vllm_config", lambda: config)
+    monkeypatch.setattr(selector, "_cached_get_attn_backend", select_backend)
     monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
+
+    def specialize_backend(backend, *, ulysses_degree):
+        specialized_backends.append((backend, ulysses_degree))
+        return backend
+
     monkeypatch.setattr(
         adapter_module.current_omni_platform,
         "get_diffusion_paged_kv_attn_backend",
-        lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
+        specialize_backend,
     )
     layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
     spec = FullAttentionSpec(
@@ -1269,7 +1288,7 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         non_causal=True,
     )
 
-    native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(
+    kwargs = dict(
         layer_name="layer-0",
         layer=layer,
         spec=spec,
@@ -1277,12 +1296,21 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
         device=torch.device("cpu"),
         ulysses_degree=2,
     )
+    if selection_fails:
+        with pytest.raises(ValueError, match="test backend unavailable"):
+            adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
+    else:
+        native_layer = adapter_module.DiffusionPagedAttentionLayerAdapter(**kwargs)
 
-    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 1)]
-    assert specialized_backends == [(native_backend, 2)]
+    assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
     assert config.attention_config.backend is None
     assert config.attention_config.backend_per_kind is original_backend_per_kind
-    assert config.parallel_config.prefill_context_parallel_size == 2
+    assert config.parallel_config.prefill_context_parallel_size == pcp_size
+    if selection_fails:
+        assert not specialized_backends
+        impl_cls.assert_not_called()
+        return
+    assert specialized_backends == [(native_backend, 2)]
     assert native_layer.num_heads == 4
     assert native_layer.num_kv_heads == 2
     assert native_layer.spec.num_kv_heads == 2
@@ -1432,7 +1460,12 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
     layer._no_parallel_strategy = object()
     layer._get_active_parallel_strategy = lambda: Strategy()
     layer._with_kv_cache_dtype = lambda metadata: metadata
-    layer._run_local_attention = lambda query, _key, _value, _metadata: events.append("dense") or query
+
+    def run_local_attention(query, _key, _value, _metadata):
+        events.append("dense")
+        return query
+
+    layer._run_local_attention = run_local_attention
     qkv = torch.zeros(1, 2, 2, 4)
 
     assert not layer.is_paged_kv_active()

tests/diffusion/diffusion_kv/test_request.py

@@ -4,8 +4,10 @@
 import pytest
 import torch
 import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
+from vllm.utils.hashing import get_hash_fn_by_name
 from vllm.v1.core.kv_cache_manager import KVCacheManager
-from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.v1.core.kv_cache_utils import BlockHash, init_none_hash
 from vllm.v1.kv_cache_interface import (
     FullAttentionSpec,
     KVCacheConfig,
@@ -49,6 +51,51 @@ def test_request_exposes_native_request_and_diffusion_semantics() -> None:
     assert request.status is RequestStatus.WAITING
 
 
+def test_request_builds_only_full_prefix_block_hashes() -> None:
+    hash_function = get_hash_fn_by_name("sha256")
+    init_none_hash(hash_function)
+    request = _request(
+        prefix_len=6,
+        target_len=2,
+        cache_token_ids=range(6),
+        mm_features=[
+            MultiModalFeatureSpec(
+                data=None,
+                modality="image",
+                identifier="tail",
+                mm_position=PlaceholderRange(offset=4, length=2),
+            )
+        ],
+    )
+
+    request.build_block_hashes(BLOCK_SIZE, hash_function)
+
+    assert len(request.block_hashes) == 1
+    assert request.skip_reading_prefix_cache is False
+
+
+def test_changed_multimodal_identity_changes_the_block_hash() -> None:
+    hash_function = get_hash_fn_by_name("sha256")
+    init_none_hash(hash_function)
+    first = _request(cache_token_ids=range(4))
+    second = _request(
+        cache_token_ids=range(4),
+        mm_features=[
+            MultiModalFeatureSpec(
+                data=None,
+                modality="image",
+                identifier="image-a",
+                mm_position=PlaceholderRange(offset=2, length=2),
+            )
+        ],
+    )
+
+    first.build_block_hashes(BLOCK_SIZE, hash_function)
+    second.build_block_hashes(BLOCK_SIZE, hash_function)
+
+    assert first.block_hashes != second.block_hashes
+
+
 @pytest.mark.parametrize(
     ("overrides", "message"),
     [

tests/diffusion/diffusion_kv/test_worker_contract.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
 from unittest.mock import Mock
@@ -111,6 +111,29 @@ def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> Non
     ] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
 
 
+def test_runner_builds_suffix_prefill_row_for_a_prefix_hit() -> None:
+    runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
+    metadata = DiffusionKVMetadata(
+        request_id="req-0",
+        allocation_generation=1,
+        sequences=(
+            DiffusionKVSequenceMetadata(
+                sequence_id=0,
+                prefix_len=8,
+                target_len=4,
+                seq_len=12,
+                block_ids=([1, 2, 3],),
+                cached_prefix_len=4,
+            ),
+        ),
+    )
+
+    attn_metadata = runner._build_paged_attention_metadata([metadata])
+
+    row = attn_metadata.prefill_rows[0]
+    assert (row.kv_start_pos, row.query_len, row.seq_len) == (4, 8, 12)
+
+
 def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
     runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
     runner.od_config.parallel_config.cfg_parallel_size = 2

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py

@@ -2,13 +2,24 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
+import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers
 from transformers.utils.generic import ModelOutput
+from vllm.lora.request import LoRARequest
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec
 
+import vllm_omni.diffusion.diffusion_engine as engine_module
+import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
 import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module
+import vllm_omni.diffusion.models.hunyuan_image3.request_layout as layout_module
+from tests.helpers.kv_layout import build_kv_cache_tensor
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
+from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVCacheManager
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
@@ -27,6 +38,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
     hunyuan_num_special_tokens,
     normalize_hunyuan_cot_text,
     prepare_hunyuan_layout,
+    prepare_hunyuan_prefix_cache,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -68,18 +80,26 @@ class _FakeTokenizerWrapper:
         seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens]
         padded_seq_len = max(seq_lens)
         positions = torch.arange(padded_seq_len)
-        joint_image_slices = [[slice(1, 3), slice(3, 5)] if has_cond_image else [] for _ in range(rows)]
+        # Like the real tokenizer, joint spans enclose VAE + separator + ViT,
+        # while all_image_slices contains the individual RoPE spans. Leave one
+        # full text block before the image so partial reuse can be tested.
+        joint_image_slices = [[slice(5, 10)] if has_cond_image else [] for _ in range(rows)]
+        cond_vae_image_slices = [[slice(5, 7)] if has_cond_image else [] for _ in range(rows)]
+        cond_vit_image_slices = [[slice(8, 10)] if has_cond_image else [] for _ in range(rows)]
         gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens]
-        all_image_slices = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)]
+        all_image_slices = [
+            vae + vit + generated
+            for vae, vit, generated in zip(cond_vae_image_slices, cond_vit_image_slices, gen_image_slices)
+        ]
         cond_vae_image_mask = (
-            torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None
+            torch.stack([(positions >= 5) & (positions < 7) for _ in range(rows)]) if has_cond_image else None
         )
         cond_vit_image_mask = (
-            torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None
+            torch.stack([(positions >= 8) & (positions < 10) for _ in range(rows)]) if has_cond_image else None
         )
         sections = []
         for _ in range(rows):
-            row_sections = []
+            row_sections: list[dict[str, object]] = []
             if has_cond_image:
                 row_sections.append(
                     {
@@ -101,6 +121,8 @@ class _FakeTokenizerWrapper:
                 ),
                 cond_vae_image_mask=cond_vae_image_mask,
                 cond_vit_image_mask=cond_vit_image_mask,
+                cond_vae_image_slices=cond_vae_image_slices,
+                cond_vit_image_slices=cond_vit_image_slices,
                 joint_image_slices=joint_image_slices,
                 gen_image_slices=gen_image_slices,
                 real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1),
@@ -155,6 +177,7 @@ def _request(
     guidance_scale: float,
     prompt="draw a cat",
     request_id: str = "req",
+    seed: int | None = None,
 ) -> OmniDiffusionRequest:
     return OmniDiffusionRequest(
         prompt=prompt,
@@ -163,6 +186,7 @@ def _request(
             width=1024,
             guidance_scale=guidance_scale,
             num_inference_steps=4,
+            seed=seed,
         ),
         request_id=request_id,
     )
@@ -205,6 +229,8 @@ def test_builds_kv_request_lengths_without_model_execution() -> None:
     assert tokenizer.calls[0]["sequence_template"] == "instruct"
     assert tokenizer.calls[0]["cfg_factor"] == 1
     assert kv_requests[0].block_hashes == []
+    assert kv_requests[0].cache_token_ids == ()
+    assert kv_requests[0].mm_features == []
     assert kv_requests[0].kv_contexts == ()
     assert kv_requests[0].skip_reading_prefix_cache is True
 
@@ -226,12 +252,14 @@ def _reference_image() -> JointImageInfo:
     return JointImageInfo(
         vae_image_info=ImageInfo(
             image_type="vae",
+            image_tensor=torch.arange(12, dtype=torch.float32).reshape(1, 3, 2, 2),
             token_width=8,
             token_height=8,
             image_token_length=64,
         ),
         vision_image_info=ImageInfo(
             image_type="siglip2",
+            image_tensor=torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2),
             token_width=4,
             token_height=4,
             image_token_length=16,
@@ -243,6 +271,342 @@ def _reference_image() -> JointImageInfo:
     )
 
 
+def _prepare_cached_requests(request, tokenizer, image_processor):
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    prepare_hunyuan_prefix_cache(request)
+    return request.diffusion_kv_requests
+
+
+def test_reference_prefix_identity_includes_image_content_and_vae_seed() -> None:
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id: str, image: JointImageInfo, seed: int):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            seed=seed,
+            prompt={
+                "prompt": "edit this image",
+                "additional_information": {"batch_cond_image_info": [image]},
+            },
+        )
+        return _prepare_cached_requests(request, tokenizer, image_processor)[0]
+
+    first = build("first", _reference_image(), 7)
+    same = build("same", _reference_image(), 7)
+    changed_seed = build("changed-seed", _reference_image(), 8)
+    changed_image_value = _reference_image()
+    changed_image_value.vae_image_info.image_tensor[0, 0, 0, 0] = 99
+    changed_image = build("changed-image", changed_image_value, 7)
+
+    assert first.mm_features == same.mm_features
+    assert [(feature.mm_position.offset, feature.mm_position.length) for feature in first.mm_features] == [
+        (5, 2),
+        (8, 2),
+    ]
+    for changed in (changed_seed, changed_image):
+        # Both subspans must carry the identity, not just the first one.
+        assert all(a.identifier != b.identifier for a, b in zip(first.mm_features, changed.mm_features, strict=True))
+
+
+def test_reference_prefix_rejects_missing_inputs_with_nested_spans() -> None:
+    tokenizer, image_processor = _components([20])
+    request = _request(
+        guidance_scale=1.0,
+        prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    request.prompt["additional_information"].pop("batch_cond_image_info")
+
+    with pytest.raises(ValueError, match="missing its canonical inputs"):
+        prepare_hunyuan_prefix_cache(request)
+
+
+def test_reference_content_is_hashed_once_across_cfg_rows(monkeypatch):
+    tokenizer, image_processor = _components([20, 20])
+    image = _reference_image()
+    hasher = Mock(wraps=layout_module.hash_prefix_cache_value)
+    monkeypatch.setattr(layout_module, "hash_prefix_cache_value", hasher)
+
+    def build(request_id):
+        request = _request(
+            guidance_scale=5.0,
+            request_id=request_id,
+            seed=7,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [image]}},
+        )
+        return _prepare_cached_requests(request, tokenizer, image_processor)
+
+    first = build("first")
+    assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 1
+    assert len(first[0].mm_features) == len(first[1].mm_features) == 2
+    assert first[0].mm_features == first[1].mm_features
+    image.vae_image_info.image_tensor[0, 0, 0, 0] += 1
+    changed = build("changed")
+    assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 2
+    assert all(a.identifier != b.identifier for a, b in zip(first[0].mm_features, changed[0].mm_features, strict=True))
+
+
+def test_failed_cfg_cache_preparation_does_not_install_partial_inputs():
+    tokenizer, image_processor = _components([20, 20])
+    request = _request(
+        guidance_scale=5.0,
+        prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    layout = _prepare(request, tokenizer, image_processor)
+    request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+    # First row prepares successfully; the second exposes inconsistent inputs.
+    layout.rope_image_info[0] = []
+    request.prompt["additional_information"].pop("batch_cond_image_info")
+    with pytest.raises(ValueError, match="missing its canonical inputs"):
+        prepare_hunyuan_prefix_cache(request)
+    assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests)
+
+
+@pytest.mark.parametrize("generator_list", [False, True])
+def test_reference_identity_uses_explicit_generator_state_before_seed(generator_list):
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id, seed, advance):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            seed=seed,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+        )
+        generator = torch.Generator(device="cpu").manual_seed(123)
+        if advance:
+            torch.rand(1, generator=generator)
+        request.sampling_params.generator = [generator] if generator_list else generator
+        before = generator.get_state().clone()
+        row = _prepare_cached_requests(request, tokenizer, image_processor)[0]
+        torch.testing.assert_close(before, generator.get_state())
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return row.block_hashes
+
+    first = build("first", 7, False)
+    assert first == build("same-generator-different-unused-seed", 8, False)
+    changed = build("advanced-generator-same-seed", 7, True)
+    assert first[:1] == changed[:1]
+    assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True))
+
+
+def test_unseeded_reference_identity_is_request_local():
+    tokenizer, image_processor = _components([20])
+
+    def build(request_id):
+        request = _request(
+            guidance_scale=1.0,
+            request_id=request_id,
+            prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+        )
+        # Real requests normally receive an automatic seed in __post_init__.
+        request.sampling_params.seed = None
+        row = _prepare_cached_requests(request, tokenizer, image_processor)[0]
+        row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return row.block_hashes
+
+    first = build("first")
+    assert first == build("first")
+    changed = build("second")
+    assert first[:1] == changed[:1]
+    assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True))
+
+
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+def test_text_prefix_hashes_include_lora_identity_and_scale(guidance_scale: float) -> None:
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+
+    def build(request_id: str, adapter_id: int | None, scale: float):
+        request = _request(guidance_scale=guidance_scale, request_id=request_id)
+        if adapter_id is not None:
+            request.sampling_params.lora_request = LoRARequest(
+                lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}"
+            )
+        request.sampling_params.lora_scale = scale
+        kv_requests = _prepare_cached_requests(request, tokenizer, image_processor)
+        for kv_request in kv_requests:
+            kv_request.build_block_hashes(4, get_hash_fn_by_name("sha256"))
+        return kv_requests
+
+    base = build("base", None, 1.0)
+    same_base = build("same-base", None, 0.5)  # Scale has no effect without an adapter.
+    first = build("first", 1, 0.5)
+    same = build("same", 1, 0.5)
+    changed_id = build("changed-id", 2, 0.5)
+    changed_scale = build("changed-scale", 1, 1.0)
+
+    for row in range(len(first)):
+        assert base[row].block_hashes == same_base[row].block_hashes
+        assert first[row].block_hashes == same[row].block_hashes
+        assert len(first[row].block_hashes) == 5
+        for changed in (base, changed_id, changed_scale):
+            assert first[row].cache_token_ids == changed[row].cache_token_ids
+            # LoRA changes even the text prefix, so every block must differ.
+            assert all(a != b for a, b in zip(first[row].block_hashes, changed[row].block_hashes, strict=True))
+
+
+@pytest.fixture
+def prefix_cache_manager(request):
+    native_kv_managers.register_all_kvcache_specs(None)
+    spec = FullAttentionSpec(block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16)
+    num_blocks = 64
+    config = KVCacheConfig(
+        num_blocks=num_blocks,
+        kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])],
+        kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)],
+    )
+    manager = DiffusionKVCacheManager(
+        config,
+        max_model_len=64,
+        scheduler_block_size=4,
+        hash_block_size=4,
+        enable_prefix_caching=getattr(request, "param", True),
+    )
+    yield manager
+    manager.close()
+
+
+@pytest.mark.parametrize("prefix_cache_manager", [False], indirect=True)
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+def test_paged_cache_off_prepares_and_allocates_without_identity_work(
+    monkeypatch, prefix_cache_manager, guidance_scale
+):
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+    config = SimpleNamespace(
+        model="fake",
+        model_class_name="HunyuanImage3ForCausalMM",
+        diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
+        enable_prefix_caching=False,
+    )
+    monkeypatch.setattr(
+        pipeline_module,
+        "get_config",
+        lambda *_args, **_kwargs: SimpleNamespace(
+            vae_downsample_factor=(8, 8),
+            patch_size=2,
+            image_base_size=1024,
+        ),
+    )
+    monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor)
+    monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer)
+    monkeypatch.setattr(pipeline_module.GenerationConfig, "from_pretrained", lambda _model: SimpleNamespace())
+    monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
+    fail = Mock(side_effect=AssertionError("disabled prefix cache must not prepare identity"))
+    monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", fail)
+    monkeypatch.setattr(layout_module, "prepare_hunyuan_prefix_cache", fail)
+    monkeypatch.setattr(layout_module, "hash_prefix_cache_value", fail)
+    monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", fail)
+    monkeypatch.setattr(kv_utils.MultiModalHasher, "hash_kwargs", fail)
+    monkeypatch.setattr(torch.Tensor, "tolist", fail)
+    monkeypatch.setattr(prefix_cache_manager, "_hash_function", fail)
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.od_config = config
+    engine._init_process_hooks(config)
+    # Guard admission too, even if someone injects a hook after initialization.
+    engine.prefix_cache_func = fail
+    request = _request(
+        guidance_scale=guidance_scale,
+        prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}},
+    )
+    request.sampling_params.seed = None
+    request.sampling_params.generator = SimpleNamespace(get_state=fail)
+    engine._prepare_request_for_admission(request)
+    assert isinstance(request.prepared_layout, HunyuanPreparedLayout)
+    assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests)
+    metadata = prefix_cache_manager.reserve_request(request.request_id, request.diffusion_kv_requests)
+    assert metadata is not None
+    assert all(row.cached_prefix_len == 0 for row in metadata.sequences)
+    prefix_cache_manager.publish_request(request.request_id)
+    prefix_cache_manager.free_request(request.request_id)
+    fail.assert_not_called()
+
+
+@pytest.mark.parametrize("guidance_scale", [1.0, 5.0])
+@pytest.mark.parametrize(
+    ("change", "expected_hit"),
+    [
+        ("same", 20),
+        ("prompt_suffix", 12),
+        ("target", 20),
+        ("vae", 4),
+        ("vit", 4),
+        ("vision_kwargs", 4),
+        ("rope", 4),
+        ("seed", 4),
+        ("lora_id", 0),
+        ("lora_scale", 0),
+        ("base", 0),
+    ],
+)
+def test_prefix_lookup_respects_reference_and_lora_identity(
+    prefix_cache_manager: DiffusionKVCacheManager, guidance_scale: float, change: str, expected_hit: int
+) -> None:
+    tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0)))
+
+    def build(request_id: str, change: str):
+        image = _reference_image()
+        if change == "vae":
+            image.vae_image_info.image_tensor[0, 0, 0, 0] = 99
+        elif change == "vit":
+            image.vision_image_info.image_tensor[0, 0, 0, 0] = 99
+        elif change == "vision_kwargs":
+            image.vision_encoder_kwargs["pixel_attention_mask"][0] = False
+        request = _request(
+            guidance_scale=guidance_scale,
+            request_id=request_id,
+            seed=8 if change == "seed" else 7,
+            prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [image]}},
+        )
+        if change != "base":
+            adapter_id = 2 if change == "lora_id" else 1
+            request.sampling_params.lora_request = LoRARequest(
+                lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}"
+            )
+        request.sampling_params.lora_scale = 1.0 if change == "lora_scale" else 0.5
+        layout = _prepare(request, tokenizer, image_processor)
+        if change == "rope":
+            for spans in layout.rope_image_info:
+                spans[0] = (slice(5, 7), (2, 1))
+        elif change == "prompt_suffix":
+            layout.tokenizer_output.tokens[:, 12] += 100
+        elif change == "target":
+            layout.tokenizer_output.tokens[:, 20:] += 100
+        request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout)
+        prepare_hunyuan_prefix_cache(request)
+        return request.diffusion_kv_requests
+
+    manager = prefix_cache_manager
+    cold = build("cold", "same")
+    cold_metadata = manager.reserve_request("cold", cold)
+    assert cold_metadata is not None
+    assert all(sequence.cached_prefix_len == 0 for sequence in cold_metadata.sequences)
+    # CPU-only: simulate successful materialization to test native cache lookup.
+    manager.publish_request("cold")
+    manager.free_request("cold")
+
+    warm = build("warm", change)
+    warm_metadata = manager.reserve_request("warm", warm)
+    assert warm_metadata is not None
+    for cold_row, warm_row, cold_sequence, warm_sequence in zip(
+        cold, warm, cold_metadata.sequences, warm_metadata.sequences, strict=True
+    ):
+        if change == "prompt_suffix":
+            assert cold_row.cache_token_ids[:12] == warm_row.cache_token_ids[:12]
+            assert cold_row.cache_token_ids[12] != warm_row.cache_token_ids[12]
+        else:
+            assert cold_row.cache_token_ids == warm_row.cache_token_ids
+        assert warm_sequence.cached_prefix_len == expected_hit
+        hit_blocks = expected_hit // 4
+        assert cold_sequence.block_ids[0][:hit_blocks] == warm_sequence.block_ids[0][:hit_blocks]
+        assert cold_row.block_hashes[:hit_blocks] == warm_row.block_hashes[:hit_blocks]
+        assert all(
+            a != b for a, b in zip(cold_row.block_hashes[hit_blocks:], warm_row.block_hashes[hit_blocks:], strict=True)
+        )
+
+
 def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None:
     tokenizer, image_processor = _components([20])
     joint_image = _reference_image()
@@ -366,13 +730,17 @@ def test_prepared_model_inputs_match_local_tokenization(
         # Preparation runs without a forward context. Coverage must come from
         # the scheduler metadata copied onto this request, including every CFG row.
         reused = pipeline.prepare_model_inputs(
-            **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(5,) * len(prefix_lens)
+            # The image span is [5, 10); native KV must cover its end, not
+            # merely reach the first image token.
+            **common_kwargs,
+            prepared_layout=prepared_layout,
+            kv_computed_tokens=(10,) * len(prefix_lens),
         )
         assert reused["cond_vae_images"] is None and reused["cond_vit_images"] is None
         partial = pipeline.prepare_model_inputs(
             **common_kwargs,
             prepared_layout=prepared_layout,
-            kv_computed_tokens=(5,) * (len(prefix_lens) - 1) + (4,),
+            kv_computed_tokens=(10,) * (len(prefix_lens) - 1) + (9,),
         )
         assert partial["cond_vae_images"] is not None and partial["cond_vit_images"] is not None
 
@@ -496,6 +864,9 @@ def test_cfg_rows_preserve_reusable_token_ids_and_native_transfer_boundary(trans
 
     tokenizer, image_processor = _components([12, 14])
     request = _request(guidance_scale=5.0)
+    # Prompt token IDs are only materialized for the native AR -> DiT
+    # transfer path; local prefix caching does not need them.
+    request.kv_transfer_params = {"num_transfer_tokens": 0}
     layout = _prepare(request, tokenizer, image_processor)
     layout.tokenizer_output.think_recaption_end_pos = [[7], [5]]
     layout.tokenizer_output.uncond_cfg_start_pos = [[None], [5]]

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -147,6 +147,80 @@ def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch):
     torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens)
 
 
+def test_slice_cached_prefix_inputs_keeps_full_kv_axis() -> None:
+    inputs_embeds = torch.arange(12, dtype=torch.float32).reshape(1, 6, 2)
+    attention_mask = torch.arange(36).reshape(1, 1, 6, 6)
+    position_ids = torch.arange(6).reshape(1, 6)
+    custom_pos_emb = (torch.arange(6).reshape(1, 6), torch.arange(10, 16).reshape(1, 6))
+    image_mask = torch.tensor([[False, False, True, True, True, True]])
+    gen_timestep_scatter_index = torch.tensor([[2]])
+
+    sliced = HunyuanImage3Pipeline._slice_cached_prefix_inputs(
+        inputs_embeds,
+        attention_mask,
+        position_ids,
+        custom_pos_emb,
+        image_mask,
+        gen_timestep_scatter_index,
+        [6],
+        [6],
+        2,
+    )
+
+    (
+        sliced_embeds,
+        sliced_mask,
+        sliced_positions,
+        sliced_rope,
+        sliced_image_mask,
+        sliced_scatter,
+        sliced_query_lens,
+        sliced_seq_len,
+    ) = sliced
+    torch.testing.assert_close(sliced_embeds, inputs_embeds[:, 2:])
+    torch.testing.assert_close(sliced_mask, attention_mask[:, :, 2:, :])
+    torch.testing.assert_close(sliced_positions, position_ids[:, 2:])
+    torch.testing.assert_close(sliced_rope[0], custom_pos_emb[0][:, 2:])
+    torch.testing.assert_close(sliced_rope[1], custom_pos_emb[1][:, 2:])
+    torch.testing.assert_close(sliced_image_mask, image_mask[:, 2:])
+    torch.testing.assert_close(sliced_scatter, torch.tensor([[0]]))
+    assert sliced_query_lens == [4]
+    assert sliced_seq_len == 4
+
+
+@pytest.mark.parametrize("local_prefix_hit", [True, False])
+def test_ar_reuse_does_not_consume_local_prefix_hits(local_prefix_hit):
+    from vllm_omni.diffusion.forward_context import set_forward_context
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+        HunyuanImage3Text2ImagePipeline,
+    )
+
+    pipe = object.__new__(HunyuanImage3Text2ImagePipeline)
+    input_ids = torch.arange(12).reshape(1, 12)
+    cond_image = torch.ones(1, 1)
+    kwargs = dict(
+        query_lens=[12],
+        attention_mask=torch.ones(1, 1, 12, 12),
+        position_ids=input_ids.clone(),
+        image_mask=torch.zeros_like(input_ids, dtype=torch.bool),
+        gen_timestep_scatter_index=torch.tensor([[8]]),
+        cond_vae_images=cond_image,
+    )
+    runtime = SimpleNamespace(metadata=SimpleNamespace(prefill_rows=[SimpleNamespace(kv_start_pos=4)]))
+    with set_forward_context(paged_kv_runtime=runtime, paged_kv_cached_prefix_len=4 if local_prefix_hit else 0):
+        output, ar_reuse_len = pipe._maybe_handle_ar_kv_reuse(input_ids, kwargs, 1, False, None, torch.device("cpu"))
+    if local_prefix_hit:
+        assert ar_reuse_len == 0
+        assert output is input_ids
+        assert kwargs["query_lens"] == [12]
+        assert kwargs["cond_vae_images"] is cond_image
+    else:
+        assert ar_reuse_len == 4
+        torch.testing.assert_close(output, input_ids[:, 4:])
+        assert kwargs["query_lens"] == [8]
+        assert "cond_vae_images" not in kwargs
+
+
 def test_hunyuan_step_group_key_ignores_step_index_for_later_steps():
     pipeline = _pipeline()
     states = [_state("req-0", 1), _state("req-1", 3)]
@@ -416,7 +490,7 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
             }
         ]
 
-    captured = {}
+    captured: dict[str, object] = {}
 
     def fake_restore_prompt_kv_cache(states_arg, row_state_indexes, row_branches):
         del states_arg
@@ -443,7 +517,9 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
     assert captured["row_state_indexes"] == [0, 1, 0, 1]
     assert captured["row_branches"] == [0, 0, 1, 1]
     assert captured["input_ids"] is None
+    assert isinstance(captured["images"], torch.Tensor)
     assert tuple(captured["images"].shape) == (4, 1)
+    assert isinstance(captured["timestep"], torch.Tensor)
     assert captured["timestep"].tolist() == [0.5, 0.0, 0.5, 0.0]
     assert captured["merged_attention_mask_shape"] == (4, 1, 2, 6)
     assert captured["merged_full_attn_spans"] == [[(4, 6)], [(4, 6)], [(4, 6)], [(4, 6)]]

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -8,7 +8,7 @@ Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP,
 from __future__ import annotations
 
 import math
-from contextlib import ExitStack, contextmanager
+from contextlib import AbstractContextManager, ExitStack, contextmanager
 from types import SimpleNamespace
 from unittest.mock import patch
 
@@ -59,7 +59,7 @@ class MockAttention(nn.Module):
 @contextmanager
 def patched_mgr_env(sp_size=1):
     target = _TRANSFORMER_MODULE
-    patches = [
+    patches: list[AbstractContextManager] = [
         patch(f"{target}.get_sequence_parallel_world_size", return_value=sp_size),
         patch(f"{target}.get_allgather_parallel_world_size", return_value=sp_size),
         patch(f"{target}.get_ulysses_parallel_world_size", return_value=1, create=True),
@@ -242,6 +242,39 @@ def test_legacy_tensor_prefix_still_uses_dense_cache_and_reuse() -> None:
     assert mgr.attn.calls[-1][1].shape == (1, prefix_len + first_len, NUM_HEADS, HEAD_DIM)
 
 
+@pytest.mark.parametrize("prefix_len", [0, 19])
+@pytest.mark.parametrize("prompt_query_len", [0, 3])
+def test_paged_sp_splits_uncached_queries(prefix_len, prompt_query_len) -> None:
+    from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
+
+    mgr = _make_cache_mgr(sp_size=2)
+    mgr.attn.paged_kv_active = True
+    shard_image_size = IMAGE_TOKEN_LEN // 2
+    query_len = prompt_query_len + shard_image_size
+    query = torch.randn(2 * query_len, NUM_HEADS, HEAD_DIM)
+    key, value = _make_known_kv(2 * query_len)
+    with set_forward_context(), override_paged_kv_adapter(object()):
+        output = mgr(
+            query,
+            key,
+            value,
+            first_step=True,
+            query_lens=[query_len] * 2,
+            seq_lens=[prefix_len + query_len] * 2,
+            shard_image_size=shard_image_size,
+            full_attn_spans=[[(prefix_len + prompt_query_len, prefix_len + query_len)]] * 2,
+        )
+
+    torch.testing.assert_close(output, query)
+    metadata = mgr.attn.paged_metadata[0]
+    torch.testing.assert_close(
+        metadata.joint_key, key.reshape(2, query_len, NUM_KV_HEADS, HEAD_DIM)[:, :prompt_query_len]
+    )
+    assert metadata.joint_query.shape[1] == prompt_query_len
+    assert mgr.attn.paged_calls[0][0].shape[1] == shard_image_size
+    assert mgr.image_kv_cache_map is None
+
+
 def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
     from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
 
@@ -391,6 +424,7 @@ def test_no_ar_kv(bs):
         gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
     )
 
+    assert mgr.image_kv_cache_map is not None
     cached_key, cached_value = mgr.image_kv_cache_map
     # 3 prompt tokens cached per batch
     assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
@@ -478,12 +512,14 @@ def test_ar_kv_no_cfg(sp_size):
         gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
     )
 
+    assert mgr.image_kv_cache_map is not None
     cached_key, cached_value = mgr.image_kv_cache_map
     if sp_size == 1:
         # cached = ar(5) + prompt(3) = 8
         expected_cached_len = 8
     else:
         # cached = seq_len - shard_image_size = 17 - 4 = 13
+        assert shard_image_size is not None
         expected_cached_len = seq_len - shard_image_size
 
     assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)

tests/diffusion/test_diffusion_config_propagation.py

@@ -5,6 +5,7 @@
 Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
 """
 
+import json
 from collections.abc import Mapping
 
 import pytest
@@ -266,6 +267,21 @@ def test_architecture_name_resolves_via_pipeline_class_fallback():
     assert hunyuan_od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
 
 
+def test_architecture_only_checkpoint_propagates_multimodal_limit(tmp_path):
+    """A config.json-only checkpoint must expose its image-input capability."""
+    (tmp_path / "config.json").write_text(
+        json.dumps({"architectures": ["HunyuanImage3ForCausalMM"]}),
+        encoding="utf-8",
+    )
+
+    od_config = OmniDiffusionConfig(model=str(tmp_path))
+    od_config.enrich_config()
+
+    assert od_config.model_class_name == "HunyuanImage3ForCausalMM"
+    assert od_config.supports_multimodal_inputs is True
+    assert od_config.max_multimodal_image_inputs == HUNYUAN_IMAGE3_MAX_INPUT_IMAGES
+
+
 def test_additional_config_roundtrip():
     additional_config = {"torchair_graph_config": {"enabled": True}}
     od = _roundtrip_diffusion_config(model="x", additional_config=additional_config)

tests/diffusion/test_diffusion_engine_dummy_run.py

@@ -149,6 +149,24 @@ def test_dense_mode_does_not_build_kv_profile_request() -> None:
     engine._make_dummy_request.assert_not_called()
 
 
+def test_explicit_profile_frames_preserve_warmup_skip(monkeypatch):
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = OmniDiffusionConfig.__new__(OmniDiffusionConfig)
+    engine.od_config.model_class_name = "HunyuanImage3ForCausalMM"
+    monkeypatch.setattr("vllm_omni.diffusion.diffusion_engine.supports_multimodal_input", lambda _: (True, False))
+    monkeypatch.setattr("vllm_omni.diffusion.diffusion_engine.image_color_format", lambda _: "RGB")
+    monkeypatch.setattr("vllm_omni.diffusion.diffusion_engine.get_dummy_run_num_frames", lambda *_: 0)
+    kwargs = dict(height=1024, width=1024, guidance_scale=5.0, num_image_inputs=3)
+
+    assert engine._make_dummy_request(**kwargs) is None
+    request = engine._make_dummy_request(**kwargs, num_frames=1)
+
+    assert request is not None
+    assert request.sampling_params.num_frames == 1
+    assert request.sampling_params.num_inference_steps == 1
+    assert len(request.prompt["multi_modal_data"]["image"]) == 3
+
+
 @pytest.mark.parametrize(
     ("execution_mode", "uses_dlo_dp", "max_num_seqs", "expected_profile_requests"),
     [

tests/diffusion/test_diffusion_model_runner.py

@@ -223,7 +223,7 @@ def test_release_captured_graphs_tolerates_a_pipeline_without_captures():
     assert not hasattr(runner, "graph_runners")
 
 
-def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summary: bool = True):
+def _make_runner(cache_backend, cache_backend_name: str | None, enable_cache_dit_summary: bool = True):
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()
     runner.device = torch.device("cpu")
@@ -253,6 +253,55 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
     return runner
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+@pytest.mark.parametrize("imported_tokens", [None, (4, 4), (4, 2)])
+def test_request_forward_separates_imported_kv_and_local_hits(monkeypatch, imported_tokens):
+    from contextlib import nullcontext
+
+    from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata, DiffusionKVSequenceMetadata
+
+    runner = _make_runner(cache_backend=None, cache_backend_name="none")
+    runner.od_config.kv_transfer_config = object() if imported_tokens else None
+    request = _make_request()
+    metadata = DiffusionKVMetadata(
+        request_id=request.request_id,
+        allocation_generation=1,
+        sequences=tuple(
+            DiffusionKVSequenceMetadata(
+                sequence_id=i,
+                prefix_len=8,
+                target_len=4,
+                seq_len=12,
+                block_ids=([1, 2, 3],),
+                cached_prefix_len=0 if imported_tokens else 4,
+                num_computed_tokens=imported_tokens[i] if imported_tokens else 4,
+            )
+            for i in range(2)
+        ),
+    )
+    captured = {}
+
+    @contextmanager
+    def capture_context(**kwargs):
+        captured.update(kwargs)
+        yield
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", capture_context)
+    runner.diffusion_kv_backend.activate_paged_attention_metadata = Mock(return_value=(object(), nullcontext()))
+    runner._execute_request_list(
+        [request],
+        od_config=runner.od_config,
+        allow_single_output=True,
+        require_request_batch_support=False,
+        record_name="test",
+        record_output_peak_memory=False,
+        diffusion_kv_metadata=[metadata],
+    )
+    assert captured["paged_kv_cached_prefix_len"] == (0 if imported_tokens else 4)
+    assert getattr(request, "kv_computed_tokens", ()) == (imported_tokens or ())
+
+
 def _make_compile_runner(
     model=None,
     *,
@@ -310,7 +359,8 @@ def test_refresh_cache_prefers_request_steps_then_schedule_then_pipeline_default
 
 @pytest.mark.core_model
 @pytest.mark.cpu
-def test_refresh_cache_without_request_or_pipeline_default_warns(caplog):
+def test_refresh_cache_without_request_or_pipeline_default_warns(caplog, monkeypatch):
+    monkeypatch.setattr(model_runner_module.logger, "handlers", [*model_runner_module.logger.handlers, caplog.handler])
     cache_backend = _EnabledCacheBackend()
     runner = _make_runner(cache_backend=cache_backend, cache_backend_name="cache_dit")
     req = _make_request()
@@ -828,7 +878,7 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
         observed_batch_rows.append((len(states), int(input_batch.latents.shape[0])))
         return original_denoise_step(input_batch, states)
 
-    runner.pipeline.denoise_step = denoise_step
+    monkeypatch.setattr(runner.pipeline, "denoise_step", denoise_step)
     runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
     forward_context_calls = []
 
@@ -1313,7 +1363,7 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
         device_config=DeviceConfig(device="cpu"),
         compilation_config=CompilationConfig(),
     )
-    calls = []
+    calls: list[tuple[str, object]] = []
 
     @contextmanager
     def _set_priority():

tests/diffusion/test_diffusion_plugin_hooks.py

@@ -20,8 +20,10 @@ from vllm_omni.diffusion.registry import (
     _DIFFUSION_MODELS,
     _DIFFUSION_POST_PROCESS_FUNCS,
     _DIFFUSION_PRE_PROCESS_FUNCS,
+    _DIFFUSION_PREFIX_CACHE_FUNCS,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
+    get_diffusion_prefix_cache_func,
     register_diffusion_model,
 )
 from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum
@@ -70,6 +72,7 @@ class TestRegisterDiffusionModel:
         original_pre = _DIFFUSION_PRE_PROCESS_FUNCS.copy()
         original_post = _DIFFUSION_POST_PROCESS_FUNCS.copy()
         original_ir_op_priority = _DIFFUSION_IR_OP_PRIORITY_FUNCS.copy()
+        original_prefix_cache = _DIFFUSION_PREFIX_CACHE_FUNCS.copy()
         yield
         _DIFFUSION_MODELS.clear()
         _DIFFUSION_MODELS.update(original_models)
@@ -79,6 +82,8 @@ class TestRegisterDiffusionModel:
         _DIFFUSION_POST_PROCESS_FUNCS.update(original_post)
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.clear()
         _DIFFUSION_IR_OP_PRIORITY_FUNCS.update(original_ir_op_priority)
+        _DIFFUSION_PREFIX_CACHE_FUNCS.clear()
+        _DIFFUSION_PREFIX_CACHE_FUNCS.update(original_prefix_cache)
 
     def test_register_new_model(self):
         """Test registering a new diffusion model with pre/post process functions."""
@@ -89,6 +94,7 @@ class TestRegisterDiffusionModel:
             pre_process_func_name="test_pre_process",
             post_process_func_name="test_post_process",
             ir_op_priority_func_name="test_ir_op_priority",
+            prefix_cache_func_name="test_prefix_cache",
         )
         assert "TestPipeline" in _DIFFUSION_MODELS
         assert _DIFFUSION_MODELS["TestPipeline"] == (
@@ -99,6 +105,29 @@ class TestRegisterDiffusionModel:
         assert _DIFFUSION_PRE_PROCESS_FUNCS["TestPipeline"] == "test_pre_process"
         assert _DIFFUSION_POST_PROCESS_FUNCS["TestPipeline"] == "test_post_process"
         assert _DIFFUSION_IR_OP_PRIORITY_FUNCS["TestPipeline"] == "test_ir_op_priority"
+        assert _DIFFUSION_PREFIX_CACHE_FUNCS["TestPipeline"] == "test_prefix_cache"
+
+    def test_registered_prefix_hook_is_loaded_from_plugin(self):
+        hook = Mock()
+        factory = Mock(return_value=hook)
+        register_diffusion_model(
+            "TestPrefixPipeline", "test_plugin.pipeline", "TestPrefixPipeline", prefix_cache_func_name="prefix_factory"
+        )
+        config = SimpleNamespace(model_class_name="TestPrefixPipeline")
+        with patch(
+            "vllm_omni.diffusion.registry.importlib.import_module", return_value=SimpleNamespace(prefix_factory=factory)
+        ) as load:
+            assert get_diffusion_prefix_cache_func(config) is hook
+        load.assert_called_once_with("test_plugin.pipeline")
+        factory.assert_called_once_with(config)
+
+    def test_builtin_hunyuan_prefix_hook(self):
+        from vllm_omni.diffusion.models.hunyuan_image3.request_layout import prepare_hunyuan_prefix_cache
+
+        assert (
+            get_diffusion_prefix_cache_func(SimpleNamespace(model_class_name="HunyuanImage3ForCausalMM"))
+            is prepare_hunyuan_prefix_cache
+        )
 
     def test_register_model_accepts_deprecated_action_postprocess_keyword(self):
         """Deprecated action hook keyword is accepted but not registered."""

tests/diffusion/test_diffusion_scheduler.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.sched import (
     Scheduler,
     StepScheduler,
 )
-from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
+from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData, SchedulerRequestState
 from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
@@ -96,6 +96,7 @@ def _initialize_paged_scheduler(
     *,
     num_blocks: int = 64,
     max_num_seqs: int = 1,
+    enable_prefix_caching: bool = False,
 ) -> None:
     native_kv_managers.register_all_kvcache_specs(None)
     spec = FullAttentionSpec(
@@ -121,22 +122,42 @@ def _initialize_paged_scheduler(
         kv_vllm_config=SimpleNamespace(
             model_config=SimpleNamespace(max_model_len=64),
             max_in_flight_tokens=64,
+            cache_config=SimpleNamespace(
+                enable_prefix_caching=enable_prefix_caching,
+                prefix_caching_hash_algo="sha256",
+            ),
         ),
     )
 
 
-def _attach_diffusion_kv(request: OmniDiffusionRequest, *, seq_len: int = 8) -> None:
+def _attach_diffusion_kv(
+    request: OmniDiffusionRequest,
+    *,
+    seq_len: int = 8,
+    prefix_len: int = 4,
+    cache_token_ids=(),
+) -> None:
     request.diffusion_kv_requests = (
         DiffusionKVRequest(
             f"{request.request_id}/diffusion-kv/0",
             sequence_id=0,
-            prefix_len=4,
+            prefix_len=prefix_len,
             target_len=4,
             seq_len=seq_len,
+            cache_token_ids=cache_token_ids,
         ),
     )
 
 
+def _make_aborted_request_output(req_id: str) -> RunnerOutput:
+    return RunnerOutput(
+        request_id=req_id,
+        step_index=None,
+        finished=True,
+        result=DiffusionOutput(output=None, aborted=True),
+    )
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("scheduler_cls", [RequestScheduler, StepScheduler])
 @pytest.mark.parametrize("failure", ["timeout", "registration"])
@@ -168,7 +189,7 @@ async def test_single_native_kv_failure_reaches_output_stream(mocker, scheduler_
     engine.abort_queue = queue.Queue()
     engine._cv = threading.Condition()
     engine.main_loop = asyncio.get_running_loop()
-    stream = asyncio.Queue()
+    stream: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
     engine._out_streams = {"failed": stream}
     engine.executor = mocker.Mock()
     engine.executor.prepare_kv_for_forward.return_value = KVConnectorOutput()
@@ -192,7 +213,7 @@ class _StubScheduler:
         self._output = output
         self.initialized_with = None
         self._request_id = request.request_id
-        self._state = None
+        self._state: SchedulerRequestState | None = None
         self._scheduled = False
         self.max_num_running_reqs = 1
 
@@ -201,7 +222,7 @@ class _StubScheduler:
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         assert request is self._request
-        self._state = SimpleNamespace(request_id=self._request_id, req=request)
+        self._state = SchedulerRequestState(request_id=self._request_id, req=request)
         return self._request_id
 
     def schedule(self):
@@ -223,6 +244,7 @@ class _StubScheduler:
     def update_from_output(self, sched_output, output) -> set[str]:
         del sched_output
         assert output is self._output
+        assert self._state is not None
         self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
         return {self._request_id}
 
@@ -553,6 +575,197 @@ class TestRequestScheduler:
         assert metadata.request_id == request.request_id
         assert len(metadata.sequences[0].block_ids[0]) == 4
 
+    def test_diffusion_kv_publishes_only_after_successful_completion(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, enable_prefix_caching=True)
+        first = _make_request("publish-success")
+        _attach_diffusion_kv(first, cache_token_ids=range(4))
+        self.scheduler.add_request(first)
+        first_schedule = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first_schedule, _make_request_output(first.request_id)) == {
+            first.request_id
+        }
+
+        warm = _make_request("publish-warm")
+        _attach_diffusion_kv(warm, cache_token_ids=range(4))
+        self.scheduler.add_request(warm)
+        warm_schedule = self.scheduler.schedule()
+        metadata = warm_schedule.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.sequences[0].cached_prefix_len == 4
+
+    @pytest.mark.parametrize("boundaries", [(8, 4), (8, 0), (0, 8), (4, 8), (8, 8), (4, 4), (0, 0)])
+    @pytest.mark.parametrize("num_branches", [1, 2])
+    def test_diffusion_kv_batches_only_matching_prefix_boundaries(self, boundaries, num_branches) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        assert manager is not None
+        pool = manager.native_manager.block_pool
+        empty_free_blocks = pool.get_num_free_blocks()
+
+        warmup = _make_request("warmup")
+        _attach_diffusion_kv(warmup, seq_len=12, prefix_len=8, cache_token_ids=range(8))
+        self.scheduler.add_request(warmup)
+        self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        for request_id, boundary in zip(("first", "second"), boundaries):
+            request = _make_request(request_id)
+            # Each request's CFG branches share a boundary, but requests may not.
+            tokens = tuple(range(boundary)) + tuple(range(100, 108 - boundary))
+            request.diffusion_kv_requests = tuple(
+                DiffusionKVRequest(
+                    f"{request_id}/diffusion-kv/{branch}",
+                    sequence_id=branch,
+                    prefix_len=8,
+                    target_len=4,
+                    seq_len=12,
+                    cache_token_ids=tokens,
+                )
+                for branch in range(num_branches)
+            )
+            self.scheduler.add_request(request)
+
+        scheduled = self.scheduler.schedule()
+        matching = boundaries[0] == boundaries[1]
+        expected_ids = ["first", "second"] if matching else ["first"]
+        assert _new_ids(scheduled) == expected_ids
+        assert {
+            sequence.cached_prefix_len
+            for req in scheduled.scheduled_new_reqs
+            for sequence in req.diffusion_kv_metadata.sequences
+        } == {boundaries[0]}
+
+        if not matching:
+            second_state = self.scheduler.get_request_state("second")
+            assert second_state.status is DiffusionRequestStatus.WAITING
+            assert not manager.has_request("second")
+            # Retrying admission cannot leak reservations or publish its suffix.
+            free_blocks = pool.get_num_free_blocks()
+            for _ in range(3):
+                retry = self.scheduler.schedule()
+                assert _new_ids(retry) == [] and _cached_ids(retry) == ["first"]
+                assert pool.get_num_free_blocks() == free_blocks
+                assert not manager.has_request("second")
+                for row in second_state.diffusion_kv_requests:
+                    assert row.num_computed_tokens == 0
+                    assert manager.native_manager.get_computed_blocks(row)[1] == boundaries[1]
+
+        self.scheduler.update_from_output(
+            scheduled,
+            BatchRunnerOutput.from_list([_make_request_output(req_id) for req_id in expected_ids]),
+        )
+        if not matching:
+            next_wave = self.scheduler.schedule()
+            assert _new_ids(next_wave) == ["second"]
+            assert next_wave.scheduled_new_reqs[0].diffusion_kv_metadata.sequences[0].cached_prefix_len == boundaries[1]
+            self.scheduler.update_from_output(next_wave, _make_request_output("second"))
+        assert not self.scheduler.has_requests()
+        assert pool.get_num_free_blocks() == empty_free_blocks
+
+    def test_deferred_allocation_is_not_registered_for_transfer(self, mocker) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2)
+        manager = self.scheduler._diffusion_kv_manager
+        connector = mocker.Mock()
+        connector.get_num_new_matched_tokens.return_value = (4, True)
+        self.scheduler._kv_connector = connector
+        # Model a compatibility decision that needs the completed lookup.
+        can_schedule = self.scheduler._can_schedule_waiting
+        mocker.patch.object(
+            self.scheduler,
+            "_can_schedule_waiting",
+            side_effect=lambda state: can_schedule(state)
+            and not (state.request_id == "deferred" and manager.has_request("deferred")),
+        )
+        for request_id in ("admitted", "deferred"):
+            request = _make_request(request_id)
+            _attach_diffusion_kv(request)
+            request.diffusion_kv_requests[0].prompt_token_ids = [1, 2, 3, 4]
+            request.kv_transfer_params = {"num_transfer_tokens": 4}
+            self.scheduler.add_request(request)
+
+        assert _new_ids(self.scheduler.schedule()) == ["admitted"]
+        assert connector.update_state_after_alloc.call_count == 1
+        assert not manager.has_request("deferred")
+        assert "deferred" not in self.scheduler._kv_request_generations
+        assert "deferred" not in self.scheduler._kv_loading_request_ids
+        assert "deferred/diffusion-kv/0" not in self.scheduler._kv_transfer_request_ids
+        assert self.scheduler.get_request_state("deferred").status is DiffusionRequestStatus.WAITING
+
+    def test_diffusion_kv_deferred_request_can_be_cancelled(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        pool = manager.native_manager.block_pool
+        empty_free_blocks = pool.get_num_free_blocks()
+        for request_id, tokens in (("warmup", range(4)), ("hit", range(4)), ("miss", range(10, 14))):
+            request = _make_request(request_id)
+            _attach_diffusion_kv(request, cache_token_ids=tokens)
+            self.scheduler.add_request(request)
+            if request_id == "warmup":
+                self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        scheduled = self.scheduler.schedule()
+        assert _new_ids(scheduled) == ["hit"]
+        self.scheduler.finish_requests("miss", DiffusionRequestStatus.FINISHED_ABORTED)
+        self.scheduler.update_from_output(scheduled, _make_request_output("hit", error="worker failed"))
+        assert not manager.has_request("miss")
+        assert not self.scheduler.has_requests()
+        assert pool.get_num_free_blocks() == empty_free_blocks
+
+    def test_diffusion_kv_preempted_request_keeps_reservation_when_deferred(self) -> None:
+        _initialize_paged_scheduler(self.scheduler, max_num_seqs=2, enable_prefix_caching=True)
+        manager = self.scheduler._diffusion_kv_manager
+        warmup = _make_request("warmup")
+        _attach_diffusion_kv(warmup, cache_token_ids=range(4))
+        self.scheduler.add_request(warmup)
+        self.scheduler.update_from_output(self.scheduler.schedule(), _make_request_output("warmup"))
+
+        hit = _make_request("hit")
+        _attach_diffusion_kv(hit, cache_token_ids=range(4))
+        self.scheduler.add_request(hit)
+        self.scheduler.schedule()
+        miss = _make_request("miss")
+        _attach_diffusion_kv(miss, cache_token_ids=range(10, 14))
+        self.scheduler.add_request(miss)
+        # Install retained state as if this request were resuming an earlier wave.
+        miss_state = self.scheduler.get_request_state("miss")
+        miss_state.status = DiffusionRequestStatus.PREEMPTED
+        miss_metadata = manager.reserve_request("miss", miss_state.diffusion_kv_requests)
+        assert _cached_ids(self.scheduler.schedule()) == ["hit"]
+        assert self.scheduler.get_request_state("miss").status is DiffusionRequestStatus.PREEMPTED
+        assert manager.get_metadata("miss") is miss_metadata
+        self.scheduler.finish_requests("hit", DiffusionRequestStatus.FINISHED_ABORTED)
+        assert _cached_ids(self.scheduler.schedule()) == ["miss"]
+        assert manager.get_metadata("miss") is miss_metadata
+        self.scheduler.close()
+
+    @pytest.mark.parametrize(
+        "terminal_output",
+        [
+            pytest.param(
+                _make_request_output("publish-error", error="worker failed"),
+                id="error",
+            ),
+            pytest.param(_make_aborted_request_output("publish-abort"), id="abort"),
+        ],
+    )
+    def test_diffusion_kv_does_not_publish_failed_or_aborted_request(
+        self,
+        terminal_output: RunnerOutput,
+    ) -> None:
+        _initialize_paged_scheduler(self.scheduler, enable_prefix_caching=True)
+        first = _make_request(terminal_output.request_id)
+        _attach_diffusion_kv(first, cache_token_ids=range(4))
+        self.scheduler.add_request(first)
+        first_schedule = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first_schedule, terminal_output) == {first.request_id}
+
+        warm = _make_request(f"{first.request_id}-warm")
+        _attach_diffusion_kv(warm, cache_token_ids=range(4))
+        self.scheduler.add_request(warm)
+        warm_schedule = self.scheduler.schedule()
+        metadata = warm_schedule.scheduled_new_reqs[0].diffusion_kv_metadata
+        assert metadata is not None
+        assert metadata.sequences[0].cached_prefix_len == 0
+
     def test_diffusion_kv_capacity_backpressures_fifo_until_blocks_are_freed(self) -> None:
         _initialize_paged_scheduler(self.scheduler, num_blocks=3, max_num_seqs=2)
         first_request = _make_request("first")

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import copy
 import gc
-import importlib
+import importlib.util
 import json
 import os
 import tempfile
@@ -118,11 +118,9 @@ HUNYUAN_COT_REF_PATH = get_asset_path(f"{ASSET_SUBDIR}/{COT_REF_NAME}")
 # Constants
 # ============================================================================
 # Test input
-PROMPT = "基于图一的logo，参考图二中冰箱贴的材质，制作一个新的冰箱贴"
-TEST_IMAGE_URLS = [
-    "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_0.png",
-    "https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_1_1.png",
-]
+_IT2I_INPUT = json.loads(get_asset_path("hunyuan_image3/it2i.jsonl").read_text(encoding="utf-8"))
+PROMPT = _IT2I_INPUT["prompt"]
+TEST_IMAGE_URLS = _IT2I_INPUT["image_files"]
 SEED = 42
 AR_TP_SIZE = len(AR_DEVICES.split(","))
 DIT_TP_SIZE = len(DIT_DEVICES.split(","))

tests/e2e/accuracy/test_hunyuan_image3_prefix_cache_accuracy.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Single-DiT IT2I accuracy against the existing official Instruct reference.
+
+Reuse the IT2I inputs and quality thresholds for cold, partial and repeated prefix
+hits. The reference comes from AR->DiT; this is a reference-quality check, not a
+claim of identical conditioning or bitwise equivalence between the pipelines.
+"""
+
+import json
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+import pytest
+from PIL import Image
+
+from benchmarks.accuracy.common import VllmOmniImageClient
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import get_asset_path
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+_REPO_ROOT = Path(__file__).resolve().parents[3]
+_BENCHMARK_CONFIG = _REPO_ROOT / "tests/dfx/perf/tests/test_hunyuan_image3_prefix_caching.json"
+_IT2I_INPUT = json.loads(get_asset_path("hunyuan_image3/it2i.jsonl").read_text(encoding="utf-8"))
+_REFERENCE_PATH = get_asset_path("hunyuan_image3/hunyuan_image_instruct_ref.png")
+_PROMPT = _IT2I_INPUT["prompt"]
+# Warm up with a different prompt, then score only the original IT2I prompt.
+# This exercises partial image-prefix hits without changing the golden's input.
+_EDIT_PROMPT = (
+    f"请制作一张清晰的产品展示图：{_PROMPT}。保留标志的形状、文字与颜色，"
+    "呈现真实的材质纹理，使用柔和的自然光和简洁背景，不添加其他物体。"
+)
+# Match the official Instruct reference thresholds in test_hunyuan_image3.py.
+_ACCURACY_THRESHOLDS = {"clip_score": 90.0, "ssim": 0.26, "psnr": 12.5}
+_PREFILL_RE = re.compile(
+    r"Diffusion prefix prefill: request_id=(\S+) sequence_id=(\d+) "
+    r"cached_prefix_len=(\d+) prefix_len=(\d+) query_len=(\d+)"
+)
+_SLICE_RE = re.compile(r"Hunyuan prefix slice: cached_prefix_len=(\d+) query_len=(\d+)")
+_REFERENCE_RE = re.compile(r"Hunyuan reference span: sequence_id=(\d+) start=(\d+) end=(\d+)")
+
+
+@dataclass(frozen=True)
+class PrefixPrefill:
+    request_id: str
+    sequence_id: int
+    cached_prefix_len: int
+    prefix_len: int
+    query_len: int
+
+
+def _assert_prefix_execution(log: str) -> list[PrefixPrefill]:
+    events = [PrefixPrefill(request_id, *map(int, values)) for request_id, *values in _PREFILL_RE.findall(log)]
+    # TP ranks execute identical slices; compare sets rather than log counts.
+    expected = {(event.cached_prefix_len, event.query_len) for event in events if event.cached_prefix_len > 0}
+    actual = {tuple(map(int, match)) for match in _SLICE_RE.findall(log)}
+    assert actual == expected, f"Scheduled/model prefix slices differ: {expected=}, {actual=}"
+    if events:
+        assert len(events) == 2 and {event.sequence_id for event in events} == {0, 1}, events
+        assert len({event.request_id for event in events}) == 1, events
+        assert len({event.cached_prefix_len for event in events}) == 1, f"CFG boundaries differ: {events}"
+    return events
+
+
+def _assert_reference_prefix_hit(events: list[PrefixPrefill], log: str) -> None:
+    # Require a hit into BOTH reference images, not only the system prompt.
+    spans = {tuple(map(int, match)) for match in _REFERENCE_RE.findall(log)}
+    for event in events:
+        row_spans = [(start, end) for row, start, end in spans if row == event.sequence_id]
+        assert len(row_spans) == len(_IT2I_INPUT["image_files"]), spans
+        assert all(0 <= start < end for start, end in row_spans), spans
+        assert max(start for start, _ in row_spans) < event.cached_prefix_len <= event.prefix_len, event
+        assert event.query_len > 0, event
+
+
+def _generate_edit(client, *, model, images, prompt, output_path, capfd, **generation):
+    # Preserve startup/previous diagnostics and capture one sequential request.
+    before = capfd.readouterr()
+    output_path.with_suffix(".before.log").write_text(before.out + before.err)
+    try:
+        image = client.generate_image_edit(model=model, prompt=prompt, images=images, **generation)
+        assert image.size == (generation["width"], generation["height"]), image.size
+        image.save(output_path)
+    finally:
+        captured = capfd.readouterr()
+        output_path.with_suffix(".log").write_text(captured.out + captured.err)
+    log = captured.out + captured.err
+    return image, _assert_prefix_execution(log), log
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+@pytest.mark.parametrize("tp_size,sp_size", [(4, 1), (2, 2)], ids=["tp4-sp1", "tp2-sp2"])
+def test_hunyuan_image3_warm_prefix_accuracy(tmp_path, capfd, tp_size, sp_size):
+    from tests.e2e.accuracy.helpers import CLIPScorer, compute_image_ssim_psnr, download_images, model_output_dir
+    from tests.helpers.runtime import OmniServer
+
+    configs = json.loads(_BENCHMARK_CONFIG.read_text())
+    model = os.environ.get("HUNYUAN_IMAGE3_MODEL", configs[0]["server_params"]["model"])
+    output_dir = model_output_dir(tmp_path, model)
+    images = download_images(_IT2I_INPUT["image_files"])
+    reference = Image.open(_REFERENCE_PATH).convert("RGB")
+    # Pin the official reference's settings independently of the shorter perf run.
+    generation = dict(
+        seed=42,
+        num_inference_steps=50,
+        guidance_scale=2.5,
+        width=1280,
+        height=720,
+        bot_task="think_recaption",
+        sys_type="en_unified",
+    )
+    results, traces = {}, {}
+    for config in configs:
+        mode = config["test_name"].removeprefix("test_hunyuan_image3_prefix_")
+        server_params = config["server_params"]
+        # Keep four workers; exercise trimmed queries with strict Ulysses,
+        # not the T2I/no-hit AllGather-KV path covered by other SP tests.
+        server_params["stage_overrides"]["0"]["parallel_config"].update(
+            tensor_parallel_size=tp_size,
+            sequence_parallel_size=sp_size,
+            ulysses_degree=sp_size,
+            ulysses_mode="strict",
+            ring_degree=1,
+            allgather_degree=1,
+            cfg_parallel_size=1,
+        )
+        server_args = [
+            "--deploy-config",
+            str(_REPO_ROOT / server_params["serve_args"]["deploy-config"]),
+            "--stage-overrides",
+            json.dumps(server_params["stage_overrides"]),
+            "--trust-remote-code",
+            "--stage-init-timeout",
+            "900",
+        ]
+        requests = [("original", _PROMPT)]
+        if mode == "paged_prefix":
+            requests = [("warmup", _EDIT_PROMPT), ("partial", _PROMPT)]
+            requests += [(f"repeat{i}", _PROMPT) for i in range(1, 4)]
+        with OmniServer(
+            model,
+            server_args,
+            env_dict={
+                "VLLM_LOGGING_LEVEL": "DEBUG",
+                "DIFFUSION_ATTENTION_BACKEND": os.environ.get("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN"),
+            },
+        ) as server:
+            client = VllmOmniImageClient(f"http://{server.host}:{server.port}")
+            for case, prompt in requests:
+                label = f"{mode}_{case}"
+                image, events, log = _generate_edit(
+                    client,
+                    model=server.model,
+                    images=images,
+                    prompt=prompt,
+                    output_path=output_dir / f"{label}.png",
+                    capfd=capfd,
+                    **generation,
+                )
+                if case != "warmup":
+                    results[mode, case] = image
+                traces[label] = [vars(event) for event in events]
+                (output_dir / "prefix_traces.json").write_text(json.dumps(traces, indent=2))
+                if mode == "dense":
+                    assert not events, events
+                    continue
+                assert len(events) == 2, f"Missing CFG prefill traces: {events}"
+                if mode == "paged_no_cache" or case == "warmup":
+                    assert all(event.cached_prefix_len == 0 for event in events), events
+                    continue
+                _assert_reference_prefix_hit(events, log)
+                if case == "partial":
+                    assert all(event.cached_prefix_len < event.prefix_len for event in events), events
+                else:
+                    partial_len = traces["paged_prefix_partial"][0]["cached_prefix_len"]
+                    assert events[0].cached_prefix_len > partial_len, events
+                    if case != "repeat1":
+                        first = traces["paged_prefix_repeat1"]
+                        assert [event.cached_prefix_len for event in events] == [
+                            event["cached_prefix_len"] for event in first
+                        ], events
+                        assert image.tobytes() == results["paged_prefix", "repeat1"].tobytes(), label
+
+    # Score every mode against the official output, not an uncached Omni output.
+    clip_scorer = CLIPScorer()
+    metrics, failures = {}, []
+    for (mode, case), image in results.items():
+        ssim, psnr = compute_image_ssim_psnr(prediction=image, reference=reference)
+        label = f"{mode}_{case} vs official_ref"
+        metrics[label] = {"clip_score": clip_scorer.image_image_score(image, reference), "ssim": ssim, "psnr": psnr}
+        for name, threshold in _ACCURACY_THRESHOLDS.items():
+            value = metrics[label][name]
+            print(f"{label} {name}: {value:.6f}, required >= {threshold}")
+            if not value >= threshold:
+                failures.append(f"{label}: {name}={value:.6f} < {threshold}")
+    (output_dir / "quality_metrics.json").write_text(json.dumps(metrics, indent=2))
+    assert not failures, "\n".join(failures)

vllm_omni/config/omni_config.py

@@ -824,6 +824,7 @@ class _DiffusionConfigProjection:
     video_output_transport: object = field(default_factory=dict)
     enable_cache_dit_summary: bool = False
     diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+    enable_prefix_caching: bool = False
     diffusion_kv_max_rows_per_request: int | None = Field(default=None, ge=1, strict=True)
     enable_prompt_embed_cache: bool = False
     prompt_embed_cache_size: int = Field(default=32, ge=1)
@@ -1003,6 +1004,11 @@ class _DiffusionConfigProjection:
             )
 
         self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode)
+        if self.enable_prefix_caching and self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
+            raise ValueError(
+                "enable_prefix_caching=True requires diffusion_kv_mode='paged_scheduler'; "
+                "set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching"
+            )
         if (
             self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
             and self.diffusion_kv_max_rows_per_request is None
@@ -1111,6 +1117,7 @@ _DIFFUSION_SHARED_CONFIG_FIELDS = frozenset(
         "dist_timeout",
         "model_config",
         "quantization_config",
+        "enable_prefix_caching",
     }
 )
 _DIFFUSION_RUNTIME_CONFIG_FIELDS = frozenset(
@@ -1898,6 +1905,7 @@ def _build_diffusion_stage_config(
         engine.diffusion,
         model=common_kwargs["model_config"].model,
         quantization_config=common_kwargs["quantization_config"],
+        enable_prefix_caching=bool(common_kwargs["cache_config"].enable_prefix_caching),
     )
     return cast(
         VllmOmniDiffusionStageConfig,
@@ -2159,8 +2167,11 @@ def _build_diffusion_config_projection(
     *,
     model: str | None,
     quantization_config: _QuantizationConfigType,
+    enable_prefix_caching: bool,
 ) -> _DiffusionConfigProjection:
     diffusion_kwargs = engine.to_kwargs()
+    # Mirror the resolved cache setting, including deploy/CLI precedence.
+    diffusion_kwargs["enable_prefix_caching"] = enable_prefix_caching
     # Match the legacy builder: topology supplies defaults, while explicit
     # deploy/CLI values (including empty tuples) retain precedence.
     diffusion_kwargs.setdefault("stage_input_payload_keys", tuple(topology.stage_input_payload_keys))

vllm_omni/diffusion/data.py

@@ -846,7 +846,7 @@ class OmniDiffusionConfig:
     parallel_config: DiffusionParallelConfig = field(default_factory=DiffusionParallelConfig)
 
     # Cache backend configuration (NEW)
-    cache_backend: str = "none"  # "tea_cache", "deep_cache", etc.
+    cache_backend: str | None = "none"  # "tea_cache", "deep_cache", etc.
     cache_config: DiffusionCacheConfig | dict[str, Any] = field(default_factory=dict)
     video_output_transport: VideoOutputTransportConfig = field(default_factory=VideoOutputTransportConfig)
     enable_cache_dit_summary: bool = False
@@ -883,6 +883,9 @@ class OmniDiffusionConfig:
 
     # Local Diffusion KV ownership and cache-layout mode.
     diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY
+    # Reuse block-aligned, immutable diffusion prefixes across requests when
+    # Scheduler-owned paged KV is active.
+    enable_prefix_caching: bool = False
     # Maximum number of native BlockTable rows one public request can own
     # (sequences plus independent contexts). The model adapter defines it.
     diffusion_kv_max_rows_per_request: int | None = None
@@ -1204,6 +1207,19 @@ 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 not isinstance(self.enable_prefix_caching, bool):
+            raise TypeError("enable_prefix_caching must be a bool")
+        if self.enable_prefix_caching and self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
+            raise ValueError(
+                "enable_prefix_caching=True requires diffusion_kv_mode='paged_scheduler'; "
+                "set diffusion_kv_mode='paged_scheduler' or disable enable_prefix_caching"
+            )
+        if self.enable_prefix_caching and self.enable_sleep_mode:
+            raise ValueError(
+                "Diffusion prefix caching cannot be combined with sleep mode: "
+                "sleep discards KV pages without invalidating cached prefixes; "
+                "disable enable_prefix_caching or enable_sleep_mode"
+            )
         if self.diffusion_kv_max_rows_per_request is not None and (
             type(self.diffusion_kv_max_rows_per_request) is not int or self.diffusion_kv_max_rows_per_request <= 0
         ):
@@ -1247,6 +1263,11 @@ class OmniDiffusionConfig:
             if self.diffusion_kv_mode is not DiffusionKVCacheMode.PAGED_SCHEDULER:
                 raise ValueError("native kv_transfer_config requires diffusion_kv_mode='paged_scheduler'")
             self.kv_transfer_config = parse_kv_transfer_config(self.kv_transfer_config)
+            if self.enable_prefix_caching and self.kv_transfer_config is not None:
+                raise ValueError(
+                    "Diffusion prefix caching cannot be combined with native kv_transfer_config; "
+                    "disable enable_prefix_caching for AR KV import or remove kv_transfer_config for local DiT reuse"
+                )
             if self.enable_sleep_mode:
                 raise ValueError("Native KV transfer does not support sleep mode: registered pages must remain mapped")
 
@@ -1722,6 +1743,7 @@ class OmniDiffusionConfig:
                         or DiffusionModelRegistry._try_load_model_cls(architecture) is not None
                     ):
                         self.model_class_name = architecture
+                    self.update_multimodal_support()
                 else:
                     raise
 

vllm_omni/diffusion/diffusion_engine.py

@@ -54,6 +54,7 @@ from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
+    get_diffusion_prefix_cache_func,
 )
 from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
 from vllm_omni.diffusion.sched import BaseScheduler, RequestScheduler, StepScheduler
@@ -130,7 +131,7 @@ __all__ = [
 
 
 def _func_accepts_parameter(func: object | None, parameter_name: str) -> bool:
-    if func is None:
+    if not callable(func):
         return False
     parameters = inspect.signature(func).parameters
     return parameter_name in parameters or any(
@@ -300,6 +301,13 @@ class DiffusionEngine:
     def _init_process_hooks(self, od_config: OmniDiffusionConfig) -> None:
         self.post_process_func = get_diffusion_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
+        self.prefix_cache_func = get_diffusion_prefix_cache_func(od_config) if self._prefix_cache_enabled() else None
+        if self._prefix_cache_enabled() and self.prefix_cache_func is None:
+            raise ValueError(
+                "enable_prefix_caching=True requires a registered prefix-cache hook for "
+                f"{od_config.model_class_name!r}; "
+                "disable enable_prefix_caching or use a supported native pipeline such as HunyuanImage3ForCausalMM"
+            )
         # Cache whether the model-specific postprocess accepts request-level
         # sampling params so step() can support both legacy and extended hooks.
         self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")
@@ -585,6 +593,7 @@ class DiffusionEngine:
         )
 
     def _busy_loop(self):
+        assert self.stop_event is not None
         while not self.stop_event.is_set():
             self._process_aborts_queue()
             self._process_rpc_queue()
@@ -696,6 +705,7 @@ class DiffusionEngine:
 
         Caller must hold ``self._cv``.
         """
+        assert self.stop_event is not None
         start = time.monotonic()
         decision = self.scheduler.get_admission_wait_decision(
             now=start,
@@ -1045,6 +1055,12 @@ class DiffusionEngine:
         engine.run_startup_warmup()
         return engine
 
+    def _prefix_cache_enabled(self) -> bool:
+        config = getattr(self, "od_config", None)
+        return getattr(config, "diffusion_kv_mode", None) is DiffusionKVCacheMode.PAGED_SCHEDULER and bool(
+            getattr(config, "enable_prefix_caching", False)
+        )
+
     def _prepare_request_for_admission(self, request: OmniDiffusionRequest) -> OmniDiffusionRequest:
         """Run model-owned preprocessing once, before entering Engine locks."""
 
@@ -1052,6 +1068,12 @@ class DiffusionEngine:
         if pre_process_func is not None:
             request = pre_process_func(request)
         self._validate_diffusion_kv_profile_limits(request)
+        # Gate cache-input preparation itself: disabled caching must not inspect
+        # tensors / RNG state or copy token IDs just to discard their hashes.
+        # Both dependency hashing and preprocessing stay outside Engine locks.
+        prefix_cache_func = getattr(self, "prefix_cache_func", None)
+        if self._prefix_cache_enabled() and prefix_cache_func is not None and request.diffusion_kv_requests:
+            prefix_cache_func(request)
         return request
 
     def _validate_diffusion_kv_profile_limits(self, request: OmniDiffusionRequest) -> None:
@@ -1192,7 +1214,7 @@ class DiffusionEngine:
                 # sync func should receive one result
                 if (
                     sched_output.scheduled_request_ids
-                    and not isinstance(runner_output, RunnerOutput)
+                    and isinstance(runner_output, BatchRunnerOutput)
                     and len(runner_output) != 1
                 ):
                     raise ValueError("Sync func should receive one result at one time")
@@ -1249,7 +1271,7 @@ class DiffusionEngine:
         num_inference_steps: int = 1,
         num_frames: int | None = None,
     ) -> OmniDiffusionRequest | None:
-        """Build a minimal model request for startup profiling or warmup."""
+        """Build a startup request; explicit frame counts bypass the warmup policy."""
         prompt = OmniTextPrompt(prompt="dummy run")
         model_class_name = self.od_config.model_class_name
         if model_class_name is None:

vllm_omni/diffusion/diffusion_kv/kv_cache_utils.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Small diffusion identity helpers alongside vLLM's native KV cache utilities.
+
+Content hashing runs before Scheduler admission, only with prefix caching on.
+Requests carry native MultiModalFeatureSpec / PlaceholderRange values; native
+generate_block_hash_extra_keys handles block intersections, with no per-token
+extra-key expansion or intermediate dependency objects.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING
+
+import numpy as np
+import torch
+from vllm.multimodal.hasher import MultiModalHasher
+
+if TYPE_CHECKING:
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def _normalize_identity(value: object) -> object:
+    # MultiModalHasher handles tensor bytes (including BF16). Its generic
+    # recursive serialization alone does not canonicalize nested dict order
+    # or distinguish every container/scalar boundary. A typed JSON envelope
+    # supplies those guarantees and rejects the upstream pickle fallback.
+    if value is None:
+        return ["none"]
+    if isinstance(value, np.generic):
+        scalar = value.item()
+        if isinstance(scalar, np.generic):
+            # Extended-precision scalars have no lossless Python scalar;
+            # item() returns another NumPy scalar instead of making progress.
+            raise TypeError(f"Unsupported prefix-cache identity scalar dtype: {value.dtype}")
+        return _normalize_identity(scalar)
+    if isinstance(value, bool):
+        return ["bool", value]
+    if isinstance(value, int):
+        return ["int", str(value)]
+    if isinstance(value, float):
+        return ["float", value.hex()]
+    if isinstance(value, str):
+        return ["str", value]
+    if isinstance(value, bytes):
+        return ["bytes", value.hex()]
+    if isinstance(value, torch.Tensor):
+        if value.layout != torch.strided or value.is_quantized:
+            raise TypeError("Prefix-cache identity requires a dense, non-quantized tensor")
+        tensor = value.detach().cpu().contiguous()
+        # A byte view also covers dtypes NumPy cannot represent directly.
+        data = tensor.reshape(-1).view(torch.uint8).numpy()
+        return ["tensor", str(tensor.dtype), list(tensor.shape), MultiModalHasher.hash_kwargs("sha256", data=data)]
+    if isinstance(value, np.ndarray):
+        if value.dtype.hasobject or value.dtype.fields is not None:
+            raise TypeError("Prefix-cache identity does not support object or structured arrays")
+        return [
+            "array",
+            value.dtype.str,
+            list(value.shape),
+            MultiModalHasher.hash_kwargs("sha256", data=np.ascontiguousarray(value)),
+        ]
+    if isinstance(value, dict):
+        if any(not isinstance(key, str) for key in value):
+            raise TypeError("Prefix-cache identity dictionaries require string keys")
+        return ["dict", [[key, _normalize_identity(value[key])] for key in sorted(value)]]
+    if isinstance(value, list | tuple):
+        return ["sequence", [_normalize_identity(item) for item in value]]
+    raise TypeError(f"Unsupported prefix-cache identity value: {type(value).__name__}")
+
+
+def hash_prefix_cache_value(value: object) -> bytes:
+    """Deterministic, type-framed identity using vLLM's multimodal hasher."""
+
+    encoded = json.dumps(_normalize_identity(value), ensure_ascii=True, separators=(",", ":"))
+    return bytes.fromhex(MultiModalHasher.hash_kwargs("sha256", identity=encoded))
+
+
+def get_cache_namespace(model_namespace: str, sampling: OmniDiffusionSamplingParams) -> str:
+    """Isolate model semantics and the adapter actually activated by Worker.
+
+    Native LoRA block keys use lora_name, but diffusion also supports a scale
+    and identifies loaded adapters by lora_int_id. Keep that difference here.
+    Seeds are NOT global identity inputs: models attach random state only to
+    the multimodal features whose computation depends on it.
+    """
+
+    if not isinstance(model_namespace, str) or not model_namespace:
+        raise ValueError("Model cache namespace must be non-empty")
+    lora = sampling.lora_request
+    adapter = None if lora is None else (lora.lora_int_id, float(sampling.lora_scale))
+    return hash_prefix_cache_value(("diffusion-prefix-identity-v2", model_namespace, adapter)).hex()

vllm_omni/diffusion/diffusion_kv/manager.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
 from collections.abc import Sequence
 
-from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.utils.hashing import get_hash_fn_by_name
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
 from vllm.v1.kv_cache_interface import KVCacheConfig
 
 from vllm_omni.core.sched.utils import free_kv_blocks_in_physical_order
@@ -31,6 +33,8 @@ class DiffusionKVCacheManager:
         scheduler_block_size: int,
         hash_block_size: int,
         max_in_flight_tokens: int | None = None,
+        enable_prefix_caching: bool = False,
+        prefix_caching_hash_algo: str = "sha256",
     ) -> None:
         if max_model_len <= 0:
             raise ValueError(f"max_model_len must be positive, got {max_model_len}")
@@ -40,15 +44,18 @@ class DiffusionKVCacheManager:
             raise ValueError("scheduler_block_size and hash_block_size must be positive")
         if max_in_flight_tokens is not None and max_in_flight_tokens <= 0:
             raise ValueError(f"max_in_flight_tokens must be positive when provided, got {max_in_flight_tokens}")
+        self._hash_function = get_hash_fn_by_name(prefix_caching_hash_algo)
         self.native_manager = KVCacheManager(
             kv_cache_config=kv_cache_config,
             max_model_len=max_model_len,
             scheduler_block_size=scheduler_block_size,
             hash_block_size=hash_block_size,
             max_in_flight_tokens=max_in_flight_tokens,
-            enable_caching=False,
+            enable_caching=enable_prefix_caching,
         )
         self.max_model_len = max_model_len
+        self.hash_block_size = hash_block_size
+        self.enable_prefix_caching = enable_prefix_caching
         # Native vLLM may reserve a null block, so an idle BlockPool does not
         # necessarily report ``kv_cache_config.num_blocks`` free blocks.
         self._empty_pool_num_free_blocks = self.native_manager.block_pool.get_num_free_blocks()
@@ -57,31 +64,87 @@ class DiffusionKVCacheManager:
         self._internal_request_ids: set[str] = set()
         self._next_allocation_generation = 1
 
-    def _get_empty_pool_required_blocks(self, requests: Sequence[DiffusionKVRequest]) -> int:
-        """Return the native full-sequence reservation for one public request.
+    def _prepare_block_hashes(self, requests: Sequence[DiffusionKVRequest]) -> None:
+        if not self.enable_prefix_caching:
+            return
+        for request in requests:
+            request.build_block_hashes(self.hash_block_size, self._hash_function)
 
-        vLLM performs this calculation inside ``allocate_slots`` when
-        ``full_sequence_must_fit`` is enabled. Diffusion CFG adds one semantic
-        requirement that native requests do not have: every execution sequence
-        must fit atomically. Sum the native per-sequence result so a request
-        that can never fit is rejected independently of current pool usage.
+    def _clip_computed_blocks(
+        self,
+        blocks: KVCacheBlocks,
+        *,
+        hit_len: int,
+        cached_prefix_len: int,
+    ) -> KVCacheBlocks:
+        if cached_prefix_len == hit_len:
+            return blocks
+        clipped_groups: list[list[KVCacheBlock]] = []
+        for group_blocks, cache_group in zip(
+            blocks.blocks,
+            self.native_manager.kv_cache_config.kv_cache_groups,
+            strict=True,
+        ):
+            block_size = int(cache_group.kv_cache_spec.block_size)
+            if cached_prefix_len % block_size:
+                raise RuntimeError(
+                    "Diffusion prefix hit is not aligned to a native cache group: "
+                    f"cached_prefix_len={cached_prefix_len}, block_size={block_size}"
+                )
+            keep_blocks = cached_prefix_len // block_size
+            if keep_blocks > len(group_blocks):
+                raise RuntimeError(
+                    "Native prefix lookup returned too few blocks for the common CFG boundary: "
+                    f"required={keep_blocks}, returned={len(group_blocks)}"
+                )
+            clipped_groups.append(list(group_blocks[:keep_blocks]))
+        return self.native_manager.create_kv_cache_blocks(tuple(clipped_groups))
+
+    @staticmethod
+    def _unique_blocks(block_sets: Sequence[KVCacheBlocks]) -> list[KVCacheBlock]:
+        unique: dict[int, KVCacheBlock] = {}
+        for blocks in block_sets:
+            for group in blocks.blocks:
+                for block in group:
+                    if not block.is_null:
+                        unique[id(block)] = block
+        return list(unique.values())
+
+    def _get_empty_pool_required_blocks(
+        self,
+        requests: Sequence[DiffusionKVRequest],
+        computed_blocks: Sequence[KVCacheBlocks],
+        cached_prefix_len: int,
+    ) -> int:
+        """Return the minimum physical capacity for one public request.
+
+        Cached CFG rows may point to the same physical blocks. Count those hit
+        blocks once, then ask the native coordinator how many additional blocks
+        each row needs for its suffix. This gives an admission bound that still
+        works for warm requests and does not over-count shared prefix pages.
         """
 
         coordinator = self.native_manager.coordinator
-        empty_blocks = self.native_manager.empty_kv_cache_blocks.blocks
-        return sum(
-            coordinator.get_num_blocks_to_allocate(
+        unique_hit_blocks = len(self._unique_blocks(computed_blocks))
+        suffix_blocks = 0
+        for request, hit in zip(requests, computed_blocks, strict=True):
+            native_required = coordinator.get_num_blocks_to_allocate(
                 request_id=request.request_id,
                 num_tokens=request.num_tokens,
-                new_computed_blocks=empty_blocks,
+                new_computed_blocks=hit.blocks,
                 num_encoder_tokens=0,
-                total_computed_tokens=0,
-                num_local_computed_tokens=0,
+                total_computed_tokens=cached_prefix_len,
+                num_local_computed_tokens=cached_prefix_len,
                 num_tokens_main_model=request.num_tokens,
                 apply_admission_cap=True,
             )
-            for request in requests
-        )
+            evictable_hits = sum(block.ref_cnt == 0 and not block.is_null for group in hit.blocks for block in group)
+            suffix_blocks += max(native_required - evictable_hits, 0)
+        # ``get_num_blocks_to_allocate`` includes evictable hit pages because
+        # the live free-capacity check must first remove them from the free
+        # queue. They are already part of the pool's physical capacity, so do
+        # not count them a second time in this empty-pool bound.
+        return unique_hit_blocks + suffix_blocks
 
     def has_request(self, public_request_id: str) -> bool:
         return public_request_id in self._requests
@@ -123,7 +186,25 @@ class DiffusionKVCacheManager:
                     f"seq_len={request.seq_len}, max_model_len={self.max_model_len}"
                 )
 
-        required_blocks = self._get_empty_pool_required_blocks(requests)
+        self._prepare_block_hashes(requests)
+        hit_results = [self.native_manager.get_computed_blocks(request) for request in requests]
+        cached_prefix_len = min(result[1] for result in hit_results)
+        computed_blocks = [
+            self._clip_computed_blocks(
+                result[0],
+                hit_len=result[1],
+                cached_prefix_len=cached_prefix_len,
+            )
+            for result in hit_results
+        ]
+        for request, result in zip(requests, hit_results, strict=True):
+            request.shared_prefix_boundary = result[2]
+
+        required_blocks = self._get_empty_pool_required_blocks(
+            requests,
+            computed_blocks,
+            cached_prefix_len,
+        )
         if required_blocks > self._empty_pool_num_free_blocks:
             raise DiffusionKVAdmissionError(
                 f"Diffusion KV request {public_request_id!r} cannot fit even when the block pool is empty: "
@@ -133,11 +214,20 @@ class DiffusionKVCacheManager:
 
         allocated: list[DiffusionKVRequest] = []
         sequence_metadata: list[DiffusionKVSequenceMetadata] = []
+        pinned_blocks = self._unique_blocks(computed_blocks)
+        if pinned_blocks:
+            # All CFG lookups happen before allocation so they can share one
+            # execution boundary. Pin every distinct hit until each native
+            # allocation has taken its own reference; otherwise an earlier
+            # branch's miss allocation could evict a later branch's hit.
+            self.native_manager.block_pool.touch(pinned_blocks)
         try:
-            for request in requests:
+            for request, request_computed_blocks in zip(requests, computed_blocks, strict=True):
                 blocks = self.native_manager.allocate_slots(
                     request,
-                    num_new_tokens=request.seq_len,
+                    num_new_tokens=request.seq_len - cached_prefix_len,
+                    num_new_computed_tokens=cached_prefix_len,
+                    new_computed_blocks=request_computed_blocks,
                     delay_cache_blocks=True,
                     full_sequence_must_fit=True,
                 )
@@ -146,18 +236,23 @@ class DiffusionKVCacheManager:
                     allocated.clear()
                     return None
                 allocated.append(request)
+                request.num_computed_tokens = cached_prefix_len
                 sequence_metadata.append(
                     DiffusionKVSequenceMetadata(
                         sequence_id=request.sequence_id,
                         prefix_len=request.prefix_len,
                         target_len=request.target_len,
                         seq_len=request.seq_len,
-                        block_ids=blocks.get_block_ids(),
+                        block_ids=self.native_manager.get_block_ids(request.request_id),
+                        cached_prefix_len=cached_prefix_len,
                     )
                 )
         except Exception:
             self._rollback(allocated)
             raise
+        finally:
+            if pinned_blocks:
+                self.native_manager.block_pool.free_blocks(reversed(pinned_blocks))
 
         metadata = DiffusionKVMetadata(
             request_id=public_request_id,
@@ -170,6 +265,23 @@ class DiffusionKVCacheManager:
         self._internal_request_ids.update(internal_ids)
         return metadata
 
+    def publish_request(self, public_request_id: str) -> None:
+        """Publish successfully materialized, block-aligned stable prefixes."""
+
+        if not self.enable_prefix_caching:
+            return
+        try:
+            requests = self._requests[public_request_id]
+        except KeyError as exc:
+            raise KeyError(f"Diffusion KV request {public_request_id!r} is not allocated") from exc
+        for request in requests:
+            publishable_prefix_len = min(
+                request.prefix_len // self.hash_block_size * self.hash_block_size,
+                len(request.block_hashes) * self.hash_block_size,
+            )
+            if publishable_prefix_len > 0:
+                self.native_manager.cache_blocks(request, publishable_prefix_len)
+
     def get_metadata(self, public_request_id: str) -> DiffusionKVMetadata:
         try:
             return self._metadata[public_request_id]
@@ -181,6 +293,8 @@ class DiffusionKVCacheManager:
         self._metadata.pop(public_request_id, None)
         for request in reversed(requests):
             free_kv_blocks_in_physical_order(self.native_manager, request)
+            request.num_computed_tokens = 0
+            request.shared_prefix_boundary = 0
             self._internal_request_ids.discard(request.request_id)
 
     def close(self) -> None:
@@ -190,3 +304,5 @@ class DiffusionKVCacheManager:
     def _rollback(self, requests: Sequence[DiffusionKVRequest]) -> None:
         for request in reversed(requests):
             free_kv_blocks_in_physical_order(self.native_manager, request)
+            request.num_computed_tokens = 0
+            request.shared_prefix_boundary = 0

vllm_omni/diffusion/diffusion_kv/metadata.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Scheduler-to-Worker allocation payloads for diffusion KV.
 
@@ -40,6 +40,7 @@ class DiffusionKVSequenceMetadata:
     target_len: int
     seq_len: int
     block_ids: tuple[list[int], ...]
+    cached_prefix_len: int = 0
     context_ids: tuple[str, ...] = ()
     num_computed_tokens: int = 0
 

vllm_omni/diffusion/diffusion_kv/model_runner_backend.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from dataclasses import dataclass
 from typing import Any
 
@@ -154,6 +154,7 @@ class DiffusionKVModelRunnerBackend:
 
         if attention_geometry is not None:
             num_heads, num_kv_heads, head_size = attention_geometry
+            assert callable(set_attention_geometry)
             try:
                 set_attention_geometry(
                     num_heads=num_heads,
@@ -396,6 +397,15 @@ class DiffusionKVModelRunnerBackend:
         for sequence in metadata.sequences:
             if type(sequence.sequence_id) is not int or sequence.sequence_id < 0:
                 raise ValueError(f"Diffusion KV sequence_id must be a non-negative integer: {sequence.sequence_id!r}")
+            if (
+                type(sequence.cached_prefix_len) is not int
+                or sequence.cached_prefix_len < 0
+                or sequence.cached_prefix_len > sequence.prefix_len
+            ):
+                raise ValueError(
+                    "Diffusion KV cached_prefix_len must lie within the stable prefix: "
+                    f"cached={sequence.cached_prefix_len!r}, prefix={sequence.prefix_len!r}"
+                )
             sequence_identity = (metadata.request_id, sequence.sequence_id, None)
             sequence_install = self._validate_row(
                 identity=sequence_identity,
@@ -410,6 +420,7 @@ class DiffusionKVModelRunnerBackend:
                     sequence.prefix_len,
                     sequence.target_len,
                     sequence.seq_len,
+                    sequence.cached_prefix_len,
                     sequence_install.block_ids,
                     sequence.context_ids,
                 )
@@ -593,7 +604,7 @@ class DiffusionKVModelRunnerBackend:
                     )
         raise RuntimeError(f"Diffusion KV request state is missing logical length for {identity!r}")
 
-    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
+    def remove_diffusion_kv_requests(self, request_ids: Sequence[str | tuple[str, int]]) -> int:
         """Retire Worker rows without logically freeing Scheduler-owned blocks."""
         if (
             not is_scheduler_paged_kv_mode(

vllm_omni/diffusion/diffusion_kv/paged_attention_adapter.py

@@ -84,14 +84,18 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
         attention_config = vllm_config.attention_config
         previous_backend = attention_config.backend
         previous_backend_per_kind = attention_config.backend_per_kind
-        previous_pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
+        parallel_config = vllm_config.parallel_config
+        previous_pcp_size = parallel_config.prefill_context_parallel_size
         try:
             # This is a portable vLLM backend request. The active platform
             # resolves it to its native implementation, such as FlashAttention
             # on CUDA or AscendAttentionBackend on NPU.
             attention_config.backend = AttentionBackendEnum.FLASH_ATTN
             attention_config.backend_per_kind = {}
-            vllm_config.parallel_config.prefill_context_parallel_size = 1
+            # Diffusion maps SP to PCP for MoE, not for attention. Omni's
+            # Ulysses hooks gather the sequence and shard heads before this
+            # native kernel, so backend selection must not request PCP.
+            parallel_config.prefill_context_parallel_size = 1
             with set_current_vllm_config(vllm_config):
                 attn_backend = get_attn_backend(
                     head_size=spec.head_size,
@@ -103,7 +107,7 @@ class DiffusionPagedAttentionLayerAdapter(AttentionLayerBase):
         finally:
             attention_config.backend = previous_backend
             attention_config.backend_per_kind = previous_backend_per_kind
-            vllm_config.parallel_config.prefill_context_parallel_size = previous_pcp_size
+            parallel_config.prefill_context_parallel_size = previous_pcp_size
         attn_backend = current_omni_platform.get_diffusion_paged_kv_attn_backend(
             attn_backend,
             ulysses_degree=ulysses_degree,

vllm_omni/diffusion/diffusion_kv/request.py

@@ -1,13 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from __future__ import annotations
 
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass
 from typing import Any
 
-from vllm.v1.core.kv_cache_utils import BlockHash
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
 from vllm.v1.request import RequestStatus
 
 
@@ -51,8 +52,10 @@ class DiffusionKVRequest:
     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.
+    With prefix caching enabled, Engine invokes model input preparation to
+    attach token IDs and native multimodal identities / positions. The Manager
+    builds one canonical hash for every cacheable full block before lookup;
+    publication requires those hashes and successfully materialized KV.
     """
 
     def __init__(
@@ -64,6 +67,9 @@ class DiffusionKVRequest:
         target_len: int,
         seq_len: int,
         block_hashes: Sequence[BlockHash] = (),
+        cache_token_ids: Sequence[int] = (),
+        mm_features: Sequence[MultiModalFeatureSpec] = (),
+        cache_namespace: str = "diffusion-primary-v1",
         kv_contexts: Sequence[DiffusionKVContext] = (),
         kv_transfer_params: dict[str, Any] | None = None,
         prompt_token_ids: list[int] | None = None,
@@ -83,6 +89,15 @@ class DiffusionKVRequest:
                 "prefix_len + target_len must not exceed seq_len: "
                 f"prefix_len={prefix_len}, target_len={target_len}, seq_len={seq_len}"
             )
+        token_ids = tuple(cache_token_ids)
+        if token_ids and len(token_ids) < prefix_len:
+            raise ValueError(
+                f"cache_token_ids must cover the reusable prefix: tokens={len(token_ids)}, prefix_len={prefix_len}"
+            )
+        if any(type(token_id) is not int for token_id in token_ids):
+            raise TypeError("cache_token_ids must contain only integers")
+        if not isinstance(cache_namespace, str) or not cache_namespace:
+            raise ValueError("cache_namespace must be a non-empty string")
 
         contexts = tuple(kv_contexts)
         if any(not isinstance(context, DiffusionKVContext) for context in contexts):
@@ -96,6 +111,15 @@ class DiffusionKVRequest:
         self.prefix_len = prefix_len
         self.target_len = target_len
         self.kv_contexts = contexts
+        self.cache_token_ids = token_ids
+        self.mm_features = list(mm_features)
+        self.cache_namespace = cache_namespace
+        # Native generate_block_hash_extra_keys consumes this Request surface.
+        # Diffusion adapter ID + scale are already isolated in cache_namespace;
+        # there are no separate LoRA-name, salt, or prompt-embedding inputs.
+        self.lora_request = None
+        self.cache_salt = None
+        self.prompt_embeds = None
 
         # Native vLLM Request surface. Keep this list intentionally small and
         # cover it with real-KVCacheManager conformance tests.
@@ -115,6 +139,70 @@ class DiffusionKVRequest:
         self.kv_transfer_params = kv_transfer_params
         self.prompt_token_ids = prompt_token_ids
 
+    def _validate_mm_features(self) -> None:
+        # Native block traversal expects sorted, non-overlapping placeholders.
+        # Do not silently accept nested/overlapping ranges: the traversal can
+        # skip a feature and incorrectly reuse KV. Combine such dependencies
+        # in the feature identifier instead.
+        previous_end = 0
+        for feature in self.mm_features:
+            if not isinstance(feature, MultiModalFeatureSpec):
+                raise TypeError("mm_features must contain native MultiModalFeatureSpec values")
+            position = feature.mm_position
+            end = position.offset + position.length
+            if position.offset < previous_end or position.length <= 0 or end > self.seq_len:
+                raise ValueError("mm_features must have valid sorted, non-overlapping positions within seq_len")
+            if not isinstance(feature.identifier, str) or not feature.identifier:
+                raise ValueError("mm_features must have non-empty content identifiers")
+            previous_end = end
+
+    def build_block_hashes(
+        self,
+        hash_block_size: int,
+        hash_function: Callable[[object], bytes],
+    ) -> None:
+        """Build native chained hashes for complete reusable-prefix blocks."""
+
+        if self.block_hashes:
+            self.skip_reading_prefix_cache = False
+            return
+        if hash_block_size <= 0:
+            raise ValueError(f"hash_block_size must be positive, got {hash_block_size}")
+        num_cache_tokens = self.prefix_len // hash_block_size * hash_block_size
+        if num_cache_tokens == 0 or not self.cache_token_ids:
+            return
+        if len(self.cache_token_ids) < num_cache_tokens:
+            raise ValueError(
+                "cache_token_ids do not cover the block-aligned reusable prefix: "
+                f"tokens={len(self.cache_token_ids)}, required={num_cache_tokens}"
+            )
+        self._validate_mm_features()
+
+        # Keep the first hash deterministic without mutating vLLM's process-wide
+        # ``NONE_HASH``.  A diffusion manager can be created more than once in
+        # one process (for example in tests or worker reinitialization); using
+        # the global random root would make requests built by different manager
+        # instances impossible to match.
+        parent_hash: BlockHash | None = BlockHash(hash_function(("diffusion-prefix-root-v1", self.cache_namespace)))
+        hashes: list[BlockHash] = []
+        mm_index = 0
+        for start in range(0, num_cache_tokens, hash_block_size):
+            end = start + hash_block_size
+            # Reuse upstream MM intersection/offset logic directly. Keep our
+            # stable-prefix bound: native get_request_block_hasher hashes the
+            # entire token axis, including diffusion's mutable target region.
+            extra_keys, mm_index = generate_block_hash_extra_keys(self, start, end, mm_index)
+            parent_hash = hash_block_tokens(
+                hash_function=hash_function,
+                parent_block_hash=parent_hash,
+                curr_block_token_ids=self.cache_token_ids[start:end],
+                extra_keys=extra_keys,
+            )
+            hashes.append(parent_hash)
+
+        self.block_hashes = hashes
+        self.skip_reading_prefix_cache = not hashes
+
     @property
     def seq_len(self) -> int:
         """Complete first-step sequence length and allocation boundary."""

vllm_omni/diffusion/executor/abstract.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
+from concurrent.futures import Future
 from dataclasses import replace
 from typing import TYPE_CHECKING, Any
 
@@ -12,7 +13,7 @@ from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
 from vllm.v1.outputs import KVConnectorOutput
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.sched.interface import CachedRequestData
 
 if TYPE_CHECKING:
@@ -132,6 +133,10 @@ class DiffusionExecutor(ABC):
         """
         return None
 
+    def wait_output_ready(self, async_output_id: str) -> Future[DiffusionOutput]:
+        """Resolve deferred output; only asynchronous executors implement this."""
+        raise NotImplementedError(f"{type(self).__name__} does not support asynchronous output")
+
     def get_kv_cache_specs(self) -> list[dict[str, KVCacheSpec]]:
         """Collect rank-local native specs after every Worker loads its model."""
 
@@ -202,7 +207,7 @@ class DiffusionExecutor(ABC):
             and not scheduler_output.kv_required_request_ids
             and not scheduler_output.kv_poll_only
         ):
-            return
+            return None
         transfer_output = replace(
             scheduler_output,
             scheduled_new_reqs=[],

vllm_omni/diffusion/forward_context.py

@@ -19,6 +19,8 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 if TYPE_CHECKING:
     import torch
 
+    from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import DiffusionPagedAttentionRuntime
+
 
 @dataclass
 class ForwardContext:
@@ -32,7 +34,10 @@ class ForwardContext:
     # Runner-owned paged execution metadata/runtime. Attention resolves the
     # active Worker adapter from it; model code must not construct BlockTable
     # rows or activate the runtime directly.
-    paged_kv_runtime: object | None = None
+    paged_kv_runtime: DiffusionPagedAttentionRuntime | None = None
+    # Block-aligned prefix already resident in Scheduler-owned pages for the
+    # active request-level prefill. Zero keeps the cold/full-prefill path.
+    paged_kv_cached_prefix_len: int = 0
     # Active Worker-side paged KV adapter.  The adapter is installed only for
     # the duration of a paged forward; dense forwards leave this as ``None``.
     # Keep the field opaque here to avoid coupling the common context module to
@@ -177,7 +182,8 @@ def create_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
-    paged_kv_runtime: object | None = None,
+    paged_kv_runtime: DiffusionPagedAttentionRuntime | None = None,
+    paged_kv_cached_prefix_len: int = 0,
     in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
@@ -187,6 +193,7 @@ def create_forward_context(
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
         paged_kv_runtime=paged_kv_runtime,
+        paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
         in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,
@@ -213,7 +220,8 @@ def set_forward_context(
     vllm_config: VllmConfig | None = None,
     omni_diffusion_config: OmniDiffusionConfig | None = None,
     attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None = None,
-    paged_kv_runtime: object | None = None,
+    paged_kv_runtime: DiffusionPagedAttentionRuntime | None = None,
+    paged_kv_cached_prefix_len: int = 0,
     in_diffusion_kv_memory_profile: bool = False,
     split_text_embed_in_sp: bool = False,
     denoise_step_idx: int | None = None,
@@ -227,6 +235,7 @@ def set_forward_context(
         omni_diffusion_config=omni_diffusion_config,
         attn_metadata=attn_metadata,
         paged_kv_runtime=paged_kv_runtime,
+        paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
         in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
         split_text_embed_in_sp=split_text_embed_in_sp,
         denoise_step_idx=denoise_step_idx,

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -76,7 +76,9 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import (
+    get_forward_context,
     get_paged_kv_computed_tokens,
+    is_forward_context_available,
     paged_kv_prefill,
     set_forward_context_denoise_step_idx,
 )
@@ -2481,6 +2483,7 @@ class HunyuanImage3Model(nn.Module):
         gen_timestep_scatter_index: torch.Tensor | None = None,
         uncond_cfg_prefill: bool = False,
         ar_kv_reuse_len: int = 0,
+        paged_kv_cached_prefix_len: int = 0,
         full_attn_spans: list[list[tuple[int, int]]] | None = None,
     ) -> tuple | BaseModelOutputWithPast:
         current_omni_platform.reset_diffusion_fused_moe_forward_context()
@@ -2539,7 +2542,9 @@ class HunyuanImage3Model(nn.Module):
             else:
                 shard_padding_size = shard_image_size * sp_world_size - num_image_tokens
             if first_step:
-                seq_lens = [prompt_size + shard_image_size + ar_kv_reuse_len for _ in seq_lens]
+                seq_lens = [
+                    prompt_size + shard_image_size + ar_kv_reuse_len + paged_kv_cached_prefix_len for _ in seq_lens
+                ]
             else:
                 seq_lens = [x - y for x, y in zip(seq_lens, query_lens)]
                 seq_lens = [seq_len + shard_image_size for seq_len in seq_lens]
@@ -3023,6 +3028,10 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         cfg_rank,
         device,
     ):
+        if is_forward_context_available() and get_forward_context().paged_kv_cached_prefix_len:
+            # Local hits are sliced after conditional-image embeddings are
+            # prepared in forward_call; the AR path must not truncate them.
+            return input_ids, 0
         ar_kv_data = model_kwargs.pop("ar_kv_data", None)
         computed_tokens = get_paged_kv_computed_tokens()
         if computed_tokens and computed_tokens[0] > 0:

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -23,7 +23,11 @@ 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.forward_context import (
+    get_forward_context,
+    is_forward_context_available,
+    set_forward_context_denoise_step_idx,
+)
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
@@ -105,6 +109,10 @@ def to_device(data, device):
         return data
 
 
+def get_hunyuan_image_3_prefix_cache_func(od_config: OmniDiffusionConfig):
+    return request_layout_utils.prepare_hunyuan_prefix_cache
+
+
 def _to_pil_image(image: Any) -> PILImage.Image:
     if isinstance(image, PILImage.Image):
         return image
@@ -1408,6 +1416,13 @@ class HunyuanImage3Pipeline(
         )
         full_attn_spans: list[list[tuple[int, int]]] = [[] for _ in range(bsz)]
         for i in range(bsz):
+            for reference_slice in tokenizer_output.joint_image_slices[i]:
+                logger.debug(
+                    "Hunyuan reference span: sequence_id=%d start=%d end=%d",
+                    i,
+                    int(reference_slice.start or 0),
+                    int(reference_slice.stop or seq_len),
+                )
             for j, image_slice in enumerate(batch_image_slices[i]):
                 attention_mask[i, image_slice, image_slice] = True
                 start = image_slice.start if image_slice.start is not None else 0
@@ -1420,6 +1435,69 @@ class HunyuanImage3Pipeline(
         model_kwargs["full_attn_spans"] = full_attn_spans
         return attention_mask
 
+    @staticmethod
+    def _slice_cached_prefix_inputs(
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.Tensor | None,
+        custom_pos_emb: tuple[torch.Tensor, ...] | None,
+        image_mask: torch.Tensor | None,
+        gen_timestep_scatter_index: torch.Tensor | None,
+        query_lens: list[int] | None,
+        seq_lens: list[int] | None,
+        cached_prefix_len: int,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor,
+        tuple[torch.Tensor, ...],
+        torch.Tensor,
+        torch.Tensor,
+        list[int],
+        int,
+    ]:
+        """Remove a cached first-step prefix while preserving full KV lengths."""
+
+        seq_len = int(inputs_embeds.shape[1])
+        if (
+            query_lens is None
+            or seq_lens is None
+            or len(query_lens) != inputs_embeds.shape[0]
+            or len(seq_lens) != len(query_lens)
+        ):
+            raise ValueError("Hunyuan paged prefix execution requires one query/sequence length per row")
+        if any(query_len != seq_len for query_len in query_lens) or any(
+            logical_len != seq_len for logical_len in seq_lens
+        ):
+            raise ValueError(
+                "Hunyuan paged prefix slicing requires an untrimmed first-step tensor: "
+                f"query_lens={query_lens}, tensor_len={seq_len}"
+            )
+        if not 0 < cached_prefix_len < seq_len:
+            raise ValueError(
+                "Hunyuan cached prefix must leave a non-empty first-step suffix: "
+                f"cached={cached_prefix_len}, seq_len={seq_len}"
+            )
+        if (
+            attention_mask is None
+            or position_ids is None
+            or custom_pos_emb is None
+            or image_mask is None
+            or gen_timestep_scatter_index is None
+        ):
+            raise ValueError("Hunyuan paged prefix slicing is missing first-step layout tensors")
+
+        return (
+            inputs_embeds[:, cached_prefix_len:],
+            attention_mask[:, :, cached_prefix_len:, :],
+            position_ids[:, cached_prefix_len:],
+            tuple(pos[:, cached_prefix_len:] for pos in custom_pos_emb),
+            image_mask[:, cached_prefix_len:],
+            gen_timestep_scatter_index - cached_prefix_len,
+            [query_len - cached_prefix_len for query_len in query_lens],
+            seq_len - cached_prefix_len,
+        )
+
     def prepare_inputs_for_generation(
         self,
         input_ids,
@@ -1747,6 +1825,40 @@ class HunyuanImage3Pipeline(
         assert inputs_embeds is not None
         bsz, seq_len, n_embd = inputs_embeds.shape
 
+        paged_kv_cached_prefix_len = 0
+        if first_step and self._uses_scheduler_paged_kv():
+            if not is_forward_context_available():
+                raise RuntimeError("Hunyuan paged prefix execution requires an Omni ForwardContext")
+            paged_kv_cached_prefix_len = int(get_forward_context().paged_kv_cached_prefix_len)
+        if paged_kv_cached_prefix_len:
+            if ar_kv_reuse_len:
+                raise ValueError("Hunyuan cannot combine imported AR KV and Scheduler prefix-cache hits")
+            (
+                inputs_embeds,
+                attention_mask,
+                position_ids,
+                custom_pos_emb,
+                image_mask,
+                gen_timestep_scatter_index,
+                query_lens,
+                seq_len,
+            ) = self._slice_cached_prefix_inputs(
+                inputs_embeds,
+                attention_mask,
+                position_ids,
+                custom_pos_emb,
+                image_mask,
+                gen_timestep_scatter_index,
+                query_lens,
+                seq_lens,
+                paged_kv_cached_prefix_len,
+            )
+            logger.debug(
+                "Hunyuan prefix slice: cached_prefix_len=%d query_len=%d",
+                paged_kv_cached_prefix_len,
+                seq_len,
+            )
+
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         from vllm.forward_context import set_forward_context
 
@@ -1770,6 +1882,7 @@ class HunyuanImage3Pipeline(
                 gen_timestep_scatter_index=gen_timestep_scatter_index,
                 uncond_cfg_prefill=uncond_cfg_prefill,
                 ar_kv_reuse_len=ar_kv_reuse_len,
+                paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
                 full_attn_spans=full_attn_spans,
             )
         hidden_states = outputs[0]

vllm_omni/diffusion/models/hunyuan_image3/request_layout.py

@@ -8,7 +8,9 @@ from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
 
+from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
 from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
@@ -357,7 +359,7 @@ def build_hunyuan_diffusion_kv_requests(
     request: OmniDiffusionRequest,
     prepared_layout: HunyuanPreparedLayout,
 ) -> tuple[DiffusionKVRequest, ...]:
-    """Build one persistent Scheduler KV request per Hunyuan execution row."""
+    """Build allocation-only KV requests, even when prefix caching is disabled."""
 
     tokenizer_output = prepared_layout.tokenizer_output
     cfg_factor = hunyuan_cfg_factor(
@@ -408,7 +410,15 @@ def build_hunyuan_diffusion_kv_requests(
             prefix_len=int(prefix_row[-1].item()),
             target_len=target_len,
             seq_len=int(valid_row[-1].item()),
-            prompt_token_ids=tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist(),
+            # Native AR -> DiT transfer needs token IDs to describe the
+            # transferred prefix.  Local paged prefix caching does not: its
+            # identity is attached later by ``prepare_hunyuan_prefix_cache``.
+            # Keep this conversion out of the disabled/local-only path.
+            prompt_token_ids=(
+                tokenizer_output.tokens[sequence_id, : reusable_lens[sequence_id]].tolist()
+                if request.kv_transfer_params is not None
+                else None
+            ),
             # 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.
@@ -418,6 +428,93 @@ def build_hunyuan_diffusion_kv_requests(
     )
 
 
+def prepare_hunyuan_prefix_cache(request: OmniDiffusionRequest) -> None:
+    """Attach native MM identities / positions when Engine enables caching.
+
+    This is model input adaptation, not a hashing framework. All reference
+    inputs are hashed once, then reused by VAE/ViT subspans and CFG rows. The
+    native block hasher consumes these ranges directly, without token extras.
+    """
+
+    layout = get_hunyuan_prepared_layout(request)
+    if layout is None or not request.diffusion_kv_requests:
+        raise ValueError("Hunyuan prefix caching requires prepared layout and KV requests")
+    sampling = request.sampling_params
+    _, _, _, images, _ = extract_hunyuan_prompt_inputs(
+        [request.prompt], sampling.extra_args or {}, request_id=request.request_id, allow_cond_image=True
+    )
+    reference_digest = None
+    if images:
+        # Match Worker: an explicitly supplied generator overrides the seed.
+        # Conditional VAE samples latents, so image bytes alone are not enough.
+        if sampling.generator is not None:
+            generators = sampling.generator if isinstance(sampling.generator, list) else [sampling.generator]
+            random_state: object = (
+                "generator-state",
+                tuple((str(generator.device), generator.get_state()) for generator in generators),
+            )
+        elif sampling.seed is not None:
+            random_state = ("seed", int(sampling.seed), str(sampling.generator_device or "worker-default"))
+        else:
+            random_state = ("request-local-random-state", request.request_id)
+        # Preserve the conservative whole-reference-set policy. Splitting
+        # independent image identities also needs the VAE RNG sequence modeled.
+        reference_digest = hash_prefix_cache_value(
+            (
+                "hunyuan-reference-v1",
+                [[joint_image_info_to_payload(image) for image in row] for row in images],
+                random_state,
+            )
+        )
+
+    namespace = get_cache_namespace("hunyuan-image3-primary-v4", sampling)
+    output = layout.tokenizer_output
+    joint_rows = output.joint_image_slices or [[] for _ in range(layout.num_branches)]
+    prepared = []
+    for row in request.diffusion_kv_requests:
+        if row.block_hashes or row.num_computed_tokens:
+            raise ValueError("Hunyuan cache inputs must be prepared before KV execution")
+        joint_spans = [(int(span.start or 0), int(span.stop or 0)) for span in joint_rows[row.sequence_id]]
+        features = []
+        for image_slice, (height, width) in layout.rope_image_info[row.sequence_id]:
+            start = int(image_slice.start or 0)
+            end = int(image_slice.stop or output.tokens.shape[1])
+            if start >= row.prefix_len or end <= 0:
+                continue
+            # Joint spans enclose VAE + separator + ViT; RoPE lists subspans.
+            is_reference = any(joint_start <= start and end <= joint_end for joint_start, joint_end in joint_spans)
+            if is_reference and reference_digest is None:
+                raise ValueError("Hunyuan reference-image prefix is missing its canonical inputs")
+            identifier = hash_prefix_cache_value(
+                (
+                    "hunyuan-image-span-v2",
+                    start,
+                    end,
+                    int(height),
+                    int(width),
+                    is_reference,
+                    reference_digest if is_reference else None,
+                )
+            ).hex()
+            features.append(
+                MultiModalFeatureSpec(
+                    data=None,
+                    modality="image",
+                    identifier=identifier,
+                    mm_position=PlaceholderRange(offset=start, length=end - start),
+                )
+            )
+        # Raw model inputs remain in request.prompt/prepared_layout; data=None
+        # here does not imply an encoder hit. These are Scheduler-only KV keys.
+        tokens = tuple(output.tokens[row.sequence_id, : row.prefix_len].tolist())
+        prepared.append((row, tokens, sorted(features, key=lambda feature: feature.mm_position.offset)))
+
+    for row, tokens, features in prepared:
+        row.cache_token_ids = tokens
+        row.mm_features = features
+        row.cache_namespace = namespace
+
+
 def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None:
     prepared_layout = getattr(source, "prepared_layout", None)
     if prepared_layout is None:

vllm_omni/diffusion/registry.py

@@ -749,6 +749,11 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
 }
 
 
+_DIFFUSION_PREFIX_CACHE_FUNCS = {
+    "HunyuanImage3ForCausalMM": "get_hunyuan_image_3_prefix_cache_func",
+}
+
+
 def register_diffusion_model(
     model_arch: str,
     module_name: str,
@@ -757,6 +762,7 @@ def register_diffusion_model(
     post_process_func_name: str | None = None,
     ir_op_priority_func_name: str | None = None,
     action_post_process_func_name: str | None = None,
+    prefix_cache_func_name: str | None = None,
 ) -> None:
     """Register a diffusion model pipeline from an out-of-tree plugin.
 
@@ -782,6 +788,9 @@ def register_diffusion_model(
             for out-of-tree plugins. Action postprocess hooks are no longer
             registered separately; move action handling into
             ``post_process_func_name`` and return a payload/metadata envelope.
+        prefix_cache_func_name: Optional factory for a CPU cache-input hook,
+            called after preprocessing only when paged prefix caching is
+            enabled. It fills existing DiffusionKVRequest cache inputs in place.
     """
     if action_post_process_func_name is not None:
         logger.warning(
@@ -810,6 +819,8 @@ def register_diffusion_model(
         _DIFFUSION_POST_PROCESS_FUNCS[model_arch] = post_process_func_name
     if ir_op_priority_func_name is not None:
         _DIFFUSION_IR_OP_PRIORITY_FUNCS[model_arch] = ir_op_priority_func_name
+    if prefix_cache_func_name is not None:
+        _DIFFUSION_PREFIX_CACHE_FUNCS[model_arch] = prefix_cache_func_name
 
     logger.info(
         "Registered diffusion model %s -> %s.%s",
@@ -860,3 +871,11 @@ def get_diffusion_pre_process_func(od_config: OmniDiffusionConfig):
         return None  # Return None if no pre-processing function is registered (for backward compatibility)
     func_name = _DIFFUSION_PRE_PROCESS_FUNCS[od_config.model_class_name]
     return _load_process_func(od_config, func_name)
+
+
+def get_diffusion_prefix_cache_func(od_config: OmniDiffusionConfig):
+    """Load optional model preparation for native multimodal KV identities."""
+    if uses_diffusers_adapter(od_config) or od_config.model_class_name is None:
+        return None
+    func_name = _DIFFUSION_PREFIX_CACHE_FUNCS.get(od_config.model_class_name)
+    return None if func_name is None else _load_process_func(od_config, func_name)

vllm_omni/diffusion/sched/base_scheduler.py

@@ -38,7 +38,7 @@ from vllm_omni.diffusion.sched.interface import (
     StepBatchSamplingParamsKey,
     _AdmissionWaitDecision,
 )
-from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
 logger = init_logger(__name__)
 
@@ -120,6 +120,12 @@ class BaseScheduler(ABC):
                 scheduler_block_size=scheduler_block_size,
                 hash_block_size=hash_block_size,
                 max_in_flight_tokens=kv_vllm_config.max_in_flight_tokens,
+                enable_prefix_caching=bool(getattr(kv_vllm_config.cache_config, "enable_prefix_caching", False)),
+                prefix_caching_hash_algo=getattr(
+                    kv_vllm_config.cache_config,
+                    "prefix_caching_hash_algo",
+                    "sha256",
+                ),
             )
         else:
             if any(
@@ -182,7 +188,8 @@ class BaseScheduler(ABC):
 
             diffusion_kv_metadata: DiffusionKVMetadata | None = None
             if self._diffusion_kv_manager is not None:
-                if self._diffusion_kv_manager.has_request(request_id):
+                already_reserved = self._diffusion_kv_manager.has_request(request_id)
+                if already_reserved:
                     diffusion_kv_metadata = self._diffusion_kv_manager.get_metadata(request_id)
                 else:
                     matched_tokens: list[int] = []
@@ -210,6 +217,11 @@ class BaseScheduler(ABC):
                         continue
                     if allocation is None:
                         break
+                    # Check the lookup result before registering destination
+                    # pages with the connector: deferred pages can be freed.
+                    if not self._can_schedule_waiting(state):
+                        self._diffusion_kv_manager.free_request(request_id)
+                        break
                     diffusion_kv_metadata = allocation
                     self._kv_request_generations[request_id] = allocation.allocation_generation
                     if self._kv_connector is not None:
@@ -397,7 +409,7 @@ class BaseScheduler(ABC):
         )
 
     @abstractmethod
-    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
+    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
         pass
 
     def has_requests(self) -> bool:
@@ -633,6 +645,11 @@ class BaseScheduler(ABC):
         # Also surface admission failures recorded while schedule() built this
         # output. Older finished ids retained only for Worker cleanup have
         # already been popped by the Engine and are deliberately ignored.
+        if self._diffusion_kv_manager is not None:
+            for request_id, status in statuses.items():
+                if status == DiffusionRequestStatus.FINISHED_COMPLETED:
+                    self._diffusion_kv_manager.publish_request(request_id)
+
         finished_req_ids = {
             request_id for request_id in sched_output.finished_req_ids if request_id in self._request_states
         }

vllm_omni/diffusion/sched/request_scheduler.py

@@ -12,11 +12,12 @@ from vllm_omni.diffusion.sched.interface import (
     DiffusionRequestStatus,
     DiffusionSchedulerOutput,
     RequestBatchSamplingParamsKey,
+    SchedulerRequestState,
     _AdmissionWaitDecision,
 )
 
 if TYPE_CHECKING:
-    from vllm_omni.diffusion.worker.utils import RunnerOutput
+    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
 # LoRA identity is derived from `sampling.lora_request`, not a same-named field
 # on sampling params, so it must be resolved separately from the bulk lookup.
@@ -32,7 +33,7 @@ def _normalize_explicit_sample_solver(value: object | None) -> str | None:
     return str(value).strip().lower()
 
 
-def _normalize_explicit_flow_shift(value: object | None) -> float | None:
+def _normalize_explicit_flow_shift(value: float | str | None) -> float | None:
     """Normalize an explicitly provided flow shift without selecting a default."""
     if value is None:
         return None
@@ -59,6 +60,18 @@ def build_request_batch_sampling_params_key(request: OmniDiffusionRequest) -> Re
 class RequestScheduler(BaseScheduler):
     """Scheduler for static request waves, including admission coalescing."""
 
+    def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool:
+        if not super()._can_schedule_waiting(state):
+            return False
+        manager = self._diffusion_kv_manager
+        if not self._running or manager is None or not manager.has_request(state.request_id):
+            return True
+        # The request-level runner uses one first-step query boundary per wave.
+        # The Manager already aligns all CFG sequences within each request.
+        current = manager.get_metadata(self._running[0]).sequences[0].cached_prefix_len
+        candidate = manager.get_metadata(state.request_id).sequences[0].cached_prefix_len
+        return candidate == current
+
     def get_admission_wait_decision(
         self,
         *,
@@ -98,7 +111,7 @@ class RequestScheduler(BaseScheduler):
     def _build_sampling_params_key(self, request: OmniDiffusionRequest) -> RequestBatchSamplingParamsKey:
         return build_request_batch_sampling_params_key(request)
 
-    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
+    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
         scheduled_request_ids = sched_output.scheduled_request_ids
         if not scheduled_request_ids and not sched_output.finished_req_ids:
             return set()

vllm_omni/diffusion/sched/step_scheduler.py

@@ -16,7 +16,7 @@ from vllm_omni.diffusion.sched.interface import (
 )
 
 if TYPE_CHECKING:
-    from vllm_omni.diffusion.worker.utils import RunnerOutput
+    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
 
 logger = init_logger(__name__)
 
@@ -62,7 +62,7 @@ class StepScheduler(BaseScheduler):
         )
         return request_id
 
-    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: RunnerOutput) -> set[str]:
+    def update_from_output(self, sched_output: DiffusionSchedulerOutput, output: BaseRunnerOutput) -> set[str]:
         scheduled_request_ids = sched_output.scheduled_request_ids
         if not scheduled_request_ids and not sched_output.finished_req_ids:
             return set()
@@ -125,6 +125,7 @@ class StepScheduler(BaseScheduler):
             return self._sequence_length(sampling.timesteps)
         if sampling.sigmas is not None:
             return len(sampling.sigmas)
+        assert sampling.num_inference_steps is not None
         return int(sampling.num_inference_steps)
 
     @staticmethod

vllm_omni/diffusion/vllm_config.py

@@ -152,6 +152,7 @@ class _DiffusionVllmModelConfig:
         return self._require_attention_geometry()[2]
 
     def get_total_num_kv_heads(self) -> int:
+        assert self.hf_text_config is not None
         return self.hf_text_config.num_key_value_heads
 
     @property
@@ -206,6 +207,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
     """Populate native model, parallel, cache, and Scheduler configuration."""
 
     parallel_config = od_config.parallel_config
+    assert parallel_config.data_parallel_size is not None
     vllm_config.parallel_config.tensor_parallel_size = parallel_config.tensor_parallel_size
     vllm_config.parallel_config.data_parallel_size = parallel_config.data_parallel_size
     if parallel_config.enable_expert_parallel and od_config.is_moe:
@@ -225,8 +227,7 @@ def configure_diffusion_vllm_config(vllm_config: VllmConfig, od_config: OmniDiff
         vllm_config.parallel_config.pipeline_parallel_size = parallel_config.pipeline_parallel_size
         vllm_config.cache_config.gpu_memory_utilization = float(getattr(od_config, "gpu_memory_utilization", 0.9))
         vllm_config.cache_config.kv_cache_memory_bytes = getattr(od_config, "kv_cache_memory_bytes", None)
-        # Prefix hashes/publication are intentionally deferred to the next PR.
-        vllm_config.cache_config.enable_prefix_caching = False
+        vllm_config.cache_config.enable_prefix_caching = bool(getattr(od_config, "enable_prefix_caching", False))
         max_rows_per_request = getattr(od_config, "diffusion_kv_max_rows_per_request", None)
         if type(max_rows_per_request) is not int or max_rows_per_request <= 0:
             raise ValueError("paged_scheduler requires a positive diffusion_kv_max_rows_per_request")

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -13,7 +13,7 @@ from __future__ import annotations
 import copy
 import gc
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Sequence
 from contextlib import AbstractContextManager, nullcontext
 from typing import TYPE_CHECKING, Any, cast
 
@@ -97,9 +97,10 @@ def _dit_any_rank_failed(local_failed: bool) -> bool:
     if not torch.distributed.is_initialized():
         return local_failed
     try:
-        from vllm_omni.diffusion.distributed.parallel_state import get_dit_group
+        from vllm_omni.diffusion.distributed import parallel_state
 
-        group = get_dit_group()
+        get_dit_group = getattr(parallel_state, "get_dit_group", None)
+        group = get_dit_group() if get_dit_group is not None else None
     except (AssertionError, ImportError):
         group = None
     if group is None:
@@ -489,6 +490,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
     def prepare_kv_for_forward(self, scheduler_output: DiffusionSchedulerOutput):
         from vllm_omni.diffusion.diffusion_kv.kv_connector import wait_for_kv_load
 
+        assert self.od_config.kv_transfer_config is not None
         timeout = self.od_config.kv_transfer_config.kv_connector_extra_config.get("transfer_timeout", 60.0)
         if self._kv_receive_progress is not None:
             return self._kv_receive_progress.prepare(self._kv_connector, scheduler_output, timeout)
@@ -505,7 +507,7 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
     ) -> int:
         return self.diffusion_kv_backend.get_diffusion_kv_row(request_id, sequence_id, context_id)
 
-    def remove_diffusion_kv_requests(self, request_ids: list[str | tuple[str, int]]) -> int:
+    def remove_diffusion_kv_requests(self, request_ids: Sequence[str | tuple[str, int]]) -> int:
         return self.diffusion_kv_backend.remove_diffusion_kv_requests(request_ids)
 
     def refresh_diffusion_kv_block_table_layout(self) -> None:
@@ -538,13 +540,19 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
                         f"request={request_metadata.request_id!r}, sequence={sequence.sequence_id}, "
                         f"active={active_seq_len}, allocated={sequence.seq_len}"
                     )
+                # Imported AR KV and local hits have separate owners.
+                kv_start_pos = (
+                    sequence.num_computed_tokens
+                    if getattr(self.od_config, "kv_transfer_config", None) is not None
+                    else sequence.cached_prefix_len
+                )
                 prefill_rows.append(
                     DiffusionPagedAttentionRow(
                         request_id=request_metadata.request_id,
                         sequence_id=sequence.sequence_id,
-                        query_len=sequence.seq_len - sequence.num_computed_tokens,
+                        kv_start_pos=kv_start_pos,
+                        query_len=sequence.seq_len - kv_start_pos,
                         seq_len=sequence.seq_len,
-                        kv_start_pos=sequence.num_computed_tokens,
                     )
                 )
                 denoise_rows.append(
@@ -800,17 +808,48 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
                         "Diffusion KV metadata count must match the request batch: "
                         f"metadata={len(diffusion_kv_metadata)}, requests={len(reqs)}"
                     )
-                for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
-                    req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
+                native_kv_transfer = getattr(self.od_config, "kv_transfer_config", None) is not None
+                if native_kv_transfer:
+                    for req, metadata in zip(reqs, diffusion_kv_metadata, strict=True):
+                        # This field is consumed by Hunyuan only for native
+                        # AR->DiT transfer. Local prefix hits must still run
+                        # their VAE/ViT conditioning path.
+                        req.kv_computed_tokens = tuple(seq.num_computed_tokens for seq in metadata.sequences)
                 paged_metadata = self._build_paged_attention_metadata(diffusion_kv_metadata)
+                paged_kv_cached_prefix_len = 0
+                if not native_kv_transfer:
+                    cached_prefix_lens = {row.kv_start_pos for row in paged_metadata.prefill_rows}
+                    if len(cached_prefix_lens) != 1:
+                        raise ValueError(
+                            "One paged request-level forward requires a uniform cached prefix boundary; "
+                            f"got {sorted(cached_prefix_lens)}"
+                        )
+                    paged_kv_cached_prefix_len = next(iter(cached_prefix_lens))
                 paged_kv_runtime, paged_kv_context = self.diffusion_kv_backend.activate_paged_attention_metadata(
                     paged_metadata
                 )
+                if is_primary:
+                    # Trace the boundary actually passed to the model, not a
+                    # speculative lookup. Useful for warm-cache regressions.
+                    for request_metadata in diffusion_kv_metadata:
+                        for sequence in request_metadata.sequences:
+                            logger.debug(
+                                "Diffusion prefix prefill: request_id=%s sequence_id=%d "
+                                "cached_prefix_len=%d prefix_len=%d query_len=%d",
+                                request_metadata.request_id,
+                                sequence.sequence_id,
+                                sequence.cached_prefix_len,
+                                sequence.prefix_len,
+                                sequence.seq_len - sequence.cached_prefix_len,
+                            )
+            else:
+                paged_kv_cached_prefix_len = 0
             with (
                 set_forward_context(
                     vllm_config=self.vllm_config,
                     omni_diffusion_config=od_config,
                     paged_kv_runtime=paged_kv_runtime,
+                    paged_kv_cached_prefix_len=paged_kv_cached_prefix_len,
                     in_diffusion_kv_memory_profile=in_diffusion_kv_memory_profile,
                 ),
                 paged_kv_context,
@@ -1054,7 +1093,10 @@ class DiffusionModelRunner(DiffusionStagePayloadMixin):
                     prepared_layout=getattr(sched_new_req.req, "prepared_layout", None),
                     external_req_id=getattr(sched_new_req.req, "external_req_id", None),
                 )
-                if sched_new_req.diffusion_kv_metadata is not None:
+                if (
+                    sched_new_req.diffusion_kv_metadata is not None
+                    and getattr(self.od_config, "kv_transfer_config", None) is not None
+                ):
                     new_state.extra["kv_computed_tokens"] = tuple(
                         seq.num_computed_tokens for seq in sched_new_req.diffusion_kv_metadata.sequences
                     )