[Bugfix][Diffusion] Keep TeaCache CFG branches separate in the shared CFG helper

tests/diffusion/cache/test_teacache_cfg_group.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext
+from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook, apply_teacache_hook
+from vllm_omni.diffusion.distributed import parallel_state
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
+from vllm_omni.diffusion.hooks import HookRegistry
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+class _TeaCacheProbe(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.computed_residuals: list[float] = []
+
+    def forward(self, hidden_states: torch.Tensor, residual: float) -> tuple[torch.Tensor]:
+        raise AssertionError("TeaCache hook must intercept this forward")
+
+
+def _extract_probe_context(module: _TeaCacheProbe, hidden_states: torch.Tensor, residual: float) -> CacheContext:
+    def run_blocks() -> tuple[torch.Tensor]:
+        module.computed_residuals.append(residual)
+        return (hidden_states + residual,)
+
+    return CacheContext(
+        modulated_input=hidden_states.clone(),
+        hidden_states=hidden_states.clone(),
+        encoder_hidden_states=None,
+        temb=torch.zeros_like(hidden_states),
+        run_transformer_blocks=run_blocks,
+        postprocess=lambda output: (output,),
+    )
+
+
+class _ProbePipeline(CFGParallelMixin):
+    def __init__(self) -> None:
+        self.transformer = _TeaCacheProbe()
+        apply_teacache_hook(
+            self.transformer,
+            TeaCacheConfig(transformer_type="_TeaCacheProbe", coefficients=[0, 0, 0, 0, 0], rel_l1_thresh=0.2),
+        )
+
+
+@pytest.fixture
+def pipeline(monkeypatch: pytest.MonkeyPatch) -> _ProbePipeline:
+    monkeypatch.setattr(parallel_state, "_CFG", None)
+    monkeypatch.setitem(EXTRACTOR_REGISTRY, "_TeaCacheProbe", _extract_probe_context)
+    return _ProbePipeline()
+
+
+@pytest.mark.parametrize("group_initialized", [False, True])
+def test_sequential_cfg_with_teacache_keeps_branch_residuals_separate(
+    pipeline: _ProbePipeline, monkeypatch: pytest.MonkeyPatch, group_initialized: bool
+) -> None:
+    """Direct pipelines without a CFG group must support real TeaCache hits."""
+    if group_initialized:
+        monkeypatch.setattr(parallel_state, "_CFG", SimpleNamespace(world_size=1, rank_in_group=0))
+
+    hidden_states = torch.ones(1, 1)
+    positive_kwargs = {"hidden_states": hidden_states, "residual": 2.0}
+    negative_kwargs = {"hidden_states": hidden_states, "residual": 10.0}
+
+    for _ in range(2):
+        prediction = 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,
+        )
+        torch.testing.assert_close(prediction, torch.full_like(hidden_states, -21.0))
+
+    # Each branch computes once, then reuses its own residual on the next step.
+    assert pipeline.transformer.computed_residuals == [2.0, 10.0]
+    hook = HookRegistry.get_or_create(pipeline.transformer).get_hook("teacache")
+    assert isinstance(hook, TeaCacheHook)
+    assert set(hook.state_manager._states) == {"teacache_positive", "teacache_negative"}
+    assert all(state.cnt == 2 for state in hook.state_manager._states.values())
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_teacache_uses_initialized_cfg_rank(
+    pipeline: _ProbePipeline, monkeypatch: pytest.MonkeyPatch, rank: int
+) -> None:
+    monkeypatch.setattr(parallel_state, "_CFG", SimpleNamespace(world_size=2, rank_in_group=rank))
+    pipeline.transformer.do_true_cfg = True
+    hidden_states = torch.ones(1, 1)
+    for _ in range(2):
+        output = pipeline.transformer(hidden_states=hidden_states, residual=2.0)[0]
+        torch.testing.assert_close(output, hidden_states + 2.0)
+
+    hook = HookRegistry.get_or_create(pipeline.transformer).get_hook("teacache")
+    assert isinstance(hook, TeaCacheHook)
+    expected_branch = "negative" if rank == 1 else "positive"
+    assert set(hook.state_manager._states) == {f"teacache_{expected_branch}"}
+    assert pipeline.transformer.computed_residuals == [2.0]

tests/diffusion/distributed/test_cfg_parallel.py

@@ -1054,3 +1054,86 @@ def test_predict_noise_with_multi_branch_cfg_parity(
         batch_size=batch_size,
         cfg_normalize=cfg_normalize,
     )
