diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 16f7d8f06..aed06bcb9 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -77,6 +77,24 @@ def parse_args(): parser.add_argument("--log-stats", action="store_true", default=False) parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.") parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.") + parser.add_argument( + "--diffusion-kv-cache-dtype", + type=str, + default=None, + help="Diffusion attention KV cache dtype, for example 'fp8'. Separate from vLLM --kv-cache-dtype.", + ) + parser.add_argument( + "--diffusion-kv-cache-skip-steps", + type=str, + default=None, + help="Denoising step selector to keep diffusion KV cache in native dtype, for example '0,1,4-6'.", + ) + parser.add_argument( + "--diffusion-kv-cache-skip-layers", + type=str, + default=None, + help="Transformer layer selector to keep diffusion KV cache in native dtype, for example '0-2,10'.", + ) parser.add_argument( "--additional-config", type=str, @@ -140,6 +158,9 @@ def main(): "init_timeout": args.init_timeout, "enforce_eager": args.enforce_eager, "mode": _MODALITY_MODE[args.modality], + "diffusion_kv_cache_dtype": args.diffusion_kv_cache_dtype, + "diffusion_kv_cache_skip_steps": args.diffusion_kv_cache_skip_steps, + "diffusion_kv_cache_skip_layers": args.diffusion_kv_cache_skip_layers, } if additional_config is not None: @@ -237,6 +258,9 @@ def main(): print(f" Inference steps: {args.steps}") print(f" Guidance scale: {args.guidance_scale}") print(f" Seed: {args.seed}") + print(f" diffusion_kv_cache_dtype: {args.diffusion_kv_cache_dtype}") + print(f" diffusion_kv_cache_skip_steps: {args.diffusion_kv_cache_skip_steps}") + print(f" diffusion_kv_cache_skip_layers: {args.diffusion_kv_cache_skip_layers}") if args.modality == "text2img": print(f" Output size: {args.width}x{args.height}") if args.image_path: 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 4edcfb6ca..c39965f67 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -72,6 +72,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelOutput, ) from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.layers.rope import RotaryEmbedding from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding @@ -868,7 +869,15 @@ class ImageKVCacheManager: Manages specialized caching and updating of KV-Cache for image tokens in multimodal models. """ - def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int, scaling: float, image_token_len: int = 4097): + def __init__( + self, + num_heads: int, + num_kv_heads: int, + head_dim: int, + scaling: float, + image_token_len: int = 4097, + prefix: str = "", + ): """ Args: image_token_len: Number of tokens per image (including special placeholders), @@ -892,6 +901,7 @@ class ImageKVCacheManager: causal=False, softmax_scale=self.scaling, num_kv_heads=self.num_kv_heads, + prefix=f"{prefix}.attn" if prefix else "", ) def _cache_prompt_kv( @@ -1650,6 +1660,7 @@ class HunYuanAttention(nn.Module): num_kv_heads=self.num_kv_heads, scaling=self.scaling, image_token_len=4097, + prefix=f"{prefix}.image_attn", ) self.image_rope2d_emb = HunYuanRotary2DEmbedder( num_heads=self.num_heads, @@ -3021,6 +3032,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): + set_forward_context_denoise_step_idx(i) if cfg_parallel_ready: # CFG parallel: each rank forwards its own branch (no batch doubling) latent_model_input = latents @@ -3109,6 +3121,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline): if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() + set_forward_context_denoise_step_idx(None) + if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor: latents = latents / self.vae.config.scaling_factor if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor: