[Feature][Hunyuan image 3.0] AR + DIT with kv reuse. (#3346)

Signed-off-by: dengyunyang <584797741@qq.com>

examples/offline_inference/hunyuan_image3/end2end.py

@@ -203,6 +203,7 @@ def main():
         if isinstance(sp, OmniDiffusionSamplingParams):
             sp.num_inference_steps = args.steps
             sp.guidance_scale = args.guidance_scale
+            sp.guidance_scale_provided = True
             if args.seed is not None:
                 sp.seed = args.seed
             if args.modality in ("text2img",):

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -0,0 +1,440 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ImageKVCacheManager.
+
+Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation.
+"""
+
+from __future__ import annotations
+
+import math
+from contextlib import ExitStack, contextmanager
+from unittest.mock import patch
+
+import pytest
+import torch
+import torch.nn as nn
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer"
+
+NUM_HEADS = 4
+NUM_KV_HEADS = 2
+HEAD_DIM = 16
+IMAGE_TOKEN_LEN = 8
+SCALING = 1.0 / math.sqrt(HEAD_DIM)
+
+
+# ============================================================
+# Mocks + helpers
+# ============================================================
+
+
+class MockAttention(nn.Module):
+    def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs):
+        super().__init__()
+
+    def forward(self, query, key, value, attn_metadata=None, **kwargs):
+        return query
+
+
+@contextmanager
+def patched_mgr_env(sp_size=1):
+    target = _TRANSFORMER_MODULE
+    patches = [
+        patch(f"{target}.get_sequence_parallel_world_size", return_value=sp_size),
+        patch(f"{target}.get_sequence_parallel_rank", return_value=0),
+        patch(f"{target}.Attention", MockAttention),
+    ]
+    with ExitStack() as stack:
+        for p in patches:
+            stack.enter_context(p)
+        yield
+
+
+def _make_cache_mgr(image_token_len=IMAGE_TOKEN_LEN, sp_size=1):
+    with patched_mgr_env(sp_size=sp_size):
+        from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+            ImageKVCacheManager,
+        )
+
+        mgr = ImageKVCacheManager(
+            num_heads=NUM_HEADS,
+            num_kv_heads=NUM_KV_HEADS,
+            head_dim=HEAD_DIM,
+            scaling=SCALING,
+            image_token_len=image_token_len,
+        )
+    return mgr
+
+
+def _make_known_kv(num_tokens, base=0.0):
+    """Create key/value with known values. Token i has all elements = base+i / base+i+0.5."""
+    k = torch.full((num_tokens, NUM_KV_HEADS, HEAD_DIM), 0.0)
+    v = torch.full((num_tokens, NUM_KV_HEADS, HEAD_DIM), 0.0)
+    for i in range(num_tokens):
+        k[i] = base + i
+        v[i] = base + i + 0.5
+    return k, v
+
+
+def _call_mgr(
+    mgr,
+    bs,
+    q_len,
+    seq_len,
+    key_flat,
+    value_flat,
+    first_step=False,
+    uncond_cfg_prefill=False,
+    num_image_tokens=IMAGE_TOKEN_LEN,
+    shard_image_size=None,
+):
+    query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
+    attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
+    mgr(
+        query,
+        key_flat,
+        value_flat,
+        attn_mask,
+        query_lens=[q_len] * bs,
+        seq_lens=[seq_len] * bs,
+        first_step=first_step,
+        uncond_cfg_prefill=uncond_cfg_prefill,
+        num_image_tokens=num_image_tokens,
+        shard_image_size=shard_image_size,
+    )
+
+
+# ============================================================
+# Test 1: No AR KV — basic cache → reuse
+# ============================================================
+
+
+@pytest.mark.parametrize("bs", [1, 2])
+def test_no_ar_kv(bs):
+    """
+    No AR KV injected. Tests the basic first_step cache → update reuse path.
+
+    Sequence layout per batch on first_step (q_len=12, IMAGE_TOKEN_LEN=8):
+        [prompt(3) | image(8) | eoi(1)]
+        image_size = IMAGE_TOKEN_LEN + 1 = 9
+        cached_prompt_len = seq_len - image_size = 12 - 9 = 3
+
+    After first_step:
+        image_kv_cache_map stores prompt tokens [0:3] for each batch (total 3*bs tokens flat).
+
+    Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + eoi(1) + image(8) = 12):
+        _reuse_prompt_kv produces [cached_prompt(3) | new_image(8) | zero_eoi(1)] per batch.
+    """
+    mgr = _make_cache_mgr()
+    assert mgr.image_kv_cache_map is None
+
+    # --- first_step ---
+    q_len = 12
+    k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
+
+    _call_mgr(mgr, bs, q_len=q_len, seq_len=q_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+
+    cached_key, cached_value = mgr.image_kv_cache_map
+    # 3 prompt tokens cached per batch
+    assert cached_key.shape[0] == 3 * bs
+    for b in range(bs):
+        flat_offset = b * q_len
+        cache_offset = b * 3
+        assert torch.allclose(cached_key[cache_offset : cache_offset + 3], k_flat[flat_offset : flat_offset + 3])
+        assert torch.allclose(cached_value[cache_offset : cache_offset + 3], v_flat[flat_offset : flat_offset + 3])
+
+    # --- update step ---
+    img_q_len = IMAGE_TOKEN_LEN
+    update_seq_len = 3 + 1 + img_q_len  # cached_prompt(3) + eoi(1) + image(8) = 12
+    new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
+
+    key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+    val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
+
+    assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
+    for b in range(bs):
+        cache_offset = b * 3
+        img_offset = b * img_q_len
+        # Cached prompt preserved
+        assert torch.allclose(result_k[b, :3], cached_key[cache_offset : cache_offset + 3])
+        # New image tokens
+        assert torch.allclose(result_k[b, 3 : 3 + img_q_len], new_img_k[img_offset : img_offset + img_q_len])
+        # Zero eoi at the end
+        assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
+
+
+# ============================================================
+# Test 2: AR KV, no CFG
+# ============================================================
+
+
+@pytest.mark.parametrize("sp_size", [1, 2])
+def test_ar_kv_no_cfg(sp_size):
+    """
+    AR KV injected, bs=1, no CFG. Tests AR prefix prepend + cache + reuse.
+
+    sp_size=1:
+        first_step input: q_len=12, ar_len=5, seq_len=17
+        Sequence layout: [ar(5) | prompt(3) | image(8) | eoi(1)]
+        image_size = IMAGE_TOKEN_LEN + 1 = 9
+        cached_prompt_len = seq_len - image_size = 17 - 9 = 8 (= ar(5) + prompt(3))
+
+        Update step: q_len=8, seq_len = cached_prompt(8) + eoi(1) + image(8) = 17
+        Result: [cached(8) | new_image(8) | zero_eoi(1)]
+
+    sp_size=2:
+        shard_image_size = 4 (passed externally, simulating SP sharding)
+        cached_prompt_len = seq_len - shard_image_size = 17 - 4 = 13 (= ar(5) + prompt(8))
+        Note: with SP, more of the sequence is treated as "prompt" because
+        image tokens are sharded across ranks.
+
+        Update step: q_len=4 (shard_image_size), seq_len = cached_prompt(13) + shard_image(4) = 17
+        SP path returns only cached prompt [bs, cached_prompt_len, ...] (no eoi concat).
+
+    Both sp_size cases test _cache_prompt_kv and _reuse_prompt_kv directly to avoid
+    needing full SP infrastructure mocks in __call__.
+    """
+    mgr = _make_cache_mgr(sp_size=sp_size)
+    ar_len = 5
+    ar_k, ar_v = _make_known_kv(ar_len, base=100.0)
+    mgr._injected_ar_kv = [(ar_k.clone(), ar_v.clone())]
+
+    # --- first_step: call _cache_prompt_kv directly ---
+    bs, q_len = 1, 12
+    seq_len = q_len + ar_len  # 17
+    k_raw, v_raw = _make_known_kv(q_len, base=1.0)
+    key_4d = k_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
+    val_4d = v_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
+
+    shard_image_size = 4 if sp_size > 1 else None
+    mgr.image_kv_cache_map = None
+    mgr._cache_prompt_kv(key_4d, val_4d, seq_len, shard_image_size)
+
+    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
+        expected_cached_len = seq_len - shard_image_size
+
+    assert cached_key.shape[0] == expected_cached_len
+    # AR KV always at the front
+    assert torch.allclose(cached_key[:ar_len], ar_k)
+    assert torch.allclose(cached_value[:ar_len], ar_v)
+    # Prompt tokens follow AR
+    prompt_cached = expected_cached_len - ar_len
+    assert torch.allclose(cached_key[ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
+    # AR KV consumed
+    assert mgr._injected_ar_kv is None
+
+    # --- update step: call _reuse_prompt_kv directly ---
+    if sp_size == 1:
+        img_q_len = IMAGE_TOKEN_LEN
+        update_seq_len = expected_cached_len + 1 + img_q_len  # 8 + 1 + 8 = 17
+        new_img_k, new_img_v = _make_known_kv(img_q_len, base=50.0)
+
+        key_input = new_img_k.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+        val_input = new_img_v.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+        result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=1)
+
+        assert result_k.shape == (1, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
+        # AR + prompt preserved
+        assert torch.allclose(result_k[0, :ar_len], ar_k)
+        assert torch.allclose(result_k[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
+        # New image tokens
+        assert torch.allclose(result_k[0, expected_cached_len : expected_cached_len + img_q_len], new_img_k)
+        # Zero eoi
+        assert torch.allclose(result_k[0, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
+    else:
+        # SP path: _reuse_prompt_kv returns only cached prompt (no image concat)
+        img_q_len = shard_image_size  # 4
+        update_seq_len = expected_cached_len + img_q_len  # 13 + 4 = 17
+        new_img_k, new_img_v = _make_known_kv(img_q_len, base=50.0)
+
+        key_input = new_img_k.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+        val_input = new_img_v.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+        result_k, result_v = mgr._reuse_prompt_kv(
+            key_input,
+            val_input,
+            update_seq_len,
+            bs=bs,
+            shard_image_size=img_q_len,
+        )
+
+        # SP returns only the cached prompt portion
+        assert result_k.shape == (1, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)
+        assert torch.allclose(result_k[0, :ar_len], ar_k)
+        assert torch.allclose(result_k[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
+
+
+# ============================================================
+# Test 3: AR KV + CFG (sequential & parallel)
+# ============================================================
+
+
+@pytest.mark.parametrize("cfg_parallel,bs", [(False, 2), (True, 1)])
+def test_ar_kv_with_cfg(cfg_parallel, bs):
+    """
+    AR KV + CFG. Tests uncond_cfg_prefill → first_step → update.
+
+    Common setup:
+        positive_reuse_len = 10, negative_reuse_len = 6, neg_uncond_cfg_q_len = 4
+        AR KV: 10 tokens (base=100)
+
+    Sequential CFG (cfg_parallel=False, bs=2):
+        uncond_cfg_prefill (bs=1):
+            Builds neg AR KV = [shared_prefix(6) from pos_ar | neg_prefill_tokens(4)]
+            → _injected_ar_kv becomes [(pos_ar(10), pos_av(10)), (neg_k(10), neg_v(10))]
+
+        first_step (bs=2, q_len=12, seq_len=22):
+            Batch 0 (pos): [pos_ar(10) | prompt(3) | image(8) | eoi(1)]
+            Batch 1 (neg): [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
+            cached_prompt_len per batch = 22 - 9 = 13
+            Total cached = 13 * 2 = 26
+
+        Update (bs=2, seq_len = 13 + 1 + 8 = 22):
+            Result per batch: [cached(13) | new_image(8) | zero_eoi(1)]
+
+    CFG Parallel (cfg_parallel=True, bs=1):
+        This rank handles only the negative branch.
+        uncond_cfg_prefill (bs=1):
+            Same as above: builds neg AR KV.
+            Then we keep only the negative entry (simulating _keep_negative_kv_only).
+            → _injected_ar_kv = [(neg_k(10), neg_v(10))]
+
+        first_step (bs=1, q_len=12, seq_len=22):
+            [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
+            cached_prompt_len = 22 - 9 = 13
+
+        Update (bs=1, seq_len = 13 + 1 + 8 = 22):
+            Result: [cached(13) | new_image(8) | zero_eoi(1)]
+    """
+    positive_reuse_len = 10
+    negative_reuse_len = 6
+    neg_uncond_cfg_q_len = positive_reuse_len - negative_reuse_len  # 4
+
+    mgr = _make_cache_mgr()
+    pos_ar_k, pos_ar_v = _make_known_kv(positive_reuse_len, base=100.0)
+    mgr._injected_ar_kv = [(pos_ar_k.clone(), pos_ar_v.clone())]
+
+    # --- uncond_cfg_prefill ---
+    neg_k, neg_v = _make_known_kv(neg_uncond_cfg_q_len, base=200.0)
+    prefill_seq_len = negative_reuse_len + neg_uncond_cfg_q_len  # 6 + 4 = 10
+
+    _call_mgr(
+        mgr,
+        bs=1,
+        q_len=neg_uncond_cfg_q_len,
+        seq_len=prefill_seq_len,
+        key_flat=neg_k,
+        value_flat=neg_v,
+        first_step=True,
+        uncond_cfg_prefill=True,
+        num_image_tokens=0,
+    )
+
+    # After prefill: _injected_ar_kv = [(pos), (neg)]
+    assert len(mgr._injected_ar_kv) == 2
+    neg_ar_k, neg_ar_v = mgr._injected_ar_kv[1]
+    # neg AR KV = [shared_prefix(6) from pos | neg_prefill(4)], total 10
+    assert neg_ar_k.shape[0] == positive_reuse_len
+    assert torch.allclose(neg_ar_k[:negative_reuse_len], pos_ar_k[:negative_reuse_len])
+    assert torch.allclose(neg_ar_k[negative_reuse_len:], neg_k)
+
+    # --- simulate cfg_parallel: keep only negative ---
+    if cfg_parallel:
+        mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
+
+    # --- first_step ---
+    q_len = 12
+    seq_len = q_len + positive_reuse_len  # 22
+    k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
+
+    _call_mgr(mgr, bs, q_len=q_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+
+    cached_key, cached_value = mgr.image_kv_cache_map
+    cached_prompt_len_per_batch = seq_len - (IMAGE_TOKEN_LEN + 1)  # 22 - 9 = 13
+    assert cached_key.shape[0] == cached_prompt_len_per_batch * bs
+    assert mgr._injected_ar_kv is None
+
+    if not cfg_parallel:
+        # bs=2: batch 0 = pos, batch 1 = neg
+        # Batch 0: pos_ar(10) + prompt(3)
+        assert torch.allclose(cached_key[:positive_reuse_len], pos_ar_k)
+        assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
+        # Batch 1: neg_ar(10) + prompt(3)
+        assert torch.allclose(cached_key[13:23], neg_ar_k)
+        assert torch.allclose(cached_key[23:26], k_flat[q_len : q_len + 3])
+    else:
+        # bs=1: only neg branch
+        assert torch.allclose(cached_key[:positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
+
+    # --- update step ---
+    img_q_len = IMAGE_TOKEN_LEN
+    update_seq_len = cached_prompt_len_per_batch + 1 + img_q_len  # 13 + 1 + 8 = 22
+    new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
+
+    key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+    val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
+    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
+
+    assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
+    for b in range(bs):
+        cache_offset = b * cached_prompt_len_per_batch
+        # Cached prompt preserved
+        assert torch.allclose(
+            result_k[b, :cached_prompt_len_per_batch],
+            cached_key[cache_offset : cache_offset + cached_prompt_len_per_batch],
+        )
+        # New image tokens
+        img_offset = b * img_q_len
+        assert torch.allclose(
+            result_k[b, cached_prompt_len_per_batch : cached_prompt_len_per_batch + img_q_len],
+            new_img_k[img_offset : img_offset + img_q_len],
+        )
+        # Zero eoi
+        assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
+
+
+# ============================================================
+# Test 4: Cross-request isolation
+# ============================================================
+
+
+def test_cross_request_isolation():
+    """
+    Verify leftover image_kv_cache_map from a previous request is NOT treated as AR KV.
+
+    Setup: mgr has stale cache from a prior request (9 tokens, base=999).
+    New request: first_step with q_len=12, no AR KV.
+
+    Expected: stale cache is overwritten. New cache = prompt tokens from current request.
+    The stale values (999.x) must NOT appear in the new cache.
+    """
+    mgr = _make_cache_mgr()
+
+    # Simulate leftover from previous request
+    leftover_k, leftover_v = _make_known_kv(9, base=999.0)
+    mgr.image_kv_cache_map = (leftover_k, leftover_v)
+    assert mgr._injected_ar_kv is None
+
+    # New request first_step
+    bs, seq_len = 1, 12
+    k_flat, v_flat = _make_known_kv(bs * seq_len, base=1.0)
+
+    _call_mgr(mgr, bs, q_len=seq_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
+
+    cached_key, cached_value = mgr.image_kv_cache_map
+    # cached_prompt_len = 12 - 9 = 3
+    assert cached_key.shape[0] == 3
+    # Must be from current request, not stale
+    assert torch.allclose(cached_key[:3], k_flat[:3])
+    assert torch.allclose(cached_value[:3], v_flat[:3])
+    # Stale values must not be present
+    assert not torch.any(cached_key >= 999.0)

tests/e2e/accuracy/helpers.py

@@ -142,6 +142,40 @@ def compute_image_ssim_psnr(
     return ssim_value, psnr_value
 
 
+class CLIPScorer:
+    def __init__(self, model_name_or_path: str = "openai/clip-vit-base-patch16"):
+        from transformers import CLIPModel, CLIPProcessor
+
+        self._model = CLIPModel.from_pretrained(model_name_or_path)
+        self._processor = CLIPProcessor.from_pretrained(model_name_or_path)
+        self._model.eval()
+
+    def score(self, image: Image.Image, text: str) -> float:
+        inputs = self._processor(text=[text], images=[image], return_tensors="pt", padding=True)
+        with torch.no_grad():
+            outputs = self._model(**inputs)
+        img_emb = outputs.image_embeds
+        txt_emb = outputs.text_embeds
+        img_emb = img_emb / img_emb.norm(p=2, dim=-1, keepdim=True)
+        txt_emb = txt_emb / txt_emb.norm(p=2, dim=-1, keepdim=True)
+        similarity = (img_emb * txt_emb).sum(dim=-1)
+        return float(similarity.item() * 100)
+
+    def assert_score(self, *, model_name: str, image: Image.Image, text: str, threshold: float) -> None:
+        value = self.score(image, text)
+        print(f"{model_name} CLIP score:")
+        print(f"  CLIPScore: value={value:.4f}, threshold>={threshold:.4f}, higher_is_better=True")
+        assert value >= threshold, (
+            f"CLIP score below threshold for {model_name}: got {value:.4f}, expected >= {threshold:.4f}."
+        )
+
+
+def _pil_to_clip_tensor(image: Image.Image) -> torch.Tensor:
+    array = np.asarray(image.convert("RGB"), dtype=np.uint8)
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor
+
+
 def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor:
     array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0
     tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+import copy
+import gc
+import tempfile
+import time
+from pathlib import Path
+
+import pytest
+import torch
+import yaml
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import CLIPScorer, assert_similarity, model_output_dir
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
+SEED = 42
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 0
+HEIGHT, WIDTH = 1024, 1024
+PSNR_THRESHOLD = 0.0
+SSIM_THRESHOLD = 0.0
+CLIP_SCORE_THRESHOLD = 20.0
+PROMPT = "A brown and white dog is running on the grass."
+
+# fmt: off
+_BASE_CONFIG = {
+    "stage_args": [
+        {
+            "stage_id": 0, "stage_type": "llm",
+            "runtime": {"process": True, "devices": "0,1", "max_batch_size": 1, "requires_multimodal_data": True},
+            "engine_args": {
+                "model_stage": "AR", "model_arch": "HunyuanImage3ForCausalMM",
+                "worker_cls": "vllm_omni.worker.gpu_ar_worker.GPUARWorker",
+                "scheduler_cls": "vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler",
+                "gpu_memory_utilization": 0.95, "enforce_eager": True, "trust_remote_code": True,
+                "engine_output_type": "latent", "enable_prefix_caching": False,
+                "max_num_batched_tokens": 32768, "tensor_parallel_size": 2, "pipeline_parallel_size": 1,
+                "hf_overrides": {"rope_parameters": {"mrope_section": [0, 32, 32], "rope_type": "default"}},
+            },
+            "is_comprehension": False, "final_output": True, "final_output_type": "text",
+            "default_sampling_params": {
+                "temperature": 0.0, "top_p": 1, "top_k": -1, "max_tokens": 8192,
+                "stop_token_ids": [128025], "detokenize": True, "skip_special_tokens": False,
+            },
+            "output_connectors": {"to_stage_1": "rdma_connector"},
+        },
+        {
+            "stage_id": 1, "stage_type": "diffusion",
+            "runtime": {"process": True, "devices": "2,3", "max_batch_size": 1, "requires_multimodal_data": True},
+            "engine_args": {
+                "model_stage": "dit", "model_arch": "HunyuanImage3ForCausalMM",
+                "enforce_eager": True, "trust_remote_code": True, "distributed_executor_backend": "mp",
+                "parallel_config": {"tensor_parallel_size": 2, "enable_expert_parallel": True},
+            },
+            "engine_input_source": [0],
+            "custom_process_input_func": "vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion",
+            "final_output": True, "final_output_type": "image",
+            "default_sampling_params": {"num_inference_steps": NUM_INFERENCE_STEPS, "guidance_scale": GUIDANCE_SCALE},
+            "input_connectors": {"from_stage_0": "rdma_connector"},
+        },
+    ],
+    "runtime": {
+        "enabled": True,
+        "connectors": {"rdma_connector": {
+            "name": "MooncakeTransferEngineConnector",
+            "extra": {"host": "auto", "zmq_port": 50051, "protocol": "rdma", "device_name": "", "memory_pool_size": 4294967296, "memory_pool_device": "cpu"},
+        }},
+        "edges": [{"from": 0, "to": 1}],
+    },
+}
+# fmt: on
+
+
+def _make_config(enable_kv_reuse: bool, path: Path) -> None:
+    config = copy.deepcopy(_BASE_CONFIG)
+    config["stage_args"][0]["engine_args"]["omni_kv_config"] = {"need_send_cache": enable_kv_reuse}
+    config["stage_args"][1]["engine_args"]["omni_kv_config"] = {"need_recv_cache": enable_kv_reuse}
+    path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+
+
+def _run(stage_config_path: str, output_path: Path) -> tuple[Image.Image, str, float]:
+    from transformers import AutoTokenizer
+
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+    from vllm_omni.platforms import current_omni_platform
+
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+    token_ids = build_prompt_tokens(PROMPT, tokenizer, task="it2i_recaption", sys_type="en_unified")
+
+    with OmniRunner(MODEL_NAME, stage_configs_path=stage_config_path) as runner:
+        params_list = list(runner.omni.default_sampling_params_list)
+        for sp in params_list:
+            if isinstance(sp, OmniDiffusionSamplingParams):
+                sp.num_inference_steps = NUM_INFERENCE_STEPS
+                sp.guidance_scale = GUIDANCE_SCALE
+                sp.seed = SEED
+                sp.generator = torch.Generator(device=current_omni_platform.device_type or "cuda").manual_seed(SEED)
+
+        prompts: list[OmniPromptType] = [
+            {
+                "prompt_token_ids": token_ids,
+                "prompt": PROMPT,
+                "use_system_prompt": "en_unified",
+                "height": HEIGHT,
+                "width": WIDTH,
+            }
+        ]
+        t0 = time.perf_counter()
+        outputs = list(runner.omni.generate(prompts=prompts, sampling_params_list=params_list))
+        elapsed = time.perf_counter() - t0
+
+    assert outputs, "Pipeline produced no outputs"
+    images = None
+    cot_text = ""
+    for out in outputs:
+        ro = getattr(out, "request_output", None)
+        if ro and getattr(ro, "outputs", None):
+            txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
+            if txt:
+                cot_text = txt
+
+        imgs = getattr(out, "images", None)
+        if not imgs and ro and hasattr(ro, "images"):
+            imgs = ro.images
+        if imgs:
+            images = imgs
+
+    assert images, "Pipeline output had no images"
+
+    image = images[0].convert("RGB")
+    image.save(output_path)
+    gc.collect()
+    if torch.accelerator.is_available():
+        torch.accelerator.empty_cache()
+    return image, cot_text, elapsed
+
+
+@pytest.mark.skipif(torch.accelerator.device_count() < 4, reason="Needs 4+ GPUs (2 AR + 2 DiT)")
+def test_text_to_image_alignment(accuracy_artifact_root: Path) -> None:
+    """KV reuse ON vs OFF: same pipeline, same seed → PSNR >= 40 dB."""
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME + "-kv-reuse")
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        tmp = Path(tmpdir)
+        _make_config(False, tmp / "off.yaml")
+        image_no_reuse_1, cot_no_reuse_1, time_no_reuse_1 = _run(
+            str(tmp / "off.yaml"), output_dir / "without_kv_reuse_1.png"
+        )
+        image_no_reuse_2, cot_no_reuse_2, time_no_reuse_2 = _run(
+            str(tmp / "off.yaml"), output_dir / "without_kv_reuse_2.png"
+        )
+        _make_config(True, tmp / "on.yaml")
+        image_reuse, cot_reuse, time_reuse = _run(str(tmp / "on.yaml"), output_dir / "with_kv_reuse.png")
+
+    print("\n--- End-to-end time ---")
+    print(f"  WITHOUT KV reuse 1: {time_no_reuse_1:.2f}s")
+    print(f"  WITHOUT KV reuse 2: {time_no_reuse_2:.2f}s")
+    print(f"  WITH KV reuse:    {time_reuse:.2f}s")
+
+    (output_dir / "cot_with_kv_reuse.txt").write_text(cot_reuse, encoding="utf-8")
+    (output_dir / "cot_without_kv_reuse_1.txt").write_text(cot_no_reuse_1, encoding="utf-8")
+    (output_dir / "cot_without_kv_reuse_2.txt").write_text(cot_no_reuse_2, encoding="utf-8")
+
+    print(f"\n--- CoT WITH KV reuse (len={len(cot_reuse)}) ---\n{cot_reuse}")
+    print(f"\n--- CoT WITHOUT KV reuse 1 (len={len(cot_no_reuse_1)}) ---\n{cot_no_reuse_1}")
+    print(f"\n--- CoT WITHOUT KV reuse 2 (len={len(cot_no_reuse_2)}) ---\n{cot_no_reuse_2}")
+
+    clip_scorer = CLIPScorer()
+    clip_scorer.assert_score(
+        model_name=f"{MODEL_NAME} with-reuse",
+        image=image_reuse,
+        text=PROMPT,
+        threshold=CLIP_SCORE_THRESHOLD,
+    )
+    clip_scorer.assert_score(
+        model_name=f"{MODEL_NAME} without-reuse 1",
+        image=image_no_reuse_1,
+        text=PROMPT,
+        threshold=CLIP_SCORE_THRESHOLD,
+    )
+    clip_scorer.assert_score(
+        model_name=f"{MODEL_NAME} without-reuse 2",
+        image=image_no_reuse_2,
+        text=PROMPT,
+        threshold=CLIP_SCORE_THRESHOLD,
+    )
+
+    assert_similarity(
+        model_name=f"{MODEL_NAME} non-reuse run1 vs run2",
+        vllm_image=image_no_reuse_1,
+        diffusers_image=image_no_reuse_2,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        width=WIDTH,
+        height=HEIGHT,
+    )
+
+    assert_similarity(
+        model_name=f"{MODEL_NAME} KV-reuse vs no-reuse",
+        vllm_image=image_reuse,
+        diffusers_image=image_no_reuse_1,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+        width=WIDTH,
+        height=HEIGHT,
+    )

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py

@@ -35,6 +35,8 @@ class TokenizerEncodeOutput(BaseOutput):
     all_image_slices: list[slice] | None = None
     cond_timestep_scatter_index: torch.Tensor | None = None
     gen_timestep_scatter_index: torch.Tensor | None = None
+    think_recaption_end_pos: list[int | None] | list[list[int | None]] | None = None
+    uncond_cfg_start_pos: list[int | None] | None = None
 
 
 class Conversation:
@@ -59,6 +61,7 @@ class TokenizerWrapper:
         self.cfg_token_id = self.tokenizer.convert_tokens_to_ids("<cfg>")
         self.end_answer_token_id = self.tokenizer.convert_tokens_to_ids("</answer>")
         self.end_recaption_token_id = self.tokenizer.convert_tokens_to_ids("</recaption>")
+        self.end_think_token_id = self.tokenizer.convert_tokens_to_ids("</think>")
         self.ratio_token_offset = self.tokenizer.convert_tokens_to_ids("<img_ratio_0>")
         self.special_token_map = self.tokenizer.added_tokens_encoder
 
@@ -310,8 +313,14 @@ class TokenizerWrapper:
             if key == "text":
                 token_seq.extend(source)  # text token sequence
                 extra_token_pos["<text>_start"].append(token_count)
+                if "<cfg>_start" not in extra_token_pos and len(source) > 0 and source[0] == self.cfg_token_id:
+                    extra_token_pos["<cfg>_start"].append(token_count)
                 token_count += len(source)
                 extra_token_pos["<text>_end"].append(token_count - 1)
+                if len(source) > 0 and source[-1] == self.end_think_token_id:
+                    extra_token_pos["<think>_end"].append(token_count - 1)
+                if len(source) > 0 and source[-1] == self.end_recaption_token_id:
+                    extra_token_pos["<recaption>_end"].append(token_count - 1)
 
             elif key == "gen_image":
                 if isinstance(source, int):
@@ -828,6 +837,12 @@ class TokenizerWrapper:
         # real_pos is the first position of the <pad> token
         real_pos = torch.tensor(extra_token_pos.get("first_pad", [full_seq_token_tensor.shape[0]]), dtype=torch.long)
 
+        # used in AR kv reuse case
+        think_end_pos = extra_token_pos.get("<think>_end", [None])[0]
+        recaption_end_pos = extra_token_pos.get("<recaption>_end", [None])[0]
+        think_recaption_end_pos = [recaption_end_pos or think_end_pos]
+        uncond_cfg_start_pos = [extra_token_pos.get("<cfg>_start", [None])[0]]
+
         return TokenizerEncodeOutput(
             tokens=full_seq_token_tensor,
             timestep_scatter_index=timestep_scatter_index,
@@ -845,6 +860,8 @@ class TokenizerWrapper:
             all_image_slices=all_image_slices,
             cond_timestep_scatter_index=cond_timestep_scatter_index,
             gen_timestep_scatter_index=gen_timestep_scatter_index,
+            think_recaption_end_pos=think_recaption_end_pos,
+            uncond_cfg_start_pos=uncond_cfg_start_pos,
         )
 
     def get_cot_sections(self, cot_text, uncond_kwargs, cot_max_length=None, drop_think=False):

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -863,7 +863,8 @@ class ImageKVCacheManager:
         self.scaling = scaling
         # cache related
         self.image_token_len: int = image_token_len
-        self.image_kv_cache: tuple[torch.Tensor, torch.Tensor] = None
+        self.image_kv_cache_map: tuple[torch.Tensor, torch.Tensor] | None = None
+        self._injected_ar_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
 
         self.sp_size = get_sequence_parallel_world_size()
         self.sp_rank = get_sequence_parallel_rank()
@@ -875,133 +876,119 @@ class ImageKVCacheManager:
             num_kv_heads=self.num_kv_heads,
         )
 
-    def _save_image_kv_caches(
-        self,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        seq_len: int,
-    ) -> None:
-        bs, q_len, num_kv_heads, head_dim = key.shape
-        assert q_len == seq_len, f"for first-step, {q_len} != {seq_len}"
-
-        key = key.reshape(-1, num_kv_heads, head_dim)
-        value = value.reshape(-1, num_kv_heads, head_dim)
-
-        cached_prompt_len = seq_len - self.image_token_len - 1
-        cached_key = [key[:cached_prompt_len], key[seq_len - 1 : seq_len]]
-        cached_value = [value[:cached_prompt_len], value[seq_len - 1 : seq_len]]
-
-        if bs > 1:
-            assert bs == 2, "for cfg case, bs must be 2"
-            cached_key.append(key[seq_len : seq_len + cached_prompt_len])
-            cached_key.append(key[-1:])
-
-            cached_value.append(value[seq_len : seq_len + cached_prompt_len])
-            cached_value.append(value[-1:])
-
-        cached_key = torch.cat(cached_key, dim=0)
-        cached_value = torch.cat(cached_value, dim=0)
-        self.image_kv_cache_map = (cached_key, cached_value)
-
-    def _update_image_kv_caches(
+    def _cache_prompt_kv(
         self,
         key: torch.Tensor,
         value: torch.Tensor,
         seq_len: int,
+        shard_image_size: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        cached_key, cached_value = self.image_kv_cache_map
+        """Cache prompt KV on first_step. Consumes _injected_ar_kv.
+
+        _injected_ar_kv is a per-batch list: [(k,v)] for bs=1,
+        [(pos_k,pos_v), (neg_k,neg_v)] for bs=2 CFG.
+        """
+
+        ar_kv_list = self._injected_ar_kv
+        self._injected_ar_kv = None
+
         bs, q_len, num_kv_heads, head_dim = key.shape
 
-        cached_prompt_len = cached_key.shape[0] // bs - 1
-        assert (cached_prompt_len + 1) == (seq_len - q_len), f"{cached_prompt_len + 1} != {seq_len - q_len}"
+        ar_kv_len = ar_kv_list[0][0].shape[0] if ar_kv_list is not None else 0
+        assert q_len + ar_kv_len == seq_len, f"q_len({q_len}) + ar_kv_len({ar_kv_len}) != seq_len({seq_len})"
 
-        key = key.reshape(-1, num_kv_heads, head_dim)
-        value = value.reshape(-1, num_kv_heads, head_dim)
+        if ar_kv_len > 0:
+            new_keys = []
+            new_values = []
+            for b in range(bs):
+                ar_k, ar_v = ar_kv_list[b]
+                ar_k = ar_k.reshape(1, ar_kv_len, num_kv_heads, head_dim)
+                ar_v = ar_v.reshape(1, ar_kv_len, num_kv_heads, head_dim)
+                k = torch.cat([ar_k, key[b : b + 1]], dim=1)
+                v = torch.cat([ar_v, value[b : b + 1]], dim=1)
+                new_keys.append(k)
+                new_values.append(v)
+            key = torch.cat(new_keys, dim=0)
+            value = torch.cat(new_values, dim=0)
 
-        new_key = [
-            cached_key[:cached_prompt_len],
-            key[:q_len],
-            cached_key[cached_prompt_len : cached_prompt_len + 1],
-        ]
-        new_value = [
-            cached_value[:cached_prompt_len],
-            value[:q_len],
-            cached_value[cached_prompt_len : cached_prompt_len + 1],
-        ]
+        image_size = shard_image_size if self.sp_size > 1 else (self.image_token_len + 1)
+        cached_prompt_len = seq_len - image_size
 
-        if bs > 1:
-            assert bs == 2, "for cfg case, bs must be 2"
-            new_key.append(cached_key[cached_prompt_len + 1 : cached_prompt_len + 1 + cached_prompt_len])
-            new_key.append(key[q_len:])
-            new_key.append(cached_key[-1:])
+        cached_key = key[:, :cached_prompt_len].reshape(-1, num_kv_heads, head_dim)
+        cached_value = value[:, :cached_prompt_len].reshape(-1, num_kv_heads, head_dim)
+        self.image_kv_cache_map = (cached_key, cached_value)
+        return key, value
 
-            new_value.append(cached_value[cached_prompt_len + 1 : cached_prompt_len + 1 + cached_prompt_len])
-            new_value.append(value[q_len:])
-            new_value.append(cached_value[-1:])
-
-        new_key = torch.cat(new_key, dim=0)
-        new_value = torch.cat(new_value, dim=0)
-        new_key = new_key.reshape(bs, seq_len, num_kv_heads, head_dim)
-        new_value = new_value.reshape(bs, seq_len, num_kv_heads, head_dim)
-
-        return new_key.contiguous(), new_value.contiguous()
-
-    def _sp_save_prompt_kv_caches(
+    def _reuse_prompt_kv(
         self,
         key: torch.Tensor,
         value: torch.Tensor,
         seq_len: int,
-        shard_image_size: int,
-    ) -> None:
-        """
-        We don't need to cached last token, since it all false and nevel attn to in non-first step.
-        With sp the key len is [text, image_shard], image_shard in rank include [image token, padding token]
-        The cache is [prompt0, prompt1], see _prepare_attention_mask_for_generation and
-        _update_model_kwargs_for_generation
-        """
-        bs, q_len, num_kv_heads, head_dim = key.shape
-        assert q_len == seq_len, f"for first-step, {q_len} != {seq_len}"
-        key = key.reshape(-1, num_kv_heads, head_dim)
-        value = value.reshape(-1, num_kv_heads, head_dim)
-        cached_prompt_len = seq_len - shard_image_size
-        if bs > 1:
-            assert bs == 2, "for cfg case, bs must be 2"
-
-        cached_key = []
-        cached_value = []
-        for b in range(bs):
-            base = b * seq_len
-            # cache text prompt
-            cached_key.append(key[base : base + cached_prompt_len])
-            cached_value.append(value[base : base + cached_prompt_len])
-
-        cached_key = torch.cat(cached_key, dim=0)
-        cached_value = torch.cat(cached_value, dim=0)
-        self.image_kv_cache_map = (cached_key, cached_value)
-
-    def _sp_get_prompt_kv_caches(
-        self,
-        key: torch.Tensor,
-        seq_len: int,
+        bs: int,
+        shard_image_size: int | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        cached_key, cached_value = self.image_kv_cache_map
-        bs, q_len, kv_head_num, head_dim = key.shape
+        """
+        Reuse cached prompt KV in subsequent denoising steps.
 
+        Non-SP: returns [cached_prompt, new_image, zero_eoi] shaped [bs, seq_len, ...].
+        SP: returns cached prompt only shaped [bs, cached_prompt_len, ...] (used as joint_text).
+        """
+        cached_key, cached_value = self.image_kv_cache_map
+        _, q_len, num_kv_heads, head_dim = key.shape
         cached_prompt_len = cached_key.shape[0] // bs
 
-        assert cached_prompt_len == seq_len - q_len
+        if shard_image_size is not None:
+            assert cached_prompt_len == seq_len - q_len
+            result_k = cached_key.reshape(bs, cached_prompt_len, num_kv_heads, head_dim)
+            result_v = cached_value.reshape(bs, cached_prompt_len, num_kv_heads, head_dim)
+            return result_k.contiguous(), result_v.contiguous()
 
-        joint_text_key = []
-        joint_text_value = []
+        assert cached_prompt_len + 1 == seq_len - q_len, f"{cached_prompt_len + 1} != {seq_len - q_len}"
+
+        key = key.reshape(-1, num_kv_heads, head_dim)
+        value = value.reshape(-1, num_kv_heads, head_dim)
+        zero_eoi = torch.zeros(1, num_kv_heads, head_dim, device=key.device, dtype=key.dtype)
+
+        new_key, new_value = [], []
         for b in range(bs):
             cache_base = b * cached_prompt_len
-            joint_text_key.append(cached_key[cache_base : cache_base + cached_prompt_len])
-            joint_text_value.append(cached_value[cache_base : cache_base + cached_prompt_len])
+            img_base = b * q_len
+            new_key.append(cached_key[cache_base : cache_base + cached_prompt_len])
+            new_key.append(key[img_base : img_base + q_len])
+            new_key.append(zero_eoi)
+            new_value.append(cached_value[cache_base : cache_base + cached_prompt_len])
+            new_value.append(value[img_base : img_base + q_len])
+            new_value.append(zero_eoi)
 
-        joint_text_key = torch.cat(joint_text_key, dim=0).reshape(bs, cached_prompt_len, kv_head_num, head_dim)
-        joint_text_value = torch.cat(joint_text_value, dim=0).reshape(bs, cached_prompt_len, kv_head_num, head_dim)
+        new_key = torch.cat(new_key).reshape(bs, seq_len, num_kv_heads, head_dim)
+        new_value = torch.cat(new_value).reshape(bs, seq_len, num_kv_heads, head_dim)
+        return new_key.contiguous(), new_value.contiguous()
 
-        return joint_text_key.contiguous(), joint_text_value.contiguous()
+    def _build_neg_ar_kv(
+        self,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        seq_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build neg AR KV from pos AR KV shared prefix + uncond text prefill tokens.
+
+        After this call, _injected_ar_kv becomes [(pos_k,pos_v), (neg_k,neg_v)].
+        """
+        bs, q_len_actual, num_kv_heads, head_dim = key.shape
+        assert bs == 1
+
+        shared_prefix_len = seq_len - q_len_actual
+        if shared_prefix_len > 0 and self._injected_ar_kv is not None:
+            pos_key, pos_value = self._injected_ar_kv[0]
+            assert shared_prefix_len <= pos_key.shape[0]
+            pfx_k = pos_key[:shared_prefix_len].reshape(1, shared_prefix_len, num_kv_heads, head_dim)
+            pfx_v = pos_value[:shared_prefix_len].reshape(1, shared_prefix_len, num_kv_heads, head_dim)
+            key = torch.cat([pfx_k, key], dim=1)
+            value = torch.cat([pfx_v, value], dim=1)
+
+        neg_kv = (key.reshape(-1, num_kv_heads, head_dim), value.reshape(-1, num_kv_heads, head_dim))
+        self._injected_ar_kv = [self._injected_ar_kv[0], neg_kv]
+        return key, value
 
     def __call__(
         self,
@@ -1013,6 +1000,7 @@ class ImageKVCacheManager:
     ) -> torch.Tensor:
         self.image_token_len = kwargs.get("num_image_tokens")
         first_step = kwargs.get("first_step")
+        uncond_cfg_prefill = kwargs.get("uncond_cfg_prefill", False)
 
         query_lens = kwargs.get("query_lens")
         seq_lens = kwargs.get("seq_lens")
@@ -1030,27 +1018,33 @@ class ImageKVCacheManager:
         query = query.reshape(bs, q_len, head_num_per_rank, head_dim)
         key = key.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
         value = value.reshape(bs, q_len, kv_head_num_per_rank, head_dim)
-        if first_step:
-            self.image_kv_cache_map = None
-            if self.sp_size <= 1:
-                self._save_image_kv_caches(key, value, seq_len)
-            else:
-                self._sp_save_prompt_kv_caches(key, value, seq_len, shard_image_size)
-                cached_prompt_len = self.image_kv_cache_map[0].shape[0] // bs
-                # joint text part
-                joint_text_query = query[:, :cached_prompt_len, :, :]
-                joint_text_key = key[:, :cached_prompt_len, :, :]
-                joint_text_value = value[:, :cached_prompt_len, :, :]
-                # image part
-                query = query[:, cached_prompt_len:, :, :]
-                key = key[:, cached_prompt_len:, :, :]
-                value = value[:, cached_prompt_len:, :, :]
+
+        if uncond_cfg_prefill:
+            key, value = self._build_neg_ar_kv(key, value, seq_len)
+            if self.sp_size > 1:
+                joint_text_query = query
+                joint_text_key = key
+                joint_text_value = value
+                query = query[:, :0, :, :]
+                key = key[:, :0, :, :]
+                value = value[:, :0, :, :]
+        elif first_step:
+            self.image_kv_cache_map = None  # reset first
+            key, value = self._cache_prompt_kv(key, value, seq_len, shard_image_size)
+            if self.sp_size > 1:
+                local_prompt_len = q_len - shard_image_size
+                joint_text_query = query[:, :local_prompt_len, :, :]
+                joint_text_key = key[:, :local_prompt_len, :, :]
+                joint_text_value = value[:, :local_prompt_len, :, :]
+                query = query[:, local_prompt_len:, :, :]
+                key = key[:, local_prompt_len:, :, :]
+                value = value[:, local_prompt_len:, :, :]
         else:
             if self.sp_size <= 1:
-                key, value = self._update_image_kv_caches(key, value, seq_len)
+                key, value = self._reuse_prompt_kv(key, value, seq_len, bs)
             else:
                 joint_text_query = query[:, :0, :, :]
-                joint_text_key, joint_text_value = self._sp_get_prompt_kv_caches(key, seq_len)
+                joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
 
         key = repeat_kv(key, repeat_num)
         value = repeat_kv(value, repeat_num)
@@ -1072,12 +1066,7 @@ class ImageKVCacheManager:
                 joint_strategy="front",
                 attn_mask=attention_mask,
             )
-        # Compute attention using unified attention layer
         attn_output = self.attn(query, key, value, attn_metadata)
-
-        # attn_output = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, dropout_p=0.0)
-
-        # attn_output = attn_output.transpose(1, 2).contiguous()  # [bs, q_len, heads, head_dim]
         attn_output = attn_output.reshape(bs * q_len, head_num_per_rank, head_dim)
         return attn_output
 
@@ -1851,9 +1840,15 @@ class HunyuanImagePreprocessor(nn.Module):
         position_ids: torch.LongTensor,
         image_token_len: int,
         is_first_step: bool = False,
+        uncond_cfg_prefill: bool = False,
     ):
         # ---------- compute prompt length ----------
