From a7fe4fba4e1a7b2b9920b9e66caac34ba39c1d86 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 22 Dec 2025 06:49:04 +0800 Subject: [PATCH] [refactor] teacache refactor (#395) Signed-off-by: zjy0516 --- tests/diffusion/cache/test_cache_backends.py | 2 - vllm_omni/diffusion/cache/base.py | 12 ++++ vllm_omni/diffusion/cache/teacache/backend.py | 4 -- vllm_omni/diffusion/cache/teacache/hook.py | 9 ++- vllm_omni/diffusion/hooks.py | 4 +- .../models/qwen_image/pipeline_qwen_image.py | 59 ++++++++----------- .../qwen_image/pipeline_qwen_image_edit.py | 55 ++++++++--------- .../pipeline_qwen_image_edit_plus.py | 52 ++++++++-------- .../qwen_image/qwen_image_transformer.py | 4 +- 9 files changed, 94 insertions(+), 107 deletions(-) diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index 501a4c38c..7bc49bdb7 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/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") diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py index 2d7c64442..d78076224 100644 --- a/vllm_omni/diffusion/cache/base.py +++ b/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) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 71f9ad653..108dd5eb2 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index 1859ae17e..b2f731665 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/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 diff --git a/vllm_omni/diffusion/hooks.py b/vllm_omni/diffusion/hooks.py index 2296aa9a0..b1e1016ce 100644 --- a/vllm_omni/diffusion/hooks.py +++ b/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 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index e5b99840d..b1b7af2e8 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/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) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index 8605e01a5..8d180e625 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/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) diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 0042484a7..6068cf0f1 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/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) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index e31bbe8ba..5b961d617 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/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.