From 06c9d9b463b4aeb936dd34938b2d87347f645998 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=B1=AA=E5=BF=97=E9=B9=8F?= Date: Sun, 25 Jan 2026 14:52:27 +0800 Subject: [PATCH] [Bagel] Support TeaCache (#848) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: princepride Signed-off-by: 汪志鹏 --- docs/user_guide/diffusion_acceleration.md | 2 +- vllm_omni/diffusion/cache/teacache/backend.py | 91 +++++++++++----- vllm_omni/diffusion/cache/teacache/config.py | 3 + .../diffusion/cache/teacache/extractors.py | 102 ++++++++++++++++++ vllm_omni/diffusion/cache/teacache/hook.py | 2 +- 5 files changed, 172 insertions(+), 28 deletions(-) diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index 4572b7115..42202d8d7 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/docs/user_guide/diffusion_acceleration.md @@ -49,7 +49,7 @@ The following table shows which models are currently supported by each accelerat | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | | **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | | **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ❌ | -| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ❌ | ✅ | ❌ | ❌ | ❌ | +| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ### VideoGen diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 9db153cc4..bf328d43d 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/vllm_omni/diffusion/cache/teacache/backend.py @@ -15,10 +15,42 @@ from vllm.logger import init_logger from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook +from vllm_omni.diffusion.data import DiffusionCacheConfig logger = init_logger(__name__) +def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """ + Enable TeaCache for Bagel model. + """ + teacache_config = TeaCacheConfig( + transformer_type="Bagel", + rel_l1_thresh=config.rel_l1_thresh, + coefficients=config.coefficients, + ) + transformer = pipeline.bagel + original_forward_flow = transformer._forward_flow + + import types + + def forward_alias(self, *args, **kwargs): + return original_forward_flow(*args, **kwargs) + + transformer.forward = types.MethodType(forward_alias, transformer) + apply_teacache_hook(transformer, teacache_config) + transformer._forward_flow = transformer.forward + pipeline.transformer = transformer + + logger.info( + f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " + f"transformer_class={teacache_config.transformer_type}" + ) + + +CUSTOM_TEACACHE_ENABLERS = {"BagelPipeline": enable_bagel_teacache} + + class TeaCacheBackend(CacheBackend): """ TeaCache implementation using hooks. @@ -51,38 +83,45 @@ class TeaCacheBackend(CacheBackend): - transformer: pipeline.transformer - transformer_type: pipeline.transformer.__class__.__name__ """ - # Extract transformer and transformer_type from pipeline - transformer = pipeline.transformer - transformer_type = transformer.__class__.__name__ + # Helper to get pipeline class name + pipeline_type = pipeline.__class__.__name__ - # Create TeaCacheConfig from DiffusionCacheConfig with transformer_type - # Access parameters via attribute access: config.rel_l1_thresh - # rel_l1_thresh already has a default value of 0.2 in DiffusionCacheConfig - try: - teacache_config = TeaCacheConfig( - transformer_type=transformer_type, - rel_l1_thresh=self.config.rel_l1_thresh, - coefficients=self.config.coefficients, - ) - 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, coefficients (optional). " - f"transformer_type is automatically extracted from pipeline.transformer.__class__.__name__." - ) + # Check for pipeline-level custom enablers + if pipeline_type in CUSTOM_TEACACHE_ENABLERS: + logger.info(f"Using custom TeaCache enabler for model: {pipeline_type}") + CUSTOM_TEACACHE_ENABLERS[pipeline_type](pipeline, self.config) + else: + transformer = pipeline.transformer + transformer_type = transformer.__class__.__name__ - # Apply hook to transformer - apply_teacache_hook(transformer, teacache_config) + # Create TeaCacheConfig from DiffusionCacheConfig with transformer_type + # Access parameters via attribute access: config.rel_l1_thresh + # rel_l1_thresh already has a default value of 0.2 in DiffusionCacheConfig + try: + teacache_config = TeaCacheConfig( + transformer_type=transformer_type, + rel_l1_thresh=self.config.rel_l1_thresh, + coefficients=self.config.coefficients, + ) + 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, coefficients (optional). " + f"transformer_type is automatically extracted from pipeline.transformer.__class__.__name__." + ) + + # 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"transformer_class={teacache_config.transformer_type}" + ) # Mark as enabled self.enabled = True - logger.info( - f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, " - f"transformer_class={teacache_config.transformer_type}" - ) - def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: """ Refresh TeaCache state for new generation. diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index bb5e7cb27..30b1745f4 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/vllm_omni/diffusion/cache/teacache/config.py @@ -25,6 +25,9 @@ _MODEL_COEFFICIENTS = { 3.20000000e00, -2.00000000e-02, ], + # Bagel transformer coefficients + # Using Qwen's coefficients as reasonable default given shared architecture + "Bagel": [1.33313129e06, -1.68644226e05, 7.95050740e03, -1.63747873e02, 1.26352397e00], # Z-Image transformer coefficients # Copied from Qwen-Image, need to be tuned specifically for Z-Image in future "ZImageTransformer2DModel": [ diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 5c35e0c8c..01d453061 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -254,6 +254,107 @@ def extract_qwen_context( ) +def extract_bagel_context( + module: nn.Module, + x_t: torch.Tensor, + timestep: torch.Tensor | float | int, + packed_vae_token_indexes: torch.LongTensor, + packed_vae_position_ids: torch.LongTensor, + packed_text_ids: torch.LongTensor, + packed_text_indexes: torch.LongTensor, + packed_indexes: torch.LongTensor, + packed_position_ids: torch.LongTensor, + packed_seqlens: torch.IntTensor, + key_values_lens: torch.IntTensor, + past_key_values: Any, + packed_key_value_indexes: torch.LongTensor, + **kwargs: Any, +) -> CacheContext: + """ + Extract cache context for Bagel model. + + Args: + module: Bagel instance + x_t: Latent image input + timestep: Current timestep + packed_vae_token_indexes: Indexes for VAE tokens in packed sequence + packed_vae_position_ids: Position IDs for VAE tokens + packed_text_ids: Text token IDs + packed_text_indexes: Indexes for text tokens in packed sequence + packed_indexes: Global indexes + packed_position_ids: Global position IDs + packed_seqlens: Sequence lengths + key_values_lens: KV cache lengths + past_key_values: KV cache + packed_key_value_indexes: KV cache indexes + **kwargs: Additional keyword arguments + + Returns: + CacheContext with all information needed for generic caching + """ + + # 1. Embed text + packed_text_embedding = module.language_model.model.embed_tokens(packed_text_ids) + packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), module.hidden_size)) + packed_sequence[packed_text_indexes] = packed_text_embedding + + # 2. Embed timestep + if not isinstance(timestep, torch.Tensor): + timestep = torch.tensor([timestep], device=x_t.device) + if timestep.dim() == 0: + timestep = timestep.unsqueeze(0) + + # 3. Embed image (x_t) + packed_pos_embed = module.latent_pos_embed(packed_vae_position_ids) + packed_timestep_embeds = module.time_embedder(timestep) + + x_t_emb = module.vae2llm(x_t) + packed_timestep_embeds + packed_pos_embed + if x_t_emb.dtype != packed_sequence.dtype: + x_t_emb = x_t_emb.to(packed_sequence.dtype) + + packed_sequence[packed_vae_token_indexes] = x_t_emb + + # Use the full packed sequence as modulated input to match hidden_states size + modulated_input = packed_sequence + + def run_transformer_blocks(): + extra_inputs = {} + if module.use_moe: + extra_inputs = { + "mode": "gen", + "packed_vae_token_indexes": packed_vae_token_indexes, + "packed_text_indexes": packed_text_indexes, + } + + output = module.language_model.forward( + packed_query_sequence=packed_sequence, + query_lens=packed_seqlens, + packed_query_position_ids=packed_position_ids, + packed_query_indexes=packed_indexes, + past_key_values=past_key_values, + key_values_lens=key_values_lens, + packed_key_value_indexes=packed_key_value_indexes, + update_past_key_values=False, + is_causal=False, + **extra_inputs, + ) + return (output.packed_query_sequence,) + + def postprocess(h): + v_t = module.llm2vae(h) + v_t = v_t[packed_vae_token_indexes] + return v_t + + return CacheContext( + modulated_input=modulated_input, + hidden_states=packed_sequence, # Use full packed sequence + encoder_hidden_states=None, + temb=packed_timestep_embeds, # Approximate + run_transformer_blocks=run_transformer_blocks, + postprocess=postprocess, + ) + + def extract_zimage_context( module: nn.Module, x: list[torch.Tensor], @@ -448,6 +549,7 @@ def extract_zimage_context( # on the transformer module and multiple pipelines can share the same transformer. EXTRACTOR_REGISTRY: dict[str, Callable] = { "QwenImageTransformer2DModel": extract_qwen_context, + "Bagel": extract_bagel_context, "ZImageTransformer2DModel": extract_zimage_context, # Future models: # "FluxTransformer2DModel": extract_flux_context, diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index d896473b4..65f764c43 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/vllm_omni/diffusion/cache/teacache/hook.py @@ -122,7 +122,7 @@ class TeaCacheHook(ModelHook): # - cfg_rank 0: positive branch # - cfg_rank > 0: negative branch # Without CFG-parallel, branches alternate within a single rank - if module.do_true_cfg: + if getattr(module, "do_true_cfg", False): cfg_parallel_size = get_classifier_free_guidance_world_size() if cfg_parallel_size > 1: cfg_rank = get_classifier_free_guidance_rank()