[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/stepcache/__init__.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+StepCache: velocity-based step skipping for DreamZero-style DiT denoising.
+
+Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the
+pipeline denoise loop when successive velocity predictions are highly aligned.
+
+Usage:
+    from vllm_omni import Omni
+
+    omni = Omni(
+        model="...",
+        cache_backend="step_cache",
+    )
+
+    # Alternative: environment variable
+    # export DIFFUSION_CACHE_BACKEND=step_cache
+"""
+
+from vllm_omni.diffusion.cache.stepcache.backend import (
+    CUSTOM_STEP_CACHE_DIT_ENABLERS,
+    CUSTOM_STEPCACHE_ENABLERS,
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_DIT_CONFIG_ATTR,
+    STEP_CACHE_DIT_STATE_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheBackend,
+    StepCacheDiTBackend,
+    StepCacheDiTConfig,
+    StepCacheDiTState,
+    enable_dreamzero_step_cache_dit,
+    enable_dreamzero_stepcache,
+    get_step_cache_dit_config,
+    get_step_cache_dit_state,
+    get_stepcache_config,
+    get_stepcache_state,
+    is_step_cache_dit_active,
+    is_stepcache_active,
+)
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
+
+__all__ = [
+    "CUSTOM_STEP_CACHE_DIT_ENABLERS",
+    "CUSTOM_STEPCACHE_ENABLERS",
+    "STEP_CACHE_CONFIG_ATTR",
+    "STEP_CACHE_DIT_CONFIG_ATTR",
+    "STEP_CACHE_DIT_STATE_ATTR",
+    "STEP_CACHE_STATE_ATTR",
+    "StepCacheBackend",
+    "StepCacheConfig",
+    "StepCacheDiTBackend",
+    "StepCacheDiTConfig",
+    "StepCacheDiTState",
+    "StepCacheState",
+    "enable_dreamzero_step_cache_dit",
+    "enable_dreamzero_stepcache",
+    "get_step_cache_dit_config",
+    "get_step_cache_dit_state",
+    "get_stepcache_config",
+    "get_stepcache_state",
+    "is_step_cache_dit_active",
+    "is_stepcache_active",
+]

vllm_omni/diffusion/cache/stepcache/backend.py

@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+StepCache backend implementation.
+
+This module provides the stepcache backend that implements the CacheBackend
+interface for DreamZero-style velocity-based step skipping in the denoise loop.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheBackend
+from vllm_omni.diffusion.cache.stepcache.config import (
+    STEP_CACHE_CONFIG_ATTR,
+    STEP_CACHE_STATE_ATTR,
+    StepCacheConfig,
+)
+from vllm_omni.diffusion.cache.stepcache.state import StepCacheState
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+logger = init_logger(__name__)
+
+# Backward-compatible aliases for older imports and pipeline attribute names.
+STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR
+STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR
+StepCacheDiTConfig = StepCacheConfig
+StepCacheDiTState = StepCacheState
+
+
+def is_stepcache_active(pipeline: Any) -> bool:
+    """Return True when stepcache is enabled on *pipeline*."""
+    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
+    return isinstance(cfg, StepCacheConfig) and cfg.enabled
+
+
+def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None:
+    cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None)
+    return cfg if isinstance(cfg, StepCacheConfig) else None
+
+
+def get_stepcache_state(pipeline: Any) -> StepCacheState | None:
+    state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None)
+    return state if isinstance(state, StepCacheState) else None
+
+
+def is_step_cache_dit_active(pipeline: Any) -> bool:
+    return is_stepcache_active(pipeline)
+
+
+def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None:
+    return get_stepcache_config(pipeline)
+
+
+def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None:
+    return get_stepcache_state(pipeline)
+
+
+def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """Enable stepcache for DreamZeroPipeline."""
+    cache_config = StepCacheConfig.from_diffusion_cache_config(config)
+    setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config)
+    setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config))
+
+    logger.info(
+        "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)",
+        cache_config.enabled,
+        cache_config.sim_thresholds,
+        cache_config.skip_countdowns,
+    )
+
+
+enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache
+
+CUSTOM_STEPCACHE_ENABLERS = {
+    "DreamZeroPipeline": enable_dreamzero_stepcache,
+}
+
+CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS
+
+
+class StepCacheBackend(CacheBackend):
+    """
+    Velocity cosine step-skipping cache backend for DreamZero.
+
+    Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported
+    pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to
+    decide whether to run ``predict_noise``.
+
+    Example:
+        >>> from vllm_omni.diffusion.data import DiffusionCacheConfig
+        >>> backend = StepCacheBackend(DiffusionCacheConfig())
+        >>> backend.enable(pipeline)
+        >>> backend.refresh(pipeline, num_inference_steps=16)
+    """
+
+    def enable(self, pipeline: Any) -> None:
+        """Enable stepcache on a supported pipeline."""
+        pipeline_type = pipeline.__class__.__name__
+
+        if pipeline_type in CUSTOM_STEPCACHE_ENABLERS:
+            logger.info("Using custom stepcache enabler for model: %s", pipeline_type)
+            CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config)
+        else:
+            raise ValueError(
+                f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}"
+            )
+
+        self.enabled = True
+
+    def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh stepcache state for a new generation."""
+        state = get_stepcache_state(pipeline)
+        if state is not None:
+            state.reset()
+            if verbose:
+                logger.debug(
+                    "stepcache state refreshed (num_inference_steps=%d)",
+                    num_inference_steps,
+                )
+        elif verbose and is_stepcache_active(pipeline):
+            logger.warning("stepcache config active but state not found on pipeline")
+
+
+StepCacheDiTBackend = StepCacheBackend

vllm_omni/diffusion/cache/stepcache/config.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Configuration for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from vllm_omni.diffusion.data import DiffusionCacheConfig
+
+# Pipeline attributes set by StepCacheBackend.enable().
+STEP_CACHE_CONFIG_ATTR = "_stepcache_config"
+STEP_CACHE_STATE_ATTR = "_stepcache_state"
+
+
+@dataclass(frozen=True)
+class StepCacheConfig:
+    """Runtime config for velocity-based step skipping.
+
+    Skips entire DiT forwards when successive video velocity predictions are
+    highly aligned (cosine similarity above configured thresholds).
+
+    Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``.
+    """
+
+    enabled: bool = True
+    min_history_steps: int = 2
+    max_history: int = 2
+    sim_thresholds: tuple[float, ...] = (0.95, 0.93)
+    skip_countdowns: tuple[int, ...] = (4, 2)
+
+    @classmethod
+    def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig:
+        enabled = config.step_cache_dit_enabled
+
+        thresholds = tuple(config.velocity_sim_thresholds)
+        countdowns = tuple(config.velocity_skip_countdowns)
+        if len(thresholds) != len(countdowns):
+            raise ValueError(
+                "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; "
+                f"got {len(thresholds)} and {len(countdowns)}."
+            )
+
+        return cls(
+            enabled=enabled,
+            min_history_steps=config.step_cache_dit_min_history,
+            max_history=config.step_cache_dit_max_history,
+            sim_thresholds=thresholds,
+            skip_countdowns=countdowns,
+        )

vllm_omni/diffusion/cache/stepcache/state.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Mutable state for step-level DiT velocity caching."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig
+
+
+class StepCacheState:
+    """Per-generation mutable state (skip countdown)."""
+
+    def __init__(self, config: StepCacheConfig) -> None:
+        self.config = config
+        self.skip_countdown = 0
+
+    def reset(self) -> None:
+        self.skip_countdown = 0
+
+    def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool:
+        """Return True when the DiT forward should execute at this scheduler step."""
+        if not self.config.enabled:
+            return True
+
+        if len(prev_predictions) < self.config.min_history_steps:
+            return True
+
+        if self.skip_countdown > 1:
+            self.skip_countdown -= 1
+            return False
+        if self.skip_countdown == 1:
+            self.skip_countdown = 0
+            return True
+
+        v_last = prev_predictions[-1][0].flatten(1).float()
+        v_prev = prev_predictions[-2][0].flatten(1).float()
+        sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean()
+
+        for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns):
+            if sim > threshold:
+                self.skip_countdown = countdown
+                return False
+
+        return True
+
+    def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None:
+        while len(prev_predictions) > self.config.max_history:
+            prev_predictions.pop(0)