-        prompt_len = hidden_states.shape[1] - image_token_len - 1 if is_first_step else 0
+        if uncond_cfg_prefill:
+            prompt_len = hidden_states.shape[1]
+        elif is_first_step:
+            prompt_len = hidden_states.shape[1] - image_token_len - 1
+        else:
+            prompt_len = 0
 
         # ---------- hidden_states ----------
         text_hidden_states = hidden_states[:, :prompt_len, :]
@@ -2225,6 +2220,7 @@ class HunyuanImage3Model(nn.Module):
         seq_lens: list[int] | None = None,
         num_image_tokens: int | None = None,
         gen_timestep_scatter_index: torch.Tensor | None = None,
+        uncond_cfg_prefill: bool = False,
     ) -> tuple | BaseModelOutputWithPast:
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (
@@ -2266,6 +2262,7 @@ class HunyuanImage3Model(nn.Module):
                 position_ids,
                 num_image_tokens,
                 first_step,
+                uncond_cfg_prefill=uncond_cfg_prefill,
             )
             assert len(set(query_lens)) == 1 and len(set(seq_lens)) == 1, (
                 "query_lens and seq_lens must be the same for sequence parallel"
@@ -2331,6 +2328,7 @@ class HunyuanImage3Model(nn.Module):
                 gen_timestep_scatter_index=gen_timestep_scatter_index,
                 shard_image_size=shard_image_size,
                 shard_padding_size=shard_padding_size,
+                uncond_cfg_prefill=uncond_cfg_prefill,
             )
 
             hidden_states = layer_outputs[0]
