Skip to content

vllm_omni.diffusion.cache.cachedit

Public API for the Cache-DiT diffusion cache backend.

Modules:

Name Description
backend

Generic Cache-DiT backend lifecycle and integration.
config

Configuration types and conversion helpers for Cache-DiT.
model_specific

Model-specific Cache-DiT adapters and enablers.
runtime

Declarative request-scoped Cache-DiT lifecycle.

CUSTOM_DIT_ENABLERS module-attribute 

CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {}

BagelCachedAdapter

Bases: CachedAdapter

Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.

collect_unified_blocks classmethod 

collect_unified_blocks(
    block_adapter: BlockAdapter, contexts_kwargs: list[dict]
) -> list[dict[str, ModuleList]]

create_context classmethod 

create_context(
    block_adapter: BlockAdapter, **context_kwargs
) -> tuple[list[str], list[dict[str, Any]]]

CacheDiTAdapterConfig dataclass

Describe how Cache-DiT wraps a model's transformer blocks.

block_forward_patterns instance-attribute 

block_forward_patterns: dict[str, ForwardPattern]

cached_adapter_cls class-attribute instance-attribute 

cached_adapter_cls: type[CachedAdapter] | None = None

check_forward_pattern class-attribute instance-attribute 

check_forward_pattern: bool = True

has_separate_cfg class-attribute instance-attribute 

has_separate_cfg: bool = False

CacheDiTBackend

Bases: CacheBackend

Manage Cache-DiT through the common diffusion cache lifecycle.

disable 

disable(pipeline: SupportsComponentDiscovery) -> None

Remove Cache-DiT hooks so later requests use native forwards.

enable 

enable(pipeline: SupportsComponentDiscovery) -> None

refresh 

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

CacheDiTConfig dataclass

Cache-DiT-specific projection of the shared diffusion cache config.

Bn_compute_blocks class-attribute instance-attribute 

Bn_compute_blocks: int = 0

Fn_compute_blocks class-attribute instance-attribute 

Fn_compute_blocks: int = 1

enable_taylorseer class-attribute instance-attribute 

enable_taylorseer: bool = False

force_refresh_step_hint class-attribute instance-attribute 

force_refresh_step_hint: int | None = None

force_refresh_step_policy class-attribute instance-attribute 

force_refresh_step_policy: str = 'once'

max_cached_steps class-attribute instance-attribute 

max_cached_steps: int = -1

max_continuous_cached_steps class-attribute instance-attribute 

max_continuous_cached_steps: int = 3

max_warmup_steps class-attribute instance-attribute 

max_warmup_steps: int = 4

residual_diff_threshold class-attribute instance-attribute 

residual_diff_threshold: float = 0.24

scm_steps_mask_policy class-attribute instance-attribute 

scm_steps_mask_policy: str | None = None

scm_steps_policy class-attribute instance-attribute 

scm_steps_policy: str = 'dynamic'

taylorseer_order class-attribute instance-attribute 

taylorseer_order: int = 1

from_diffusion_config classmethod 

from_diffusion_config(
    config: DiffusionCacheConfig,
) -> CacheDiTConfig

Project the shared config onto fields consumed by Cache-DiT.

to_calibrator_config 

to_calibrator_config() -> CalibratorConfig | None

Build the optional TaylorSeer calibrator configuration.

to_db_cache_config 

to_db_cache_config() -> DBCacheConfig

Build the third-party cache configuration used during installation.

CacheDiTEnableResult dataclass

Custom-enabler result with the exact targets needed for teardown.

refresh instance-attribute 

refresh: RefreshCacheContextFunc

targets instance-attribute 

targets: tuple[Any, ...]

CacheDiTRequestSpec dataclass

Desired Cache-DiT state for one request.

installation_key identifies configurations that require reinstalling hooks when switching between them. Step-count changes only require a context refresh.

cache_config instance-attribute 

cache_config: DiffusionCacheConfig

installation_key instance-attribute 

installation_key: str

num_inference_steps instance-attribute 

num_inference_steps: int

RequestScopedCacheDiTRuntime

Reconcile declarative request state with installed Cache-DiT hooks.

is_enabled property 

is_enabled: bool

adopt 

adopt(
    backend: CacheDiTBackend, *, installation_key: str
) -> None

Assume ownership of an already-enabled startup backend.

disable 

disable() -> None

prepare 

prepare(spec: CacheDiTRequestSpec | None) -> None

Apply the minimal transition needed before the request denoise.

None disables installed hooks. An unchanged installation key only refreshes request context; a changed key disables the previous hooks before installing the requested profile.

SensenovaCachedAdapter

Bases: CachedAdapter

Custom CachedAdapter for SenseNova-U1 that uses SensenovaCachedBlocks.

collect_unified_blocks classmethod 

collect_unified_blocks(
    block_adapter: BlockAdapter, contexts_kwargs: list[dict]
) -> list[dict[str, ModuleList]]

cache_summary 

cache_summary(
    pipeline: SupportsComponentDiscovery,
    details: bool = True,
) -> None

Log Cache-DiT statistics for every transformer on the pipeline.

enable_cache_for_dit 

enable_cache_for_dit(
    pipeline: SupportsComponentDiscovery,
    cache_config: DiffusionCacheConfig,
    block_adapter: BlockAdapter | None = None,
    adapter_cls: type[CachedAdapter] | None = None,
    get_pipeline_transformer: Callable[
        [SupportsComponentDiscovery], Module
    ] = _default_get_pipeline_transformer,
) -> RefreshCacheContextFunc

Enable Cache-DiT for a standard single-transformer DiT pipeline.