[Bagel] Support TeaCache (#848)

Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>

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
 

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.

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": [

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,

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()