[Bugfix][Diffusion] Stamp the CFG branch for TeaCache instead of inferring it from parity

tests/diffusion/cache/test_teacache_cfg_branch.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""TeaCacheHook takes the CFG branch from the caller's stamp (#8482).
+
+Under step execution the TeaCache hook is not refreshed between scheduler
+waves, so a CFG wave followed by a non-CFG wave with an odd number of forwards
+used to flip the forward-count parity: the next CFG step's positive call landed
+in the negative state. CFGParallelMixin now stamps ``transformer.cfg_branch`` on
+every call and the hook prefers it.
+"""
+
+from typing import Any
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.cache.teacache.hook as teacache_hook
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_POS, _NEG = "teacache_positive", "teacache_negative"
+
+
+class _BranchProbeDiT(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.do_true_cfg = False
+        self.cfg_branch: str | None = None
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor]:  # replaced by the hook
+        return (x,)
+
+
+def _extract(module: torch.nn.Module, x: torch.Tensor, **_: Any) -> CacheContext:
+    return CacheContext(
+        modulated_input=x,
+        hidden_states=x,
+        encoder_hidden_states=None,
+        temb=x,
+        run_transformer_blocks=lambda: (x + 1.0,),
+        postprocess=lambda h: (h,),
+    )
+
+
+class _Pipeline(CFGParallelMixin):
+    def __init__(self, transformer: torch.nn.Module) -> None:
+        self.transformer = transformer
+
+
+@pytest.fixture
+def hooked(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "_BranchProbeDiT", _extract)
+    monkeypatch.setattr(teacache_hook, "get_classifier_free_guidance_world_size", lambda: 1)
+    module = _BranchProbeDiT()
+    teacache_hook.apply_teacache_hook(
+        module,
+        TeaCacheConfig(transformer_type="_BranchProbeDiT", coefficients=[0.0, 0.0, 0.0, 1.0, 0.0], rel_l1_thresh=0.2),
+    )
+    hook = module._hook_registry.get_hook(teacache_hook.TeaCacheHook._HOOK_NAME)
+    contexts: list[str] = []
+    set_context = hook.state_manager.set_context
+
+    def record(name: str) -> None:
+        contexts.append(name)
+        set_context(name)
+
+    monkeypatch.setattr(hook.state_manager, "set_context", record)
+    return module, contexts
+
+
+def _step(pipeline: _Pipeline, x: torch.Tensor, *, cfg: bool) -> None:
+    with torch.no_grad():
+        pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=cfg,
+            true_cfg_scale=4.0,
+            positive_kwargs={"x": x},
+            negative_kwargs={"x": x} if cfg else None,
+            cfg_normalize=False,
+        )
+
+
+def test_interleaved_non_cfg_wave_keeps_branch_states_aligned(hooked) -> None:
+    module, contexts = hooked
+    pipeline = _Pipeline(module)
+    x = torch.ones(1, 4)
+
+    _step(pipeline, x, cfg=True)  # CFG wave: positive, negative
+    _step(pipeline, x, cfg=False)  # non-CFG wave: one forward, odd count
+    _step(pipeline, x, cfg=True)  # next CFG wave must start positive again
+
+    assert contexts == [_POS, _NEG, _POS, _POS, _NEG]
+    assert module.cfg_branch is None
+
+
+def test_unstamped_callers_keep_the_parity_fallback(hooked) -> None:
+    """Callers outside CFGParallelMixin are unchanged: branch inferred from parity."""
+    module, contexts = hooked
+    x = torch.ones(1, 4)
+
+    with torch.no_grad():
+        for cfg in (True, True, False, True, True):
+            module.do_true_cfg = cfg
+            module(x)
+
+    assert contexts == [_POS, _NEG, _POS, _NEG, _POS]
+
+
+def _extract_with_hint(module: torch.nn.Module, x: torch.Tensor, teacache_branch: str | None = None, **_: Any):
+    ctx = _extract(module, x)
+    ctx.extra_states = {"teacache_branch": teacache_branch} if teacache_branch is not None else None
+    return ctx
+
+
+def test_extractor_hint_takes_precedence_over_stamp(hooked, monkeypatch: pytest.MonkeyPatch) -> None:
+    """An extractor's per-call hint (e.g. MammothModa2) is more specific than the mixin stamp."""
+    module, contexts = hooked
+    hook = module._hook_registry.get_hook(teacache_hook.TeaCacheHook._HOOK_NAME)
+    monkeypatch.setattr(hook, "extractor_fn", _extract_with_hint)
+    x = torch.ones(1, 4)
+
+    with torch.no_grad():
+        module.cfg_branch = "negative"
+        module(x, teacache_branch="positive")
+        module(x)  # no hint: the stamp applies
+    module.cfg_branch = None
+
+    assert contexts == [_POS, _NEG]
+
+
+def test_invalid_stamp_is_rejected(hooked) -> None:
+    module, _ = hooked
+    module.cfg_branch = "uncond"
+
+    with pytest.raises(ValueError, match="cfg_branch"), torch.no_grad():
+        module(torch.ones(1, 4))