+
+
+class _FlagRecordingTransformer(SimpleTransformer):
+    """Records ``do_true_cfg`` at each forward call."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.seen_flags: list[Any] = []
+
+    def forward(self, x: torch.Tensor, **kwargs) -> tuple[torch.Tensor]:
+        self.seen_flags.append(getattr(self, "do_true_cfg", None))
+        return super().forward(x, **kwargs)
+
+
+class _FlagRecordingPipeline(CFGParallelMixin):
+    def __init__(self) -> None:
+        _set_random_seeds(0)
+        self.transformer = _FlagRecordingTransformer(in_channels=4, hidden_dim=32)
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_sequential_cfg_marks_transformer_for_per_branch_caches():
+    """``predict_noise_maybe_with_cfg`` sets ``transformer.do_true_cfg`` per call and clears it without CFG."""
+    pipeline = _FlagRecordingPipeline()
+    positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+        batch_size=1, channels=4, height=4, width=4, dtype=torch.float32, device=torch.device("cpu"), input_seed=1
+    )
+
+    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_flags == [True, True]
+
+        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_flags == [True, True, False]
+    assert pipeline.transformer.do_true_cfg is False
+
+
+class _NoTransformerPipeline(CFGParallelMixin):
+    """Pipelines may route predict_noise to a module not named ``transformer``."""
+
+    def __init__(self) -> None:
+        _set_random_seeds(0)
+        self.denoiser = SimpleTransformer(in_channels=4, hidden_dim=32)
+
+    def predict_noise(self, *args: Any, **kwargs: Any) -> torch.Tensor:
+        return self.denoiser(*args, **kwargs)[0]
+
+
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_sequential_cfg_without_transformer_attribute():
+    pipeline = _NoTransformerPipeline()
+    positive_kwargs, negative_kwargs = _make_two_branch_inputs(
+        batch_size=1, channels=4, height=4, width=4, dtype=torch.float32, device=torch.device("cpu"), input_seed=2
+    )
+
+    with torch.no_grad():
+        out = 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 out.shape == (1, 4, 4, 4)
+    assert not hasattr(pipeline, "transformer")

tests/diffusion/distributed/test_pipeline_parallel.py

@@ -172,6 +172,66 @@ def predict_noise(self, x=None, intermediate_tensors=None, **_kwargs) -> torch.T
         return x
 
 
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@pytest.mark.cpu
+def test_pp_predict_noise_sets_transformer_do_true_cfg(monkeypatch: pytest.MonkeyPatch) -> None:
+    """PP>1 path sets ``transformer.do_true_cfg`` before each ``predict_noise`` call."""
+
+    class _StubPPGroup:
+        is_first_rank = True
+        is_last_rank = True
+
+        def irecv_tensor_dict(self):
+            raise AssertionError("unexpected recv in stub last rank")
+
+        def isend_tensor_dict(self, *args, **kwargs):
+            return []
+
+    class _FlagTransformer(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.seen_at_predict: list[bool | None] = []
+
+    class _MiniPPPipeline(PipelineParallelMixin, CFGParallelMixin):
+        def __init__(self) -> None:
+            self.transformer = _FlagTransformer()
+
+        def predict_noise(self, **kwargs) -> torch.Tensor:
+            self.transformer.seen_at_predict.append(getattr(self.transformer, "do_true_cfg", None))
+            return torch.ones(2, 4)
+
+    monkeypatch.setattr(pp_module, "get_pipeline_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(pp_module, "get_pp_group", lambda: _StubPPGroup())
+    monkeypatch.setattr(pp_module, "get_classifier_free_guidance_world_size", lambda: 1)
+
+    pipeline = _MiniPPPipeline()
+    positive_kwargs = {"x": torch.zeros(2, 4)}
+    negative_kwargs = {"x": torch.ones(2, 4)}
+
+    with torch.inference_mode():
+        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_at_predict == [True, True]
+        assert pipeline.transformer.do_true_cfg is True
+
+        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_at_predict[-1] is False
+    assert pipeline.transformer.do_true_cfg is False
+
+
 # ---------------------------------------------------------------------------
 # 1.  AsyncLatents – unit tests (no distributed env required)
 # ---------------------------------------------------------------------------

vllm_omni/diffusion/cache/teacache/hook.py

@@ -24,6 +24,7 @@
     get_classifier_free_guidance_rank,
     get_classifier_free_guidance_world_size,
     get_sp_group,
+    is_cfg_group_initialized,
     model_parallel_is_initialized,
 )
 from vllm_omni.diffusion.hooks import HookRegistry, ModelHook, StateManager
@@ -154,7 +155,8 @@ def new_forward(self, module: torch.nn.Module, *args: Any, **kwargs: Any) -> Any
                 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()
+            # 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"

vllm_omni/diffusion/distributed/cfg_parallel.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 pipeline class for Diffusion models with shared CFG functionality.
@@ -88,6 +88,11 @@ class CFGParallelMixin(metaclass=ABCMeta):
         and set self.scheduler to a composite scheduler that handles tuples.
     """
 
