[diffusion] Cleanup TeaCache

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -146,6 +146,7 @@ class SamplingParams:
 
     # TeaCache parameters
     enable_teacache: bool = False
+    calibrate_cache: bool = False
 
     # Profiling
     profile: bool = False
@@ -592,6 +593,12 @@ def add_cli_args(parser: Any) -> Any:
             action="store_true",
             default=SamplingParams.enable_teacache,
         )
+        parser.add_argument(
+            "--calibrate-cache",
+            action="store_true",
+            default=SamplingParams.calibrate_cache,
+            help="Run in calibration mode: collect magnitude ratio statistics instead of skipping steps.",
+        )
 
         # profiling
         parser.add_argument(

python/sglang/multimodal_gen/configs/sample/teacache.py

@@ -10,14 +10,18 @@
 class TeaCacheParams(CacheParams):
     cache_type: str = "teacache"
     teacache_thresh: float = 0.0
+    skip_start_step: int = 5
+    skip_end_step: int = 0
     coefficients: list[float] = field(default_factory=list)
 
 
 @dataclass
 class WanTeaCacheParams(CacheParams):
     # Unfortunately, TeaCache is very different for Wan than other models
     cache_type: str = "teacache"
-    teacache_thresh: float = 0.0
+    teacache_thresh: float = 0.08
+    skip_start_step: int = 5
+    skip_end_step: int = 0
     use_ret_steps: bool = True
     ret_steps_coeffs: list[float] = field(default_factory=list)
     non_ret_steps_coeffs: list[float] = field(default_factory=list)
@@ -28,16 +32,3 @@ def coefficients(self) -> list[float]:
             return self.ret_steps_coeffs
         else:
             return self.non_ret_steps_coeffs
-
-    @property
-    def ret_steps(self) -> int:
-        if self.use_ret_steps:
-            return 5 * 2
-        else:
-            return 1 * 2
-
-    def get_cutoff_steps(self, num_inference_steps: int) -> int:
-        if self.use_ret_steps:
-            return num_inference_steps * 2
-        else:
-            return num_inference_steps * 2 - 2

python/sglang/multimodal_gen/runtime/cache/__init__.py

@@ -10,18 +10,26 @@
 
 """
 
+from sglang.multimodal_gen.runtime.cache.base import DiffusionCache
 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.cache.teacache import TeaCacheContext, TeaCacheMixin
+from sglang.multimodal_gen.runtime.cache.teacache import (
+    TeaCacheContext,
+    TeaCacheState,
+    TeaCacheStrategy,
+)
 
 __all__ = [
+    # Base
+    "DiffusionCache",
     # TeaCache (always available)
     "TeaCacheContext",
-    "TeaCacheMixin",
+    "TeaCacheState",
+    "TeaCacheStrategy",
     # cache-dit integration (lazy-loaded, requires cache-dit package)
     "CacheDitConfig",
     "enable_cache_on_transformer",

python/sglang/multimodal_gen/runtime/cache/base.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Base class for diffusion model cache strategies (TeaCache, MagCache, etc.).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.multimodal_gen.runtime.cache.magcache import MagCacheState
+    from sglang.multimodal_gen.runtime.cache.teacache import TeaCacheState
+
+
+class DiffusionCache:
+    """
+    Base class for diffusion model timestep caching strategies (TeaCache, MagCache, etc.).
+
+    Note: this is *timestep-level* caching — skipping the full transformer forward
+    pass for certain denoising timesteps. It is unrelated to Cache-DiT's block-step
+    caching, which operates at the transformer-block level.
+
+    Each subclass owns its own state (positive + negative CFG branch) and
+    context extraction logic. CachableDiT holds a single
+    `self.cache: DiffusionCache | None` and delegates all decisions here.
+
+    Subclasses must implement: reset, get_context, should_skip.
+    maybe_cache, retrieve, and calibrate have default implementations.
+
+    Subclasses set self.state / self.state_neg to strategy-specific state
+    objects (positive and negative CFG branches respectively). state_neg is
+    None when CFG negative-branch caching is disabled.
+
+    Typical forward pass usage in CachableDiT:
+
+        ctx = self.cache.get_context(self.cnt)
+        if ctx and self.cache.should_skip(ctx, timestep_proj=..., temb=...):
+            hidden_states = self.cache.retrieve(hidden_states, ctx)
+        else:
+            original_hidden_states = hidden_states.clone()
+            # ... run transformer blocks ...
+            if calibrate_cache:
+                self.cache.calibrate(hidden_states, original_hidden_states, ctx)
+            else:
+                self.cache.maybe_cache(hidden_states, original_hidden_states, ctx)
+    """
+
+    def __init__(self) -> None:
+        self.state: MagCacheState | TeaCacheState | None = None
+        self.state_neg: MagCacheState | TeaCacheState | None = None
+
+    def reset(self) -> None:
+        """Reset all state at the start of a new generation."""
+        raise NotImplementedError
+
+    def get_context(self, cnt: int):
+        """
+        Read the global forward_context / forward_batch and return a
+        strategy-specific context dataclass, or None to bypass caching.
+
+        cnt is the monotonically increasing forward-call index owned by the
+        model (model.cnt), incremented on every call regardless of whether
+        the forward pass was skipped.
+        """
+        raise NotImplementedError
+
+    def should_skip(self, ctx, **kwargs) -> bool:
+        """
+        Decide whether to skip the transformer forward pass and reuse the
+        cached residual. kwargs carries model-specific tensors (e.g.
+        timestep_proj, temb) needed by some strategies.
+        """
+        raise NotImplementedError
+
+    def maybe_cache(
+        self,
+        hidden_states: torch.Tensor,
+        original_hidden_states: torch.Tensor,
+        ctx,
+    ) -> None:
+        """Store residual after a full forward pass for future reuse."""
+        state = (
+            self.state_neg
+            if (ctx.is_cfg_negative and self.state_neg is not None)
+            else self.state
+        )
+        state.previous_residual = hidden_states.squeeze(0) - original_hidden_states
+
+    def retrieve(self, hidden_states: torch.Tensor, ctx) -> torch.Tensor:
+        """Reconstruct output from cached residual."""
+        state = (
+            self.state_neg
+            if (ctx.is_cfg_negative and self.state_neg is not None)
+            else self.state
+        )
+        return hidden_states + state.previous_residual
+
+    def calibrate(
+        self,
+        hidden_states: torch.Tensor,
+        original_hidden_states: torch.Tensor,
+        ctx,
+    ) -> None:
+        """Calibrate the cache. No-op by default."""
+        pass