[Feature] teacache integration (#179)

Signed-off-by: Prajwal A <prajwalanagani@gmail.com>
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Signed-off-by: Jiangyun Zhu <riverclouds.zhu@qq.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Samit <285365963@qq.com>

.gitignore

@@ -179,6 +179,7 @@ checkpoints/
 
 # Cache directories
 cache/
+!vllm_omni/diffusion/cache/
 .cache/
 diffusion_cache/
 kv_cache/

docs/.nav.yml

@@ -20,6 +20,8 @@ nav:
   - Configuration:
     - configuration/README.md
     - configuration/*
+  - Diffusion Acceleration:
+     - TeaCache: user_guide/teacache.md
   - Models:
     - models/supported_models.md
 - Developer Guide:

docs/api/README.md

@@ -96,6 +96,7 @@ Model execution components.
 Configuration classes.
 
 - [vllm_omni.config.model.OmniModelConfig][]
+- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 
 ## Workers
 

docs/configuration/README.md

@@ -9,3 +9,7 @@ Currently, the main options are maintained by stage configs for each model.
 For specific example, please refer to [Qwen2.5-omni stage config](stage_configs/qwen2_5_omni.yaml)
 
 For introduction, please check [Introduction for stage config](./stage_configs.md)
+
+## Optimization Features
+
+- **[TeaCache Configuration](../user_guide/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/user_guide/teacache.md

@@ -0,0 +1,99 @@
+# TeaCache Configuration Guide
+
+TeaCache speeds up diffusion model inference by caching transformer computations when consecutive timesteps are similar. This typically provides **1.5x-2.0x speedup** with minimal quality loss.
+
+## Quick Start
+
+Enable TeaCache by setting `cache_adapter` to `"tea_cache"`:
+
+```python
+from vllm_omni import Omni
+
+# Simple configuration - model_class_name and model_type are auto-detected
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_adapter="tea_cache",
+    cache_config={
+        "rel_l1_thresh": 0.2  # Optional, defaults to 0.2
+    }
+)
+outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+```
+
+### Using Environment Variable
+
+You can also enable TeaCache via environment variable:
+
+```bash
+export DIFFUSION_CACHE_ADAPTER=tea_cache
+```
+
+Then initialize without explicitly setting `cache_adapter`:
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_config={"rel_l1_thresh": 0.2}  # Optional
+)
+```
+
+## Configuration Parameters
+
+### `rel_l1_thresh` (float, default: `0.2`)
+
+Controls the balance between speed and quality. Lower values prioritize quality, higher values prioritize speed.
+
+**Recommended values:**
+
+- `0.2` - **~1.5x speedup** with minimal quality loss (recommended)
+- `0.4` - **~1.8x speedup** with slight quality loss
+- `0.6` - **~2.0x speedup** with noticeable quality loss
+- `0.8` - **~2.25x speedup** with significant quality loss
+
+## Examples
+
+### Python API
+
+```python
+from vllm_omni import Omni
+
+omni = Omni(
+    model="Qwen/Qwen-Image",
+    cache_adapter="tea_cache",
+    cache_config={"rel_l1_thresh": 0.2}
+)
+outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+```
+
+## Performance Tuning
+
+Start with the default `rel_l1_thresh=0.2` and adjust based on your needs:
+
+- **Maximum quality**: Use `0.1-0.2`
+- **Balanced**: Use `0.2-0.4` (recommended)
+- **Maximum speed**: Use `0.6-0.8` (may reduce quality)
+
+## Troubleshooting
+
+### Quality Degradation
+
+If you notice quality issues, lower the threshold:
+
+```python
+cache_config={"rel_l1_thresh": 0.1}  # More conservative caching
+```
+
+## Supported Models
+
+Currently supported models:
+
+- **QwenImagePipeline**
+  - Model identifier: `"Qwen/Qwen-Image"`
+  - Example: `model="Qwen/Qwen-Image"`
+
+Future support (extractors in development):
+
+- **FluxPipeline** - Coming soon
+- **CogVideoXPipeline** - Coming soon

vllm_omni/diffusion/cache/__init__.py

@@ -0,0 +1,41 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Cache module for diffusion model inference acceleration.
+
+This module provides a unified cache adapter system for different caching strategies:
+- TeaCache: Timestep Embedding Aware Cache for adaptive transformer caching
+- (Future: DeepCache, DistriFusion, etc.)
+
+The cache adapter system uses a registry pattern where cache types can be
+dynamically registered and instantiated. Cache is configured via OmniDiffusionConfig
+or DIFFUSION_CACHE_ADAPTER environment variable.
+"""
+
+from vllm_omni.diffusion.cache.apply import setup_cache
+from vllm_omni.diffusion.cache.base import CacheAdapter
+from vllm_omni.diffusion.cache.registry import (
+    CacheType,
+    get_cache_adapter,
+    register_cache_adapter,
+)
+from vllm_omni.diffusion.cache.teacache import (
+    CacheContext,
+    TeaCacheConfig,
+    apply_teacache_hook,
+)
+
+# Import teacache adapter to trigger registration
+from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter  # noqa: F401
+
+__all__ = [
+    "CacheAdapter",
+    "CacheType",
+    "TeaCacheConfig",
+    "CacheContext",
+    "setup_cache",
+    "get_cache_adapter",
+    "register_cache_adapter",
+    "apply_teacache_hook",
+]

vllm_omni/diffusion/cache/apply.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unified cache setup function for diffusion models.
+
+This module provides the main entry point for setting up cache adapters
+on transformer models. It handles cache type selection, adapter instantiation,
+and application to the model.
+"""
+
+from typing import Any, Optional
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheAdapter
+from vllm_omni.diffusion.cache.registry import get_cache_adapter
+
+logger = init_logger(__name__)
+
+
+def setup_cache(
+    transformer: torch.nn.Module,
+    cache_type: str = "none",
+    cache_config: Optional[dict[str, Any]] = None,
+) -> Optional[CacheAdapter]:
+    """
+    Setup cache adapter for transformer (one-time initialization).
+
+    This is the main entry point for applying caching to diffusion transformers.
+    It's called once during pipeline initialization and applies the appropriate
+    cache adapter based on configuration.
+
+    Args:
+        transformer: Transformer module to apply cache to
+        cache_type: Type of cache to apply. Options:
+            - "none": No caching (default)
+            - "tea_cache": TeaCache adaptive caching
+            - Additional types can be registered via register_cache_adapter()
+        cache_config: Cache-specific configuration dictionary. Contents depend
+            on cache type:
+            - For "tea_cache": {"rel_l1_thresh": 0.2, "model_type": "QwenImagePipeline"}
+            - For other types: see their respective adapter documentation
+
+    Returns:
+        CacheAdapter instance for state management, or None if cache_type="none"
+
+    Example:
+        >>> # Setup TeaCache during pipeline initialization
+        >>> adapter = setup_cache(
+        ...     transformer,
+        ...     cache_type="tea_cache",
+        ...     cache_config={"rel_l1_thresh": 0.2}
+        ... )
+        >>>
+        >>> # Reset cache state before each generation
+        >>> if adapter is not None:
+        ...     adapter.reset(transformer)
+    """
+    # No caching requested
+    if cache_type == "none":
+        logger.debug("No cache adapter requested (cache_type='none')")
+        return None
+
+    # Get adapter instance from registry
+    try:
+        adapter = get_cache_adapter(cache_type, cache_config or {})
+    except ValueError as e:
+        logger.error(f"Failed to get cache adapter: {e}")
+        raise
+
+    # Apply cache to transformer
+    logger.info(f"Applying {cache_type} cache to transformer")
+    try:
+        adapter.apply(transformer)
+    except Exception as e:
+        logger.error(f"Failed to apply cache adapter: {e}")
+        raise
+
+    logger.info(f"Cache adapter '{cache_type}' successfully applied")
+    return adapter

vllm_omni/diffusion/cache/base.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Base cache adapter interface for diffusion models.
+
+This module defines the abstract base class that all cache adapters must implement.
+Cache adapters provide a unified interface for applying different caching strategies
+(TeaCache, DeepCache, etc.) to transformer models using hooks.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+import torch
+
+
+class CacheAdapter(ABC):
+    """
+    Abstract base class for cache adapters.
+
+    Cache adapters apply caching strategies to transformer models to accelerate
+    inference. Each adapter type (TeaCache, DeepCache, etc.) implements the
+    apply() and reset() methods to manage cache lifecycle.
+
+    Attributes:
+        config: Dictionary containing cache-specific configuration parameters
+    """
+
+    def __init__(self, config: dict[str, Any]):
+        """
+        Initialize cache adapter with configuration.
+
+        Args:
+            config: Cache-specific configuration dictionary
+        """
+        self.config = config
+
+    @abstractmethod
+    def apply(self, transformer: torch.nn.Module) -> None:
+        """
+        Apply cache to transformer using hooks.
+
+        This method should register the appropriate hooks on the transformer
+        to enable caching during inference. Called once during pipeline
+        initialization.
+
+        Args:
+            transformer: Transformer module to apply cache to
+        """
+        raise NotImplementedError("Subclasses must implement apply()")
+
+    @abstractmethod
+    def reset(self, transformer: torch.nn.Module) -> None:
+        """
+        Reset cache state for new generation.
+
+        This method should clear any cached values and reset counters/accumulators.
+        Called at the start of each generation to ensure clean state.
+
+        Args:
+            transformer: Transformer module to reset cache on
+        """
+        raise NotImplementedError("Subclasses must implement reset()")
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(config={self.config})"

vllm_omni/diffusion/cache/registry.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Cache adapter registry for diffusion models.
+
+This module provides a registry pattern for cache adapters, allowing dynamic
+registration and instantiation of different cache types (TeaCache, DeepCache, etc.).
+"""
+
+from enum import Enum
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheAdapter
+
+logger = init_logger(__name__)
+
+
+class CacheType(Enum):
+    """Supported cache adapter types."""
+
+    NONE = "none"
+    TEA_CACHE = "tea_cache"
+    # Future cache types can be added here:
+    # DEEP_CACHE = "deep_cache"
+    # DISTRI_FUSION = "distri_fusion"
+
+
+# Global registry mapping cache types to adapter classes
+CACHE_ADAPTER_REGISTRY: dict[CacheType, type[CacheAdapter]] = {}
+
+
+def register_cache_adapter(cache_type: CacheType, adapter_class: type[CacheAdapter]) -> None:
+    """
+    Register a cache adapter class for a given cache type.
+
+    This allows extending the system with new cache types without modifying
+    the core cache infrastructure.
+
+    Args:
+        cache_type: CacheType enum value
+        adapter_class: CacheAdapter subclass to register
+
+    Example:
+        >>> register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter)
+    """
+    if not issubclass(adapter_class, CacheAdapter):
+        raise TypeError(f"{adapter_class} must be a subclass of CacheAdapter")
+
+    CACHE_ADAPTER_REGISTRY[cache_type] = adapter_class
+    logger.debug(f"Registered cache adapter: {cache_type.value} -> {adapter_class.__name__}")
+
+
+def get_cache_adapter(cache_type: str, config: dict[str, Any]) -> CacheAdapter:
+    """
+    Factory function to get cache adapter instance.
+
+    Converts string cache type to enum, looks up in registry, and instantiates
+    the appropriate adapter class with the provided configuration.
+
+    Args:
+        cache_type: String name of cache type ("tea_cache", "deep_cache", etc.)
+        config: Configuration dictionary to pass to adapter constructor
+
+    Returns:
+        Instantiated CacheAdapter subclass
+
+    Raises:
+        ValueError: If cache_type is unknown or not registered
+
+    Example:
+        >>> adapter = get_cache_adapter("tea_cache", {"rel_l1_thresh": 0.2})
+        >>> adapter.apply(transformer)
+    """
+    # Normalize cache type string
+    cache_type_str = cache_type.lower().strip()
+
+    # Convert string to enum
+    try:
+        cache_enum = CacheType(cache_type_str)
+    except ValueError:
+        available = [ct.value for ct in CacheType if ct != CacheType.NONE]
+        raise ValueError(f"Unknown cache type: '{cache_type}'. Available types: {available}")
+
+    # Check if it's the special "none" case
+    if cache_enum == CacheType.NONE:
+        raise ValueError("Cannot instantiate adapter for cache_type='none'. Use setup_cache() which handles this case.")
+
+    # Lookup in registry
+    if cache_enum not in CACHE_ADAPTER_REGISTRY:
+        raise ValueError(
+            f"Cache type '{cache_type}' is not registered. Registered types: {list(CACHE_ADAPTER_REGISTRY.keys())}"
+        )
+
+    adapter_class = CACHE_ADAPTER_REGISTRY[cache_enum]
+
+    # Instantiate and return
+    return adapter_class(config)

vllm_omni/diffusion/cache/teacache/__init__.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+TeaCache: Timestep Embedding Aware Cache for diffusion model acceleration.
+
+TeaCache speeds up diffusion inference by reusing transformer block computations
+when consecutive timestep embeddings are similar.
+
+This implementation uses a hooks-based approach that requires zero changes to
+model code. Model developers only need to add an extractor function to support
+new models.
+
+Usage:
+    from vllm_omni import Omni
+
+    omni = Omni(
+        model="Qwen/Qwen-Image",
+        cache_adapter="tea_cache",
+        cache_config={"rel_l1_thresh": 0.2}
+    )
+    images = omni.generate("a cat")
+
+    # Alternative: Using environment variable
+    # export DIFFUSION_CACHE_ADAPTER=tea_cache
+"""
+
+from vllm_omni.diffusion.cache.teacache.adapter import TeaCacheAdapter
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import (
+    CacheContext,
+    register_extractor,
+)
+from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook
+from vllm_omni.diffusion.cache.teacache.state import TeaCacheState
+
+__all__ = [
+    "TeaCacheAdapter",
+    "TeaCacheConfig",
+    "TeaCacheState",
+    "TeaCacheHook",
+    "apply_teacache_hook",
+    "register_extractor",
+    "CacheContext",
+]

vllm_omni/diffusion/cache/teacache/adapter.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+TeaCache adapter implementation.
+
+This module provides the TeaCache adapter that implements the CacheAdapter
+interface using the hooks-based TeaCache system.
+"""
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.cache.base import CacheAdapter
+from vllm_omni.diffusion.cache.registry import CacheType, register_cache_adapter
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook
+
+logger = init_logger(__name__)
+
+
+class TeaCacheAdapter(CacheAdapter):
+    """
+    TeaCache implementation using hooks.
+
+    TeaCache (Timestep Embedding Aware Cache) is an adaptive caching technique
+    that speeds up diffusion inference by reusing transformer block computations
+    when consecutive timestep embeddings are similar.
+
+    The adapter applies TeaCache hooks to the transformer which intercept the
+    forward pass and implement the caching logic transparently.
+
+    Example:
+        >>> adapter = TeaCacheAdapter({"rel_l1_thresh": 0.2})
+        >>> adapter.apply(transformer)
+        >>> # Generate with cache enabled
+        >>> adapter.reset(transformer)  # Reset before each generation
+    """
+
+    def apply(self, transformer: torch.nn.Module) -> None:
+        """
+        Apply TeaCache to transformer using hooks.
+
+        This creates a TeaCacheConfig from the adapter's config dict and
+        applies the TeaCache hook to the transformer.
+
+        Args:
+            transformer: Transformer module to apply TeaCache to
+        """
+        # Validate that model_type is provided
+        if "model_type" not in self.config:
+            raise ValueError(
+                "model_type must be provided in cache_config. "
+                "Should be set to OmniDiffusionConfig.model_class_name (e.g., 'QwenImagePipeline')"
+            )
+
+        # Create TeaCacheConfig from dict
+        try:
+            teacache_config = TeaCacheConfig(**self.config)
+        except Exception as e:
+            logger.error(f"Failed to create TeaCacheConfig: {e}")
+            raise ValueError(
+                f"Invalid TeaCache configuration: {e}. "
+                f"Expected keys: rel_l1_thresh, model_type (pipeline class name), coefficients (optional)"
+            )
+
+        # Apply hook to transformer
+        apply_teacache_hook(transformer, teacache_config)
+        logger.info(
+            f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
+            f"model_type={teacache_config.model_type}"
+        )
+
+    def reset(self, transformer: torch.nn.Module) -> None:
+        """
+        Reset TeaCache state for new generation.
+
+        Clears all cached residuals and resets counters/accumulators.
+        Should be called before each generation to ensure clean state.
+
+        Args:
+            transformer: Transformer module to reset cache on
+        """
+        if hasattr(transformer, "_hook_registry"):
+            hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME)
+            if hook is not None:
+                transformer._hook_registry.reset_hook(TeaCacheHook._HOOK_NAME)
+                logger.debug("TeaCache state reset")
+            else:
+                logger.warning("TeaCache hook not found, nothing to reset")
+        else:
+            logger.warning("Transformer has no hook registry, TeaCache may not be applied")
+
+
+# Register TeaCache adapter in the global registry
+register_cache_adapter(CacheType.TEA_CACHE, TeaCacheAdapter)

vllm_omni/diffusion/cache/teacache/config.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from dataclasses import dataclass
+from typing import Optional
+
+# Model-specific polynomial coefficients for rescaling L1 distances
+# These coefficients account for model-specific characteristics in how embeddings change
+# Source: TeaCache paper and ComfyUI-TeaCache empirical tuning
+_MODEL_COEFFICIENTS = {
+    # FLUX model coefficients from TeaCache paper
+    "FluxPipeline": [
+        4.98651651e02,
+        -2.83781631e02,
+        5.58554382e01,
+        -3.82021401e00,
+        2.64230861e-01,
+    ],
+    # Qwen-Image model coefficients from ComfyUI-TeaCache
+    # Tuned specifically for Qwen's dual-stream transformer architecture
+    "QwenImagePipeline": [
+        -4.50000000e02,
+        2.80000000e02,
+        -4.50000000e01,
+        3.20000000e00,
+        -2.00000000e-02,
+    ],
+}
+
+
+@dataclass
+class TeaCacheConfig:
+    """
+    Configuration for TeaCache applied to transformer models.
+
+    TeaCache (Timestep Embedding Aware Cache) is an adaptive caching technique that speeds up
+    diffusion model inference by reusing transformer block computations when consecutive
+    timestep embeddings are similar.
+
+    Args:
+        rel_l1_thresh: Threshold for accumulated relative L1 distance. When below threshold,
+            cached residual is reused. Values in [0.1, 0.3] work best:
+            - 0.2: ~1.5x speedup with minimal quality loss
+            - 0.4: ~1.8x speedup with slight quality loss
+            - 0.6: ~2.0x speedup with noticeable quality loss
+        coefficients: Polynomial coefficients for rescaling L1 distance. If None, uses
+            model-specific defaults based on model_type.
+        model_type: Pipeline class name (e.g., "QwenImagePipeline", "FluxPipeline").
+            Must match OmniDiffusionConfig.model_class_name for proper extractor lookup.
+            Defaults to "QwenImagePipeline".
+    """
+
+    rel_l1_thresh: float = 0.2
+    coefficients: Optional[list[float]] = None
+    model_type: str = "QwenImagePipeline"
+
+    def __post_init__(self) -> None:
+        """Validate and set default coefficients."""
+        if self.rel_l1_thresh <= 0:
+            raise ValueError(f"rel_l1_thresh must be positive, got {self.rel_l1_thresh}")
+
+        if self.coefficients is None:
+            # Use model-specific coefficients, fallback to FluxPipeline if not found
+            self.coefficients = _MODEL_COEFFICIENTS.get(self.model_type, _MODEL_COEFFICIENTS["FluxPipeline"])
+
+        if len(self.coefficients) != 5:
+            raise ValueError(f"coefficients must contain exactly 5 elements, got {len(self.coefficients)}")

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -0,0 +1,334 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Model-specific extractors for TeaCache.
+
+This module provides a registry of extractor functions that know how to extract
+modulated inputs from different transformer architectures. Adding support for
+a new model requires only adding a new extractor function to the registry.
+
+With Option B enhancement, extractors now return a CacheContext object containing
+all model-specific information needed for generic caching, including preprocessing,
+transformer execution, and postprocessing logic.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Callable, Union
+
+import torch
+import torch.nn as nn
+
+
+@dataclass
+class CacheContext:
+    """
+    Context object containing all model-specific information for caching.
+
+    This allows the TeaCacheHook to remain completely generic - all model-specific
+    logic is encapsulated in the extractor that returns this context.
+
+    Attributes:
+        modulated_input: Tensor used for cache decision (similarity comparison).
+            Must be a torch.Tensor extracted from the first transformer block,
+            typically after applying normalization and modulation.
+
+        hidden_states: Current hidden states (will be modified by caching).
+            Must be a torch.Tensor representing the main image/latent states
+            after preprocessing but before transformer blocks.
+
+        encoder_hidden_states: Optional encoder states (for dual-stream models).
+            Set to None for single-stream models (e.g., Flux).
+            For dual-stream models (e.g., Qwen), contains text encoder outputs.
+
+        temb: Timestep embedding tensor.
+            Must be a torch.Tensor containing the timestep conditioning.
+
+        run_transformer_blocks: Callable that executes model-specific transformer blocks.
+            Signature: () -> tuple[torch.Tensor, ...]
+
+            Returns:
+                tuple containing:
+                - [0]: processed hidden_states (required)
+                - [1]: processed encoder_hidden_states (optional, only for dual-stream)
+
+            Example for single-stream:
+                def run_blocks():
+                    h = hidden_states
+                    for block in module.transformer_blocks:
+                        h = block(h, temb=temb)
+                    return (h,)
+
+            Example for dual-stream:
+                def run_blocks():
+                    h, e = hidden_states, encoder_hidden_states
+                    for block in module.transformer_blocks:
+                        e, h = block(h, e, temb=temb)
+                    return (h, e)
+
+        postprocess: Callable that does model-specific output postprocessing.
+            Signature: (torch.Tensor) -> Union[torch.Tensor, Transformer2DModelOutput, tuple]
+
+            Takes the processed hidden_states and applies final transformations
+            (normalization, projection) to produce the model output.
+
+            Example:
+                def postprocess(h):
+                    h = module.norm_out(h, temb)
+                    output = module.proj_out(h)
+                    return Transformer2DModelOutput(sample=output)
+
+        extra_states: Optional dict for additional model-specific state.
+            Use this for models that need to pass additional context beyond
+            the standard fields.
+    """
+
+    modulated_input: torch.Tensor
+    hidden_states: torch.Tensor
+    encoder_hidden_states: torch.Tensor | None
+    temb: torch.Tensor
+    run_transformer_blocks: Callable[[], tuple[torch.Tensor, ...]]
+    postprocess: Callable[[torch.Tensor], Any]
+    extra_states: dict[str, Any] | None = None
+
+    def validate(self) -> None:
+        """
+        Validate that the CacheContext contains valid data.
+
+        Raises:
+            TypeError: If fields have wrong types
+            ValueError: If tensors have invalid properties
+            RuntimeError: If callables fail basic invocation tests
+
+        This method should be called after creating a CacheContext to catch
+        common developer errors early with clear error messages.
+        """
+        # Validate tensor fields
+        if not isinstance(self.modulated_input, torch.Tensor):
+            raise TypeError(f"modulated_input must be torch.Tensor, got {type(self.modulated_input)}")
+
+        if not isinstance(self.hidden_states, torch.Tensor):
+            raise TypeError(f"hidden_states must be torch.Tensor, got {type(self.hidden_states)}")
+
+        if self.encoder_hidden_states is not None and not isinstance(self.encoder_hidden_states, torch.Tensor):
+            raise TypeError(
+                f"encoder_hidden_states must be torch.Tensor or None, got {type(self.encoder_hidden_states)}"
+            )
+
+        if not isinstance(self.temb, torch.Tensor):
+            raise TypeError(f"temb must be torch.Tensor, got {type(self.temb)}")
+
+        # Validate callables
+        if not callable(self.run_transformer_blocks):
+            raise TypeError(f"run_transformer_blocks must be callable, got {type(self.run_transformer_blocks)}")
+
+        if not callable(self.postprocess):
+            raise TypeError(f"postprocess must be callable, got {type(self.postprocess)}")
+
+        # Validate tensor shapes are compatible
+        if self.modulated_input.shape[0] != self.hidden_states.shape[0]:
+            raise ValueError(
+                f"Batch size mismatch: modulated_input has batch size "
+                f"{self.modulated_input.shape[0]}, but hidden_states has "
+                f"{self.hidden_states.shape[0]}"
+            )
+
+        # Validate devices match
+        if self.modulated_input.device != self.hidden_states.device:
+            raise ValueError(
+                f"Device mismatch: modulated_input on {self.modulated_input.device}, "
+                f"hidden_states on {self.hidden_states.device}"
+            )
+
+
+def extract_qwen_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_hidden_states_mask: torch.Tensor,
+    timestep: Union[torch.Tensor, float, int],
+    img_shapes: torch.Tensor,
+    txt_seq_lens: torch.Tensor,
+    guidance: torch.Tensor | None = None,
+    attention_kwargs: dict[str, Any] | None = None,
+    **kwargs: Any,
+) -> CacheContext:
+    """
+    Extract cache context for QwenImageTransformer2DModel.
+
+    This is the ONLY Qwen-specific code needed for TeaCache support.
+    It encapsulates preprocessing, modulated input extraction, transformer execution,
+    and postprocessing logic.
+
+    Args:
+        module: QwenImageTransformer2DModel instance
+        hidden_states: Input hidden states tensor
+        encoder_hidden_states: Text encoder outputs
+        encoder_hidden_states_mask: Mask for text encoder
+        timestep: Current diffusion timestep
+        img_shapes: Image shapes for position embedding
+        txt_seq_lens: Text sequence lengths
+        guidance: Optional guidance scale for CFG
+        attention_kwargs: Additional attention arguments
+        **kwargs: Additional keyword arguments ignored by this extractor
+
+    Returns:
+        CacheContext with all information needed for generic caching
+    """
+    from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+    if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
+        raise ValueError("Module must have transformer_blocks")
+
+    # ============================================================================
+    # PREPROCESSING (Qwen-specific)
+    # ============================================================================
+    hidden_states = module.img_in(hidden_states)
+    timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype)
+    encoder_hidden_states = module.txt_norm(encoder_hidden_states)
+    encoder_hidden_states = module.txt_in(encoder_hidden_states)
+
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+
+    temb = (
+        module.time_text_embed(timestep, hidden_states)
+        if guidance is None
+        else module.time_text_embed(timestep, guidance, hidden_states)
+    )
+
+    image_rotary_emb = module.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
+
+    # ============================================================================
+    # EXTRACT MODULATED INPUT (for cache decision)
+    # ============================================================================
+    block = module.transformer_blocks[0]
+    img_mod_params = block.img_mod(temb)
+    img_mod1, _ = img_mod_params.chunk(2, dim=-1)
+    img_normed = block.img_norm1(hidden_states)
+    img_modulated, _ = block._modulate(img_normed, img_mod1)
+
+    # ============================================================================
+    # DEFINE TRANSFORMER EXECUTION (Qwen-specific)
+    # ============================================================================
+    def run_transformer_blocks():
+        """Execute all Qwen transformer blocks."""
+        h = hidden_states
+        e = encoder_hidden_states
+        for block in module.transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                encoder_hidden_states_mask=encoder_hidden_states_mask,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=attention_kwargs,
+            )
+        return (h, e)
+
+    # ============================================================================
+    # DEFINE POSTPROCESSING (Qwen-specific)
+    # ============================================================================
+    return_dict = kwargs.get("return_dict", True)
+
+    def postprocess(h):
+        """Apply Qwen-specific output postprocessing."""
+        h = module.norm_out(h, temb)
+        output = module.proj_out(h)
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    # ============================================================================
+    # RETURN CONTEXT
+    # ============================================================================
+    return CacheContext(
+        modulated_input=img_modulated,
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+
+
+# Registry for model-specific extractors
+# Key: Pipeline class name (from OmniDiffusionConfig.model_class_name)
+# Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
+#
+# Note: Use the exact pipeline class name as specified in OmniDiffusionConfig.model_class_name
+# to ensure consistent lookup without substring matching.
+EXTRACTOR_REGISTRY: dict[str, Callable] = {
+    "QwenImagePipeline": extract_qwen_context,
+    # Future models:
+    # "FluxPipeline": extract_flux_context,
+    # "CogVideoXPipeline": extract_cogvideox_context,
+}
+
+
+def register_extractor(model_identifier: str, extractor_fn: Callable) -> None:
+    """
+    Register a new extractor function for a model type.
+
+    This allows extending TeaCache support to new models without modifying
+    the core TeaCache code.
+
+    Args:
+        model_identifier: Model type identifier (class name or type string)
+        extractor_fn: Function with signature (module, *args, **kwargs) -> CacheContext
+
+    Example:
+        >>> def extract_flux_context(module, hidden_states, timestep, guidance=None, **kwargs):
+        ...     # Preprocessing
+        ...     temb = module.time_text_embed(timestep, guidance)
+        ...     # Extract modulated input
+        ...     modulated = module.transformer_blocks[0].norm1(hidden_states, emb=temb)
+        ...     # Define execution
+        ...     def run_blocks():
+        ...         h = hidden_states
+        ...         for block in module.transformer_blocks:
+        ...             h = block(h, temb=temb)
+        ...         return (h,)
+        ...     # Define postprocessing
+        ...     def postprocess(h):
+        ...         return module.proj_out(module.norm_out(h, temb))
+        ...     # Return context
+        ...     return CacheContext(modulated, hidden_states, None, temb, run_blocks, postprocess)
+        >>> register_extractor("FluxTransformer2DModel", extract_flux_context)
+    """
+    EXTRACTOR_REGISTRY[model_identifier] = extractor_fn
+
+
+def get_extractor(model_type: str) -> Callable:
+    """
+    Get extractor function for given model type.
+
+    This function looks up the extractor based on the exact model_type string,
+    which should match the pipeline class name from OmniDiffusionConfig.model_class_name.
+
+    Args:
+        model_type: Model type string (e.g., "QwenImagePipeline", "FluxPipeline")
+                   Must exactly match a key in EXTRACTOR_REGISTRY.
+
+    Returns:
+        Extractor function with signature (module, *args, **kwargs) -> CacheContext
+
+    Raises:
+        ValueError: If model type not found in registry
+
+    Example:
+        >>> # Get extractor for Qwen model
+        >>> extractor = get_extractor("QwenImagePipeline")
+        >>> ctx = extractor(transformer, hidden_states, encoder_hidden_states, timestep, ...)
+    """
+    # Direct lookup - no substring matching
+    if model_type in EXTRACTOR_REGISTRY:
+        return EXTRACTOR_REGISTRY[model_type]
+
+    # No match found
+    available_types = list(EXTRACTOR_REGISTRY.keys())
+    raise ValueError(
+        f"Unknown model type: '{model_type}'. "
+        f"Available types: {available_types}\n"
+        f"To add support for a new model, use register_extractor() or add to EXTRACTOR_REGISTRY."
+    )

vllm_omni/diffusion/cache/teacache/hook.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Hook-based TeaCache implementation for vLLM-Omni.
+
+This module implements a diffusers-style hook system that completely intercepts
+the transformer forward pass, eliminating the need for any TeaCache-specific
+code in model definitions. Model developers only need to add an extractor function
+to support new models.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import get_extractor
+from vllm_omni.diffusion.cache.teacache.state import TeaCacheState
+from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager
+
+
+class TeaCacheHook(ModelHook):
+    """
+    ModelHook implementing TeaCache for transformer models.
+
+    This hook completely intercepts the transformer's forward pass and implements
+    adaptive caching based on timestep embedding similarity. It's model-agnostic
+    and supports multiple model types through extractor functions.
+
+    Key features:
+    - Zero changes to model code
+    - CFG-aware with separate states for positive/negative branches
+    - Model-specific polynomial rescaling
+    - Auto-detection of model types
+
+    Attributes:
+        config: TeaCache configuration with thresholds and callbacks
+        rescale_func: Polynomial function for rescaling L1 distances
+        state_manager: Manages TeaCacheState across forward passes
+        extractor_fn: Model-specific function to extract modulated input
+    """
+
+    _HOOK_NAME = "teacache"
+
+    def __init__(self, config: TeaCacheConfig):
+        """
+        Initialize TeaCacheHook.
+
+        Args:
+            config: TeaCache configuration object.
+        """
+        super().__init__()
+        self.config = config
+        self.rescale_func = np.poly1d(config.coefficients)
+        self.state_manager = StateManager(TeaCacheState)
+        self.extractor_fn = None
+
+    def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
+        """
+        Initialize hook with extractor from config model type.
+
+        Args:
+            module: The module to initialize the hook for.
+
+        Returns:
+            The initialized module.
+        """
+        # Get extractor function based on model_type from config
+        # model_type should be the pipeline class name (e.g., "QwenImagePipeline")
+        self.extractor_fn = get_extractor(self.config.model_type)
+
+        # Set default context
+        self.state_manager.set_context("teacache")
+
+        return module
+
+    def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any:
+        """
+        Generic forward handler that works for ANY model.
+
+        This method is completely model-agnostic. All model-specific logic
+        is encapsulated in the extractor function that returns a CacheContext.
+
+        The extractor does:
+        - Model-specific preprocessing
+        - Extraction of modulated input for cache decision
+        - Providing transformer execution callable
+        - Providing postprocessing callable
+
+        This hook does:
+        - CFG-aware state management
+        - Cache decision logic (generic)
+        - Residual caching and reuse
+
+        Args:
+            module: Transformer module (any architecture)
+            *args: Positional arguments for model forward
+            **kwargs: Keyword arguments for model forward
+
+        Returns:
+            Model output (format depends on model)
+        """
+        # Get model-specific context from extractor
+        # The extractor encapsulates ALL model-specific logic
+        ctx = self.extractor_fn(module, *args, **kwargs)
+
+        # ============================================================================
+        # GENERIC CACHING LOGIC (works for all models)
+        # ============================================================================
+        # Set context based on CFG branch for separate state tracking
+        cache_branch = kwargs.get("cache_branch", "default")
+        context_name = f"teacache_{cache_branch}"
+        self.state_manager.set_context(context_name)
+        state = self.state_manager.get_state()
+
+        # Decide whether to compute or cache based on modulated input similarity
+        should_compute = self._should_compute_full_transformer(state, ctx.modulated_input)
+
+        if not should_compute and state.previous_residual is not None:
+            # ============================================================================
+            # FAST PATH: Reuse cached residuals
+            # ============================================================================
+            ctx.hidden_states = ctx.hidden_states + state.previous_residual
+            if state.previous_residual_encoder is not None and ctx.encoder_hidden_states is not None:
+                ctx.encoder_hidden_states = ctx.encoder_hidden_states + state.previous_residual_encoder
+            output = ctx.hidden_states
+        else:
+            # ============================================================================
+            # SLOW PATH: Full transformer computation
+            # ============================================================================
+            ori_hidden_states = ctx.hidden_states.clone()
+            ori_encoder_hidden_states = (
+                ctx.encoder_hidden_states.clone() if ctx.encoder_hidden_states is not None else None
+            )
+
+            # Run transformer blocks using model-specific callable
+            outputs = ctx.run_transformer_blocks()
+
+            # Update context with outputs
+            ctx.hidden_states = outputs[0]
+            if len(outputs) > 1 and ctx.encoder_hidden_states is not None:
+                ctx.encoder_hidden_states = outputs[1]
+
+            # Cache residuals for next timestep
+            state.previous_residual = (ctx.hidden_states - ori_hidden_states).detach()
+            if ori_encoder_hidden_states is not None:
+                state.previous_residual_encoder = (ctx.encoder_hidden_states - ori_encoder_hidden_states).detach()
+
+            output = ctx.hidden_states
+
+        # Update state
+        state.previous_modulated_input = ctx.modulated_input.detach()
+        state.cnt += 1
+
+        # ============================================================================
+        # POSTPROCESSING (model-specific, via callable)
+        # ============================================================================
+        return ctx.postprocess(output)
+
+    def _should_compute_full_transformer(self, state: TeaCacheState, modulated_inp: torch.Tensor) -> bool:
+        """
+        Determine whether to compute full transformer or reuse cached residual.
+
+        This implements the core TeaCache algorithm:
+        1. Always compute first timestep
+        2. For intermediate steps:
+           - Compute relative L1 distance between current and previous modulated inputs
+           - Apply polynomial rescaling with model-specific coefficients
+           - Accumulate rescaled distances
+           - Compare to threshold: below = cache, above = compute
+
+        Args:
+            state: Current TeaCacheState containing counters and cached values
+            modulated_inp: Modulated input extracted from first transformer block
+
+        Returns:
+            True to compute full transformer, False to reuse cached residual
+        """
+        # First timestep: always compute
+        if state.cnt == 0:
+            state.accumulated_rel_l1_distance = 0.0
+            return True
+
+        # Need previous input for comparison
+        if state.previous_modulated_input is None:
+            return True
+
+        # Compute relative L1 distance between consecutive modulated inputs
+        rel_distance = (
+            (
+                (modulated_inp - state.previous_modulated_input).abs().mean()
+                / (state.previous_modulated_input.abs().mean() + 1e-8)
+            )
+            .cpu()
+            .item()
+        )
+
+        # Apply model-specific polynomial rescaling
+        rescaled_distance = float(self.rescale_func(rel_distance))
+        state.accumulated_rel_l1_distance += abs(rescaled_distance)
+
+        # Decision: below threshold = cache, above = compute
+        if state.accumulated_rel_l1_distance < self.config.rel_l1_thresh:
+            return False  # Use cache
+        else:
+            state.accumulated_rel_l1_distance = 0.0  # Reset accumulator
+            return True  # Compute
+
+    def reset_state(self, module: torch.nn.Module) -> torch.nn.Module:
+        """
+        Reset all cached states for a new inference run.
+
+        Args:
+            module: The module to reset state for.
+
+        Returns:
+            The module with reset state.
+        """
+        self.state_manager.reset()
+        return module
+
+
+def apply_teacache_hook(module: torch.nn.Module, config: TeaCacheConfig) -> None:
+    """
+    Apply TeaCache optimization to a transformer module.
+
+    This function registers a TeaCacheHook that completely intercepts the
+    module's forward pass, implementing adaptive caching without any changes
+    to the model code.
+
+    Args:
+        module: Transformer model to optimize (e.g., QwenImageTransformer2DModel)
+        config: TeaCacheConfig specifying caching parameters
+
+    Example:
+        >>> config = TeaCacheConfig(rel_l1_thresh=0.2, model_type="Qwen")
+        >>> apply_teacache_hook(transformer, config)
+        >>> # Model now uses TeaCache automatically, no code changes needed!
+    """
+    registry = HookRegistry.get_or_create(module)
+    hook = TeaCacheHook(config)
+    registry.register_hook(TeaCacheHook._HOOK_NAME, hook)

vllm_omni/diffusion/cache/teacache/state.py

@@ -0,0 +1,40 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+TeaCache state management.
+
+This module manages the state for TeaCache hooks across diffusion timesteps.
+"""
+
+from typing import Optional
+
+import torch
+
+
+class TeaCacheState:
+    """
+    State management for TeaCache hook.
+
+    Tracks caching state across diffusion timesteps, managing counters,
+    accumulated distances, and cached residuals for the TeaCache algorithm.
+    """
+
+    def __init__(self):
+        """Initialize empty TeaCache state."""
+        # Timestep tracking
+        self.cnt = 0
+
+        # Caching state
+        self.accumulated_rel_l1_distance = 0.0
+        self.previous_modulated_input: Optional[torch.Tensor] = None
+        self.previous_residual: Optional[torch.Tensor] = None
+        self.previous_residual_encoder: Optional[torch.Tensor] = None
+
+    def reset(self) -> None:
+        """Reset all state variables for a new inference run."""
+        self.cnt = 0
+        self.accumulated_rel_l1_distance = 0.0
+        self.previous_modulated_input = None
+        self.previous_residual = None
+        self.previous_residual_encoder = None

vllm_omni/diffusion/data.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import enum
+import os
 import random
 from dataclasses import dataclass, field
 from typing import Any, Callable
@@ -60,9 +61,13 @@ class OmniDiffusionConfig:
     # Workload type
     # workload_type: WorkloadType = WorkloadType.T2V
 
-    # Cache strategy
+    # Cache strategy (legacy)
     cache_strategy: str = "none"
 
+    # Cache adapter configuration (NEW)
+    cache_adapter: str = "none"  # "tea_cache", "deep_cache", etc.
+    cache_config: dict[str, Any] = field(default_factory=dict)
+
     # Distributed executor backend
     distributed_executor_backend: str = "mp"
     nccl_port: int | None = None
@@ -208,8 +213,17 @@ class OmniDiffusionConfig:
         initial_master_port = (self.master_port or 30005) + random.randint(0, 100)
         self.master_port = self.settle_port(initial_master_port, 37)
 
+        # Automatically inject model_class_name into cache_config if not present
+        if self.cache_adapter != "none" and self.model_class_name:
+            if "model_type" not in self.cache_config:
+                self.cache_config["model_type"] = self.model_class_name
+                logger.debug(f"Auto-injected model_type='{self.model_class_name}' into cache_config")
+
     @classmethod
     def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig":
+        # Check environment variable as fallback for cache_adapter
+        if "cache_adapter" not in kwargs:
+            kwargs["cache_adapter"] = os.environ.get("DIFFUSION_CACHE_ADAPTER", "none").lower()
         return cls(**kwargs)
 
 

vllm_omni/diffusion/hooks.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import torch.nn as nn
+
+
+class BaseState:
+    """Base class for hook state containers."""
+
+    def reset(self) -> None:  # pragma: no cover - default is no-op
+        pass
+
+
+class StateManager:
+    """Manage per-context hook state instances."""
+
+    def __init__(self, state_cls: Callable[[], BaseState]):
+        self._state_cls = state_cls
+        self._states: dict[str, BaseState] = {}
+        self._context: str = "default"
+
+    def set_context(self, name: str) -> None:
+        self._context = name or "default"
+
+    def get_state(self) -> BaseState:
+        if self._context not in self._states:
+            self._states[self._context] = self._state_cls()
+        return self._states[self._context]
+
+    def reset(self) -> None:
+        self._states.clear()
+
+
+class ModelHook:
+    """Base class for model hooks that can override a module's forward."""
+
+    def initialize_hook(self, module: nn.Module) -> nn.Module:
+        return module
+
+    def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any):
+        raise NotImplementedError
+
+    def reset_state(self, module: nn.Module) -> nn.Module:
+        return module
+
+
+@dataclass
+class _WrappedForward:
+    module: nn.Module
+
+    def __call__(self, *args: Any, **kwargs: Any):
+        registry: HookRegistry | None = getattr(self.module, "_hook_registry", None)
+        if registry is None or not registry._hooks:
+            return self.module._original_forward(*args, **kwargs)
+        return registry.dispatch(*args, **kwargs)
+
+
+class HookRegistry:
+    """Registry of hooks attached to a module."""
+
+    def __init__(self, module: nn.Module):
+        self.module = module
+        self._hooks: dict[str, ModelHook] = {}
+
+    @classmethod
+    def get_or_create(cls, module: nn.Module) -> HookRegistry:
+        registry: HookRegistry | None = getattr(module, "_hook_registry", None)
+        if registry is None:
+            registry = cls(module)
+            setattr(module, "_hook_registry", registry)
+
+            # Wrap module.forward once so hooks can intercept calls.
+            if not hasattr(module, "_original_forward"):
+                module._original_forward = module.forward  # type: ignore[attr-defined]
+                module.forward = _WrappedForward(module)  # type: ignore[assignment]
+
+        return registry
+
+    def register_hook(self, name: str, hook: ModelHook) -> None:
+        hook.initialize_hook(self.module)
+        self._hooks[name] = hook
+
+    def get_hook(self, name: str) -> ModelHook | None:
+        return self._hooks.get(name)
+
+    def dispatch(self, *args: Any, **kwargs: Any):
+        # For now we support a single active hook and call it directly.
+        # This can be extended to a chain if needed.
+        if not self._hooks:
+            return self.module._original_forward(*args, **kwargs)  # type: ignore[attr-defined]
+        # Deterministic order: sort by name.
+        name = sorted(self._hooks.keys())[0]
+        hook = self._hooks[name]
+        return hook.new_forward(self.module, *args, **kwargs)
+
+    def reset_hook(self, name: str) -> None:
+        hook = self._hooks.get(name)
+        if hook is not None:
+            hook.reset_state(self.module)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -270,6 +270,9 @@ class QwenImagePipeline(
         self.transformer = QwenImageTransformer2DModel(od_config=od_config)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
+        # Initialize cache adapter to None (will be set by worker or setup_cache if needed)
+        self._cache_adapter = None
+
         self.stage = None
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
@@ -546,31 +549,46 @@ class QwenImagePipeline(
             if self.interrupt:
                 continue
             self._current_timestep = t
-            # broadcast to batch dimension and place on same device/dtype as latents
+
+            # Broadcast timestep to match batch size
             timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
-            noise_pred = self.transformer(
-                hidden_states=latents,
-                timestep=timestep / 1000,
-                guidance=guidance,
-                encoder_hidden_states_mask=prompt_embeds_mask,
-                encoder_hidden_states=prompt_embeds,
-                img_shapes=img_shapes,
-                txt_seq_lens=txt_seq_lens,
-                attention_kwargs=self.attention_kwargs,
-                return_dict=False,
-            )[0]
+
+            # Forward pass for positive prompt (or unconditional if no CFG)
+            # cache_branch is passed to hook for CFG-aware state management
+            transformer_kwargs = {
+                "hidden_states": latents,
+                "timestep": timestep / 1000,
+                "guidance": guidance,
+                "encoder_hidden_states_mask": prompt_embeds_mask,
+                "encoder_hidden_states": prompt_embeds,
+                "img_shapes": img_shapes,
+                "txt_seq_lens": txt_seq_lens,
+                "attention_kwargs": self.attention_kwargs,
+                "return_dict": False,
+            }
+            if self._cache_adapter is not None:
+                transformer_kwargs["cache_branch"] = "positive"
+
+            noise_pred = self.transformer(**transformer_kwargs)[0]
+
+            # Forward pass for negative prompt (CFG)
+            # cache_branch is passed to hook for CFG-aware state management
             if do_true_cfg:
-                neg_noise_pred = self.transformer(
-                    hidden_states=latents,
-                    timestep=timestep / 1000,
-                    guidance=guidance,
-                    encoder_hidden_states_mask=negative_prompt_embeds_mask,
-                    encoder_hidden_states=negative_prompt_embeds,
-                    img_shapes=img_shapes,
-                    txt_seq_lens=negative_txt_seq_lens,
-                    attention_kwargs=self.attention_kwargs,
-                    return_dict=False,
-                )[0]
+                neg_transformer_kwargs = {
+                    "hidden_states": latents,
+                    "timestep": timestep / 1000,
+                    "guidance": guidance,
+                    "encoder_hidden_states_mask": negative_prompt_embeds_mask,
+                    "encoder_hidden_states": negative_prompt_embeds,
+                    "img_shapes": img_shapes,
+                    "txt_seq_lens": negative_txt_seq_lens,
+                    "attention_kwargs": self.attention_kwargs,
+                    "return_dict": False,
+                }
+                if self._cache_adapter is not None:
+                    neg_transformer_kwargs["cache_branch"] = "negative"
+
+                neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
                 comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
                 cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
                 noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)

vllm_omni/diffusion/request.py

@@ -124,10 +124,6 @@ class OmniDiffusionRequest:
     save_output: bool = True
     return_frames: bool = False
 
-    # TeaCache parameters
-    enable_teacache: bool = False
-    # teacache_params: TeaCacheParams | WanTeaCacheParams | None = None
-
     # STA parameters
     STA_param: list | None = None
     is_cfg_negative: bool = False

vllm_omni/diffusion/worker/gpu_worker.py

@@ -84,12 +84,29 @@ class GPUWorker:
         )
         logger.info(f"Worker {self.rank}: Model loaded successfully.")
 
+        # Apply cache adapter (model_type is auto-injected in OmniDiffusionConfig.__post_init__)
+        from vllm_omni.diffusion.cache.apply import setup_cache
+
+        self.od_config.cache_config["model_type"] = self.od_config.model_class_name
+
+        self.pipeline._cache_adapter = setup_cache(
+            self.pipeline.transformer,
+            cache_type=self.od_config.cache_adapter,
+            cache_config=self.od_config.cache_config,
+        )
+
+    def maybe_reset_cache(self) -> None:
+        """Reset cache state before each generation if applicable."""
+        if self.pipeline._cache_adapter is not None:
+            self.pipeline._cache_adapter.reset(self.pipeline.transformer)
+
     @torch.inference_mode()
     def execute_model(self, reqs: list[OmniDiffusionRequest], od_config: OmniDiffusionConfig) -> DiffusionOutput:
         """
         Execute a forward pass.
         """
         assert self.pipeline is not None
+        self.maybe_reset_cache()
         # TODO: dealing with first req for now
         req = reqs[0]
         output = self.pipeline.forward(req)

vllm_omni/entrypoints/stage_utils.py

@@ -64,7 +64,7 @@ def set_stage_devices(
         get_device_name_fn = torch.npu.get_device_name
         device_type_label = "NPU"
     elif device_type == "cuda":
-        import torch  # noqa: WPS433
+        import torch
 
         is_available_fn = torch.cuda.is_available
         set_device_fn = torch.cuda.set_device
@@ -145,7 +145,7 @@ def set_stage_devices(
             logger.debug("[Stage-%s] Set %s to single device %s (fallback)", stage_id, env_var, selected_physical)
 
         try:
-            import torch  # noqa: WPS433
+            import torch
 
             if is_available_fn():
                 try: