diff --git a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py b/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py index 342103350..8a35d8aaa 100644 --- a/tests/diffusion/models/hunyuan_image3/test_image_kv_cache_manager.py +++ b/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) diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py index 81394297d..7ed3c85a0 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3.py +++ b/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." diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index eaf0c2afe..80bb1d7b6 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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 diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 677d67796..b4a588df9 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ 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: , 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})" ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py b/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py index b178b021f..f054fef26 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/prompt_utils.py +++ b/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: