# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Generic Cache-DiT backend lifecycle and integration.""" from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass from operator import attrgetter from typing import TYPE_CHECKING, Any, TypeAlias, cast import cache_dit import torch.nn as nn from cache_dit import BlockAdapter, DBCacheConfig from cache_dit.caching.cache_adapters.cache_adapter import CachedAdapter from vllm.logger import init_logger from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.cachedit.config import ( CacheDiTAdapterConfig, CacheDiTConfig, ) from vllm_omni.diffusion.data import DiffusionCacheConfig if TYPE_CHECKING: from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery logger = init_logger(__name__) RefreshCacheContextFunc: TypeAlias = Callable[[Any, int, bool], None] @dataclass(frozen=True) class CacheDiTEnableResult: """Custom-enabler result with the exact targets needed for teardown.""" refresh: RefreshCacheContextFunc targets: tuple[Any, ...] CacheDiTEnabler: TypeAlias = Callable[ [Any, DiffusionCacheConfig], RefreshCacheContextFunc | CacheDiTEnableResult, ] # Model-specific implementations register themselves when the package loads. CUSTOM_DIT_ENABLERS: dict[str, CacheDiTEnabler] = {} def _dit_module_names(pipeline: SupportsComponentDiscovery) -> tuple[str, ...]: """Return the pipeline DiT attributes that are present at runtime.""" names = getattr(pipeline, "_dit_modules", None) if not isinstance(names, (list, tuple)): names = ("transformer",) resolved_names = [] for name in names: if not isinstance(name, str): continue try: module = attrgetter(name)(pipeline) except AttributeError: continue if module is not None: resolved_names.append(name) return tuple(resolved_names) def cache_summary(pipeline: SupportsComponentDiscovery, details: bool = True) -> None: """Log Cache-DiT statistics for every transformer on the pipeline.""" targets = getattr(pipeline, "_cache_dit_targets", None) if targets is None: targets = tuple(attrgetter(name)(pipeline) for name in _dit_module_names(pipeline)) for target in targets: if not BlockAdapter.is_cached(target): continue cache_dit.summary(target, details=details) if not targets: logger.warning("CacheDiT summary failed; this pipeline has no defined transformer attribute") def _default_get_pipeline_transformer( pipeline: SupportsComponentDiscovery, ) -> nn.Module: return cast(nn.Module, getattr(pipeline, "transformer")) def _make_pipeline_transformer_getter( name: str, ) -> Callable[[SupportsComponentDiscovery], nn.Module]: def get_pipeline_transformer( pipeline: SupportsComponentDiscovery, ) -> nn.Module: return cast(nn.Module, attrgetter(name)(pipeline)) return get_pipeline_transformer def _build_cache_context_refresh( cache_config: DiffusionCacheConfig, get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, ) -> RefreshCacheContextFunc: """Build the cache context refresh callback for one transformer.""" projected_config = CacheDiTConfig.from_diffusion_config(cache_config) def refresh_cache_context( pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True ) -> None: transformer = get_pipeline_transformer(pipeline) # Cache-DiT has no predefined SCM mask for these small step counts. scm_supported_steps = num_inference_steps >= 8 or num_inference_steps in (4, 6) if projected_config.scm_steps_mask_policy is None or not scm_supported_steps: if projected_config.force_refresh_step_hint is not None: # Cache-DiT's ``once`` policy clears the hint after it fires. # Reapply it at every request boundary so a repeated request # with the same configured policy receives the same behavior. cache_dit.refresh_context( transformer, cache_config=DBCacheConfig().reset( num_inference_steps=num_inference_steps, force_refresh_step_hint=projected_config.force_refresh_step_hint, force_refresh_step_policy=projected_config.force_refresh_step_policy, ), verbose=verbose, ) return cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose) return refresh_config = DBCacheConfig().reset( num_inference_steps=num_inference_steps, steps_computation_mask=cache_dit.steps_mask( mask_policy=projected_config.scm_steps_mask_policy, total_steps=num_inference_steps, ), steps_computation_policy=projected_config.scm_steps_policy, ) if projected_config.force_refresh_step_hint is not None: refresh_config.force_refresh_step_hint = projected_config.force_refresh_step_hint refresh_config.force_refresh_step_policy = projected_config.force_refresh_step_policy cache_dit.refresh_context( transformer, cache_config=refresh_config, verbose=verbose, ) return refresh_cache_context def 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], nn.Module] = _default_get_pipeline_transformer, ) -> RefreshCacheContextFunc: """Enable Cache-DiT for a standard single-transformer DiT pipeline.""" projected_config = CacheDiTConfig.from_diffusion_config(cache_config) db_cache_config = projected_config.to_db_cache_config() calibrator_config = projected_config.to_calibrator_config() logger.info( "Enabling cache-dit on transformer: Fn=%s, Bn=%s, W=%s", db_cache_config.Fn_compute_blocks, db_cache_config.Bn_compute_blocks, db_cache_config.max_warmup_steps, ) transformer = get_pipeline_transformer(pipeline) cache_target = transformer if block_adapter is None else block_adapter if adapter_cls is not None: adapter_cls.apply( cache_target, cache_config=db_cache_config, calibrator_config=calibrator_config, ) elif block_adapter is None: try: cache_dit.enable_cache( cache_target, cache_config=db_cache_config, calibrator_config=calibrator_config, ) except ValueError as exc: raise ValueError( "Failed to enable Cache-DiT for pipeline " f"{type(pipeline).__name__} with transformer " f"{type(transformer).__name__}: no model-declared " "_cache_dit_adapter_config or compatible Cache-DiT built-in " "adapter was found." ) from exc else: cache_dit.enable_cache( cache_target, cache_config=db_cache_config, calibrator_config=calibrator_config, ) return _build_cache_context_refresh(cache_config, get_pipeline_transformer) def _maybe_build_block_adapter( pipeline: SupportsComponentDiscovery, get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, ) -> BlockAdapter | None: """Build the model-declared block adapter, when one is configured.""" transformer = get_pipeline_transformer(pipeline) adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) if adapter_config is None: logger.info( "Transformer %s does not declare _cache_dit_adapter_config; " "falling back to Cache-DiT's built-in adapter registry.", type(transformer).__name__, ) return None block_attributes, forward_patterns = zip(*adapter_config.block_forward_patterns.items()) missing_attributes = [ block_attribute for block_attribute in block_attributes if not hasattr(transformer, block_attribute) ] if missing_attributes: raise AttributeError(f"Missing Cache-DiT block attributes: {missing_attributes}") block_adapter = BlockAdapter( transformer=transformer, blocks=[getattr(transformer, block_attribute) for block_attribute in block_attributes], forward_pattern=list(forward_patterns), has_separate_cfg=adapter_config.has_separate_cfg, check_forward_pattern=adapter_config.check_forward_pattern, ) # cache_dit marks the wrapper pipe's class as cached, and for bare-transformer # adapters that class is shared process-wide; a marker left by a previous # engine would make this enable skip cache-context installation. pipe_cls = type(block_adapter.pipe) if "_is_cached" in vars(pipe_cls): del pipe_cls._is_cached return block_adapter def _maybe_get_cached_adapter_cls( pipeline: SupportsComponentDiscovery, get_pipeline_transformer: Callable[[SupportsComponentDiscovery], nn.Module] = _default_get_pipeline_transformer, ) -> type[CachedAdapter] | None: """Return the custom cached adapter declared by the transformer.""" transformer = get_pipeline_transformer(pipeline) adapter_config: CacheDiTAdapterConfig | None = getattr(transformer, "_cache_dit_adapter_config", None) return None if adapter_config is None else adapter_config.cached_adapter_cls class CacheDiTBackend(CacheBackend): """Manage Cache-DiT through the common diffusion cache lifecycle.""" def __init__(self, cache_config: Any = None): if cache_config is None: config = DiffusionCacheConfig() elif isinstance(cache_config, dict): config = DiffusionCacheConfig.from_dict(cache_config) else: config = cache_config super().__init__(config) self._refresh_funcs: list[RefreshCacheContextFunc] = [] self._cache_targets: list[Any] = [] def enable(self, pipeline: SupportsComponentDiscovery) -> None: pipeline_name = type(pipeline).__name__ custom_enabler = CUSTOM_DIT_ENABLERS.get(pipeline_name) self._refresh_funcs = [] self._cache_targets = [] if custom_enabler is not None: logger.info("Using custom cache-dit enabler for model: %s", pipeline_name) result = custom_enabler(pipeline, self.config) if isinstance(result, CacheDiTEnableResult): self._refresh_funcs = [result.refresh] self._cache_targets = list(result.targets) else: self._refresh_funcs = [result] self._cache_targets = [_default_get_pipeline_transformer(pipeline)] else: for name in _dit_module_names(pipeline): get_transformer = _make_pipeline_transformer_getter(name) block_adapter = _maybe_build_block_adapter(pipeline, get_transformer) adapter_cls = _maybe_get_cached_adapter_cls(pipeline, get_transformer) cache_target = get_transformer(pipeline) if block_adapter is None else block_adapter self._refresh_funcs.append( enable_cache_for_dit( pipeline, self.config, block_adapter, adapter_cls, get_transformer, ) ) self._cache_targets.append(cache_target) if not self._refresh_funcs: raise ValueError(f"Pipeline {pipeline_name} has no declared DiT modules for Cache-DiT") pipeline._cache_dit_targets = tuple(self._cache_targets) self.enabled = True logger.info("Cache-dit enabled successfully on %s", pipeline_name) def disable(self, pipeline: SupportsComponentDiscovery) -> None: """Remove Cache-DiT hooks so later requests use native forwards.""" if not self.enabled: return logger.info( "Disabling cache-dit on %d DiT target(s) for %s", len(self._cache_targets), type(pipeline).__name__, ) try: for target in self._cache_targets: cache_dit.disable_cache(target) finally: self._refresh_funcs = [] self._cache_targets = [] if hasattr(pipeline, "_cache_dit_targets"): del pipeline._cache_dit_targets self.enabled = False def refresh(self, pipeline: SupportsComponentDiscovery, num_inference_steps: int, verbose: bool = True) -> None: if not self.enabled or not self._refresh_funcs: logger.warning("Cache-dit is not enabled. Cannot refresh cache context.") return if verbose: logger.info( "Refreshing cache context for transformer with num_inference_steps: %s", num_inference_steps, ) for refresh_func in self._refresh_funcs: refresh_func(pipeline, num_inference_steps, verbose) __all__ = [ "CUSTOM_DIT_ENABLERS", "CacheDiTEnableResult", "CacheDiTBackend", "cache_summary", "enable_cache_for_dit", ]