[Perf][DreamZero]dreamzero supports CUDA graph, torch.compile, and DiT Caching. (#4213)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/diffusion/cache/test_stepcache.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Unit tests for StepCacheBackend and StepCacheState."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.selector import get_cache_backend
+from vllm_omni.diffusion.cache.stepcache import (
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheBackend,
+    StepCacheConfig,
+    StepCacheState,
+    is_stepcache_active,
+)
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]:
+    v_prev = torch.randn(1, dim)
+    if sim >= 0.99:
+        v_last = v_prev.clone()
+    else:
+        v_last = torch.randn(1, dim)
+    return [(v_prev,), (v_last,)]
+
+
+class TestStepCacheConfig:
+    def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self):
+        enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True))
+        disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False))
+        assert enabled.enabled is True
+        assert disabled.enabled is False
+
+
+class TestStepCacheState:
+    def test_disabled_always_runs(self):
+        state = StepCacheState(StepCacheConfig(enabled=False))
+        assert state.should_run_step(_make_prev_predictions(0.99)) is True
+
+    def test_insufficient_history_always_runs(self):
+        state = StepCacheState(StepCacheConfig(min_history_steps=2))
+        assert state.should_run_step([(torch.randn(1, 4),)]) is True
+
+    def test_high_similarity_skips_with_countdown(self):
+        state = StepCacheState(
+            StepCacheConfig(
+                sim_thresholds=(0.95,),
+                skip_countdowns=(2,),
+            )
+        )
+        preds = _make_prev_predictions(0.99)
+        assert state.should_run_step(preds) is False
+        assert state.skip_countdown == 2
+        assert state.should_run_step(preds) is False
+        assert state.should_run_step(preds) is True
+
+    def test_low_similarity_runs(self):
+        state = StepCacheState(StepCacheConfig())
+        preds = _make_prev_predictions(0.5)
+        assert state.should_run_step(preds) is True
+
+    def test_trim_history(self):
+        state = StepCacheState(StepCacheConfig(max_history=2))
+        history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)]
+        state.trim_history(history)
+        assert len(history) == 2
+
+
+class TestStepCacheBackend:
+    def test_enable_dreamzero_pipeline(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
+
+        backend = StepCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR)
+        assert isinstance(cfg, StepCacheConfig)
+        assert is_stepcache_active(mock_pipeline)
+        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
+        assert isinstance(state, StepCacheState)
+
+    def test_refresh_resets_state(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "DreamZeroPipeline"
+        backend = StepCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+        state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR)
+        state.skip_countdown = 3
+        backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False)
+        assert state.skip_countdown == 0
+
+    def test_enable_unsupported_pipeline_raises(self):
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "QwenImagePipeline"
+        backend = StepCacheBackend(DiffusionCacheConfig())
+
+        with pytest.raises(ValueError, match="step_cache backend does not support"):
+            backend.enable(mock_pipeline)
+
+    def test_selector_returns_backend(self):
+        backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]})
+        assert isinstance(backend, StepCacheBackend)
+        assert backend.config.velocity_sim_thresholds == [0.9]
+
+    def test_legacy_backend_aliases(self):
+        assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend)
+        assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend)