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

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

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)