[diffusion] refactor: refactor and simplify teacache for cachabledit and wanvideo (#16396)

Co-authored-by: Brain97 <Brain97@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: blahblah <blahblah>

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
+> **Note**: This is one of two caching strategies available in SGLang.
-Transformers (DiT), to achieve up to **7.4x inference speedup** with minimal quality loss.
+> 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
 

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)

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)

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.
+These variables configure caching acceleration for Diffusion Transformer (DiT) models.
-See [cache-dit documentation](cache_dit.md) for details.
+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                              |
 |-------------------------------------|---------|------------------------------------------|

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}")

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")

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._init_teacache_state()
-        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")

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:
+        """Cache residual with CFG positive/negative separation."""
-            self.previous_residual_even = (
+        residual = hidden_states.squeeze(0) - original_hidden_states
-                hidden_states.squeeze(0) - original_hidden_states
+        if not self.is_cfg_negative:
-            )
+            self.previous_residual = residual
         else:
-            self.previous_residual_odd = (
+            self.previous_residual_negative = residual
-                hidden_states.squeeze(0) - original_hidden_states
-            )
 
     def should_skip_forward_for_cached_states(self, **kwargs) -> bool:
-
+        if not self.enable_teacache:
-        forward_context = get_forward_context()
-        forward_batch = forward_context.forward_batch
-        if forward_batch is None or not forward_batch.enable_teacache:
             return False
-        teacache_params = forward_batch.teacache_params
+        ctx = self._get_teacache_context()
-        assert teacache_params is not None, "teacache_params is not initialized"
+        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
+        # Initialize Wan-specific parameters
-        coefficients = teacache_params.coefficients
         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:
+        # Adjust ret_steps and cutoff_steps for non-CFG mode
-            self.cnt = 0
+        # (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_cfg_negative = ctx.is_cfg_negative
-            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()
 
-        else:  # odd -> unconditon
+        # Wan uses ret_steps/cutoff_steps for boundary detection
-            self.is_even = False
+        is_boundary_step = self.cnt < ret_steps or self.cnt >= cutoff_steps
-            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
 
-        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:
+        """Retrieve cached residual with CFG positive/negative separation."""
-            return hidden_states + self.previous_residual_even
+        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

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 (
+        from sglang.multimodal_gen.runtime.cache.cache_dit_integration import (
-            get_sp_group,
-            get_tp_group,
-            get_world_size,
-        )
-        from sglang.multimodal_gen.runtime.utils.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

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,

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,

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)