diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index a9312f4b1..ecb993945 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/tests/diffusion/cache/test_cache_backends.py @@ -16,6 +16,7 @@ from unittest.mock import Mock, patch import pytest from vllm_omni.diffusion.cache.cache_dit_backend import ( + CUSTOM_DIT_ENABLERS, CacheDiTBackend, ) from vllm_omni.diffusion.cache.selector import get_cache_backend @@ -122,6 +123,52 @@ class TestCacheDiTBackend: assert "cache_config" in call_args[1] assert backend._last_num_inference_steps == 100 + def test_hunyuan_custom_enabler_registered(self): + """Test HunyuanImage3 custom cache-dit enabler is registered.""" + assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS + + @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter") + @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit") + def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter): + """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh.""" + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline" + mock_pipeline.model = Mock() + mock_pipeline.model.layers = Mock() + + mock_cache_dit.enable_cache = Mock() + mock_cache_dit.refresh_context = Mock() + + backend = CacheDiTBackend({"Fn_compute_blocks": 2}) + backend.enable(mock_pipeline) + + assert backend.enabled is True + assert backend._refresh_func is not None + mock_block_adapter.assert_called_once() + adapter_kwargs = mock_block_adapter.call_args.kwargs + assert adapter_kwargs["transformer"] is mock_pipeline.model + assert adapter_kwargs["blocks"] is mock_pipeline.model.layers + assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4 + assert len(adapter_kwargs["params_modifiers"]) == 1 + mock_cache_dit.enable_cache.assert_called_once() + + backend.refresh(mock_pipeline, num_inference_steps=12) + mock_cache_dit.refresh_context.assert_called_once() + call_args = mock_cache_dit.refresh_context.call_args + assert call_args[0][0] is mock_pipeline.model + assert call_args[1]["num_inference_steps"] == 12 + + def test_enable_hunyuan_pipeline_requires_model_layers(self): + """Test HunyuanImage3 enabler fails with a formatted pipeline class name.""" + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline" + mock_pipeline.model = Mock(spec=[]) + + backend = CacheDiTBackend({"Fn_compute_blocks": 2}) + + with pytest.raises(ValueError, match="HunyuanImage3Pipeline"): + backend.enable(mock_pipeline) + class TestTeaCacheBackend: """Test TeaCacheBackend implementation.""" diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index d7165eda2..32c0c8ba0 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/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, diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml new file mode 100644 index 000000000..eac7f2d63 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml @@ -0,0 +1,38 @@ +# Stage config for running Hunyuan-Image3.0 DiT. +# The following config has been verified on 4x H20 GPU. + +# Stage 0: Diffusion (DiT + VAE) +# This stage receives noise and timesteps and performs denoising + VAE decode +stage_args: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0,1,2,3" + max_batch_size: 1 + engine_args: + model_stage: dit + gpu_memory_utilization: 0.65 + enforce_eager: true + trust_remote_code: true + engine_output_type: image + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + parallel_config: + tensor_parallel_size: 4 + enable_expert_parallel: true + omni_kv_config: + need_recv_cache: true + + final_output: true + final_output_type: image + is_comprehension: false + default_sampling_params: + seed: 42 + +# Runtime edges +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 diff --git a/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml b/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml new file mode 100644 index 000000000..270e0c718 --- /dev/null +++ b/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml @@ -0,0 +1,38 @@ +# Stage config for running Hunyuan-Image3.0 DiT on NPU. +# The following config has been verified on 8x A3-64G NPUs. + +# Stage 0: Diffusion (DiT + VAE) +# This stage receives noise and timesteps and performs denoising + VAE decode. +stage_args: + - stage_id: 0 + stage_type: diffusion + runtime: + devices: "0,1,2,3,4,5,6,7" + max_batch_size: 1 + engine_args: + model_stage: dit + gpu_memory_utilization: 0.65 + enforce_eager: true + trust_remote_code: true + engine_output_type: image + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + parallel_config: + tensor_parallel_size: 8 + enable_expert_parallel: true + omni_kv_config: + need_recv_cache: true + + final_output: true + final_output_type: image + is_comprehension: false + default_sampling_params: + seed: 42 + +# Runtime defaults +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1