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

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

vllm_omni/diffusion/cache/selector.py

@@ -3,6 +3,7 @@ from typing import Any
 from vllm_omni.diffusion.cache.base import CacheBackend
 from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend
 from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend
+from vllm_omni.diffusion.cache.stepcache import StepCacheBackend
 from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
@@ -14,14 +15,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
     - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface
     - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface
     - mag_cache: Uses MagCacheBackend with enable()/refresh() interface
+    - step_cache: Uses StepCacheBackend for DreamZero velocity step skip
 
     Args:
-        cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None).
+        cache_backend: Cache backend name ("cache_dit", "tea_cache",
+            "mag_cache", "step_cache", or None).
         cache_config: Cache configuration (dict or DiffusionCacheConfig instance).
 
     Returns:
-        Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend)
-        if cache_backend is set, None otherwise.
+        Cache backend instance if cache_backend is set, None otherwise.
 
     Raises:
         ValueError: If cache_backend is unsupported.
@@ -38,7 +40,10 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack
         return TeaCacheBackend(cache_config)
     elif cache_backend == "mag_cache":
         return MagCacheBackend(cache_config)
+    elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"):
+        return StepCacheBackend(cache_config)
     else:
         raise ValueError(
-            f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'"
+            f"Unsupported cache backend: {cache_backend}. "
+            "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'"
         )