[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/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.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()

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

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