From f81d8005f832bd41bb14b6f88e4096dd430c7770 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 30 Jun 2026 22:12:18 +0000 Subject: [PATCH] Fix PAB cond/uncond cache cross-contamination via StateManager PyramidAttentionBroadcastHook stored iteration and cache on a single hook-level state object that ignored cache_context. Pipelines running separate cond and uncond forwards (CogView4, Flux true CFG, etc.) would reuse conditional attention outputs during the uncond pass when PAB skipped computation, silently corrupting guidance. Migrate PAB to StateManager so cond/uncond contexts maintain isolated state, matching MagCache and FirstBlockCache. Add a regression test and update pipeline PAB tests to use state_manager. Co-authored-by: Simon Lynch --- .../hooks/pyramid_attention_broadcast.py | 27 +++---- .../hooks/test_pyramid_attention_broadcast.py | 71 +++++++++++++++++++ tests/pipelines/test_pipelines_common.py | 20 +++--- utils/check_copies.py | 3 +- 4 files changed, 98 insertions(+), 23 deletions(-) create mode 100644 tests/hooks/test_pyramid_attention_broadcast.py diff --git a/src/diffusers/hooks/pyramid_attention_broadcast.py b/src/diffusers/hooks/pyramid_attention_broadcast.py index ed5bd24..de77321 100644 --- a/src/diffusers/hooks/pyramid_attention_broadcast.py +++ b/src/diffusers/hooks/pyramid_attention_broadcast.py @@ -27,7 +27,7 @@ _SPATIAL_TRANSFORMER_BLOCK_IDENTIFIERS, _TEMPORAL_TRANSFORMER_BLOCK_IDENTIFIERS, ) -from .hooks import HookRegistry, ModelHook +from .hooks import BaseState, HookRegistry, ModelHook, StateManager logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -105,7 +105,7 @@ def __repr__(self) -> str: ) -class PyramidAttentionBroadcastState: +class PyramidAttentionBroadcastState(BaseState): r""" State for Pyramid Attention Broadcast. @@ -148,33 +148,34 @@ def __init__( self.timestep_skip_range = timestep_skip_range self.block_skip_range = block_skip_range self.current_timestep_callback = current_timestep_callback - - def initialize_hook(self, module): - self.state = PyramidAttentionBroadcastState() - return module + self.state_manager = StateManager(PyramidAttentionBroadcastState, (), {}) def new_forward(self, module: torch.nn.Module, *args, **kwargs) -> Any: + if self.state_manager._current_context is None: + self.state_manager.set_context("inference") + + state = self.state_manager.get_state() is_within_timestep_range = ( self.timestep_skip_range[0] < self.current_timestep_callback() < self.timestep_skip_range[1] ) should_compute_attention = ( - self.state.cache is None - or self.state.iteration == 0 + state.cache is None + or state.iteration == 0 or not is_within_timestep_range - or self.state.iteration % self.block_skip_range == 0 + or state.iteration % self.block_skip_range == 0 ) if should_compute_attention: output = self.fn_ref.original_forward(*args, **kwargs) else: - output = self.state.cache + output = state.cache - self.state.cache = output - self.state.iteration += 1 + state.cache = output + state.iteration += 1 return output def reset_state(self, module: torch.nn.Module) -> None: - self.state.reset() + self.state_manager.reset() return module diff --git a/tests/hooks/test_pyramid_attention_broadcast.py b/tests/hooks/test_pyramid_attention_broadcast.py new file mode 100644 index 0000000..938d12c --- /dev/null +++ b/tests/hooks/test_pyramid_attention_broadcast.py @@ -0,0 +1,71 @@ +# Copyright 2025 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers.hooks.pyramid_attention_broadcast import _apply_pyramid_attention_broadcast_hook +from diffusers.models.cache_utils import CacheMixin +from diffusers.models.modeling_utils import ModelMixin + +from ..testing_utils import torch_device + + +class CountingBlock(torch.nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.proj = torch.nn.Linear(dim, dim) + self.call_count = 0 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + self.call_count += 1 + return self.proj(hidden_states) + + +class DummyTransformer(ModelMixin, CacheMixin, torch.nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.transformer_blocks = torch.nn.ModuleList([CountingBlock(dim)]) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for block in self.transformer_blocks: + hidden_states = block(hidden_states) + return hidden_states + + +def test_pab_isolates_state_between_cond_and_uncond_contexts(): + """PAB must not reuse cond attention cache during the uncond forward pass.""" + dim = 4 + model = DummyTransformer(dim).to(torch_device) + block = model.transformer_blocks[0] + + _apply_pyramid_attention_broadcast_hook( + block, + timestep_skip_range=(100, 800), + block_skip_range=2, + current_timestep_callback=lambda: 500, + ) + + x = torch.randn(1, dim, device=torch_device) + + with model.cache_context("cond"): + model(x) + + assert block.call_count == 1 + + with model.cache_context("uncond"): + model(x) + + assert block.call_count == 2, ( + "Uncond pass must compute attention instead of reusing cond cache from a shared hook state." + ) diff --git a/tests/pipelines/test_pipelines_common.py b/tests/pipelines/test_pipelines_common.py index fcd8ab2..30aa825 100644 --- a/tests/pipelines/test_pipelines_common.py +++ b/tests/pipelines/test_pipelines_common.py @@ -2691,7 +2691,11 @@ def test_pyramid_attention_broadcast_layers(self): isinstance(hook, PyramidAttentionBroadcastHook), "Hook should be of type PyramidAttentionBroadcastHook.", ) - self.assertTrue(hook.state.cache is None, "Cache should be None at initialization.") + self.assertEqual( + len(hook.state_manager._state_cache), + 0, + "Cache should be None at initialization.", + ) self.assertEqual(count, expected_hooks, "Number of hooks should match the expected number.") # Perform dummy inference step to ensure state is updated @@ -2701,12 +2705,13 @@ def pab_state_check_callback(pipe, i, t, kwargs): hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") if hook is None: continue + state = hook.state_manager.get_state() self.assertTrue( - hook.state.cache is not None, + state.cache is not None, "Cache should have updated during inference.", ) self.assertTrue( - hook.state.iteration == i + 1, + state.iteration == i + 1, "Hook iteration state should have updated during inference.", ) return {} @@ -2722,14 +2727,11 @@ def pab_state_check_callback(pipe, i, t, kwargs): hook = module._diffusers_hook.get_hook("pyramid_attention_broadcast") if hook is None: continue - self.assertTrue( - hook.state.cache is None, + self.assertEqual( + len(hook.state_manager._state_cache), + 0, "Cache should be reset to None after inference.", ) - self.assertTrue( - hook.state.iteration == 0, - "Iteration should be reset to 0 after inference.", - ) def test_pyramid_attention_broadcast_inference(self, expected_atol: float = 0.2): # We need to use higher tolerance because we are using a random model. With a converged/trained diff --git a/utils/check_copies.py b/utils/check_copies.py index 001366c..338a25d 100644 --- a/utils/check_copies.py +++ b/utils/check_copies.py @@ -18,6 +18,7 @@ import os import re import subprocess +import sys # All paths are set with the intent you should run this script from the root of the repo with the command @@ -94,7 +95,7 @@ def get_indent(code): def run_ruff(code): - command = ["ruff", "format", "-", "--config", "pyproject.toml", "--silent"] + command = [sys.executable, "-m", "ruff", "format", "-", "--config", "pyproject.toml", "--silent"] process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE) stdout, _ = process.communicate(input=code.encode()) return stdout.decode()