[Bugfix]Fix Hunyuan Image3 denoise flow alignment (#3857)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py

@@ -23,6 +23,7 @@ NUM_HEADS = 4
 NUM_KV_HEADS = 2
 HEAD_DIM = 16
 IMAGE_TOKEN_LEN = 8
+SUFFIX_TOKEN_LEN = 3
 SCALING = 1.0 / math.sqrt(HEAD_DIM)
 
 
@@ -79,6 +80,10 @@ def _make_known_kv(num_tokens, base=0.0):
     return k, v
 
 
+def _gen_timestep_index(bs: int, current_start: int) -> torch.Tensor:
+    return torch.full((bs, 1), current_start, dtype=torch.long)
+
+
 def _call_mgr(
     mgr,
     bs,
@@ -90,6 +95,8 @@ def _call_mgr(
     uncond_cfg_prefill=False,
     num_image_tokens=IMAGE_TOKEN_LEN,
     shard_image_size=None,
+    gen_timestep_scatter_index=None,
+    position_ids=None,
 ):
     query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
     attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
@@ -104,6 +111,8 @@ def _call_mgr(
         uncond_cfg_prefill=uncond_cfg_prefill,
         num_image_tokens=num_image_tokens,
         shard_image_size=shard_image_size,
+        gen_timestep_scatter_index=gen_timestep_scatter_index,
+        position_ids=position_ids,
     )
 
 
@@ -117,54 +126,65 @@ 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
+    Sequence layout per batch on first_step (q_len=14, IMAGE_TOKEN_LEN=8):
+        [prompt(3) | current timestep/image(8) | suffix(3)]
+        gen_timestep_scatter_index = 3
+        cached_prompt_len = 3
 
     After first_step:
-        image_kv_cache_map stores prompt tokens [0:3] for each batch (total 3*bs tokens flat).
+        image_kv_cache_map stores prompt tokens [0:3] for each batch.
 
-    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.
+    Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + current(8) = 11):
+        _reuse_prompt_kv produces [cached_prompt(3) | new_current(8)] per batch.
     """
     mgr = _make_cache_mgr()
     assert mgr.image_kv_cache_map is None
 
     # --- first_step ---
-    q_len = 12
+    prompt_len = 3
+    q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     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)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=q_len,
+        seq_len=q_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
     # 3 prompt tokens cached per batch
-    assert cached_key.shape[0] == 3 * bs
+    assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
+    assert torch.equal(mgr.image_kv_cache_lens, torch.full((bs,), prompt_len))
     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])
+        assert torch.allclose(cached_key[b, :prompt_len], k_flat[flat_offset : flat_offset + prompt_len])
+        assert torch.allclose(cached_value[b, :prompt_len], v_flat[flat_offset : flat_offset + prompt_len])
 
     # --- update step ---
     img_q_len = IMAGE_TOKEN_LEN
-    update_seq_len = 3 + 1 + img_q_len  # cached_prompt(3) + eoi(1) + image(8) = 12
+    update_seq_len = prompt_len + img_q_len  # cached_prompt(3) + current(8) = 11
     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)
+    position_ids = torch.arange(prompt_len, prompt_len + img_q_len).repeat(bs, 1)
+    result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs, position_ids=position_ids)
 
     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])
+        assert torch.allclose(result_k[b, :prompt_len], cached_key[b, :prompt_len])
         # 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))
+        assert torch.allclose(
+            result_k[b, prompt_len : prompt_len + img_q_len],
+            new_img_k[img_offset : img_offset + img_q_len],
+        )
 
 
 # ============================================================
@@ -178,16 +198,17 @@ 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))
+        first_step input: q_len=14, ar_len=5, seq_len=19
+        Sequence layout: [ar(5) | prompt(3) | current(8) | suffix(3)]
+        gen_timestep_scatter_index = 3
+        cached_prompt_len = ar(5) + prompt(3) = 8
 
-        Update step: q_len=8, seq_len = cached_prompt(8) + eoi(1) + image(8) = 17
-        Result: [cached(8) | new_image(8) | zero_eoi(1)]
+        Update step: q_len=8, seq_len = cached_prompt(8) + current(8) = 16
+        Result: [cached(8) | new_current(8)]
 
     sp_size=2:
         shard_image_size = 4 (passed externally, simulating SP sharding)
+        gen_timestep_scatter_index = 8
         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.
@@ -204,15 +225,23 @@ def test_ar_kv_no_cfg(sp_size):
     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
+    bs = 1
+    q_len = 12 if sp_size > 1 else 14
+    seq_len = q_len + ar_len
     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
+    current_start = 8 if sp_size > 1 else 3
     mgr.image_kv_cache_map = None
-    mgr._cache_prompt_kv(key_4d, val_4d, seq_len, shard_image_size)
+    mgr._cache_prompt_kv(
+        key_4d,
+        val_4d,
+        seq_len,
+        shard_image_size,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
     if sp_size == 1:
@@ -222,25 +251,32 @@ def test_ar_kv_no_cfg(sp_size):
         # 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
+    assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)
     # AR KV always at the front
-    assert torch.allclose(cached_key[:ar_len], ar_k)
-    assert torch.allclose(cached_value[:ar_len], ar_v)
+    assert torch.allclose(cached_key[0, :ar_len], ar_k)
+    assert torch.allclose(cached_value[0, :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])
+    assert torch.allclose(cached_key[0, 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
+        update_seq_len = expected_cached_len + img_q_len  # 8 + 8 = 16
         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)
+        position_ids = torch.arange(expected_cached_len, expected_cached_len + img_q_len).reshape(1, img_q_len)
+        result_k, result_v = mgr._reuse_prompt_kv(
+            key_input,
+            val_input,
+            update_seq_len,
+            bs=1,
+            position_ids=position_ids,
+        )
 
         assert result_k.shape == (1, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
         # AR + prompt preserved
@@ -248,8 +284,6 @@ def test_ar_kv_no_cfg(sp_size):
         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
@@ -292,13 +326,12 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             → _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
+            Batch 0 (pos): [pos_ar(10) | prompt(3) | current(8) | suffix(3)]
+            Batch 1 (neg): [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
+            cached_prompt_len per batch = 10 + 3 = 13
 
-        Update (bs=2, seq_len = 13 + 1 + 8 = 22):
-            Result per batch: [cached(13) | new_image(8) | zero_eoi(1)]
+        Update (bs=2, seq_len = 13 + 8 = 21):
+            Result per batch: [cached(13) | new_current(8)]
 
     CFG Parallel (cfg_parallel=True, bs=1):
         This rank handles only the negative branch.
@@ -308,11 +341,11 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             → _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
+            [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
+            cached_prompt_len = 13
 
-        Update (bs=1, seq_len = 13 + 1 + 8 = 22):
-            Result: [cached(13) | new_image(8) | zero_eoi(1)]
+        Update (bs=1, seq_len = 13 + 8 = 21):
+            Result: [cached(13) | new_current(8)]
     """
     positive_reuse_len = 10
     negative_reuse_len = 6
@@ -351,46 +384,62 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
         mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
 
     # --- first_step ---
-    q_len = 12
+    prompt_len = 3
+    q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     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)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=q_len,
+        seq_len=seq_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     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
+    cached_prompt_len_per_batch = positive_reuse_len + prompt_len
+    assert cached_key.shape == (bs, cached_prompt_len_per_batch, NUM_KV_HEADS, HEAD_DIM)
     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])