+    def _set_transformer_do_true_cfg(self, do_true_cfg: bool) -> None:
+        transformer = getattr(self, "transformer", None)
+        if transformer is not None:
+            transformer.do_true_cfg = do_true_cfg
+
     def predict_noise_maybe_with_cfg(
         self,
         do_true_cfg: bool,
@@ -118,7 +123,14 @@ def predict_noise_maybe_with_cfg(
             For multi-output models (e.g., video + audio where predict_noise
             returns a tuple), override combine_cfg_noise() for per-element CFG
             logic and set self.scheduler to a composite scheduler.
+
+            Sets ``self.transformer.do_true_cfg`` before predicting. TeaCache
+            reads it to keep one cache state per CFG branch; without it the
+            positive and negative calls of a step share a state and reuse each
+            other's residual.
         """
+        self._set_transformer_do_true_cfg(do_true_cfg)
+
         if do_true_cfg:
             # Automatically detect CFG parallel configuration
             cfg_parallel_ready = _get_cfg_world_size_or_one() > 1

vllm_omni/diffusion/distributed/pipeline_parallel.py

@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 from functools import wraps
 from typing import Any
 
 import torch
 from vllm.v1.worker.gpu_worker import AsyncIntermediateTensors
 
+from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
     get_classifier_free_guidance_rank,
@@ -196,6 +197,9 @@ def predict_noise_maybe_with_cfg(
             CFGParallelMixin.predict_noise_maybe_with_cfg exactly.
           - PP only, no CFG: cond branch only.
 
+        Also sets ``self.transformer.do_true_cfg`` (via CFGParallelMixin) before
+        predicting, matching the non-PP CFG helper path for TeaCache.
+
         Returns:
             noise_pred on the last PP rank (all CFG ranks when CFG-parallel is active).
             None on all other ranks.
@@ -205,6 +209,8 @@ def predict_noise_maybe_with_cfg(
                 do_true_cfg, true_cfg_scale, positive_kwargs, negative_kwargs, cfg_normalize, output_slice
             )
 
+        CFGParallelMixin._set_transformer_do_true_cfg(self, do_true_cfg)
+
         self._sync_pp_send()
 
         pp_group = get_pp_group()

vllm_omni/diffusion/models/flux/pipeline_flux.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
 
 import json
 import logging
@@ -456,7 +456,6 @@ def diffuse(
     ) -> torch.Tensor:
         """Diffusion loop with optional image conditioning."""
         self.scheduler.set_begin_index(0)
-        self.transformer.do_true_cfg = do_true_cfg
         for i, t in enumerate(timesteps):
             if self.interrupt:
                 continue

vllm_omni/diffusion/models/flux/pipeline_flux_kontext.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
 
 # This implementation is heavily inspired by the diffusers project.
 # Original implementation: https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/flux/pipeline_flux_kontext.py
@@ -515,7 +515,6 @@ def diffuse(
         callback_on_step_end_tensor_inputs: list[str] = ["latents"],
     ) -> torch.Tensor:
         self.scheduler.set_begin_index(0)
-        self.transformer.do_true_cfg = do_true_cfg
         with self.progress_bar(total=len(timesteps)) as pbar:
             for i, t in enumerate(timesteps):
                 if self.interrupt:

vllm_omni/diffusion/models/qwen_image/cfg_parallel.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
 """CFG Parallel Mixin for Qwen Image series
 Shared by
 - QwenImagePipeline
@@ -67,7 +67,6 @@ def diffuse(
             Denoised latents
         """
         self.scheduler.set_begin_index(0)
-        self.transformer.do_true_cfg = do_true_cfg
         additional_transformer_kwargs = additional_transformer_kwargs or {}
 
         with self.progress_bar(total=len(timesteps)) as pbar:

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -939,7 +939,6 @@ def denoise_step(
 
         t = input_batch.timesteps
         self._current_timestep = t
-        self.transformer.do_true_cfg = input_batch.do_true_cfg
 
         positive_kwargs, negative_kwargs, output_slice = self._build_denoise_kwargs(
             latents=input_batch.latents,