[Perf] Support TeaCache for MammothModa2 (#5357)

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

vllm_omni/diffusion/cache/teacache/hook.py

@@ -145,7 +145,15 @@ class TeaCacheHook(ModelHook):
         #   - cfg_rank 0: positive branch
         #   - cfg_rank > 0: negative branch
         # Without CFG-parallel, branches alternate within a single rank
-        if getattr(module, "do_true_cfg", False):
+        branch_hint = None
+        if getattr(ctx, "extra_states", None):
+            branch_hint = ctx.extra_states.get("teacache_branch")
+
+        if branch_hint is not None:
+            if branch_hint not in ("positive", "negative"):
+                raise ValueError(f"Invalid teacache_branch={branch_hint!r}; expected 'positive' or 'negative'.")
+            cache_branch = branch_hint
+        elif 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()