tests/diffusion/distributed/test_cfg_parallel.py

@@ -1207,3 +1207,110 @@ def test_sequential_cfg_without_transformer_attribute():
 
     assert out.shape == (1, 4, 4, 4)
     assert not hasattr(pipeline, "transformer")
+
+
+class _BranchRecordingTransformer(SimpleTransformer):
+    """Records ``cfg_branch`` as seen by a hook on the transformer at call time."""
+
+    def __init__(self, *args: Any, fail: bool = False, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.cfg_branch: str | None = None
+        self.seen_branches: list[str | None] = []
+        self.fail = fail
+
+    def forward(self, x: torch.Tensor, **kwargs) -> tuple[torch.Tensor]:
+        self.seen_branches.append(self.cfg_branch)
+        if self.fail:
+            raise RuntimeError("forward failed")
+        return super().forward(x, **kwargs)
+
+
+class _BranchRecordingPipeline(CFGParallelMixin):
+    def __init__(self, fail: bool = False) -> None:
+        _set_random_seeds(0)
+        self.transformer = _BranchRecordingTransformer(in_channels=4, hidden_dim=32, fail=fail)
+
+
+def _branch_inputs(seed: int = 3):
+    return _make_two_branch_inputs(
+        batch_size=1, channels=4, height=4, width=4, dtype=torch.float32, device=torch.device("cpu"), input_seed=seed
+    )
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_helper_stamps_branch_per_call_and_restores_it():
+    """#8482: TeaCache reads the branch from the caller instead of forward-count parity."""
+    pipeline = _BranchRecordingPipeline()
+    positive_kwargs, negative_kwargs = _branch_inputs()
+
+    with torch.no_grad():
+        pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=4.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+        pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=False,
+            true_cfg_scale=1.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=None,
+            cfg_normalize=False,
+        )
+
+    assert pipeline.transformer.seen_branches == ["positive", "negative", "positive"]
+    assert pipeline.transformer.cfg_branch is None
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+@pytest.mark.parametrize(("cfg_rank", "expected"), [(0, "positive"), (1, "negative")])
+def test_cfg_parallel_rank_stamps_its_branch(monkeypatch: pytest.MonkeyPatch, cfg_rank: int, expected: str):
+    pipeline = _BranchRecordingPipeline()
+    positive_kwargs, negative_kwargs = _branch_inputs()
+    with torch.no_grad():
+        pos = _wrap(pipeline.predict_noise(**positive_kwargs))[0]
+        neg = _wrap(pipeline.predict_noise(**negative_kwargs))[0]
+    pipeline.transformer.seen_branches.clear()
+    fake_group = FakeCfgGroup(world_size=2, rank_in_group=cfg_rank, rank_tensors=[pos, neg])
+    _patch_cfg_state(monkeypatch, world_size=2, rank=cfg_rank, fake_group=fake_group)
+    # _get_cfg_world_size_or_one() reports 1 unless a CFG group is initialized.
+    from vllm_omni.diffusion.distributed import cfg_parallel as cfg_parallel_mod
+
+    monkeypatch.setattr(cfg_parallel_mod, "is_cfg_group_initialized", lambda: True)
+
+    with torch.no_grad():
+        pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=4.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    assert pipeline.transformer.seen_branches == [expected]
+    assert pipeline.transformer.cfg_branch is None
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_cfg_branch_stamp_is_restored_when_predict_noise_raises():
+    pipeline = _BranchRecordingPipeline(fail=True)
+    positive_kwargs, negative_kwargs = _branch_inputs()
+
+    with pytest.raises(RuntimeError, match="forward failed"), torch.no_grad():
+        pipeline.predict_noise_maybe_with_cfg(
+            do_true_cfg=True,
+            true_cfg_scale=4.0,
+            positive_kwargs=positive_kwargs,
+            negative_kwargs=negative_kwargs,
+            cfg_normalize=False,
+        )
+
+    assert pipeline.transformer.seen_branches == ["positive"]
+    assert pipeline.transformer.cfg_branch is None

