[Feature] Add cache-dit support for HunyuanImage3 (#1848)

Signed-off-by: KexiongYu <yukexiong1@huawei.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -606,6 +606,76 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
     return refresh_cache_context
 
 
+def enable_cache_for_hunyuan_image3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for HunyuanImage3 pipeline.
+
+    HunyuanImage3 stores its main transformer stack at ``pipeline.model`` with
+    decoder blocks in ``pipeline.model.layers``.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    transformer = getattr(pipeline, "model", None)
+    if transformer is None or not hasattr(transformer, "layers"):
+        raise ValueError(
+            f"HunyuanImage3 cache-dit enabler expects pipeline.model.layers, got pipeline={pipeline.__class__.__name__}"
+        )
+
+    logger.info(
+        f"Enabling cache-dit on HunyuanImage3 model: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+    # HunyuanImage3 decoder layers are single-stream `hidden_states` blocks
+    # (`HunyuanImage3DecoderLayer.forward(hidden_states, ...)`) while caller
+    # still reads outputs as tuple (`layer_outputs[0]`). Pattern_4 matches this:
+    # hidden-only input with hidden-first tuple output.
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_4,
+            params_modifiers=[modifier],
+            check_forward_pattern=False,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = getattr(pipeline, "model", None)
+        if transformer is None:
+            raise ValueError("HunyuanImage3 cache-dit refresh expects pipeline.model to exist.")
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 class BagelCachedContextManager(CachedContextManager):
     """
     Custom CachedContextManager for Bagel that safely handles NaiveCache objects
@@ -1064,6 +1134,7 @@ CUSTOM_DIT_ENABLERS.update(
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
         "Wan22TI2VPipeline": enable_cache_for_wan22,
+        "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
         "FluxPipeline": enable_cache_for_flux,
         "Flux2KleinPipeline": enable_cache_for_flux2_klein,
         "LongCatImagePipeline": enable_cache_for_longcat_image,