diff --git a/docs/source/en/api/cache.md b/docs/source/en/api/cache.md index a5ed8751118d..193a4f61696f 100644 --- a/docs/source/en/api/cache.md +++ b/docs/source/en/api/cache.md @@ -41,6 +41,12 @@ Cache methods speedup diffusion transformers by storing and reusing intermediate [[autodoc]] apply_taylorseer_cache +## ChebyshevCacheConfig + +[[autodoc]] ChebyshevCacheConfig + +[[autodoc]] apply_chebyshev_cache + ## MagCacheConfig [[autodoc]] MagCacheConfig diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index da77fa67df52..946aa39241ae 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -175,6 +175,7 @@ ) _import_structure["hooks"].extend( [ + "ChebyshevCacheConfig", "FasterCacheConfig", "FirstBlockCacheConfig", "HookRegistry", @@ -184,6 +185,7 @@ "SmoothedEnergyGuidanceConfig", "TaylorSeerCacheConfig", "TextKVCacheConfig", + "apply_chebyshev_cache", "apply_faster_cache", "apply_first_block_cache", "apply_layer_skip", @@ -1036,6 +1038,7 @@ TangentialClassifierFreeGuidance, ) from .hooks import ( + ChebyshevCacheConfig, FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, @@ -1045,6 +1048,7 @@ SmoothedEnergyGuidanceConfig, TaylorSeerCacheConfig, TextKVCacheConfig, + apply_chebyshev_cache, apply_faster_cache, apply_first_block_cache, apply_layer_skip, diff --git a/src/diffusers/hooks/__init__.py b/src/diffusers/hooks/__init__.py index 2a9aa81608e7..935f284c78a4 100644 --- a/src/diffusers/hooks/__init__.py +++ b/src/diffusers/hooks/__init__.py @@ -16,6 +16,7 @@ if is_torch_available(): + from .chebyshev_cache import ChebyshevCacheConfig, apply_chebyshev_cache from .context_parallel import apply_context_parallel from .faster_cache import FasterCacheConfig, apply_faster_cache from .first_block_cache import FirstBlockCacheConfig, apply_first_block_cache diff --git a/src/diffusers/hooks/chebyshev_cache.py b/src/diffusers/hooks/chebyshev_cache.py new file mode 100644 index 000000000000..c124babf57b1 --- /dev/null +++ b/src/diffusers/hooks/chebyshev_cache.py @@ -0,0 +1,304 @@ +# Copyright 2025 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +from dataclasses import dataclass + +import torch +import torch.nn as nn + +from ..utils import logging +from .hooks import HookRegistry, ModelHook, StateManager + + +logger = logging.get_logger(__name__) +_CHEBYSHEV_CACHE_HOOK = "chebyshev_cache" +_SPATIAL_ATTENTION_BLOCK_IDENTIFIERS = ( + "^blocks.*attn", + "^transformer_blocks.*attn", + "^single_transformer_blocks.*attn", +) +_TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS = ("^temporal_transformer_blocks.*attn",) +_TRANSFORMER_BLOCK_IDENTIFIERS = _SPATIAL_ATTENTION_BLOCK_IDENTIFIERS + _TEMPORAL_ATTENTION_BLOCK_IDENTIFIERS + + +@dataclass +class ChebyshevCacheConfig: + """ + Configuration for Chebyshev-extrapolation cache. Adapted from "ChebBooster: A Training-Free Approach for Efficient + Diffusion Transformer Inference via Chebyshev-Inspired Extrapolation" (https://huggingface.co/papers/2608.23429). + + Like TaylorSeer, this hook reuses expensive module outputs across denoising steps by extrapolating from previously + computed values. Unlike TaylorSeer's Taylor-series (divided-difference) extrapolation — which is prone to Runge + oscillations over long cache intervals — this hook evaluates a barycentric interpolant with Chebyshev-Lobatto + weights, which stays numerically stable as the extrapolation order grows. + + Attributes: + cache_interval (`int`, defaults to `4`): + The interval between full computation steps. After a full computation, the extrapolated outputs are reused + for this many subsequent denoising steps before refreshing with a new full forward pass. + + disable_cache_before_step (`int`, defaults to `3`): + The denoising step index before which caching is disabled. The initial steps run full computations to seed + the history buffer used by the barycentric extrapolant. + + disable_cache_after_step (`int`, *optional*, defaults to `None`): + The denoising step index after which caching is disabled. For steps `>=` this value all modules run full + computations, restoring accuracy in the final refinement steps. + + max_order (`int`, defaults to `3`): + The polynomial degree of the barycentric extrapolant. The hook keeps the last `max_order + 1` computed + outputs per module as interpolation nodes. Higher orders capture more curvature; the Chebyshev weighting + keeps them stable where a Taylor expansion of the same order would oscillate. + + factors_dtype (`torch.dtype`, defaults to `torch.float32`): + Data type used for storing the history buffer and computing the barycentric extrapolation. `float32` + preserves the conditioning of the barycentric formula. + + skip_predict_identifiers (`list[str]`, *optional*, defaults to `None`): + Regex patterns (using `re.fullmatch`) for module names to place in "skip" mode, where the module returns a + zero tensor (matching the recorded shape) during prediction steps to skip computation cheaply. + + cache_identifiers (`list[str]`, *optional*, defaults to `None`): + Regex patterns (using `re.fullmatch`) for module names to place in Chebyshev-extrapolation caching mode. If + neither this nor `skip_predict_identifiers` is provided, all attention-like modules are hooked by default. + + Notes: + - Patterns are matched using `re.fullmatch` on the module name. + - The barycentric weights depend only on the number of history nodes, so they are precomputed once (the paper's + "offline weight precomputation" stage) and reused on every prediction step (the "online application" stage). + """ + + cache_interval: int = 4 + disable_cache_before_step: int = 3 + disable_cache_after_step: int | None = None + max_order: int = 3 + factors_dtype: torch.dtype | None = torch.float32 + skip_predict_identifiers: list[str] | None = None + cache_identifiers: list[str] | None = None + + def __repr__(self) -> str: + return ( + "ChebyshevCacheConfig(" + f"cache_interval={self.cache_interval}, " + f"disable_cache_before_step={self.disable_cache_before_step}, " + f"disable_cache_after_step={self.disable_cache_after_step}, " + f"max_order={self.max_order}, " + f"factors_dtype={self.factors_dtype}, " + f"skip_predict_identifiers={self.skip_predict_identifiers}, " + f"cache_identifiers={self.cache_identifiers})" + ) + + +def chebyshev_barycentric_weights(num_nodes: int) -> list[float]: + """ + Barycentric weights for the second (Chebyshev-Lobatto) barycentric form. + + These weights depend only on the node count, not the node positions, so they can be precomputed offline. Applied to + the (roughly equispaced) history of compute steps they yield Berrut's pole-free rational interpolant, which — unlike + a Taylor expansion of the same order — does not exhibit Runge oscillations when extrapolating. + """ + if num_nodes <= 1: + return [1.0] + weights = [] + for j in range(num_nodes): + half_endpoint = 0.5 if (j == 0 or j == num_nodes - 1) else 1.0 + weights.append(((-1.0) ** j) * half_endpoint) + return weights + + +class ChebyshevCacheState: + def __init__( + self, + factors_dtype: torch.dtype | None = torch.float32, + max_order: int = 3, + is_inactive: bool = False, + ): + self.factors_dtype = factors_dtype + self.num_history = max(1, max_order + 1) + self.is_inactive = is_inactive + + self.module_dtypes: tuple[torch.dtype, ...] = () + self.device: torch.device | None = None + self.current_step: int = -1 + # Ring buffers of interpolation nodes: the step indices and the outputs computed at those steps. + self.history_steps: list[int] = [] + self.history_outputs: list[tuple[torch.Tensor, ...]] = [] + self.inactive_shapes: tuple[tuple[int, ...], ...] | None = None + + def reset(self) -> None: + self.current_step = -1 + self.device = None + self.history_steps = [] + self.history_outputs = [] + self.inactive_shapes = None + + def update(self, outputs: tuple[torch.Tensor, ...]) -> None: + self.module_dtypes = tuple(output.dtype for output in outputs) + self.device = outputs[0].device + + if self.is_inactive: + self.inactive_shapes = tuple(output.shape for output in outputs) + return + + self.history_steps.append(self.current_step) + self.history_outputs.append(tuple(output.to(self.factors_dtype) for output in outputs)) + # Keep only the most recent `num_history` nodes as the extrapolation window. + if len(self.history_steps) > self.num_history: + self.history_steps = self.history_steps[-self.num_history :] + self.history_outputs = self.history_outputs[-self.num_history :] + + @torch.compiler.disable + def predict(self) -> list[torch.Tensor]: + if self.is_inactive: + if self.inactive_shapes is None: + raise ValueError("Inactive shapes not set during prediction.") + return [ + torch.zeros(shape, dtype=self.module_dtypes[i], device=self.device) + for i, shape in enumerate(self.inactive_shapes) + ] + + if not self.history_outputs: + raise ValueError("History buffer empty during prediction.") + + num_outputs = len(self.history_outputs[-1]) + nodes = self.history_steps + num_nodes = len(nodes) + + # Single node -> constant extrapolation; no barycentric evaluation needed. + if num_nodes == 1: + return [self.history_outputs[0][i].to(self.module_dtypes[i]) for i in range(num_outputs)] + + weights = chebyshev_barycentric_weights(num_nodes) + # Second barycentric form: p(x) = sum_j (w_j / (x - x_j)) f_j / sum_j (w_j / (x - x_j)). + # `current_step` is a prediction step, so it never coincides with a compute node -> no division by zero. + coeffs = [weights[j] / (self.current_step - nodes[j]) for j in range(num_nodes)] + denom = sum(coeffs) + + outputs = [] + for i in range(num_outputs): + acc = torch.zeros_like(self.history_outputs[-1][i]) + for j in range(num_nodes): + acc = acc + self.history_outputs[j][i] * coeffs[j] + outputs.append((acc / denom).to(self.module_dtypes[i])) + return outputs + + +class ChebyshevCacheHook(ModelHook): + _is_stateful = True + + def __init__( + self, + cache_interval: int, + disable_cache_before_step: int, + state_manager: StateManager, + disable_cache_after_step: int | None = None, + ): + super().__init__() + self.cache_interval = cache_interval + self.disable_cache_before_step = disable_cache_before_step + self.disable_cache_after_step = disable_cache_after_step + self.state_manager = state_manager + + def initialize_hook(self, module: torch.nn.Module): + return module + + def reset_state(self, module: torch.nn.Module) -> None: + self.state_manager.reset() + + @torch.compiler.disable + def _measure_should_compute(self): + state: ChebyshevCacheState = self.state_manager.get_state() + state.current_step += 1 + current_step = state.current_step + is_warmup_phase = current_step < self.disable_cache_before_step + is_compute_interval = (current_step - self.disable_cache_before_step - 1) % self.cache_interval == 0 + is_cooldown_phase = self.disable_cache_after_step is not None and current_step >= self.disable_cache_after_step + should_compute = is_warmup_phase or is_compute_interval or is_cooldown_phase + return should_compute, state + + def new_forward(self, module: torch.nn.Module, *args, **kwargs): + should_compute, state = self._measure_should_compute() + if should_compute: + outputs = self.fn_ref.original_forward(*args, **kwargs) + wrapped_outputs = (outputs,) if isinstance(outputs, torch.Tensor) else outputs + state.update(wrapped_outputs) + return outputs + + outputs_list = state.predict() + return outputs_list[0] if len(outputs_list) == 1 else tuple(outputs_list) + + +def _resolve_patterns(config: ChebyshevCacheConfig) -> tuple[list[str], list[str]]: + inactive_patterns = config.skip_predict_identifiers or [] + active_patterns = config.cache_identifiers or [] + return inactive_patterns, active_patterns + + +def apply_chebyshev_cache(module: torch.nn.Module, config: ChebyshevCacheConfig): + """ + Applies the Chebyshev-extrapolation cache to a model subtree (typically the transformer / UNet). + + Selected modules are hooked to extrapolate their outputs across denoising steps using a numerically stable + barycentric Chebyshev interpolant, reducing redundant computation in the diffusion loop. + + Args: + module (torch.nn.Module): The model subtree to apply the hooks to. + config (ChebyshevCacheConfig): Configuration for the cache. + + Example: + ```python + >>> import torch + >>> from diffusers import PixArtSigmaPipeline, ChebyshevCacheConfig + + >>> pipe = PixArtSigmaPipeline.from_pretrained( + ... "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16 + ... ) + >>> pipe.to("cuda") + + >>> config = ChebyshevCacheConfig(cache_interval=4, max_order=3, disable_cache_before_step=3) + >>> pipe.transformer.enable_cache(config) + ``` + """ + inactive_patterns, active_patterns = _resolve_patterns(config) + active_patterns = active_patterns or list(_TRANSFORMER_BLOCK_IDENTIFIERS) + + for name, submodule in module.named_modules(): + matches_inactive = any(re.fullmatch(pattern, name) for pattern in inactive_patterns) + matches_active = any(re.fullmatch(pattern, name) for pattern in active_patterns) + if not (matches_inactive or matches_active): + continue + _apply_chebyshev_cache_hook(module=submodule, config=config, is_inactive=matches_inactive) + + +def _apply_chebyshev_cache_hook(module: nn.Module, config: ChebyshevCacheConfig, is_inactive: bool): + state_manager = StateManager( + ChebyshevCacheState, + init_kwargs={ + "factors_dtype": config.factors_dtype, + "max_order": config.max_order, + "is_inactive": is_inactive, + }, + ) + + registry = HookRegistry.check_if_exists_or_initialize(module) + + hook = ChebyshevCacheHook( + cache_interval=config.cache_interval, + disable_cache_before_step=config.disable_cache_before_step, + disable_cache_after_step=config.disable_cache_after_step, + state_manager=state_manager, + ) + + registry.register_hook(hook, _CHEBYSHEV_CACHE_HOOK) diff --git a/src/diffusers/models/cache_utils.py b/src/diffusers/models/cache_utils.py index 161fcf426f21..de85e60e79fa 100644 --- a/src/diffusers/models/cache_utils.py +++ b/src/diffusers/models/cache_utils.py @@ -67,12 +67,14 @@ def enable_cache(self, config) -> None: """ from ..hooks import ( + ChebyshevCacheConfig, FasterCacheConfig, FirstBlockCacheConfig, MagCacheConfig, PyramidAttentionBroadcastConfig, TaylorSeerCacheConfig, TextKVCacheConfig, + apply_chebyshev_cache, apply_faster_cache, apply_first_block_cache, apply_mag_cache, @@ -98,6 +100,8 @@ def enable_cache(self, config) -> None: apply_pyramid_attention_broadcast(self, config) elif isinstance(config, TaylorSeerCacheConfig): apply_taylorseer_cache(self, config) + elif isinstance(config, ChebyshevCacheConfig): + apply_chebyshev_cache(self, config) else: raise ValueError(f"Cache config {type(config)} is not supported.") @@ -105,6 +109,7 @@ def enable_cache(self, config) -> None: def disable_cache(self) -> None: from ..hooks import ( + ChebyshevCacheConfig, FasterCacheConfig, FirstBlockCacheConfig, HookRegistry, @@ -113,6 +118,7 @@ def disable_cache(self) -> None: TaylorSeerCacheConfig, TextKVCacheConfig, ) + from ..hooks.chebyshev_cache import _CHEBYSHEV_CACHE_HOOK from ..hooks.faster_cache import _FASTER_CACHE_BLOCK_HOOK, _FASTER_CACHE_DENOISER_HOOK from ..hooks.first_block_cache import _FBC_BLOCK_HOOK, _FBC_LEADER_BLOCK_HOOK from ..hooks.mag_cache import _MAG_CACHE_BLOCK_HOOK, _MAG_CACHE_LEADER_BLOCK_HOOK @@ -141,6 +147,8 @@ def disable_cache(self) -> None: registry.remove_hook(_TEXT_KV_CACHE_BLOCK_HOOK, recurse=True) elif isinstance(self._cache_config, TaylorSeerCacheConfig): registry.remove_hook(_TAYLORSEER_CACHE_HOOK, recurse=True) + elif isinstance(self._cache_config, ChebyshevCacheConfig): + registry.remove_hook(_CHEBYSHEV_CACHE_HOOK, recurse=True) else: raise ValueError(f"Cache config {type(self._cache_config)} is not supported.") diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 8439a2b93371..3c5143d77c94 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -167,6 +167,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ChebyshevCacheConfig(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class FasterCacheConfig(metaclass=DummyObject): _backends = ["torch"] @@ -302,6 +317,10 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +def apply_chebyshev_cache(*args, **kwargs): + requires_backends(apply_chebyshev_cache, ["torch"]) + + def apply_faster_cache(*args, **kwargs): requires_backends(apply_faster_cache, ["torch"]) diff --git a/tests/hooks/test_chebyshev_cache.py b/tests/hooks/test_chebyshev_cache.py new file mode 100644 index 000000000000..41eb4e510ca8 --- /dev/null +++ b/tests/hooks/test_chebyshev_cache.py @@ -0,0 +1,118 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import ChebyshevCacheConfig, apply_chebyshev_cache +from diffusers.hooks.chebyshev_cache import chebyshev_barycentric_weights +from diffusers.models import ModelMixin +from diffusers.models.cache_utils import CacheMixin + + +class DummyBlock(torch.nn.Module): + def forward(self, hidden_states, encoder_hidden_states=None, **kwargs): + # Output is double the input, so a cached (extrapolated) step is + # distinguishable from a freshly computed one. + return hidden_states * 2.0 + + +class DummyTransformer(ModelMixin, CacheMixin): + def __init__(self): + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([DummyBlock()]) + + def forward(self, hidden_states, encoder_hidden_states=None): + for block in self.transformer_blocks: + hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states) + return hidden_states + + +def _set_context(model, context_name): + """Helper to set the state-manager context on all hooks in the model.""" + for module in model.modules(): + if hasattr(module, "_diffusers_hook"): + module._diffusers_hook._set_context(context_name) + + +def test_chebyshev_barycentric_weights(): + # Single node -> constant extrapolation. + assert chebyshev_barycentric_weights(1) == [1.0] + # For two nodes the Chebyshev-Lobatto weights coincide (up to a global + # scale that cancels in the barycentric quotient) with the exact linear + # barycentric weights, so extrapolation is exact for affine data. + assert chebyshev_barycentric_weights(2) == [0.5, -0.5] + # Alternating signs with halved endpoints for higher orders. + assert chebyshev_barycentric_weights(4) == [0.5, -1.0, 1.0, -0.5] + + +def test_chebyshev_cache_linear_extrapolation(): + """Two seeded nodes -> the cached step must be the exact linear extrapolant, not a recompute.""" + model = DummyTransformer() + config = ChebyshevCacheConfig( + cache_interval=100, + disable_cache_before_step=2, + max_order=1, + factors_dtype=torch.float32, + cache_identifiers=[r"transformer_blocks\.0"], + ) + apply_chebyshev_cache(model, config) + _set_context(model, "test_context") + + # Step 0 (warmup, compute): 2 * 1.0 -> 2.0, seeds node (step=0, value=2.0). + out0 = model(torch.tensor([[[1.0]]])) + assert torch.allclose(out0, torch.tensor([[[2.0]]])), f"Step 0 should compute, got {out0.item()}" + + # Step 1 (warmup, compute): 2 * 2.0 -> 4.0, seeds node (step=1, value=4.0). + out1 = model(torch.tensor([[[2.0]]])) + assert torch.allclose(out1, torch.tensor([[[4.0]]])), f"Step 1 should compute, got {out1.item()}" + + # Step 2 (cache): linear extrapolation of (0, 2.0) and (1, 4.0) at x=2 -> 6.0. + # A recompute would instead give 2 * 100.0 = 200.0, so this asserts the skip path. + out2 = model(torch.tensor([[[100.0]]])) + assert torch.allclose(out2, torch.tensor([[[6.0]]])), f"Step 2 should extrapolate to 6.0, got {out2.item()}" + + +def test_chebyshev_cache_constant_reproduction(): + """Berrut's rational interpolant reproduces constants exactly, at any order.""" + model = DummyTransformer() + config = ChebyshevCacheConfig( + cache_interval=100, + disable_cache_before_step=3, + max_order=2, + factors_dtype=torch.float32, + cache_identifiers=[r"transformer_blocks\.0"], + ) + apply_chebyshev_cache(model, config) + _set_context(model, "test_context") + + const_in = torch.tensor([[[5.0]]]) + for _ in range(3): # warmup computes -> history is constant 10.0 + model(const_in) + + # Cached step with a different input: must return the constant 10.0, not 2 * 999. + out = model(torch.tensor([[[999.0]]])) + assert torch.allclose(out, torch.tensor([[[10.0]]])), f"Constant extrapolation failed, got {out.item()}" + + +def test_chebyshev_cache_enable_disable_dispatch(): + """enable_cache / disable_cache must route ChebyshevCacheConfig through the CacheMixin dispatcher.""" + model = DummyTransformer() + config = ChebyshevCacheConfig(cache_identifiers=[r"transformer_blocks\.0"]) + + model.enable_cache(config) + assert model.is_cache_enabled + assert any(hasattr(m, "_diffusers_hook") for m in model.modules()) + + model.disable_cache() + assert not model.is_cache_enabled