@@ -2602,6 +2600,191 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                 k: v[s.start : s.stop] if isinstance(v, list) else v[s] for k, v in model_kwargs["vit_kwargs"].items()
             }
 
+    # ==========================================================
+    # Positive/ Negative Reuse Length
+    # ==========================================================
+    def _get_kv_reuse_len(self, model_kwargs, batch_size=1):
+        tokenizer_output = model_kwargs["tokenizer_output"]
+        # from positive prompt
+        think_recaption_end_pos = tokenizer_output.think_recaption_end_pos
+        if not think_recaption_end_pos or think_recaption_end_pos[0][0] is None:
+            return 0, None
+        pos_reuse_kv_len = think_recaption_end_pos[0][0]  # not reuse the last token </think> or <recaption>
+        # from negative prompt
+        if len(tokenizer_output.uncond_cfg_start_pos) > batch_size:
+            neg_reuse_kv_len = tokenizer_output.uncond_cfg_start_pos[batch_size][0]
+        else:  # no negative prompt in non-cfg case.
+            neg_reuse_kv_len = None
+        return pos_reuse_kv_len, neg_reuse_kv_len
+
+    # ==========================================================
+    # Negative CFG Prefill
+    # ==========================================================
+    def _maybe_run_negative_cfg_prefill(
+        self,
+        input_ids,
+        model_kwargs,
+        batch_size,
+        positive_reuse_len,
+        negative_reuse_len,
+        cfg_parallel_ready,
+        cfg_rank,
+        device,
+    ):
+        if cfg_parallel_ready and cfg_rank != 1:
+            return
+
+        assert negative_reuse_len is not None and negative_reuse_len > 0, (
+            "negative_reuse_len should be greater than 0 for running negative CFG prefill"
+        )
+
+        prefill_inputs = self._build_negative_cfg_prefill_inputs(
+            input_ids=input_ids,
+            model_kwargs=model_kwargs,
+            batch_size=batch_size,
+            negative_reuse_len=negative_reuse_len,
+            positive_reuse_len=positive_reuse_len,
+            cfg_parallel_ready=cfg_parallel_ready,
+        )
+
+        if logger.isEnabledFor(logging.DEBUG):
+            import time
+
+            t0 = time.perf_counter()
+        with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=True):
+            self.model.forward_call(**prefill_inputs)
+        if logger.isEnabledFor(logging.DEBUG):
+            torch.accelerator.synchronize()
+            logger.debug("[CFG] Uncond prefill took %.3fs", time.perf_counter() - t0)
+
+        if cfg_parallel_ready:
+            self._keep_negative_kv_only()
+
+    # ==========================================================
+    # Build Negative CFG Prefill Inputs
+    # ==========================================================
+    def _build_negative_cfg_prefill_inputs(
+        self,
+        input_ids,
+        model_kwargs,
+        batch_size,
+        negative_reuse_len,
+        positive_reuse_len,
+        cfg_parallel_ready,
+    ):
+        assert batch_size == 1
+        seq_slice = slice(negative_reuse_len, positive_reuse_len)
+        prefill_seq_len = positive_reuse_len
+        prefill_query_len = positive_reuse_len - negative_reuse_len
+        assert prefill_query_len > 0, "prefill_query_len should be greater than 0"
+
+        if cfg_parallel_ready:
+            batch_slice = slice(None)
+            custom_pos_emb = model_kwargs["custom_pos_emb"]
+        else:
+            batch_slice = slice(batch_size, batch_size * 2)
+            custom_pos_emb = (
+                model_kwargs["custom_pos_emb"][0][batch_slice],
+                model_kwargs["custom_pos_emb"][1][batch_slice],
+            )
+
+        return dict(
+            input_ids=input_ids[batch_slice, seq_slice],
+            attention_mask=model_kwargs["attention_mask"][batch_slice, :, seq_slice, :prefill_seq_len],
+            position_ids=model_kwargs["position_ids"][batch_slice, seq_slice],
+            custom_pos_emb=custom_pos_emb,
+            mode="gen_image",
+            first_step=True,
+            uncond_cfg_prefill=True,
+            query_lens=[prefill_query_len],
+            seq_lens=[prefill_seq_len],
+            num_image_tokens=0,
+        )
+
+    # ==========================================================
+    # Keep Negative KV Only
+    # ==========================================================
+    def _keep_negative_kv_only(self):
+        for layer in self.model.model.layers:
+            mgr = layer.self_attn.image_attn
+            mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
+
+    # ==========================================================
+    # Truncate Prefix
+    # ==========================================================
+    def _truncate_reused_prefix(
+        self,
+        input_ids,
+        model_kwargs,
+        positive_reuse_len,
+    ):
+        input_ids = input_ids[:, positive_reuse_len:]
+        model_kwargs["query_lens"] = [q - positive_reuse_len for q in model_kwargs["query_lens"]]
+        model_kwargs["attention_mask"] = model_kwargs["attention_mask"][:, :, positive_reuse_len:, :]
+        model_kwargs["position_ids"] = model_kwargs["position_ids"][:, positive_reuse_len:]
+
+        # Shift image_mask and gen_timestep_scatter_index to match truncated sequence
+        model_kwargs["image_mask"] = model_kwargs["image_mask"][:, positive_reuse_len:]
+        model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"] - positive_reuse_len
+        model_kwargs["ar_kv_reuse_offset"] = positive_reuse_len
+
+        # cond-image have computed in ar, we may skip it by index in the future.
+        model_kwargs.pop("cond_vae_images", None)
+        model_kwargs.pop("cond_vae_image_mask", None)
+        model_kwargs.pop("cond_vit_image_mask", None)
+        model_kwargs.pop("cond_timestep_scatter_index", None)
+        model_kwargs.pop("cond_timestep", None)
+        model_kwargs.pop("cond_vit_images", None)
+
+        return input_ids
+
+    def _maybe_handle_ar_kv_reuse(
+        self,
+        input_ids,
+        model_kwargs,
+        batch_size,
+        cfg_parallel_ready,
+        cfg_rank,
+        device,
+    ):
+        ar_kv_data = model_kwargs.pop("ar_kv_data", None)
+        if ar_kv_data is None:
+            return input_ids
+
+        # 1. positive prefix len
+        positive_reuse_len, negative_reuse_len = self._get_kv_reuse_len(model_kwargs, batch_size)
+        logger.info(
+            f"Handling AR KV reuse with positive_reuse_len={positive_reuse_len}, "
+            f"negative_reuse_len={negative_reuse_len}"
+        )
+        if positive_reuse_len <= 0:
+            return input_ids
+
+        # 2. inject positive kv
+        self.model.inject_ar_kv_into_layers(ar_kv_data, positive_reuse_len)
+
+        # 3. negative cfg prefill
+        if self.do_classifier_free_guidance:
+            self._maybe_run_negative_cfg_prefill(
+                input_ids=input_ids,
+                model_kwargs=model_kwargs,
+                batch_size=batch_size,
+                positive_reuse_len=positive_reuse_len,
+                negative_reuse_len=negative_reuse_len,
+                cfg_parallel_ready=cfg_parallel_ready,
+                cfg_rank=cfg_rank,
+                device=device,
+            )
+
+        # 4. truncate reused prefix
+        input_ids = self._truncate_reused_prefix(
+            input_ids=input_ids,
+            model_kwargs=model_kwargs,
+            positive_reuse_len=positive_reuse_len,
+        )
+
+        return input_ids
+
     @torch.no_grad()
     def __call__(
         self,
@@ -2743,6 +2926,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
             self._split_model_kwargs_for_cfg_parallel(model_kwargs, batch_size, cfg_rank)
         else:
             cfg_factor = 1 + self.do_classifier_free_guidance
+            cfg_rank = None
 
         b, _, q_len1, seq_len = attention_mask.shape
         query_lens = [q_len1] * b
@@ -2751,6 +2935,12 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
         model_kwargs["seq_lens"] = seq_lens
         model_kwargs["attention_mask"] = attention_mask.to(latents.device)
 
+        # Attempt to reuse KV cache from the AR stage.
+        # Note: the reusable KV length may differ between positive and negative prompts.
+        input_ids = self._maybe_handle_ar_kv_reuse(
+            input_ids, model_kwargs, batch_size, cfg_parallel_ready, cfg_rank, device
+        )
+
         # Sampling loop
         num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
         self._num_timesteps = len(timesteps)
@@ -2824,12 +3014,15 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
                 # Scheduler step (all ranks compute locally in CFG parallel)
                 latents = self.scheduler.step(pred, t, latents, **_scheduler_step_extra_kwargs, return_dict=False)[0]
                 if i != len(timesteps) - 1 and should_compute:
+                    offset = model_kwargs.get("ar_kv_reuse_offset", 0)
                     model_kwargs = self.model._update_model_kwargs_for_generation(  # noqa
                         model_output,
                         model_kwargs,
                     )
                     if input_ids.shape[1] != model_kwargs["position_ids"].shape[1]:
-                        input_ids = torch.gather(input_ids, 1, index=model_kwargs["position_ids"])
+                        # Select using relative positions
+                        index = model_kwargs["position_ids"] - offset
+                        input_ids = torch.gather(input_ids, 1, index=index)
                     attention_mask = model_kwargs.get("attention_mask")
                     b, _, q_len1, seq_len = attention_mask.shape
                     query_lens = [q_len1] * b

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -723,6 +723,26 @@ class HunyuanImage3Pipeline(
                     f"Each message should be a list of dicts or a dict, but got {type(message)}."
                 )
 
+    @staticmethod
+    def _normalize_cot_text(cot: str | None) -> str | None:
+        """
+        Ensure cot_text starts with the appropriate opening tag.
+
+        AR-generated text may omit the leading <think> or <recaption> tag
+        (since it was used as a generation trigger). This normalizes the text
+        so downstream parsing in get_cot_sections works correctly.
+        """
+        if not cot:
+            return cot
+
+        if "</think>" in cot and not cot.startswith("<think>"):
+            cot = "<think>" + cot
+            return cot
+        if "</recaption>" in cot and not cot.startswith("<recaption>"):
+            cot = "<recaption>" + cot
+            return cot
+        return cot
+
     def prepare_model_inputs(
         self,
         prompt=None,
@@ -1030,19 +1050,21 @@ class HunyuanImage3Pipeline(
                 # position ids
                 image_mask = model_kwargs["image_mask"]
                 bsz, seq_len = image_mask.shape
-                index = torch.arange(seq_len, device=image_mask.device).unsqueeze(0).repeat(bsz, 1)
+                offset = model_kwargs.get("ar_kv_reuse_offset", 0)  # should be an absolute position.
+                index = torch.arange(offset, offset + seq_len, device=image_mask.device).unsqueeze(0).repeat(bsz, 1)
                 position_ids = index.masked_select(image_mask.bool()).reshape(bsz, -1)
                 timestep_position_ids = index[
                     torch.arange(bsz), model_kwargs["gen_timestep_scatter_index"][:, -1]
                 ].unsqueeze(-1)
                 updated_model_kwargs["position_ids"] = torch.cat([timestep_position_ids, position_ids], dim=1)
-
                 # attention mask
                 mask_list = []
                 for attention_mask_i, position_ids_i in zip(
                     model_kwargs["attention_mask"], updated_model_kwargs["position_ids"]
                 ):
-                    mask_list.append(torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1)))
+                    mask_list.append(
+                        torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1) - offset)
+                    )
                 attention_mask = torch.stack(mask_list, dim=0)
                 updated_model_kwargs["attention_mask"] = attention_mask
                 updated_model_kwargs["gen_timestep_scatter_index"] = model_kwargs["gen_timestep_scatter_index"]
