[HY-Imgae3.0] support hunyuan image3 dit fa-fp8 on npu (#3540)

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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:

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: