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

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

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