[refactor] teacache refactor (#395)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

tests/diffusion/cache/test_cache_backends.py

@@ -146,8 +146,6 @@ class TestTeaCacheBackend:
 
         # Verify hook was applied
         assert backend.enabled is True
-        assert hasattr(mock_pipeline, "_cache_backend")
-        assert mock_pipeline._cache_backend is backend
         mock_apply_hook.assert_called_once()
 
     @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")

vllm_omni/diffusion/cache/base.py

@@ -23,6 +23,8 @@ All backends implement the same interface:
 from abc import ABC, abstractmethod
 from typing import Any
 
+import torch.nn as nn
+
 from vllm_omni.diffusion.data import DiffusionCacheConfig
 
 
@@ -98,3 +100,13 @@ class CacheBackend(ABC):
 
     def __repr__(self) -> str:
         return f"{self.__class__.__name__}(config={self.config})"
+
+
+class CachedTransformer(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.do_true_cfg = False
+
+    def __init_subclass__(cls, enable_separate_cfg: bool = True, **kwargs):
+        cls.enable_separate_cfg = enable_separate_cfg
+        super().__init_subclass__(**kwargs)

vllm_omni/diffusion/cache/teacache/backend.py

@@ -75,10 +75,6 @@ class TeaCacheBackend(CacheBackend):
         # Apply hook to transformer
         apply_teacache_hook(transformer, teacache_config)
 
-        # Store backend reference on pipeline for pipeline's internal use
-        # (e.g., checking if cache is enabled for cache_branch parameter)
-        pipeline._cache_backend = self
-
         # Mark as enabled
         self.enabled = True
 

vllm_omni/diffusion/cache/teacache/hook.py

@@ -58,6 +58,7 @@ class TeaCacheHook(ModelHook):
         self.rescale_func = np.poly1d(config.coefficients)
         self.state_manager = StateManager(TeaCacheState)
         self.extractor_fn = None
+        self._forward_cnt = 0
 
     def initialize_hook(self, module: torch.nn.Module) -> torch.nn.Module:
         """
@@ -112,7 +113,11 @@ class TeaCacheHook(ModelHook):
         # GENERIC CACHING LOGIC (works for all models)
         # ============================================================================
         # Set context based on CFG branch for separate state tracking
-        cache_branch = kwargs.get("cache_branch", "default")
+        if module.do_true_cfg and self._forward_cnt % 2 == 1:
+            cache_branch = "negative"
+        else:
+            cache_branch = "positive"
+
         context_name = f"teacache_{cache_branch}"
         self.state_manager.set_context(context_name)
         state = self.state_manager.get_state()
@@ -155,6 +160,7 @@ class TeaCacheHook(ModelHook):
         # Update state
         state.previous_modulated_input = ctx.modulated_input.detach()
         state.cnt += 1
+        self._forward_cnt += 1
 
         # ============================================================================
         # POSTPROCESSING (model-specific, via callable)
@@ -221,6 +227,7 @@ class TeaCacheHook(ModelHook):
             The module with reset state.
         """
         self.state_manager.reset()
+        self._forward_cnt = 0
         return module
 
 

vllm_omni/diffusion/hooks.py

@@ -38,13 +38,13 @@ 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
+        raise NotImplementedError
 
     def new_forward(self, module: nn.Module, *args: Any, **kwargs: Any):
         raise NotImplementedError
 
     def reset_state(self, module: nn.Module) -> nn.Module:
-        return module
+        raise NotImplementedError
 
 
 @dataclass

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -271,9 +271,6 @@ class QwenImagePipeline(
 
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
-        # Initialize cache backend to None (will be set by worker if needed)
-        self._cache_backend = None
-
         self.stage = None
 
         self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
@@ -555,41 +552,31 @@ class QwenImagePipeline(
             timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
 
             # 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_backend is not None:
-                transformer_kwargs["cache_branch"] = "positive"
-
-            noise_pred = self.transformer(**transformer_kwargs)[0]
-
+            self.transformer.do_true_cfg = do_true_cfg
+            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 negative prompt (CFG)
-            # cache_branch is passed to hook for CFG-aware state management
             if do_true_cfg:
-                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_backend is not None:
-                    neg_transformer_kwargs["cache_branch"] = "negative"
-
-                neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
+                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]
                 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/models/qwen_image/pipeline_qwen_image_edit.py

@@ -604,43 +604,34 @@ class QwenImageEditPipeline(
             if image_latents is not None:
                 latent_model_input = torch.cat([latents, image_latents], dim=1)
 
+            self.transformer.do_true_cfg = do_true_cfg  # used in teacache hook
             # 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": latent_model_input,
-                "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_backend is not None:
-                transformer_kwargs["cache_branch"] = "positive"
-
-            noise_pred = self.transformer(**transformer_kwargs)[0]
+            noise_pred = self.transformer(
+                hidden_states=latent_model_input,
+                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]
             noise_pred = noise_pred[:, : latents.size(1)]
 
             # Forward pass for negative prompt (CFG)
-            # cache_branch is passed to hook for CFG-aware state management
             if do_true_cfg:
-                neg_transformer_kwargs = {
-                    "hidden_states": latent_model_input,
-                    "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_backend is not None:
-                    neg_transformer_kwargs["cache_branch"] = "negative"
-
-                neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
+                neg_noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    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_noise_pred = neg_noise_pred[:, : latents.size(1)]
                 comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
                 cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -544,39 +544,33 @@ class QwenImageEditPlusPipeline(
             if image_latents is not None:
                 latent_model_input = torch.cat([latents, image_latents], dim=1)
 
-            transformer_kwargs = {
-                "hidden_states": latent_model_input,
-                "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_backend is not None:
-                transformer_kwargs["cache_branch"] = "positive"
+            self.transformer.do_true_cfg = do_true_cfg
 
-            noise_pred = self.transformer(**transformer_kwargs)[0]
+            noise_pred = self.transformer(
+                hidden_states=latent_model_input,
+                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]
             noise_pred = noise_pred[:, : latents.size(1)]
 
             if do_true_cfg:
-                neg_transformer_kwargs = {
-                    "hidden_states": latent_model_input,
-                    "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_backend is not None:
-                    neg_transformer_kwargs["cache_branch"] = "negative"
-
-                neg_noise_pred = self.transformer(**neg_transformer_kwargs)[0]
+                neg_noise_pred = self.transformer(
+                    hidden_states=latent_model_input,
+                    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_noise_pred = neg_noise_pred[:, : latents.size(1)]
                 comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
                 cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)

vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinea
 from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.cache.base import CachedTransformer
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_sequence_parallel_rank,
@@ -577,7 +578,8 @@ class QwenImageTransformerBlock(nn.Module):
         return encoder_hidden_states, hidden_states
 
 
-class QwenImageTransformer2DModel(nn.Module):
+# Note: inheriting from CachedTransformer only when we support caching
+class QwenImageTransformer2DModel(CachedTransformer):
     """
     The Transformer model introduced in Qwen.