From b105dad5dac56d149febea9a23d57b07f56b666a Mon Sep 17 00:00:00 2001
From: Mick <mickjagger19@icloud.com>
Date: Sun, 25 Jan 2026 22:43:22 +0800
Subject: [PATCH] [diffusion] refactor: remove useless lazy-import cache-dit
 codes (#17659)

---
 .../multimodal_gen/runtime/cache/__init__.py  | 41 ++++---------------
 .../multimodal_gen/runtime/cache/teacache.py  |  3 ++
 2 files changed, 10 insertions(+), 34 deletions(-)

diff --git a/python/sglang/multimodal_gen/runtime/cache/__init__.py b/python/sglang/multimodal_gen/runtime/cache/__init__.py
index 895f0c594..f01f41cc0 100644
--- a/python/sglang/multimodal_gen/runtime/cache/__init__.py
+++ b/python/sglang/multimodal_gen/runtime/cache/__init__.py
@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 """
-Cache acceleration module for SGLang multimodal generation.
+Cache acceleration module for SGLang-diffusion
 
 This module provides various caching strategies to accelerate
 diffusion transformer (DiT) inference:
@@ -8,26 +8,13 @@ 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.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__ = [
@@ -40,17 +27,3 @@ __all__ = [
     "enable_cache_on_dual_transformer",
     "get_scm_mask",
 ]
-
-
-# Lazy imports for cache-dit integration (requires cache-dit package)
-def __getattr__(name):
-    if name in (
-        "CacheDitConfig",
-        "enable_cache_on_transformer",
-        "enable_cache_on_dual_transformer",
-        "get_scm_mask",
-    ):
-        from sglang.multimodal_gen.runtime.cache import cache_dit_integration
-
-        return getattr(cache_dit_integration, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
diff --git a/python/sglang/multimodal_gen/runtime/cache/teacache.py b/python/sglang/multimodal_gen/runtime/cache/teacache.py
index 40f9123ef..5cdafd08b 100644
--- a/python/sglang/multimodal_gen/runtime/cache/teacache.py
+++ b/python/sglang/multimodal_gen/runtime/cache/teacache.py
@@ -23,6 +23,8 @@ 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
 
@@ -125,6 +127,7 @@ class TeaCacheMixin:
     # 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__."""