+        assert torch.allclose(cached_key[0, :positive_reuse_len], pos_ar_k)
+        assert torch.allclose(cached_key[0, 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])
+        assert torch.allclose(cached_key[1, :positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[1, positive_reuse_len:13], 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])
+        assert torch.allclose(cached_key[0, :positive_reuse_len], neg_ar_k)
+        assert torch.allclose(cached_key[0, 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
+    update_seq_len = cached_prompt_len_per_batch + img_q_len  # 13 + 8 = 21
     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)
+    position_ids = torch.arange(cached_prompt_len_per_batch, cached_prompt_len_per_batch + img_q_len).repeat(bs, 1)
+    result_k, result_v = mgr._reuse_prompt_kv(
+        key_input,
+        val_input,
+        update_seq_len,
+        bs=bs,
+        position_ids=position_ids,
+    )
 
     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],
+            cached_key[b, :cached_prompt_len_per_batch],
         )
         # New image tokens
         img_offset = b * img_q_len
@@ -398,8 +447,6 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
             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))
 
 
 # ============================================================
@@ -412,7 +459,7 @@ 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.
+    New request: first_step with q_len=14, 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.
@@ -421,20 +468,33 @@ def test_cross_request_isolation():
 
     # 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)
+    mgr.image_kv_cache_map = (
+        leftover_k.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
+        leftover_v.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
+    )
     assert mgr._injected_ar_kv is None
 
     # New request first_step
