diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 5a01858a4..f8f92944f 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/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",): 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 new file mode 100644 index 000000000..342103350 --- /dev/null +++ b/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) diff --git a/tests/e2e/accuracy/helpers.py b/tests/e2e/accuracy/helpers.py index 62d3e0a69..517fd9e13 100644 --- a/tests/e2e/accuracy/helpers.py +++ b/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) diff --git a/tests/e2e/accuracy/test_hunyuan_image3.py b/tests/e2e/accuracy/test_hunyuan_image3.py new file mode 100644 index 000000000..2c96ee479 --- /dev/null +++ b/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, + ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py index 4a29e9df9..751bfb21a 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py +++ b/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("") self.end_answer_token_id = self.tokenizer.convert_tokens_to_ids("") self.end_recaption_token_id = self.tokenizer.convert_tokens_to_ids("") + self.end_think_token_id = self.tokenizer.convert_tokens_to_ids("") self.ratio_token_offset = self.tokenizer.convert_tokens_to_ids("") 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["_start"].append(token_count) + if "_start" not in extra_token_pos and len(source) > 0 and source[0] == self.cfg_token_id: + extra_token_pos["_start"].append(token_count) token_count += len(source) extra_token_pos["_end"].append(token_count - 1) + if len(source) > 0 and source[-1] == self.end_think_token_id: + extra_token_pos["_end"].append(token_count - 1) + if len(source) > 0 and source[-1] == self.end_recaption_token_id: + extra_token_pos["_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 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("_end", [None])[0] + recaption_end_pos = extra_token_pos.get("_end", [None])[0] + think_recaption_end_pos = [recaption_end_pos or think_end_pos] + uncond_cfg_start_pos = [extra_token_pos.get("_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): 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 f2cfa3ed2..a4ba514d9 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/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 or + # 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 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 96849a135..3a4d0d2c6 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ 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 or 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 "" in cot and not cot.startswith(""): + cot = "" + cot + return cot + if "" in cot and not cot.startswith(""): + cot = "" + 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], diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i_kv_reuse.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i_kv_reuse.yaml new file mode 100644 index 000000000..23c3bbbb2 --- /dev/null +++ b/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] # + 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