@@ -1133,11 +1155,12 @@ class HunyuanImage3Pipeline(
         query_lens: list[int] | None = None,
         seq_lens: list[int] | None = None,
         num_image_tokens: int | None = None,
+        uncond_cfg_prefill: bool = False,
     ) -> tuple | CausalMMOutputWithPast:
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         # Sanity Check of Inputs
         self._check_inputs(
-            mode == "gen_image",
+            mode == "gen_image" and not uncond_cfg_prefill,
             "in `gen_image` mode",
             [
                 ("images", images),
@@ -1146,7 +1169,7 @@ class HunyuanImage3Pipeline(
             ],
         )
         self._check_inputs(
-            mode == "gen_image" and first_step,
+            mode == "gen_image" and first_step and not uncond_cfg_prefill,
             "in `gen_image` mode at the first step",
             [
                 ("image_mask", image_mask),
@@ -1169,7 +1192,6 @@ class HunyuanImage3Pipeline(
                 ("vit_kwargs", vit_kwargs),
             ],
         )
-
         custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids)
 
         inputs_embeds = self.model.embed_tokens(input_ids)
@@ -1181,6 +1203,8 @@ class HunyuanImage3Pipeline(
             # For gen_text, make sure gen_timestep_scatter_index is None
             gen_timestep_scatter_index = None
             token_h, token_w = None, None
+        elif uncond_cfg_prefill:
+            token_h, token_w = None, None
         else:
             if first_step:
                 inputs_embeds, token_h, token_w = self.instantiate_vae_image_tokens(
@@ -1226,6 +1250,7 @@ class HunyuanImage3Pipeline(
                 seq_lens=seq_lens,
                 num_image_tokens=num_image_tokens,
                 gen_timestep_scatter_index=gen_timestep_scatter_index,
+                uncond_cfg_prefill=uncond_cfg_prefill,
             )
         hidden_states = outputs[0]
 
@@ -1234,6 +1259,9 @@ class HunyuanImage3Pipeline(
             logits = self.lm_head(hidden_states)
             logits = logits.float()
             diffusion_prediction = None
+        elif uncond_cfg_prefill:
+            logits = None
+            diffusion_prediction = None
         else:
             logits = None
             hidden_states = hidden_states.to(input_ids.device)
@@ -1259,6 +1287,45 @@ class HunyuanImage3Pipeline(
 
         return output
 
+    def inject_ar_kv_into_layers(
+        self,
+        ar_kv_data: dict[int, dict[str, torch.Tensor]],
+        positive_reuse_len: int,
+    ) -> None:
+        """Inject AR-stage KV cache into each layer's ImageKVCacheManager.
+
+        Truncates to positive_reuse_len and sets image_kv_cache_map directly.
+        """
+        for layer in self.model.layers:
+            layer_idx = layer.layer_idx
+            if layer_idx not in ar_kv_data:
+                continue
+            kv = ar_kv_data[layer_idx]
+            cache_mgr = layer.self_attn.image_attn
+            k, v = kv["key"], kv["value"]
+            cache_mgr._injected_ar_kv = [(k[:positive_reuse_len], v[:positive_reuse_len])]
+
+    def _extract_ar_kv_from_request(self, req) -> dict[str, Any]:
+        kv = getattr(getattr(req, "sampling_params", None), "past_key_values", None)
+        if kv is None:
+            return {}
+        key_cache = getattr(kv, "key_cache", None)
+        value_cache = getattr(kv, "value_cache", None)
+        if not key_cache or not value_cache:
+            return {}
+        ar_kv_data = {
+            i: {"key": k, "value": v}
+            for i, (k, v) in enumerate(zip(key_cache, value_cache))
+            if k is not None and v is not None
+        }
+        if not ar_kv_data:
+            return {}
+        logger.info(
+            f"[AR KV Reuse] Extracted {len(ar_kv_data)} layers of AR KV, "
+            f"each with length: {next(iter(ar_kv_data.values()))['key'].shape}"
+        )
+        return {"ar_kv_data": ar_kv_data}
+
     def forward(
         self,
         req: OmniDiffusionRequest,
@@ -1294,7 +1361,9 @@ class HunyuanImage3Pipeline(
         for p in req.prompts:
             extra = p.get("extra", {}) if isinstance(p, dict) else {}
             cot_text_list.append(extra.get("ar_generated_text") or None)
-        cot_text = cot_text_list if any(t is not None for t in cot_text_list) else None
+        cot_text = (
+            [self._normalize_cot_text(t) for t in cot_text_list] if any(t is not None for t in cot_text_list) else None
+        )
 
         batch_cond_image_info: list[list[JointImageInfo]] | None = None
         if any(not isinstance(p, str) for p in req.prompts):
@@ -1329,6 +1398,9 @@ class HunyuanImage3Pipeline(
             logger.info("HunyuanImage3.0 runs without classifier-free guidance when guidance_scale <= 1.0.")
         image_size = (height, width)
 
+        # ---- AR KV Reuse: extract injected KV from request ----
+        ar_kv_kwargs = self._extract_ar_kv_from_request(req)
+
         model_inputs = self.prepare_model_inputs(
             prompt=prompt,
             cot_text=cot_text,
@@ -1340,6 +1412,9 @@ class HunyuanImage3Pipeline(
             guidance_scale=guidance_scale,
             batch_cond_image_info=batch_cond_image_info,
         )
+
+        model_inputs.update(ar_kv_kwargs)
+
         outputs = self._generate(**model_inputs, **kwargs)
         return DiffusionOutput(
             output=outputs[0],

vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i_kv_reuse.yaml

@@ -0,0 +1,89 @@
+stage_args:
+  # Stage 0: AR Model
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      process: true
+      devices: "0,1"
+      max_batch_size: 1
+      requires_multimodal_data: true  # AR needs the original image
+    engine_args:
+      model_stage: AR
+      model_arch: HunyuanImage3ForCausalMM
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.95
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # AR outputs latent for DiT
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      tensor_parallel_size: 2
+      pipeline_parallel_size: 1
+      hf_overrides:
+        rope_parameters:
+          mrope_section: [0, 32, 32]
+          rope_type: default
+      omni_kv_config:
+        need_send_cache: true
+    is_comprehension: false  # Generation task, not comprehension
+    final_output: true
+    final_output_type: text
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1
+      top_k: -1
+      max_tokens: 8192
+      stop_token_ids: [128025]  # <answer>
+      detokenize: true  # DiT bridge consumes ar_generated_text; let the AR engine produce it
+      skip_special_tokens: False
+    output_connectors:
+      to_stage_1: rdma_connector
+
+  # Stage 1: Diffusion (DiT + VAE)
+  # Receives latents from AR stage, performs denoising + VAE decode
+  - stage_id: 1
+    stage_type: diffusion
+    runtime:
+      process: true
+      devices: "2,3"
+      max_batch_size: 1
+      requires_multimodal_data: true  # May need condition images
+    engine_args:
+      model_stage: dit
+      model_arch: HunyuanImage3ForCausalMM
+      enforce_eager: true
+      trust_remote_code: true
+      distributed_executor_backend: "mp"
+      omni_kv_config:
+        need_recv_cache: true  # Receive AR KV cache from stage 0
+      parallel_config:
+        tensor_parallel_size: 2
+        enable_expert_parallel: true
+    engine_input_source: [0]  # Input from AR stage
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion
+    final_output: true
+    final_output_type: image
+    default_sampling_params:
+      num_inference_steps: 50
+      guidance_scale: 0
+    input_connectors:
+      from_stage_0: rdma_connector
+
+
+# Top-level runtime config
+runtime:
+  enabled: true
+  connectors:
+    rdma_connector:
+      name: MooncakeTransferEngineConnector
+      extra:
+        host: "auto"
+        zmq_port: 50051
+        protocol: "rdma"
+        device_name: ""
+        memory_pool_size: 4294967296
+        memory_pool_device: "cpu"
+  edges:
+    - from: 0  # AR → Diffusion
+      to: 1