From 0a7a2017a02b88f6678804d6d024b100e809b02d Mon Sep 17 00:00:00 2001 From: blahblah <28567807+Brain97@users.noreply.github.com> Date: Wed, 21 Jan 2026 15:42:45 +0800 Subject: [PATCH] [diffusion] refactor: refactor and simplify teacache for cachabledit and wanvideo (#16396) Co-authored-by: Brain97 Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Mick Co-authored-by: blahblah --- .../docs/{ => cache}/cache_dit.md | 7 +- .../multimodal_gen/docs/cache/caching.md | 60 ++++ .../multimodal_gen/docs/cache/teacache.md | 84 +++++ .../docs/environment_variables.md | 10 +- .../multimodal_gen/runtime/cache/__init__.py | 56 ++++ .../{utils => cache}/cache_dit_integration.py | 0 .../multimodal_gen/runtime/cache/teacache.py | 313 ++++++++++++++++++ .../runtime/models/dits/base.py | 45 +-- .../runtime/models/dits/wanvideo.py | 142 +++----- .../pipelines_core/stages/denoising.py | 12 +- .../test/server/perf_baselines.json | 67 ++++ .../test/server/test_server_utils.py | 36 +- .../test/server/testcase_configs.py | 17 + 13 files changed, 702 insertions(+), 147 deletions(-) rename python/sglang/multimodal_gen/docs/{ => cache}/cache_dit.md (95%) create mode 100644 python/sglang/multimodal_gen/docs/cache/caching.md create mode 100644 python/sglang/multimodal_gen/docs/cache/teacache.md create mode 100644 python/sglang/multimodal_gen/runtime/cache/__init__.py rename python/sglang/multimodal_gen/runtime/{utils => cache}/cache_dit_integration.py (100%) create mode 100644 python/sglang/multimodal_gen/runtime/cache/teacache.py diff --git a/python/sglang/multimodal_gen/docs/cache_dit.md b/python/sglang/multimodal_gen/docs/cache/cache_dit.md similarity index 95% rename from python/sglang/multimodal_gen/docs/cache_dit.md rename to python/sglang/multimodal_gen/docs/cache/cache_dit.md index fa4a948d3..01038b6ca 100644 --- a/python/sglang/multimodal_gen/docs/cache_dit.md +++ b/python/sglang/multimodal_gen/docs/cache/cache_dit.md @@ -1,7 +1,10 @@ # Cache-DiT Acceleration -SGLang integrates [Cache-DiT](https://github.com/vipshop/cache-dit), a caching acceleration engine for Diffusion -Transformers (DiT), to achieve up to **7.4x inference speedup** with minimal quality loss. +> **Note**: This is one of two caching strategies available in SGLang. +> For an overview of all caching options, see [caching.md](caching.md). +> For TeaCache documentation, see [teacache.md](teacache.md). + +SGLang integrates [Cache-DiT](https://github.com/vipshop/cache-dit), a caching acceleration engine for Diffusion Transformers (DiT), to achieve up to **1.69x inference speedup** with minimal quality loss. ## Overview diff --git a/python/sglang/multimodal_gen/docs/cache/caching.md b/python/sglang/multimodal_gen/docs/cache/caching.md new file mode 100644 index 000000000..d1f8e61d0 --- /dev/null +++ b/python/sglang/multimodal_gen/docs/cache/caching.md @@ -0,0 +1,60 @@ +# Caching Acceleration for Diffusion Models + +SGLang provides multiple caching acceleration strategies for Diffusion Transformer (DiT) models. These strategies can significantly reduce inference time by skipping redundant computation. + +## Overview + +SGLang supports two complementary caching approaches: + +| Strategy | Scope | Mechanism | Best For | +|----------|-------|-----------|----------| +| **Cache-DiT** | Block-level | Skip individual transformer blocks dynamically | Advanced, higher speedup | +| **TeaCache** | Timestep-level | Skip entire denoising steps based on L1 similarity | Simple, built-in | + + + +## Cache-DiT + +[Cache-DiT](https://github.com/vipshop/cache-dit) provides block-level caching with +advanced strategies like DBCache and TaylorSeer. It can achieve up to **1.69x speedup**. + +See [cache_dit.md](cache_dit.md) for detailed configuration. + +### Quick Start + +```bash +SGLANG_CACHE_DIT_ENABLED=true \ +sglang generate --model-path Qwen/Qwen-Image \ + --prompt "A beautiful sunset over the mountains" +``` + +### Key Features + +- **DBCache**: Dynamic block-level caching based on residual differences +- **TaylorSeer**: Taylor expansion-based calibration for optimized caching +- **SCM**: Step-level computation masking for additional speedup + +## TeaCache + +TeaCache (Temporal similarity-based caching) accelerates diffusion inference by detecting when consecutive denoising steps are similar enough to skip computation entirely. + +See [teacache.md](teacache.md) for detailed documentation. + +### Quick Overview + +- Tracks L1 distance between modulated inputs across timesteps +- When accumulated distance is below threshold, reuses cached residual +- Supports CFG with separate positive/negative caches + +### Supported Models + +- Wan (wan2.1, wan2.2) +- Hunyuan (HunyuanVideo) +- Z-Image + +For Flux and Qwen models, TeaCache is automatically disabled when CFG is enabled. + +## References + +- [Cache-DiT Repository](https://github.com/vipshop/cache-dit) +- [TeaCache Paper](https://arxiv.org/abs/2411.14324) diff --git a/python/sglang/multimodal_gen/docs/cache/teacache.md b/python/sglang/multimodal_gen/docs/cache/teacache.md new file mode 100644 index 000000000..5eb0b6c19 --- /dev/null +++ b/python/sglang/multimodal_gen/docs/cache/teacache.md @@ -0,0 +1,84 @@ +# TeaCache Acceleration + +> **Note**: This is one of two caching strategies available in SGLang. +> For an overview of all caching options, see [caching.md](caching.md). + +TeaCache (Temporal similarity-based caching) accelerates diffusion inference by detecting when consecutive denoising steps are similar enough to skip computation entirely. + +## Overview + +TeaCache works by: +1. Tracking the L1 distance between modulated inputs across consecutive timesteps +2. Accumulating the rescaled L1 distance over steps +3. When accumulated distance is below a threshold, reusing the cached residual +4. Supporting CFG (Classifier-Free Guidance) with separate positive/negative caches + +## How It Works + +### L1 Distance Tracking + +At each denoising step, TeaCache computes the relative L1 distance between the current and previous modulated inputs: + +``` +rel_l1 = |current - previous|.mean() / |previous|.mean() +``` + +This distance is then rescaled using polynomial coefficients and accumulated: + +``` +accumulated += poly(coefficients)(rel_l1) +``` + +### Cache Decision + +- If `accumulated >= threshold`: Force computation, reset accumulator +- If `accumulated < threshold`: Skip computation, use cached residual + +### CFG Support + +For models that support CFG cache separation (Wan, Hunyuan, Z-Image), TeaCache maintains separate caches for positive and negative branches: +- `previous_modulated_input` / `previous_residual` for positive branch +- `previous_modulated_input_negative` / `previous_residual_negative` for negative branch + +For models that don't support CFG separation (Flux, Qwen), TeaCache is automatically disabled when CFG is enabled. + +## Configuration + +TeaCache is configured via `TeaCacheParams` in the sampling parameters: + +```python +from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams + +params = TeaCacheParams( + teacache_thresh=0.1, # Threshold for accumulated L1 distance + coefficients=[1.0, 0.0, 0.0], # Polynomial coefficients for L1 rescaling +) +``` + +### Parameters + +| Parameter | Type | Description | +|-----------|------|-------------| +| `teacache_thresh` | float | Threshold for accumulated L1 distance. Lower = more caching, faster but potentially lower quality | +| `coefficients` | list[float] | Polynomial coefficients for L1 rescaling. Model-specific tuning | + +### Model-Specific Configurations + +Different models may have different optimal configurations. The coefficients are typically tuned per-model to balance speed and quality. + +## Supported Models + +TeaCache is built into the following model families: + +| Model Family | CFG Cache Separation | Notes | +|--------------|---------------------|-------| +| Wan (wan2.1, wan2.2) | Yes | Full support | +| Hunyuan (HunyuanVideo) | Yes | To be supported | +| Z-Image | Yes | To be supported | +| Flux | No | To be supported | +| Qwen | No | To be supported | + + +## References + +- [TeaCache: Accelerating Diffusion Models with Temporal Similarity](https://arxiv.org/abs/2411.14324) diff --git a/python/sglang/multimodal_gen/docs/environment_variables.md b/python/sglang/multimodal_gen/docs/environment_variables.md index 3e4d64354..2c07a3aec 100644 --- a/python/sglang/multimodal_gen/docs/environment_variables.md +++ b/python/sglang/multimodal_gen/docs/environment_variables.md @@ -1,7 +1,11 @@ -## Cache-DiT Acceleration +## Caching Acceleration -These variables configure cache-dit caching acceleration for Diffusion Transformer (DiT) models. -See [cache-dit documentation](cache_dit.md) for details. +These variables configure caching acceleration for Diffusion Transformer (DiT) models. +SGLang supports multiple caching strategies - see [caching documentation](cache/caching.md) for an overview. + +### Cache-DiT Configuration + +See [cache-dit documentation](cache/cache_dit.md) for detailed configuration. | Environment Variable | Default | Description | |-------------------------------------|---------|------------------------------------------| diff --git a/python/sglang/multimodal_gen/runtime/cache/__init__.py b/python/sglang/multimodal_gen/runtime/cache/__init__.py new file mode 100644 index 000000000..895f0c594 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/cache/__init__.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Cache acceleration module for SGLang multimodal generation. + +This module provides various caching strategies to accelerate +diffusion transformer (DiT) inference: + +- TeaCache: Temporal similarity-based caching for diffusion models +- cache-dit integration: Block-level caching with DBCache and TaylorSeer + +Usage: + # TeaCache (built into DiT models via TeaCacheMixin) + from sglang.multimodal_gen.runtime.cache import TeaCacheMixin, TeaCacheContext + + # cache-dit integration (requires cache-dit package installed) + # These are lazy-loaded and will raise ImportError if cache-dit is not available + from sglang.multimodal_gen.runtime.cache import ( + CacheDitConfig, + enable_cache_on_transformer, + enable_cache_on_dual_transformer, + get_scm_mask, + ) + + # Or import directly from the submodule: + from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( + CacheDitConfig, + enable_cache_on_transformer, + ) +""" + +from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext, TeaCacheMixin + +__all__ = [ + # TeaCache (always available) + "TeaCacheContext", + "TeaCacheMixin", + # cache-dit integration (lazy-loaded, requires cache-dit package) + "CacheDitConfig", + "enable_cache_on_transformer", + "enable_cache_on_dual_transformer", + "get_scm_mask", +] + + +# Lazy imports for cache-dit integration (requires cache-dit package) +def __getattr__(name): + if name in ( + "CacheDitConfig", + "enable_cache_on_transformer", + "enable_cache_on_dual_transformer", + "get_scm_mask", + ): + from sglang.multimodal_gen.runtime.cache import cache_dit_integration + + return getattr(cache_dit_integration, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py b/python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/utils/cache_dit_integration.py rename to python/sglang/multimodal_gen/runtime/cache/cache_dit_integration.py diff --git a/python/sglang/multimodal_gen/runtime/cache/teacache.py b/python/sglang/multimodal_gen/runtime/cache/teacache.py new file mode 100644 index 000000000..40f9123ef --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/cache/teacache.py @@ -0,0 +1,313 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +TeaCache: Temporal similarity-based caching for diffusion models. + +TeaCache accelerates diffusion inference by selectively skipping redundant +computation when consecutive diffusion steps are similar enough. This is +achieved by tracking the L1 distance between modulated inputs across timesteps. + +Key concepts: +- Modulated input: The input to transformer blocks after timestep conditioning +- L1 distance: Measures how different consecutive timesteps are +- Threshold: When accumulated L1 distance exceeds threshold, force computation +- CFG support: Separate caches for positive and negative branches + +References: +- TeaCache: Accelerating Diffusion Models with Temporal Similarity + https://arxiv.org/abs/2411.14324 +""" + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +if TYPE_CHECKING: + from sglang.multimodal_gen.configs.sample.teacache import TeaCacheParams + + +@dataclass +class TeaCacheContext: + """Common context extracted for TeaCache skip decision. + + This context is populated from the forward_batch and forward_context + during each denoising step, providing all information needed to make + cache decisions. + + Attributes: + current_timestep: Current denoising timestep index (0-indexed). + num_inference_steps: Total number of inference steps. + do_cfg: Whether classifier-free guidance is enabled. + is_cfg_negative: True if currently processing negative CFG branch. + teacache_thresh: Threshold for accumulated L1 distance. + coefficients: Polynomial coefficients for L1 rescaling. + teacache_params: Full TeaCacheParams for model-specific access. + """ + + current_timestep: int + num_inference_steps: int + do_cfg: bool + is_cfg_negative: bool # For CFG branch selection + teacache_thresh: float + coefficients: list[float] + teacache_params: "TeaCacheParams" # Full params for model-specific access + + +class TeaCacheMixin: + """ + Mixin class providing TeaCache optimization functionality. + + TeaCache accelerates diffusion inference by selectively skipping redundant + computation when consecutive diffusion steps are similar enough. + + This mixin should be inherited by DiT model classes that want to support + TeaCache optimization. It provides: + - State management for tracking L1 distances + - CFG-aware caching (separate caches for positive/negative branches) + - Decision logic for when to compute vs. use cache + + Example usage in a DiT model: + class MyDiT(TeaCacheMixin, BaseDiT): + def __init__(self, config, **kwargs): + super().__init__(config, **kwargs) + self._init_teacache_state() + + def forward(self, hidden_states, timestep, ...): + ctx = self._get_teacache_context() + if ctx is not None: + # Compute modulated input (model-specific, e.g., after timestep embedding) + modulated_input = self._compute_modulated_input(hidden_states, timestep) + is_boundary = (ctx.current_timestep == 0 or + ctx.current_timestep >= ctx.num_inference_steps - 1) + + should_calc = self._compute_teacache_decision( + modulated_inp=modulated_input, + is_boundary_step=is_boundary, + coefficients=ctx.coefficients, + teacache_thresh=ctx.teacache_thresh, + ) + + if not should_calc: + # Use cached residual (must implement retrieve_cached_states) + return self.retrieve_cached_states(hidden_states) + + # Normal forward pass... + output = self._transformer_forward(hidden_states, timestep, ...) + + # Cache states for next step + if ctx is not None: + self.maybe_cache_states(output, hidden_states) + + return output + + Subclass implementation notes: + - `_compute_modulated_input()`: Model-specific method to compute the input + after timestep conditioning (used for L1 distance calculation) + - `retrieve_cached_states()`: Must be overridden to return cached output + - `maybe_cache_states()`: Override to store states for cache retrieval + + Attributes: + cnt: Counter for tracking steps. + enable_teacache: Whether TeaCache is enabled. + previous_modulated_input: Cached modulated input for positive branch. + previous_residual: Cached residual for positive branch. + accumulated_rel_l1_distance: Accumulated L1 distance for positive branch. + is_cfg_negative: Whether currently processing negative CFG branch. + _supports_cfg_cache: Whether this model supports CFG cache separation. + + CFG-specific attributes (only when _supports_cfg_cache is True): + previous_modulated_input_negative: Cached input for negative branch. + previous_residual_negative: Cached residual for negative branch. + accumulated_rel_l1_distance_negative: L1 distance for negative branch. + """ + + # Models that support CFG cache separation (wan/hunyuan/zimage) + # Models not in this set (flux/qwen) auto-disable TeaCache when CFG is enabled + _CFG_SUPPORTED_PREFIXES: set[str] = {"wan", "hunyuan", "zimage"} + + def _init_teacache_state(self) -> None: + """Initialize TeaCache state. Call this in subclass __init__.""" + # Common TeaCache state + self.cnt = 0 + self.enable_teacache = True + # Flag indicating if this model supports CFG cache separation + self._supports_cfg_cache = ( + self.config.prefix.lower() in self._CFG_SUPPORTED_PREFIXES + ) + + # Always initialize positive cache fields (used in all modes) + self.previous_modulated_input: torch.Tensor | None = None + self.previous_residual: torch.Tensor | None = None + self.accumulated_rel_l1_distance: float = 0.0 + + self.is_cfg_negative = False + # CFG-specific fields initialized to None (created when CFG is used) + # These are only used when _supports_cfg_cache is True AND do_cfg is True + if self._supports_cfg_cache: + self.previous_modulated_input_negative: torch.Tensor | None = None + self.previous_residual_negative: torch.Tensor | None = None + self.accumulated_rel_l1_distance_negative: float = 0.0 + + def reset_teacache_state(self) -> None: + """Reset all TeaCache state at the start of each generation task.""" + self.cnt = 0 + + # Primary cache fields (always present) + self.previous_modulated_input = None + self.previous_residual = None + self.accumulated_rel_l1_distance = 0.0 + self.is_cfg_negative = False + self.enable_teacache = True + # CFG negative cache fields (always reset, may be unused) + if self._supports_cfg_cache: + self.previous_modulated_input_negative = None + self.previous_residual_negative = None + self.accumulated_rel_l1_distance_negative = 0.0 + + def _compute_l1_and_decide( + self, + modulated_inp: torch.Tensor, + coefficients: list[float], + teacache_thresh: float, + ) -> tuple[float, bool]: + """ + Compute L1 distance and decide whether to calculate or use cache. + + Args: + modulated_inp: Current timestep's modulated input. + coefficients: Polynomial coefficients for L1 rescaling. + teacache_thresh: Threshold for cache decision. + + Returns: + Tuple of (new_accumulated_distance, should_calc). + """ + prev_modulated_inp = ( + self.previous_modulated_input_negative + if self.is_cfg_negative + else self.previous_modulated_input + ) + + # Defensive check: if previous input is not set, force calculation + if prev_modulated_inp is None: + return 0.0, True + + # Compute relative L1 distance + diff = modulated_inp - prev_modulated_inp + rel_l1 = (diff.abs().mean() / prev_modulated_inp.abs().mean()).cpu().item() + + # Apply polynomial rescaling + rescale_func = np.poly1d(coefficients) + + accumulated_rel_l1_distance = ( + self.accumulated_rel_l1_distance_negative + if self.is_cfg_negative + else self.accumulated_rel_l1_distance + ) + accumulated_rel_l1_distance = accumulated_rel_l1_distance + rescale_func(rel_l1) + + if accumulated_rel_l1_distance >= teacache_thresh: + # Threshold exceeded: force compute and reset accumulator + return 0.0, True + # Cache hit: keep accumulated distance + return accumulated_rel_l1_distance, False + + def _compute_teacache_decision( + self, + modulated_inp: torch.Tensor, + is_boundary_step: bool, + coefficients: list[float], + teacache_thresh: float, + ) -> bool: + """ + Compute cache decision for TeaCache. + + Args: + modulated_inp: Current timestep's modulated input. + is_boundary_step: True for boundary timesteps that always compute. + coefficients: Polynomial coefficients for L1 rescaling. + teacache_thresh: Threshold for cache decision. + + Returns: + True if forward computation is needed, False to use cache. + """ + if not self.enable_teacache: + return True + + if is_boundary_step: + new_accum, should_calc = 0.0, True + else: + new_accum, should_calc = self._compute_l1_and_decide( + modulated_inp=modulated_inp, + coefficients=coefficients, + teacache_thresh=teacache_thresh, + ) + + # Advance baseline and accumulator for the active branch + if not self.is_cfg_negative: + self.previous_modulated_input = modulated_inp.clone() + self.accumulated_rel_l1_distance = new_accum + elif self._supports_cfg_cache: + self.previous_modulated_input_negative = modulated_inp.clone() + self.accumulated_rel_l1_distance_negative = new_accum + + return should_calc + + def _get_teacache_context(self) -> TeaCacheContext | None: + """ + Check TeaCache preconditions and extract common context. + + Returns: + TeaCacheContext if TeaCache is enabled and properly configured, + None if should skip TeaCache logic entirely. + """ + from sglang.multimodal_gen.runtime.managers.forward_context import ( + get_forward_context, + ) + + forward_context = get_forward_context() + forward_batch = forward_context.forward_batch + + # Early return checks + if ( + forward_batch is None + or not forward_batch.enable_teacache + or forward_batch.teacache_params is None + ): + return None + + teacache_params = forward_batch.teacache_params + + # Extract common values + current_timestep = forward_context.current_timestep + num_inference_steps = forward_batch.num_inference_steps + do_cfg = forward_batch.do_classifier_free_guidance + is_cfg_negative = forward_batch.is_cfg_negative + + # Reset at first timestep + if current_timestep == 0 and not self.is_cfg_negative: + self.reset_teacache_state() + + return TeaCacheContext( + current_timestep=current_timestep, + num_inference_steps=num_inference_steps, + do_cfg=do_cfg, + is_cfg_negative=is_cfg_negative, + teacache_thresh=teacache_params.teacache_thresh, + coefficients=teacache_params.coefficients, + teacache_params=teacache_params, + ) + + def maybe_cache_states( + self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor + ) -> None: + """Cache states for later retrieval. Override in subclass if needed.""" + pass + + def should_skip_forward_for_cached_states(self, **kwargs: dict[str, Any]) -> bool: + """Check if forward can be skipped using cached states.""" + return False + + def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Retrieve cached states. Must be implemented by subclass.""" + raise NotImplementedError("retrieve_cached_states is not implemented") diff --git a/python/sglang/multimodal_gen/runtime/models/dits/base.py b/python/sglang/multimodal_gen/runtime/models/dits/base.py index 886a6a331..e3cc4a22e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/base.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/base.py @@ -8,6 +8,12 @@ import torch from torch import nn from sglang.multimodal_gen.configs.models import DiTConfig + +# NOTE: TeaCacheContext and TeaCacheMixin have been moved to +# sglang.multimodal_gen.runtime.cache.teacache +# For backwards compatibility, re-export from the new location +from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheContext # noqa: F401 +from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheMixin from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum @@ -77,11 +83,11 @@ class BaseDiT(nn.Module, ABC): return next(self.parameters()).device -class CachableDiT(BaseDiT): +class CachableDiT(TeaCacheMixin, BaseDiT): """ An intermediate base class that adds TeaCache optimization functionality to DiT models. - TeaCache accelerates inference by selectively skipping redundant computation when consecutive - diffusion steps are similar enough. + + Inherits TeaCacheMixin for cache logic and BaseDiT for core DiT functionality. """ # These are required class attributes that should be overridden by concrete implementations @@ -100,35 +106,4 @@ class CachableDiT(BaseDiT): def __init__(self, config: DiTConfig, **kwargs) -> None: super().__init__(config, **kwargs) - - self.cnt = 0 - self.teacache_thresh = 0 - self.coefficients: list[float] = [] - - # NOTE(will): Only wan2.1 needs these, so we are hardcoding it here - if self.config.prefix == "wan": - self.use_ret_steps = self.config.cache_config.use_ret_steps - self.is_even = False - self.previous_residual_even: torch.Tensor | None = None - self.previous_residual_odd: torch.Tensor | None = None - self.accumulated_rel_l1_distance_even = 0 - self.accumulated_rel_l1_distance_odd = 0 - self.should_calc_even = True - self.should_calc_odd = True - else: - self.accumulated_rel_l1_distance = 0 - self.previous_modulated_input = None - self.previous_resiual = None - self.previous_e0_even: torch.Tensor | None = None - self.previous_e0_odd: torch.Tensor | None = None - - def maybe_cache_states( - self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor - ) -> None: - pass - - def should_skip_forward_for_cached_states(self, **kwargs: dict[str, Any]) -> bool: - return False - - def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: - raise NotImplementedError("maybe_retrieve_cached_states is not implemented") + self._init_teacache_state() diff --git a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py index 3dcb8185b..31ed3c945 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/wanvideo.py @@ -5,7 +5,6 @@ import math from typing import Any -import numpy as np import torch import torch.nn as nn @@ -752,15 +751,7 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): ) # For type checking - self.previous_e0_even = None - self.previous_e0_odd = None - self.previous_residual_even = None - self.previous_residual_odd = None - self.is_even = True - self.should_calc_even = True - self.should_calc_odd = True - self.accumulated_rel_l1_distance_even = 0 - self.accumulated_rel_l1_distance_odd = 0 + self.cnt = 0 self.__post_init__() @@ -793,7 +784,9 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): **kwargs, ) -> torch.Tensor: forward_batch = get_forward_context().forward_batch - enable_teacache = forward_batch is not None and forward_batch.enable_teacache + self.enable_teacache = ( + forward_batch is not None and forward_batch.enable_teacache + ) orig_dtype = hidden_states.dtype if not isinstance(encoder_hidden_states, torch.Tensor): @@ -878,7 +871,7 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): hidden_states = self.retrieve_cached_states(hidden_states) else: # if teacache is enabled, we need to cache the original hidden states - if enable_teacache: + if self.enable_teacache: original_hidden_states = hidden_states.clone() for block in self.blocks: @@ -886,8 +879,9 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): hidden_states, encoder_hidden_states, timestep_proj, freqs_cis ) # if teacache is enabled, we need to cache the original hidden states - if enable_teacache: + if self.enable_teacache: self.maybe_cache_states(hidden_states, original_hidden_states) + self.cnt += 1 # 5. Output norm, projection & unpatchify if temb.dim() == 3: # batch_size, seq_len, inner_dim (wan 2.2 ti2v) @@ -921,114 +915,62 @@ class WanTransformer3DModel(CachableDiT, OffloadableDiTMixin): def maybe_cache_states( self, hidden_states: torch.Tensor, original_hidden_states: torch.Tensor ) -> None: - if self.is_even: - self.previous_residual_even = ( - hidden_states.squeeze(0) - original_hidden_states - ) + """Cache residual with CFG positive/negative separation.""" + residual = hidden_states.squeeze(0) - original_hidden_states + if not self.is_cfg_negative: + self.previous_residual = residual else: - self.previous_residual_odd = ( - hidden_states.squeeze(0) - original_hidden_states - ) + self.previous_residual_negative = residual def should_skip_forward_for_cached_states(self, **kwargs) -> bool: - - forward_context = get_forward_context() - forward_batch = forward_context.forward_batch - if forward_batch is None or not forward_batch.enable_teacache: + if not self.enable_teacache: return False - teacache_params = forward_batch.teacache_params - assert teacache_params is not None, "teacache_params is not initialized" + ctx = self._get_teacache_context() + if ctx is None: + return False + + # Wan uses WanTeaCacheParams with additional fields + teacache_params = ctx.teacache_params assert isinstance( teacache_params, WanTeaCacheParams ), "teacache_params is not a WanTeaCacheParams" - current_timestep = forward_context.current_timestep - num_inference_steps = forward_batch.num_inference_steps - # initialize the coefficients, cutoff_steps, and ret_steps - coefficients = teacache_params.coefficients + # Initialize Wan-specific parameters use_ret_steps = teacache_params.use_ret_steps - cutoff_steps = teacache_params.get_cutoff_steps(num_inference_steps) + cutoff_steps = teacache_params.get_cutoff_steps(ctx.num_inference_steps) ret_steps = teacache_params.ret_steps - teacache_thresh = teacache_params.teacache_thresh - if current_timestep == 0: - self.cnt = 0 + # Adjust ret_steps and cutoff_steps for non-CFG mode + # (WanTeaCacheParams uses *2 factor assuming CFG) + if not ctx.do_cfg: + ret_steps = ret_steps // 2 + cutoff_steps = cutoff_steps // 2 timestep_proj = kwargs["timestep_proj"] temb = kwargs["temb"] modulated_inp = timestep_proj if use_ret_steps else temb - if self.cnt % 2 == 0: # even -> condition - self.is_even = True - if self.cnt < ret_steps or self.cnt >= cutoff_steps: - self.should_calc_even = True - self.accumulated_rel_l1_distance_even = 0 - else: - assert ( - self.previous_e0_even is not None - ), "previous_e0_even is not initialized" - assert ( - self.accumulated_rel_l1_distance_even is not None - ), "accumulated_rel_l1_distance_even is not initialized" - rescale_func = np.poly1d(coefficients) - self.accumulated_rel_l1_distance_even += rescale_func( - ( - (modulated_inp - self.previous_e0_even).abs().mean() - / self.previous_e0_even.abs().mean() - ) - .cpu() - .item() - ) - if self.accumulated_rel_l1_distance_even < teacache_thresh: - self.should_calc_even = False - else: - self.should_calc_even = True - self.accumulated_rel_l1_distance_even = 0 - self.previous_e0_even = modulated_inp.clone() + self.is_cfg_negative = ctx.is_cfg_negative - else: # odd -> unconditon - self.is_even = False - if self.cnt < ret_steps or self.cnt >= cutoff_steps: - self.should_calc_odd = True - self.accumulated_rel_l1_distance_odd = 0 - else: - assert ( - self.previous_e0_odd is not None - ), "previous_e0_odd is not initialized" - assert ( - self.accumulated_rel_l1_distance_odd is not None - ), "accumulated_rel_l1_distance_odd is not initialized" - rescale_func = np.poly1d(coefficients) - self.accumulated_rel_l1_distance_odd += rescale_func( - ( - (modulated_inp - self.previous_e0_odd).abs().mean() - / self.previous_e0_odd.abs().mean() - ) - .cpu() - .item() - ) - if self.accumulated_rel_l1_distance_odd < teacache_thresh: - self.should_calc_odd = False - else: - self.should_calc_odd = True - self.accumulated_rel_l1_distance_odd = 0 - self.previous_e0_odd = modulated_inp.clone() - self.cnt += 1 - should_skip_forward = False - if self.is_even: - if not self.should_calc_even: - should_skip_forward = True - else: - if not self.should_calc_odd: - should_skip_forward = True + # Wan uses ret_steps/cutoff_steps for boundary detection + is_boundary_step = self.cnt < ret_steps or self.cnt >= cutoff_steps - return should_skip_forward + # Use shared helper to compute cache decision + should_calc = self._compute_teacache_decision( + modulated_inp=modulated_inp, + is_boundary_step=is_boundary_step, + coefficients=ctx.coefficients, + teacache_thresh=ctx.teacache_thresh, + ) + + return not should_calc def retrieve_cached_states(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.is_even: - return hidden_states + self.previous_residual_even + """Retrieve cached residual with CFG positive/negative separation.""" + if not self.is_cfg_negative: + return hidden_states + self.previous_residual else: - return hidden_states + self.previous_residual_odd + return hidden_states + self.previous_residual_negative EntryClass = WanTransformer3DModel diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index dd0a721b4..47bd3c321 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -161,17 +161,17 @@ class DenoisingStage(PipelineStage): if not envs.SGLANG_CACHE_DIT_ENABLED or batch.is_warmup: return - from sglang.multimodal_gen.runtime.distributed import ( - get_sp_group, - get_tp_group, - get_world_size, - ) - from sglang.multimodal_gen.runtime.utils.cache_dit_integration import ( + from sglang.multimodal_gen.runtime.cache.cache_dit_integration import ( CacheDitConfig, enable_cache_on_dual_transformer, enable_cache_on_transformer, get_scm_mask, ) + from sglang.multimodal_gen.runtime.distributed import ( + get_sp_group, + get_tp_group, + get_world_size, + ) world_size = get_world_size() parallelized = world_size > 1 diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 5de7aadbe..e0187c4b5 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -758,6 +758,73 @@ "expected_avg_denoise_ms": 104.84, "expected_median_denoise_ms": 102.01 }, + "wan2_1_t2v_1.3b_teacache_enabled": { + "stages_ms": { + "InputValidationStage": 0.07, + "TextEncodingStage": 2237.78, + "ConditioningStage": 0.01, + "TimestepPreparationStage": 2.1, + "LatentPreparationStage": 0.84, + "DenoisingStage": 13041.23, + "DecodingStage": 1274.63, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 879.71, + "1": 248.13, + "2": 246.48, + "3": 247.87, + "4": 249.38, + "5": 246.76, + "6": 250.42, + "7": 250.81, + "8": 250.98, + "9": 249.9, + "10": 246.72, + "11": 249.79, + "12": 250.46, + "13": 249.19, + "14": 247.55, + "15": 250.12, + "16": 247.57, + "17": 247.21, + "18": 247.32, + "19": 247.42, + "20": 248.21, + "21": 247.19, + "22": 247.72, + "23": 247.45, + "24": 247.9, + "25": 247.87, + "26": 247.18, + "27": 247.65, + "28": 246.91, + "29": 248.26, + "30": 247.82, + "31": 247.73, + "32": 247.38, + "33": 247.84, + "34": 247.46, + "35": 247.52, + "36": 247.94, + "37": 248.76, + "38": 248.01, + "39": 247.45, + "40": 247.84, + "41": 248.33, + "42": 247.41, + "43": 248.16, + "44": 248.18, + "45": 248.44, + "46": 248.65, + "47": 247.73, + "48": 247.48, + "49": 247.54 + }, + "expected_e2e_ms": 18382.19, + "expected_avg_denoise_ms": 260.76, + "expected_median_denoise_ms": 247.84 + }, "wan2_1_t2v_1.3b": { "stages_ms": { "InputValidationStage": 0.07, diff --git a/python/sglang/multimodal_gen/test/server/test_server_utils.py b/python/sglang/multimodal_gen/test/server/test_server_utils.py index 4bdcb7fc9..f20e9e9e0 100644 --- a/python/sglang/multimodal_gen/test/server/test_server_utils.py +++ b/python/sglang/multimodal_gen/test/server/test_server_utils.py @@ -843,12 +843,18 @@ def get_generate_fn( req_output_format = None # Not specified in current request req_background = None # Not specified in current request + # Build extra_body for optional features + extra_body = {} + if sampling_params.enable_teacache: + extra_body["enable_teacache"] = True + response = client.images.with_raw_response.generate( model=model_path, prompt=sampling_params.prompt, n=n, size=output_size, response_format="b64_json", + extra_body=extra_body if extra_body else None, ) result = response.parse() validate_image(result.data[0].b64_json) @@ -911,6 +917,11 @@ def get_generate_fn( ) # Not specified in current request req_background = None # Not specified in current request + # Build extra_body for optional features + extra_body = {"num_frames": sampling_params.num_frames} + if sampling_params.enable_teacache: + extra_body["enable_teacache"] = True + images = [open(image_path, "rb") for image_path in image_paths] try: response = client.images.with_raw_response.edit( @@ -921,7 +932,7 @@ def get_generate_fn( size=output_size, response_format="b64_json", output_format=req_output_format, - extra_body={"num_frames": sampling_params.num_frames}, + extra_body=extra_body, ) finally: for img in images: @@ -1036,6 +1047,11 @@ def get_generate_fn( if not sampling_params.prompt: pytest.skip(f"{id}: no text prompt configured") + # Build extra_body for optional features + extra_body = {} + if sampling_params.enable_teacache: + extra_body["enable_teacache"] = True + return _create_and_download_video( client, case_id, @@ -1043,6 +1059,7 @@ def get_generate_fn( prompt=sampling_params.prompt, size=output_size, seconds=video_seconds, + extra_body=extra_body if extra_body else None, ) def generate_image_to_video(case_id, client) -> str: @@ -1057,6 +1074,11 @@ def get_generate_fn( if not image_path.exists(): pytest.skip(f"{id}: file missing: {image_path}") + # Build extra_body for optional features + extra_body = {} + if sampling_params.enable_teacache: + extra_body["enable_teacache"] = True + with image_path.open("rb") as fh: return _create_and_download_video( client, @@ -1066,11 +1088,18 @@ def get_generate_fn( size=output_size, seconds=video_seconds, input_reference=fh, + extra_body=extra_body if extra_body else None, ) def generate_text_url_image_to_video(case_id, client) -> str: if not sampling_params.prompt or not sampling_params.image_path: pytest.skip(f"{id}: no edit config") + + # Build extra_body for optional features + extra_body = {"reference_url": sampling_params.image_path} + if sampling_params.enable_teacache: + extra_body["enable_teacache"] = True + return _create_and_download_video( client, case_id, @@ -1097,6 +1126,11 @@ def get_generate_fn( if not image_path.exists(): pytest.skip(f"{id}: file missing: {image_path}") + # Build extra_body for optional features + extra_body = {} + if sampling_params.enable_teacache: + extra_body["enable_teacache"] = True + with image_path.open("rb") as fh: return _create_and_download_video( client, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 8f9b7d4b5..e93ffa7f6 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -191,6 +191,9 @@ class DiffusionSamplingParams: num_outputs_per_prompt: int = 1 + # TeaCache acceleration + enable_teacache: bool = False + @dataclass(frozen=True) class DiffusionTestCase: @@ -457,6 +460,20 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [ prompt=T2V_PROMPT, ), ), + # TeaCache acceleration test for Wan video model + DiffusionTestCase( + "wan2_1_t2v_1.3b_teacache_enabled", + DiffusionServerArgs( + model_path="Wan-AI/Wan2.1-T2V-1.3B-Diffusers", + modality="video", + warmup=0, + custom_validator="video", + ), + DiffusionSamplingParams( + prompt=T2V_PROMPT, + enable_teacache=True, + ), + ), # LoRA test case for single transformer + merge/unmerge API test # Note: Uses dynamic_lora_path instead of lora_path to test LayerwiseOffload + set_lora interaction # Server starts WITHOUT LoRA, then set_lora is called after startup (Wan models auto-enable layerwise offload)