Skip to content

vllm_omni.diffusion.cache.stepcache

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

Modules:

Name Description
backend

StepCache backend implementation.
config

Configuration for step-level DiT velocity caching.
state

Mutable state for step-level DiT velocity caching.

CUSTOM_STEPCACHE_ENABLERS module-attribute 

CUSTOM_STEPCACHE_ENABLERS = {
    "DreamZeroPipeline": enable_dreamzero_stepcache
}

CUSTOM_STEP_CACHE_DIT_ENABLERS module-attribute 

CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS

STEP_CACHE_CONFIG_ATTR module-attribute 

STEP_CACHE_CONFIG_ATTR = '_stepcache_config'

STEP_CACHE_DIT_CONFIG_ATTR module-attribute 

STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR

STEP_CACHE_DIT_STATE_ATTR module-attribute 

STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR

STEP_CACHE_STATE_ATTR module-attribute 

STEP_CACHE_STATE_ATTR = '_stepcache_state'

StepCacheDiTBackend module-attribute 

StepCacheDiTBackend = StepCacheBackend

StepCacheDiTConfig module-attribute 

StepCacheDiTConfig = StepCacheConfig

StepCacheDiTState module-attribute 

StepCacheDiTState = StepCacheState

enable_dreamzero_step_cache_dit module-attribute 

enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache

StepCacheBackend

Bases: 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)

enable 

enable(pipeline: Any) -> None

Enable stepcache on a supported pipeline.

refresh 

refresh(
    pipeline: Any,
    num_inference_steps: int,
    verbose: bool = True,
) -> None

Refresh stepcache state for a new generation.

StepCacheConfig dataclass

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 class-attribute instance-attribute 

enabled: bool = True

max_history class-attribute instance-attribute 

max_history: int = 2

min_history_steps class-attribute instance-attribute 

min_history_steps: int = 2

sim_thresholds class-attribute instance-attribute 

sim_thresholds: tuple[float, ...] = (0.95, 0.93)

skip_countdowns class-attribute instance-attribute 

skip_countdowns: tuple[int, ...] = (4, 2)

from_diffusion_cache_config classmethod 

from_diffusion_cache_config(
    config: DiffusionCacheConfig,
) -> StepCacheConfig

StepCacheState

Per-generation mutable state (skip countdown).

config instance-attribute 

config = config

skip_countdown instance-attribute 

skip_countdown = 0

reset 

reset() -> None

should_run_step 

should_run_step(
    prev_predictions: list[tuple[Tensor, ...]],
) -> bool

Return True when the DiT forward should execute at this scheduler step.

trim_history 

trim_history(
    prev_predictions: list[tuple[Tensor, ...]],
) -> None

enable_dreamzero_stepcache 

enable_dreamzero_stepcache(
    pipeline: Any, config: DiffusionCacheConfig
) -> None

Enable stepcache for DreamZeroPipeline.

get_step_cache_dit_config 

get_step_cache_dit_config(
    pipeline: Any,
) -> StepCacheConfig | None

get_step_cache_dit_state 

get_step_cache_dit_state(
    pipeline: Any,
) -> StepCacheState | None

get_stepcache_config 

get_stepcache_config(
    pipeline: Any,
) -> StepCacheConfig | None

get_stepcache_state 

get_stepcache_state(pipeline: Any) -> StepCacheState | None

is_step_cache_dit_active 

is_step_cache_dit_active(pipeline: Any) -> bool

is_stepcache_active 

is_stepcache_active(pipeline: Any) -> bool

Return True when stepcache is enabled on pipeline.