[Feat] Support teacache for SenseNova-U1 (#4164)

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

examples/offline_inference/sensenova_u1/README.md

@@ -38,9 +38,9 @@ python end2end.py --modality text2text \
 python end2end.py --prompt "A futuristic cityscape at sunset" \
                   --width 2048 --height 1024 --think
 
-# Text-to-image with Cache-DiT acceleration
+# Text-to-image with cache acceleration (cache_dit or tea_cache)
 python end2end.py --prompt "A cute cat" \
-                  --cache-backend cache_dit \
+                  --cache-backend tea_cache
 ```
 
 > **Note**: Default configuration works on a single **NVIDIA A100 (80GB)** or **H100** GPU.
@@ -139,7 +139,7 @@ python end2end.py \
 | `--tensor-parallel-size` | 1 | Number of GPUs for tensor parallelism |
 | `--enforce-eager` | False | Disable torch.compile |
 | `--enable-cpu-offload` | False | Enable module-wise (sequential) CPU offload to reduce peak VRAM |
-| `--cache-backend` | None | Set to `cache_dit` for Cache-DiT acceleration |
+| `--cache-backend` | None | Set to `cache_dit` or `tea_cache` for image generation acceleration |
 | `--enable-cache-dit-summary` | False | Print Cache-DiT cache statistics after generation |
 
 ## Reducing GPU Memory Usage
@@ -211,4 +211,4 @@ The img2img counterpart lives at
 ## Online Serving
 
 For OpenAI-compatible API serving, see
-[`examples/online_serving/sensenova_u1/`](../../examples/online_serving/sensenova_u1/).
+[`examples/online_serving/sensenova_u1/`](../../online_serving/sensenova_u1/README.md).

examples/offline_inference/sensenova_u1/end2end.py

@@ -164,7 +164,7 @@ def parse_args():
         "--cache-backend",
         type=str,
         default=None,
-        choices=["cache_dit"],
+        choices=["cache_dit", "tea_cache"],
         help="Cache backend to use for image generation acceleration.",
     )
     parser.add_argument(

examples/online_serving/sensenova_u1/README.md

@@ -6,11 +6,11 @@
 vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091
 ```
 
-With Cache-DiT acceleration:
+With cache acceleration (`cache_dit` or `tea_cache`):
 
 ```bash
 vllm serve SenseNova/SenseNova-U1-8B-MoT --omni --port 8091 \
-    --cache-backend cache_dit
+    --cache-backend tea_cache
 ```
 
 Or use the convenience script:
@@ -19,8 +19,8 @@ Or use the convenience script:
 cd examples/online_serving/sensenova_u1
 bash run_server.sh
 
-# Cache-DiT
-CACHE_BACKEND=cache_dit bash run_server.sh
+# Cache acceleration (cache_dit or tea_cache)
+CACHE_BACKEND=tea_cache bash run_server.sh
 ```
 
 ### Tensor Parallelism (TP)

vllm_omni/diffusion/cache/teacache/backend.py

@@ -57,6 +57,22 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     )
 
 
+def enable_sensenova_u1_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """Enable TeaCache for SenseNova-U1 denoising forwards."""
+    teacache_config = TeaCacheConfig(
+        transformer_type="SenseNovaU1ForCausalLM",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    transformer = pipeline.denoising_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}"
+    )
+
+
 def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """
     Enable TeaCache for Flux2 Klein model.
@@ -80,6 +96,7 @@ CUSTOM_TEACACHE_ENABLERS = {
     "BagelPipeline": enable_bagel_teacache,
     "Flux2KleinPipeline": enable_flux2_klein_teacache,
     "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
+    "SenseNovaU1Pipeline": enable_sensenova_u1_teacache,
 }
 
 
@@ -180,6 +197,8 @@ class TeaCacheBackend(CacheBackend):
 
         # Extract transformer from pipeline
         transformer = pipeline.transformer
+        if not hasattr(transformer, "_hook_registry") and hasattr(pipeline, "denoising_transformer"):
+            transformer = pipeline.denoising_transformer
 
         if hasattr(transformer, "_hook_registry"):
             hook = transformer._hook_registry.get_hook(TeaCacheHook._HOOK_NAME)

vllm_omni/diffusion/cache/teacache/config.py

@@ -37,6 +37,14 @@ _MODEL_COEFFICIENTS = {
     # Bagel transformer coefficients
     # Using Qwen's coefficients as reasonable default given shared architecture
     "Bagel": [1.33313129e06, -1.68644226e05, 7.95050740e03, -1.63747873e02, 1.26352397e00],
+    # SenseNova-U1 transformer coefficients
+    "SenseNovaU1ForCausalLM": [
+        9.07281930e04,
+        -2.17699186e04,
+        1.83940990e03,
+        -6.30339273e01,
+        7.61309272e-01,
+    ],
     # 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

@@ -1187,6 +1187,68 @@ def extract_flux_context(
     )
 
 
+def extract_sensenova_u1_context(
+    module: nn.Module,
+    input_ids: torch.Tensor | None = None,
+    indexes: torch.Tensor | None = None,
+    attention_mask: dict[str, torch.Tensor] | torch.Tensor | None = None,
+    past_key_values: Any | None = None,
+    inputs_embeds: torch.Tensor | None = None,
+    use_cache: bool | None = None,
+    embed_only: bool = False,
+    compute_logits: bool = True,
+    **kwargs: Any,
+) -> CacheContext:
+    """Extract cache context for SenseNovaU1ForCausalLM denoising forwards."""
+    from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
+        SenseNovaU1CausalLMOutput,
+    )
+
+    layer_kwargs = dict(kwargs)
+    layer_kwargs.pop("cache_dit_skip", None)
+    image_gen_indicators = layer_kwargs.pop("image_gen_indicators", None)
+    exist_und = (~image_gen_indicators).any().item()
+    exist_gen = image_gen_indicators.any().item()
+    causal_mask_mapping = attention_mask
+
+    first_layer = module.model.layers[0]
+    modulated_input = first_layer.input_layernorm_mot_gen(inputs_embeds)
+
+    def run_transformer_blocks():
+        h = inputs_embeds
+        for layer in module.model.layers:
+            h = layer(
+                h,
+                image_gen_indicators=image_gen_indicators,
+                exist_und=exist_und,
+                exist_gen=exist_gen,
+                indexes=indexes,
+                attention_mask=causal_mask_mapping,
+                past_key_values=past_key_values,
+                **layer_kwargs,
+            )
+        return (h,)
+
+    def postprocess(h: torch.Tensor) -> SenseNovaU1CausalLMOutput:
+        h = module.model.norm_mot_gen(h)
+
+        logits = module.logits_processor(module.lm_head, h) if compute_logits else None
+        return SenseNovaU1CausalLMOutput(
+            logits=logits,
+            past_key_values=past_key_values if use_cache else None,
+            hidden_states=h,
+        )
+
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=inputs_embeds,
+        encoder_hidden_states=None,
+        temb=modulated_input,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+
+
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -1202,6 +1264,7 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "Flux2Transformer2DModel": extract_flux2_context,
     "LongCatImageTransformer2DModel": extract_longcat_context,
     "FluxTransformer2DModel": extract_flux_context,
+    "SenseNovaU1ForCausalLM": extract_sensenova_u1_context,
     # Future models:
     # "CogVideoXTransformer3DModel": extract_cogvideox_context,
 }

vllm_omni/diffusion/models/sensenova_u1/pipeline_sensenova_u1.py

@@ -490,6 +490,30 @@ def _optimized_scale(positive_flat, negative_flat):
 # ---------------------------------------------------------------------------
 
 
+class SenseNovaU1DenoisingAdapter(nn.Module):
+    """Denoising-only entry point used by cache backends."""
+
+    def __init__(self, language_model: SenseNovaU1ForCausalLM):
+        super().__init__()
+        object.__setattr__(self, "language_model", language_model)
+        self.do_true_cfg = True
+
+    @property
+    def model(self):
+        return self.language_model.model
+
+    @property
+    def lm_head(self):
+        return self.language_model.lm_head
+
+    @property
+    def logits_processor(self):
+        return self.language_model.logits_processor
+
+    def forward(self, *args, **kwargs):
+        return self.language_model(*args, **kwargs)
+
+
 class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProfilerMixin, CFGParallelMixin):
     """SenseNova-U1 text-to-image and image-to-image pipeline for vllm-omni.
 
@@ -536,8 +560,12 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
             self.llm_cfg,
             prefix="language_model",
         )
-        # We define this for Cache DiT compatibility
+        # Cache-DiT hooks pipeline.transformer(.blocks), so it must point at the
+        # real decoder module (exposes .blocks and real parameters).
         self.transformer = self.language_model.model
+        # TeaCache intercepts the ForCausalLM-level denoising forward; route it
+        # through a dedicated adapter so it does not collide with Cache-DiT.
+        self.denoising_transformer = SenseNovaU1DenoisingAdapter(self.language_model)
 
         # Vision model (understanding branch)
         self.vision_model = NEOVisionModel(self.vis_cfg)
@@ -704,7 +732,8 @@ class SenseNovaU1Pipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipeli
         **_kw,
     ):
         B, L = z.shape[0], z.shape[1]
-        outputs = self.language_model(
+        denoising_model = self.language_model if cache_dit_skip else self.denoising_transformer
+        outputs = denoising_model(
             inputs_embeds=input_embeds,
             image_gen_indicators=torch.ones(
                 (input_embeds.shape[0], input_embeds.shape[1]), dtype=torch.bool, device=input_embeds.device