chore: import upstream snapshot with attribution

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

@@ -0,0 +1,30 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Cache acceleration module for SGLang-diffusion
+
+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
+
+"""
+
+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
+
+__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",
+]

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

@@ -0,0 +1,683 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+cache-dit integration module for SGLang DiT pipelines.
+
+This module provides helper functions to enable cache-dit acceleration
+on transformer modules in SGLang's modular pipeline architecture.
+"""
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+import torch.distributed as dist
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_ring_parallel_world_size,
+    get_tp_world_size,
+    get_ulysses_parallel_world_size,
+)
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+import cache_dit
+from cache_dit import (
+    BlockAdapter,
+    DBCacheConfig,
+    ForwardPattern,
+    ParamsModifier,
+    TaylorSeerCalibratorConfig,
+    steps_mask,
+)
+from cache_dit.caching.block_adapters import BlockAdapterRegister
+from cache_dit.parallelism import ParallelismBackend, ParallelismConfig
+
+from sglang.multimodal_gen.runtime.distributed.parallel_state import get_dit_group
+
+_original_similarity = None
+
+
+def _patch_cache_dit_similarity():
+    from cache_dit.caching.cache_contexts import cache_manager
+
+    global _original_similarity
+    if _original_similarity is not None:
+        return
+
+    _original_similarity = cache_manager.CachedContextManager.similarity
+
+    def patched_similarity(self, t1, t2, *, threshold, parallelized=False, prefix="Fn"):
+        if not parallelized:
+            return _original_similarity(
+                self,
+                t1,
+                t2,
+                threshold=threshold,
+                parallelized=parallelized,
+                prefix=prefix,
+            )
+
+        sp_group = getattr(self, "_sglang_sp_group", None)
+        tp_group = getattr(self, "_sglang_tp_group", None)
+        tp_sp_group = getattr(self, "_sglang_tp_sp_group", None)
+        target_group = tp_sp_group or sp_group or tp_group
+
+        if target_group is None:
+            return _original_similarity(
+                self,
+                t1,
+                t2,
+                threshold=threshold,
+                parallelized=parallelized,
+                prefix=prefix,
+            )
+
+        # Adapted from https://github.com/vipshop/cache-dit/blob/main/src/cache_dit/caching/cache_contexts/cache_manager.py#L495-L523
+        condition_thresh = self.get_important_condition_threshold()
+        if condition_thresh > 0.0:
+            raw_diff = (t1 - t2).abs()
+            token_m_df = raw_diff.mean(dim=-1)
+            token_m_t1 = t1.abs().mean(dim=-1)
+            token_diff = token_m_df / token_m_t1
+            condition = token_diff > condition_thresh
+            if condition.sum() > 0:
+                condition = condition.unsqueeze(-1).expand_as(raw_diff)
+                mean_diff = raw_diff[condition].mean()
+                mean_t1 = t1[condition].abs().mean()
+            else:
+                mean_diff = (t1 - t2).abs().mean()
+                mean_t1 = t1.abs().mean()
+        else:
+            mean_diff = (t1 - t2).abs().mean()
+            mean_t1 = t1.abs().mean()
+
+        dist.all_reduce(mean_diff, op=dist.ReduceOp.AVG, group=target_group)
+        dist.all_reduce(mean_t1, op=dist.ReduceOp.AVG, group=target_group)
+
+        diff = (mean_diff / mean_t1).item()
+        self.add_residual_diff(diff)
+        return diff < threshold
+
+    cache_manager.CachedContextManager.similarity = patched_similarity
+
+
+def _build_parallelism_config(
+    sp_group: Optional[torch.distributed.ProcessGroup],
+    tp_group: Optional[torch.distributed.ProcessGroup],
+):
+    if sp_group is None and tp_group is None:
+        return None
+
+    ulysses_size = None
+    ring_size = None
+    if sp_group is not None:
+        ulysses_size = get_ulysses_parallel_world_size()
+        ring_size = get_ring_parallel_world_size()
+
+    tp_size = None
+    if tp_group is not None:
+        tp_size = get_tp_world_size()
+
+    return ParallelismConfig(
+        backend=ParallelismBackend.AUTO,
+        ulysses_size=ulysses_size,
+        ring_size=ring_size,
+        tp_size=tp_size,
+    )
+
+
+def _mark_transformer_parallelized(transformer, config, sp_group, tp_group):
+    if config is None:
+        return
+
+    transformer._is_parallelized = True
+    transformer._parallelism_config = config
+
+
+def get_scm_mask(
+    preset: str,
+    num_inference_steps: int,
+    compute_bins: Optional[List[int]] = None,
+    cache_bins: Optional[List[int]] = None,
+) -> Optional[List[int]]:
+    """
+    Get SCM mask using cache-dit's steps_mask().
+
+    This is a thin wrapper that delegates to cache-dit's built-in
+    steps_mask() function which handles all presets and scaling logic.
+
+    Args:
+        preset: Preset name ("none", "slow", "medium", "fast", "ultra").
+        compute_bins: Custom compute bins (overrides preset).
+        cache_bins: Custom cache bins (overrides preset).
+
+    Returns:
+        SCM mask list (1=compute, 0=cache), or None if disabled.
+    """
+    if preset == "none" and not (compute_bins and cache_bins):
+        return None
+
+    # Use cache-dit's steps_mask() directly
+    mask = steps_mask(
+        compute_bins=compute_bins,
+        cache_bins=cache_bins,
+        total_steps=num_inference_steps,
+        mask_policy=preset if preset != "none" else "medium",
+    )
+
+    compute_count = sum(mask)
+    cache_count = len(mask) - compute_count
+    logger.info(
+        "SCM: generated mask with %d compute steps, %d cache steps (preset=%s)",
+        compute_count,
+        cache_count,
+        preset,
+    )
+
+    return mask
+
+
+@dataclass
+class CacheDitConfig:
+    """Configuration for cache-dit integration.
+
+    Attributes:
+        enabled: Whether to enable cache-dit acceleration.
+        Fn_compute_blocks: Number of first blocks to always compute (DBCache F).
+        Bn_compute_blocks: Number of last blocks to always compute (DBCache B).
+        max_warmup_steps: Number of warmup steps before caching starts (DBCache W).
+        residual_diff_threshold: Threshold for residual difference (DBCache R).
+        max_continuous_cached_steps: Maximum consecutive cached steps (DBCache MC).
+        enable_taylorseer: Whether to enable TaylorSeer calibrator.
+        taylorseer_order: Order of Taylor expansion (1 or 2).
+        num_inference_steps: Total number of inference steps (required for transformer-only mode).
+        steps_computation_mask: Binary mask for step-level caching (1=compute, 0=cache).
+            Generated by get_scm_mask() (wrapper around cache_dit.steps_mask()).
+        steps_computation_policy: Caching policy for SCM ("dynamic" or "static").
+    """
+
+    enabled: bool = False
+    Fn_compute_blocks: int = 1
+    Bn_compute_blocks: int = 0
+    # Use 4 as default warmup steps instead of 8 in cache-dit, thus making
+    # DBCache work for few steps distilled models, e.g., Z-Image w/ 8-steps.
+    max_warmup_steps: int = 4
+    # Use a relatively higher residual diff threshold (namely, 0.24) as default
+    # to allow more aggressive caching due to we have already applied max continuous
+    # cached steps limit, otherwise, we should use a lower threshold here like 0.12.
+    residual_diff_threshold: float = 0.24
+    max_continuous_cached_steps: int = 3
+    # TaylorSeer is not suitable for few steps distilled models, so, we choose
+    # to disable it by default. Reference:
+    # - From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers,
+    #   https://arxiv.org/pdf/2503.06923
+    # - FoCa: Forecast then Calibrate: Feature Caching as ODE for Efficient
+    #   Diffusion Transformers, https://arxiv.org/pdf/2508.16211
+    enable_taylorseer: bool = False
+    taylorseer_order: int = 1
+    num_inference_steps: Optional[int] = None
+    # SCM fields (generated by _maybe_enable_cache_dit from env configuration)
+    steps_computation_mask: Optional[List[int]] = None
+    steps_computation_policy: str = "dynamic"
+
+
+@dataclass(frozen=True)
+class DualTransformerBlockAdapterSpec:
+    """BlockAdapter metadata for dual-transformer DiT pipelines.
+
+    This describes the cache-dit-facing structure of a pair of transformers.
+    The denoising loop semantics live in DenoisingStage; this spec only covers
+    how cache-dit should find blocks and interpret each block's forward.
+    """
+
+    blocks_attr: tuple[str, str]
+    blocks_name: Optional[List[str]]
+    forward_pattern: List[ForwardPattern]
+    check_forward_pattern: bool
+    check_num_outputs: bool
+    has_separate_cfg: bool
+
+
+DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS: dict[str, DualTransformerBlockAdapterSpec] = {
+    "wan2.2": DualTransformerBlockAdapterSpec(
+        blocks_attr=("blocks", "blocks"),
+        blocks_name=None,
+        forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
+        check_forward_pattern=True,
+        check_num_outputs=False,
+        has_separate_cfg=True,
+    ),
+    "ideogram4": DualTransformerBlockAdapterSpec(
+        blocks_attr=("layers", "layers"),
+        blocks_name=["layers", "layers"],
+        forward_pattern=[ForwardPattern.Pattern_3, ForwardPattern.Pattern_3],
+        check_forward_pattern=False,
+        check_num_outputs=False,
+        has_separate_cfg=False,
+    ),
+}
+
+
+# Custom BlockAdapter for DiT models absent from cache-dit's BlockAdapterRegister.
+# Value: (blocks attr, forward_pattern). forward_pattern must
+# match the block's forward signature (see cache_dit.ForwardPattern; e.g., ERNIE
+# uses Pattern_3). has_separate_cfg follows the run (passed by
+# enable_cache_on_transformer); cache-dit auto-resolves the remaining
+# fields.
+_CUSTOM_BLOCK_ADAPTER_SPECS: dict[str, tuple[str, ForwardPattern]] = {
+    "ErnieImageTransformer2DModel": ("layers", ForwardPattern.Pattern_3),
+    "Krea2Transformer2DModel": ("transformer_blocks", ForwardPattern.Pattern_3),
+}
+
+
+def _build_custom_block_adapter(
+    transformer: torch.nn.Module,
+    has_separate_cfg: bool = False,
+) -> Optional[BlockAdapter]:
+    """Build a manual BlockAdapter for a model absent from cache-dit's registry,
+    or None if the class is unknown."""
+    spec = _CUSTOM_BLOCK_ADAPTER_SPECS.get(transformer.__class__.__name__)
+    if spec is None:
+        return None
+    blocks_attr, forward_pattern = spec
+    blocks = getattr(transformer, blocks_attr, None)
+    if blocks is None:
+        raise ValueError(
+            f"Transformer {transformer.__class__.__name__} has no attribute "
+            f"{blocks_attr!r} for cache-dit blocks."
+        )
+    return BlockAdapter(
+        transformer=transformer,
+        blocks=blocks,
+        forward_pattern=forward_pattern,
+        has_separate_cfg=has_separate_cfg,
+    )
+
+
+def enable_cache_on_transformer(
+    transformer: torch.nn.Module,
+    config: CacheDitConfig,
+    model_name: str = "transformer",
+    sp_group: Optional[torch.distributed.ProcessGroup] = None,
+    tp_group: Optional[torch.distributed.ProcessGroup] = None,
+    has_separate_cfg: bool = False,
+) -> torch.nn.Module:
+    """Enable cache-dit on a transformer module, by wrapping the module with cache-dit
+
+    This function enables cache-dit acceleration using the BlockAdapterRegister
+    for pre-registered models
+
+    Args:
+        model_name: Name of the model for logging purposes.
+        sp_group: Sequence parallel process group (for Ulysses/Ring).
+        tp_group: Tensor parallel process group.
+        has_separate_cfg: Whether the run issues separate conditional/unconditional
+            passes per step (CFG). Used by custom adapters (ERNIE, Krea-2); a
+            mismatch only disables caching, never corrupts output.
+
+    """
+    if not config.enabled:
+        return transformer
+
+    if config.num_inference_steps is None:
+        raise ValueError(
+            "num_inference_steps is required for transformer-only mode. "
+            "Please provide it in CacheDitConfig."
+        )
+
+    # Prefer the standard path (transformer pre-registered in cache-dit). For
+    # models absent from the registry, fall back to a manual BlockAdapter (see
+    # _build_custom_block_adapter).
+    custom_adapter = None
+    if not BlockAdapterRegister.is_supported(transformer):
+        custom_adapter = _build_custom_block_adapter(
+            transformer, has_separate_cfg=has_separate_cfg
+        )
+        if custom_adapter is None:
+            transformer_cls_name = transformer.__class__.__name__
+            raise ValueError(
+                f"{transformer_cls_name} is not officially supported by cache-dit. "
+                "Supported cache-dit DiT families include Flux, QwenImage, HunyuanDiT, "
+                "HunyuanVideo, Wan, CogVideoX, Mochi, and others. "
+                "Please ensure your transformer belongs to one of these families or "
+                "define a custom BlockAdapter."
+            )
+
+    # Build cache config (including SCM fields if provided)
+    cache_config = DBCacheConfig(
+        num_inference_steps=config.num_inference_steps,
+        Fn_compute_blocks=config.Fn_compute_blocks,
+        Bn_compute_blocks=config.Bn_compute_blocks,
+        max_warmup_steps=config.max_warmup_steps,
+        residual_diff_threshold=config.residual_diff_threshold,
+        max_continuous_cached_steps=config.max_continuous_cached_steps,
+        # SCM fields
+        steps_computation_mask=config.steps_computation_mask,
+        steps_computation_policy=config.steps_computation_policy,
+    )
+
+    # Build calibrator config if TaylorSeer is enabled
+    calibrator_config = None
+    if config.enable_taylorseer:
+        calibrator_config = TaylorSeerCalibratorConfig(
+            taylorseer_order=config.taylorseer_order,
+        )
+
+    # Enable cache-dit on the transformer
+    logger.info(
+        "Enabling cache-dit on %s with config: Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, "
+        "TaylorSeer=%s (order=%d), steps=%d",
+        model_name,
+        config.Fn_compute_blocks,
+        config.Bn_compute_blocks,
+        config.max_warmup_steps,
+        config.residual_diff_threshold,
+        config.max_continuous_cached_steps,
+        config.enable_taylorseer,
+        config.taylorseer_order,
+        config.num_inference_steps,
+    )
+
+    # Log SCM configuration if enabled
+    if config.steps_computation_mask:
+        compute_steps = sum(config.steps_computation_mask)
+        cache_steps = len(config.steps_computation_mask) - compute_steps
+        logger.info(
+            "SCM enabled: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            config.steps_computation_policy,
+        )
+
+    parallelism_config = _build_parallelism_config(sp_group, tp_group)
+    if parallelism_config is not None:
+        _patch_cache_dit_similarity()
+
+    _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group)
+
+    # Custom path: pass a pre-built BlockAdapter, bypassing the registry.
+    # Standard path: let enable_cache discover the registered adapter.
+    target = transformer
+    if custom_adapter is not None:
+        target = custom_adapter
+        logger.info(
+            "Enabling cache-dit on %s via custom BlockAdapter (%s).",
+            model_name,
+            custom_adapter.forward_pattern,
+        )
+    cache_dit.enable_cache(
+        target,
+        cache_config=cache_config,
+        calibrator_config=calibrator_config,
+        parallelism_config=None,
+    )
+
+    if parallelism_config is not None:
+        context_manager = getattr(transformer, "_context_manager", None)
+        if context_manager is not None:
+            context_manager._sglang_sp_group = sp_group
+            context_manager._sglang_tp_group = tp_group
+            # In mixed TP + SP (Ulysses/Ring) mode, cache-dit decisions must be consistent
+            # across the full TP×SP model-parallel slice. Prefer using SGLang's DIT group
+            # as a conservative superset group; fallback to None.
+            tp_sp_group = None
+            if sp_group is not None and tp_group is not None:
+                tp_sp_group = get_dit_group()
+
+            context_manager._sglang_tp_sp_group = tp_sp_group
+
+    return transformer
+
+
+def enable_cache_on_dual_transformer(
+    transformer: torch.nn.Module,
+    transformer_2: torch.nn.Module,
+    primary_config: CacheDitConfig,
+    secondary_config: CacheDitConfig,
+    model_name: str = "wan2.2",
+    sp_group: Optional[torch.distributed.ProcessGroup] = None,
+    tp_group: Optional[torch.distributed.ProcessGroup] = None,
+) -> tuple[torch.nn.Module, torch.nn.Module]:
+    """Enable cache-dit on dual transformers using BlockAdapter.
+
+    For models with two transformers, cache-dit requires enabling cache on both
+    simultaneously via BlockAdapter. The two transformers may be split by denoising
+    range, or run as paired conditional/unconditional branches. This cannot be done
+    by calling enable_cache separately on each transformer.
+
+    Args:
+        primary_config: CacheDitConfig for primary transformer.
+        secondary_config: CacheDitConfig for secondary transformer.
+        sp_group: Sequence parallel process group (for Ulysses/Ring).
+        tp_group: Tensor parallel process group.
+    """
+    adapter_spec = DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS.get(model_name)
+    if adapter_spec is None:
+        raise ValueError(
+            f"Dual-transformer cache-dit is only supported for "
+            f"{sorted(DUAL_TRANSFORMER_BLOCK_ADAPTER_SPECS)}, got {model_name}."
+        )
+
+    if not primary_config.enabled:
+        return transformer, transformer_2
+
+    if primary_config.num_inference_steps is None:
+        raise ValueError(
+            "num_inference_steps is required for dual-transformer mode. "
+            "Please provide it in CacheDitConfig."
+        )
+
+    # Build DBCacheConfig for primary transformer
+    primary_cache_config = DBCacheConfig(
+        num_inference_steps=primary_config.num_inference_steps,
+        Fn_compute_blocks=primary_config.Fn_compute_blocks,
+        Bn_compute_blocks=primary_config.Bn_compute_blocks,
+        max_warmup_steps=primary_config.max_warmup_steps,
+        residual_diff_threshold=primary_config.residual_diff_threshold,
+        max_continuous_cached_steps=primary_config.max_continuous_cached_steps,
+        steps_computation_mask=primary_config.steps_computation_mask,
+        steps_computation_policy=primary_config.steps_computation_policy,
+    )
+
+    # Build DBCacheConfig for secondary transformer
+    secondary_cache_config = DBCacheConfig(
+        num_inference_steps=secondary_config.num_inference_steps,
+        Fn_compute_blocks=secondary_config.Fn_compute_blocks,
+        Bn_compute_blocks=secondary_config.Bn_compute_blocks,
+        max_warmup_steps=secondary_config.max_warmup_steps,
+        residual_diff_threshold=secondary_config.residual_diff_threshold,
+        max_continuous_cached_steps=secondary_config.max_continuous_cached_steps,
+        steps_computation_mask=secondary_config.steps_computation_mask,
+        steps_computation_policy=secondary_config.steps_computation_policy,
+    )
+
+    # Build calibrator configs if TaylorSeer is enabled
+    primary_calibrator = None
+    if primary_config.enable_taylorseer:
+        primary_calibrator = TaylorSeerCalibratorConfig(
+            taylorseer_order=primary_config.taylorseer_order,
+        )
+
+    secondary_calibrator = None
+    if secondary_config.enable_taylorseer:
+        secondary_calibrator = TaylorSeerCalibratorConfig(
+            taylorseer_order=secondary_config.taylorseer_order,
+        )
+
+    # Build ParamsModifier for each transformer
+    primary_modifier = ParamsModifier(
+        cache_config=primary_cache_config,
+        calibrator_config=primary_calibrator,
+    )
+    secondary_modifier = ParamsModifier(
+        cache_config=secondary_cache_config,
+        calibrator_config=secondary_calibrator,
+    )
+
+    # Log configuration
+    logger.info(
+        "Enabling cache-dit on %s dual transformers with BlockAdapter",
+        model_name,
+    )
+    logger.info(
+        "  Primary (transformer): Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, TaylorSeer=%s",
+        primary_config.Fn_compute_blocks,
+        primary_config.Bn_compute_blocks,
+        primary_config.max_warmup_steps,
+        primary_config.residual_diff_threshold,
+        primary_config.max_continuous_cached_steps,
+        primary_config.enable_taylorseer,
+    )
+    logger.info(
+        "  Secondary transformer: Fn=%d, Bn=%d, W=%d, R=%.2f, MC=%d, TaylorSeer=%s",
+        secondary_config.Fn_compute_blocks,
+        secondary_config.Bn_compute_blocks,
+        secondary_config.max_warmup_steps,
+        secondary_config.residual_diff_threshold,
+        secondary_config.max_continuous_cached_steps,
+        secondary_config.enable_taylorseer,
+    )
+
+    # Log SCM configuration if enabled
+    if primary_config.steps_computation_mask:
+        compute_steps = sum(primary_config.steps_computation_mask)
+        cache_steps = len(primary_config.steps_computation_mask) - compute_steps
+        logger.info(
+            "  SCM enabled for primary transformer: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            primary_config.steps_computation_policy,
+        )
+    if secondary_config.steps_computation_mask:
+        compute_steps = sum(secondary_config.steps_computation_mask)
+        cache_steps = len(secondary_config.steps_computation_mask) - compute_steps
+        logger.info(
+            "  SCM enabled for secondary transformer: %d compute steps, %d cache steps, policy=%s",
+            compute_steps,
+            cache_steps,
+            secondary_config.steps_computation_policy,
+        )
+
+    parallelism_config = _build_parallelism_config(sp_group, tp_group)
+    if parallelism_config is not None:
+        _patch_cache_dit_similarity()
+
+    _mark_transformer_parallelized(transformer, parallelism_config, sp_group, tp_group)
+    _mark_transformer_parallelized(
+        transformer_2, parallelism_config, sp_group, tp_group
+    )
+
+    transformer_blocks_attr, transformer_2_blocks_attr = adapter_spec.blocks_attr
+    transformer_blocks = getattr(transformer, transformer_blocks_attr, None)
+    transformer_2_blocks = getattr(transformer_2, transformer_2_blocks_attr, None)
+    if transformer_blocks is None or transformer_2_blocks is None:
+        raise ValueError(
+            f"Dual transformers for {model_name} must expose cache-dit block "
+            f"attributes {adapter_spec.blocks_attr}. "
+            f"transformer has {transformer_blocks_attr}: "
+            f"{transformer_blocks is not None}, secondary transformer has "
+            f"{transformer_2_blocks_attr}: {transformer_2_blocks is not None}"
+        )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=[transformer, transformer_2],
+            blocks=[transformer_blocks, transformer_2_blocks],
+            blocks_name=adapter_spec.blocks_name,
+            forward_pattern=adapter_spec.forward_pattern,
+            params_modifiers=[primary_modifier, secondary_modifier],
+            check_forward_pattern=adapter_spec.check_forward_pattern,
+            check_num_outputs=adapter_spec.check_num_outputs,
+            has_separate_cfg=adapter_spec.has_separate_cfg,
+        ),
+        parallelism_config=None,
+    )
+
+    if parallelism_config is not None:
+        for t in [transformer, transformer_2]:
+            context_manager = getattr(t, "_context_manager", None)
+            if context_manager is not None:
+                context_manager._sglang_sp_group = sp_group
+                context_manager._sglang_tp_group = tp_group
+                tp_sp_group = None
+                if sp_group is not None and tp_group is not None:
+                    try:
+                        tp_sp_group = get_dit_group()
+                    except Exception:
+                        tp_sp_group = None
+                context_manager._sglang_tp_sp_group = tp_sp_group
+
+    return transformer, transformer_2
+
+
+def refresh_context_on_transformer(
+    transformer: torch.nn.Module,
+    num_inference_steps: int,
+    scm_preset: str | None = None,
+    verbose: bool = False,
+) -> None:
+    """Refresh cache-dit context for transformer."""
+    steps_computation_mask = None
+    if scm_preset is not None:
+        steps_computation_mask = cache_dit.steps_mask(
+            mask_policy=scm_preset, total_steps=num_inference_steps
+        )
+    cache_dit.refresh_context(
+        transformer,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_inference_steps,
+            steps_computation_mask=steps_computation_mask,
+            steps_computation_policy=scm_preset,
+        ),
+        verbose=verbose,
+    )
+    logger.debug(f"cache-dit refreshed on transformer (steps={num_inference_steps})")
+
+
+def refresh_context_on_dual_transformer(
+    transformer: torch.nn.Module,
+    transformer_2: torch.nn.Module,
+    num_high_noise_steps: int,
+    num_low_noise_steps: int,
+    scm_preset: str | None = None,
+    verbose: bool = False,
+    steps_computation_mask: Optional[List[int]] = None,
+    steps_computation_mask_2: Optional[List[int]] = None,
+    steps_computation_policy: str | None = None,
+) -> None:
+    """Refresh cache-dit context for dual transformers."""
+    high_noise_steps_computation_mask = steps_computation_mask
+    low_noise_steps_computation_mask = steps_computation_mask_2
+    if high_noise_steps_computation_mask is None and scm_preset is not None:
+        high_noise_steps_computation_mask = cache_dit.steps_mask(
+            mask_policy=scm_preset, total_steps=num_high_noise_steps
+        )
+    if low_noise_steps_computation_mask is None and scm_preset is not None:
+        low_noise_steps_computation_mask = cache_dit.steps_mask(
+            mask_policy=scm_preset, total_steps=num_low_noise_steps
+        )
+    policy = (
+        steps_computation_policy if steps_computation_policy is not None else scm_preset
+    )
+    cache_dit.refresh_context(
+        transformer,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_high_noise_steps,
+            steps_computation_mask=high_noise_steps_computation_mask,
+            steps_computation_policy=policy,
+        ),
+        verbose=verbose,
+    )
+    cache_dit.refresh_context(
+        transformer_2,
+        cache_config=DBCacheConfig().reset(
+            num_inference_steps=num_low_noise_steps,
+            steps_computation_mask=low_noise_steps_computation_mask,
+            steps_computation_policy=policy,
+        ),
+        verbose=verbose,
+    )
+    logger.debug(
+        f"cache-dit refreshed on dual transformers (steps={num_high_noise_steps}, {num_low_noise_steps})"
+    )

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

@@ -0,0 +1,316 @@
+# 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
+
+from sglang.multimodal_gen.configs.models import DiTConfig
+
+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"}
+    config: DiTConfig
+
+    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.get_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")