-    bs, seq_len = 1, 12
+    bs = 1
+    prompt_len = 3
+    seq_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
     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)
+    _call_mgr(
+        mgr,
+        bs,
+        q_len=seq_len,
+        seq_len=seq_len,
+        key_flat=k_flat,
+        value_flat=v_flat,
+        first_step=True,
+        gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
+    )
 
     cached_key, cached_value = mgr.image_kv_cache_map
-    # cached_prompt_len = 12 - 9 = 3
-    assert cached_key.shape[0] == 3
+    assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
     # 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])
+    assert torch.allclose(cached_key[0, :prompt_len], k_flat[:prompt_len])
+    assert torch.allclose(cached_value[0, :prompt_len], v_flat[:prompt_len])
     # Stale values must not be present
     assert not torch.any(cached_key >= 999.0)

tests/e2e/accuracy/test_hunyuan_image3.py

@@ -77,9 +77,9 @@ THRESHOLDS = {
     "text_prefix_match": 10,  # First 10 characters must match exactly
     "cot_semantic_sim": 0.9,  # Full CoT semantic similarity
     # Image comparison
-    "clip_score": 85,  # CLIP image semantic similarity
-    "ssim": 0.20,  # Structural similarity
-    "psnr": 11.0,  # Peak signal-to-noise ratio (dB)
+    "clip_score": 90,  # CLIP image semantic similarity
+    "ssim": 0.26,  # Structural similarity
+    "psnr": 12.5,  # Peak signal-to-noise ratio (dB)
 }
 
 QUANT_PROMPT = "A brown and white dog is running on the grass."

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -891,6 +891,7 @@ class ImageKVCacheManager:
         # cache related
         self.image_token_len: int = image_token_len
         self.image_kv_cache_map: tuple[torch.Tensor, torch.Tensor] | None = None
+        self.image_kv_cache_lens: torch.Tensor | None = None
         self._injected_ar_kv: list[tuple[torch.Tensor, torch.Tensor]] | None = None
 
         self.sp_size = get_sequence_parallel_world_size()
@@ -904,12 +905,23 @@ class ImageKVCacheManager:
             prefix=f"{prefix}.attn" if prefix else "",
         )
 
