[Perf] Support TeaCache for MammothModa2 (#5357)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -743,7 +743,11 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         ar_image_hidden_states: torch.Tensor | None = None,
         ar_image_attention_mask: torch.Tensor | None = None,
         return_dict: bool = False,
+        teacache_branch: str | None = None,
     ) -> torch.Tensor:
+        # Consumed by TeaCacheHook/extractor before this plain forward runs.
+        del teacache_branch
+
         batch_size, height, width = self._validate_inputs(
             hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict
         )

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -807,6 +807,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 ar_image_hidden_states=ar_image_embeds,
                 ar_image_attention_mask=ar_image_attention_mask,
                 freqs_cis=self.gen_freqs_cis,
+                teacache_branch="positive",
             )
             run_uncond = (any_active_per_step is not None and any_active_per_step[i]) if needs_uncond else False
             if requires_paired_cfg and negative_prompt_embeds is not None:
@@ -819,6 +820,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                     text_attention_mask=negative_prompt_attention_mask,
                     ref_image_hidden_states=None,
                     freqs_cis=self.gen_freqs_cis,
+                    teacache_branch="negative",
                 )
                 # Fused CFG blend: torch.lerp fuses (uncond + scale * (cond - uncond)) into 1 kernel
                 blended = torch.lerp(model_pred_uncond, model_pred, scale_vec)