vllm_omni/diffusion/cache/base.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Base cache backend interface for diffusion models.
@@ -108,6 +108,9 @@ class CachedTransformer(nn.Module):
     def __init__(self, **kwargs):
         super().__init__()
         self.do_true_cfg = False
+        # Branch of the current call ("positive" / "negative"), stamped by
+        # CFGParallelMixin; None outside its calls.
+        self.cfg_branch: str | None = None
 
     def __init_subclass__(cls, enable_separate_cfg: bool = True, **kwargs):
         cls.enable_separate_cfg = enable_separate_cfg

vllm_omni/diffusion/cache/teacache/hook.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """
 Hook-based TeaCache implementation for vLLM-Omni.
@@ -141,30 +141,19 @@ def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any
         # ============================================================================
         # GENERIC CACHING LOGIC (works for all models)
         # ============================================================================
-        # Set context based on CFG branch for separate state tracking
-        # With CFG-parallel, each rank processes only one branch:
-        #   - cfg_rank 0: positive branch
-        #   - cfg_rank > 0: negative branch
-        # Without CFG-parallel, branches alternate within a single rank
-        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):
-            # Direct pipeline calls may not initialize a CFG process group.
-            cfg_parallel_size = get_classifier_free_guidance_world_size() if is_cfg_group_initialized() else 1
-            if cfg_parallel_size > 1:
-                cfg_rank = get_classifier_free_guidance_rank()
-                cache_branch = "negative" if cfg_rank > 0 else "positive"
-            else:
-                # No CFG-parallel: use forward counter to alternate branches
-                cache_branch = "negative" if self._forward_cnt % 2 == 1 else "positive"
-        else:
-            cache_branch = "positive"
+        # Set context based on CFG branch for separate state tracking.
+        # Explicit sources win, most specific first:
+        #   1. the extractor's per-call hint, ctx.extra_states["teacache_branch"]
+        #      (MammothModa2 passes it through its forward kwargs);
+        #   2. module.cfg_branch, stamped by CFGParallelMixin on every call it
+        #      makes, which stays correct when non-CFG forwards are interleaved
+        #      (step execution).
+        # Callers that provide neither fall back to the inferred branch:
+        #   - CFG-parallel: cfg_rank 0 is positive, cfg_rank > 0 negative
+        #   - otherwise branches are assumed to alternate on this rank
+        cache_branch = self._explicit_cfg_branch(module, ctx)
+        if cache_branch is None:
+            cache_branch = self._infer_cfg_branch(module)
 
         context_name = f"teacache_{cache_branch}"
         self.state_manager.set_context(context_name)
@@ -230,6 +219,28 @@ def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any
         # ============================================================================
         return ctx.postprocess(output)
 
+    @staticmethod
+    def _explicit_cfg_branch(module: torch.nn.Module, ctx: Any) -> str | None:
+        """Branch named by the extractor hint or the caller's stamp, if any."""
+        extra_states = getattr(ctx, "extra_states", None) or {}
+        source, branch = "teacache_branch", extra_states.get("teacache_branch")
+        if branch is None:
+            source, branch = "cfg_branch", getattr(module, "cfg_branch", None)
+        if branch is not None and branch not in ("positive", "negative"):
+            raise ValueError(f"Invalid {source}={branch!r}; expected 'positive' or 'negative'.")
+        return branch
+
+    def _infer_cfg_branch(self, module: torch.nn.Module) -> str:
+        """Fallback branch for callers that don't stamp ``module.cfg_branch``."""
+        if not getattr(module, "do_true_cfg", False):
+            return "positive"
+        # Direct pipeline calls may not initialize a CFG process group.
+        cfg_parallel_size = get_classifier_free_guidance_world_size() if is_cfg_group_initialized() else 1
+        if cfg_parallel_size > 1:
+            return "negative" if get_classifier_free_guidance_rank() > 0 else "positive"
+        # Sequential CFG without a stamp: assume calls alternate on this rank.
+        return "negative" if self._forward_cnt % 2 == 1 else "positive"
+
     def _should_compute_full_transformer(self, state: TeaCacheState, modulated_inp: torch.Tensor) -> bool:
         """
         Determine whether to compute full transformer or reuse cached residual.