+    @staticmethod
+    def _get_current_starts(
+        gen_timestep_scatter_index: torch.Tensor | None,
+        ar_kv_len: int = 0,
+    ) -> torch.Tensor:
+        assert gen_timestep_scatter_index is not None, (
+            "`gen_timestep_scatter_index` is required to locate the generated image timestep token for image KV reuse."
+        )
+        return gen_timestep_scatter_index[:, -1].to(dtype=torch.long) + ar_kv_len
+
     def _cache_prompt_kv(
         self,
         key: torch.Tensor,
         value: torch.Tensor,
         seq_len: int,
         shard_image_size: int | None = None,
+        gen_timestep_scatter_index: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """Cache prompt KV on first_step. Consumes _injected_ar_kv.
 
@@ -939,12 +951,26 @@ class ImageKVCacheManager:
             key = torch.cat(new_keys, dim=0)
             value = torch.cat(new_values, dim=0)
 
-        image_size = shard_image_size if self.sp_size > 1 else (self.image_token_len + 1)
-        cached_prompt_len = seq_len - image_size
+        cached_prompt_lens = self._get_current_starts(gen_timestep_scatter_index, ar_kv_len)
+        assert torch.all(cached_prompt_lens <= key.shape[1]), (
+            f"cached_prompt_lens({cached_prompt_lens.tolist()}) must be <= key length({key.shape[1]})"
+        )
+        if shard_image_size is not None:
+            expected_prompt_len = seq_len - shard_image_size
+            assert torch.all(cached_prompt_lens == expected_prompt_len), (
+                f"cached_prompt_lens({cached_prompt_lens.tolist()}) != "
+                f"seq_len({seq_len}) - shard_image_size({shard_image_size})"
+            )
 
-        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)
+        max_cached_prompt_len = int(cached_prompt_lens.max().item())
+        cached_key = key.new_zeros(bs, max_cached_prompt_len, num_kv_heads, head_dim)
+        cached_value = value.new_zeros(bs, max_cached_prompt_len, num_kv_heads, head_dim)
+        for b in range(bs):
+            cached_prompt_len = int(cached_prompt_lens[b].item())
+            cached_key[b, :cached_prompt_len] = key[b, :cached_prompt_len]
+            cached_value[b, :cached_prompt_len] = value[b, :cached_prompt_len]
         self.image_kv_cache_map = (cached_key, cached_value)
+        self.image_kv_cache_lens = cached_prompt_lens
         return key, value
 
     def _reuse_prompt_kv(
@@ -954,42 +980,36 @@ class ImageKVCacheManager:
         seq_len: int,
         bs: int,
         shard_image_size: int | None = None,
+        position_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         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).
+        Non-SP: concatenates cached prompt KV with current image KV.
+        SP: returns cached prompt KV only (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
+        _, q_len, _, _ = key.shape
+        assert cached_key.dim() == 4 and cached_key.shape[0] == bs
+        max_cached_prompt_len = cached_key.shape[1]
 
         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()
+            assert max_cached_prompt_len + q_len == seq_len, (
+                f"max_cached_prompt_len({max_cached_prompt_len}) + q_len({q_len}) != seq_len({seq_len})"
+            )
+            return cached_key.contiguous(), cached_value.contiguous()
 
-        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
-            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)
-
-        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)
+        assert max_cached_prompt_len + q_len == seq_len, (
+            f"max_cached_prompt_len({max_cached_prompt_len}) + q_len({q_len}) != seq_len({seq_len})"
+        )
+        assert self.image_kv_cache_lens is not None
+        if position_ids is not None:
+            assert position_ids.shape == (bs, q_len)
+            assert torch.all(position_ids[:, 0] == self.image_kv_cache_lens), (
+                "The first current position must immediately follow each sample's cached prompt KV."
+            )
+        new_key = torch.cat([cached_key, key], dim=1)
+        new_value = torch.cat([cached_value, value], dim=1)
         return new_key.contiguous(), new_value.contiguous()
 
     def _build_neg_ar_kv(
@@ -1032,7 +1052,7 @@ class ImageKVCacheManager:
 
         query_lens = kwargs.get("query_lens")
         seq_lens = kwargs.get("seq_lens")
-        shard_image_size = kwargs.get("shard_image_size")
+        shard_image_size = kwargs.get("shard_image_size") if self.sp_size > 1 else None
         bs = len(query_lens)
         q_len = query_lens[0]
         seq_len = seq_lens[0]
@@ -1058,7 +1078,14 @@ class ImageKVCacheManager:
                 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)
+            self.image_kv_cache_lens = None
+            key, value = self._cache_prompt_kv(
+                key,
+                value,
+                seq_len,
+                shard_image_size,
+                kwargs.get("gen_timestep_scatter_index"),
+            )
             if self.sp_size > 1:
                 local_prompt_len = seq_len - shard_image_size
                 join_query_len = query.shape[1] - shard_image_size
@@ -1070,7 +1097,7 @@ class ImageKVCacheManager:
                 value = value[:, local_prompt_len:, :, :]
         else:
             if self.sp_size <= 1:
-                key, value = self._reuse_prompt_kv(key, value, seq_len, bs)
+                key, value = self._reuse_prompt_kv(key, value, seq_len, bs, position_ids=kwargs.get("position_ids"))
             else:
                 joint_text_query = query[:, :0, :, :]
                 joint_text_key, joint_text_value = self._reuse_prompt_kv(key, value, seq_len, bs, shard_image_size)
@@ -1814,6 +1841,7 @@ class HunyuanImage3DecoderLayer(nn.Module):
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
             attention_mask=attention_mask,
+            position_ids=position_ids,
             custom_pos_emb=custom_pos_emb,
             past_key_value=past_key_value,
             output_attentions=output_attentions,
@@ -1878,12 +1906,21 @@ class HunyuanImagePreprocessor(nn.Module):
         image_token_len: int,
         is_first_step: bool = False,
         uncond_cfg_prefill: bool = False,
+        gen_timestep_scatter_index: torch.Tensor | None = None,
     ):
         # ---------- compute prompt length ----------
         if uncond_cfg_prefill:
             prompt_len = hidden_states.shape[1]
         elif is_first_step:
-            prompt_len = hidden_states.shape[1] - image_token_len - 1
+            assert gen_timestep_scatter_index is not None, (
+                "`gen_timestep_scatter_index` is required to split prompt and "
+                "generated image tokens for sequence parallel image generation."
+            )
+            prompt_lens = gen_timestep_scatter_index[:, -1]
+            assert torch.all(prompt_lens == prompt_lens[0]), (
+                "Sequence parallel image generation requires the generated timestep position to match across the batch."
+            )
+            prompt_len = int(prompt_lens[0].item())
         else:
             prompt_len = 0
 
@@ -2327,6 +2364,7 @@ class HunyuanImage3Model(nn.Module):
                 num_image_tokens,
                 first_step,
                 uncond_cfg_prefill=uncond_cfg_prefill,
+                gen_timestep_scatter_index=gen_timestep_scatter_index,
             )
             assert len(set(query_lens)) == 1 and len(set(seq_lens)) == 1, (
                 "query_lens and seq_lens must be the same for sequence parallel"
@@ -2334,15 +2372,15 @@ class HunyuanImage3Model(nn.Module):
             prompt_size = text_hidden_states.shape[1]
             shard_image_size = image_hidden_states.shape[1]
 
-            if first_step:  # [image tokens, last token]
-                shard_padding_size = shard_image_size * sp_world_size - (num_image_tokens + 1)
-            else:  # [image tokens, last token]
+            if first_step:
+                shard_padding_size = shard_image_size * sp_world_size - (origin_query_len - prompt_size)
+            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]
             else:
                 seq_lens = [x - y for x, y in zip(seq_lens, query_lens)]
-                seq_lens = [seq_len + shard_image_size - 1 for seq_len in seq_lens]
+                seq_lens = [seq_len + shard_image_size for seq_len in seq_lens]
             query_lens = [prompt_size + shard_image_size for _ in query_lens]
 
             if prompt_size > 0:
@@ -2357,10 +2395,6 @@ class HunyuanImage3Model(nn.Module):
                 position_ids = image_position_ids
                 custom_pos_emb = (image_custom_pos_emb_sin, image_custom_pos_emb_cos)
 
-            if not first_step:
-                assert torch.all(~attention_mask[..., -1]), "The last token should not be attended to"
-                attention_mask = attention_mask[..., :-1]  # we attention without last token
-
             if shard_padding_size > 0:
                 B, H, Q, K = attention_mask.shape
                 pad = shard_padding_size
@@ -3091,15 +3125,11 @@ 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]:
-                        # Select using relative positions
-                        index = model_kwargs["position_ids"] - offset
-                        input_ids = torch.gather(input_ids, 1, index=index)
+                    input_ids = None
                     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

@@ -1001,7 +1001,7 @@ class HunyuanImage3Pipeline(
         if inputs_embeds is not None and past_key_values is None:
             model_inputs = {"inputs_embeds": inputs_embeds}
         else:
-            if input_ids.shape[1] != kwargs["position_ids"].shape[1]:  # in decode steps
+            if input_ids is not None and input_ids.shape[1] != kwargs["position_ids"].shape[1]:  # in decode steps
                 input_ids = torch.gather(input_ids, dim=1, index=kwargs["position_ids"])
             model_inputs = {"input_ids": input_ids}
 
@@ -1082,12 +1082,27 @@ class HunyuanImage3Pipeline(
 
                 # attention mask
                 mask_list = []
-                for attention_mask_i, position_ids_i in zip(
-                    model_kwargs["attention_mask"], updated_model_kwargs["position_ids"]
+                current_starts = timestep_position_ids.reshape(-1)
+                max_current_start = int(current_starts.max().item())
+                for attention_mask_i, position_ids_i, current_start_i in zip(
+                    model_kwargs["attention_mask"], updated_model_kwargs["position_ids"], current_starts
                 ):
-                    mask_list.append(
-                        torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1) - offset)
+                    query_mask = torch.index_select(attention_mask_i, dim=1, index=position_ids_i.reshape(-1) - offset)
+                    current_start = int(current_start_i.item())
+                    prefix_mask = torch.index_select(
+                        query_mask,
+                        dim=2,
+                        index=torch.arange(current_start, device=image_mask.device),
                     )
+                    if current_start < max_current_start:
+                        prefix_pad = query_mask.new_zeros(
+                            query_mask.shape[0],
+                            query_mask.shape[1],
+                            max_current_start - current_start,
+                        )
+                        prefix_mask = torch.cat([prefix_mask, prefix_pad], dim=2)
+                    current_mask = torch.index_select(query_mask, dim=2, index=position_ids_i.reshape(-1))
+                    mask_list.append(torch.cat([prefix_mask, current_mask], dim=2))
                 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"]
@@ -1219,9 +1234,14 @@ class HunyuanImage3Pipeline(
         )
         custom_pos_emb = self.get_pos_emb(custom_pos_emb, position_ids)
 
-        inputs_embeds = self.model.embed_tokens(input_ids)
-
-        bsz, seq_len, n_embd = inputs_embeds.shape
+        if input_ids is not None:
+            inputs_embeds = self.model.embed_tokens(input_ids)
+            bsz, seq_len, n_embd = inputs_embeds.shape
+        else:
+            inputs_embeds = None
+            bsz = images.shape[0] if isinstance(images, torch.Tensor) else len(images)
+            seq_len = 0
+            n_embd = self.config.hidden_size
 
         # Instantiate placeholder tokens: <timestep>, <img> for the gen image
         if mode == "gen_text":
@@ -1232,6 +1252,7 @@ class HunyuanImage3Pipeline(
             token_h, token_w = None, None
         else:
             if first_step:
+                assert inputs_embeds is not None
                 inputs_embeds, token_h, token_w = self.instantiate_vae_image_tokens(
                     inputs_embeds, images, timestep, image_mask
                 )
@@ -1253,6 +1274,8 @@ class HunyuanImage3Pipeline(
             inputs_embeds = self.instantiate_vit_image_tokens(
                 inputs_embeds, cond_vit_images, cond_vit_image_mask, vit_kwargs
             )
+        assert inputs_embeds is not None
+        bsz, seq_len, n_embd = inputs_embeds.shape
 
         # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
         from vllm.forward_context import set_forward_context
@@ -1291,7 +1314,7 @@ class HunyuanImage3Pipeline(
             diffusion_prediction = None
         else:
             logits = None
-            hidden_states = hidden_states.to(input_ids.device)
+            hidden_states = hidden_states.to(inputs_embeds.device)
             assert hidden_states.numel() == bsz * seq_len * n_embd, (
                 f"Shape mismatch: {hidden_states.shape} cannot reshape to ({bsz}, {seq_len}, {n_embd})"
             )

vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py

@@ -219,7 +219,7 @@ def build_prompt(
     effective_sys_type = sys_type or preset_sys_type
 
     system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
-    sys_text = system_prompt or ""
+    sys_text = system_prompt.strip() if system_prompt is not None else ""
 
     has_image_input = task in ("i2t", "it2i")
     if has_image_input:
@@ -288,7 +288,7 @@ def build_prompt_tokens(
         )
 
     system_prompt = get_system_prompt(effective_sys_type, bot_task, custom_system_prompt)
-    sys_text = system_prompt or ""
+    sys_text = system_prompt.strip() if system_prompt is not None else ""
 
     ids: list[int] = [bos_id]
     if sys_text: