diff --git a/tests/diffusion/cache/test_magcache_hook.py b/tests/diffusion/cache/test_magcache_hook.py new file mode 100644 index 000000000..4a6fa0a76 --- /dev/null +++ b/tests/diffusion/cache/test_magcache_hook.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Unit tests for the MagCache hook forward numerics. + +The cached MagCache residual is the residual of the whole block stack +(tail output - head input). On a skipped step the head hook applies it +once; every other hooked block must pass its inputs through unchanged, +matching the reference implementation in diffusers.hooks.mag_cache. +""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.diffusion.cache.magcache.config import MagCacheConfig +from vllm_omni.diffusion.cache.magcache.hook import ( + _MAG_CACHE_LEADER_BLOCK_HOOK, + apply_mag_cache_hook, +) +from vllm_omni.diffusion.hooks.base import HookRegistry + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _CountingBlock(nn.Module): + """Toy dual-stream block: adds a constant delta to hidden_states. + + Returns (encoder_hidden_states, hidden_states) to match the default + TransformerBlockMetadata(return_hidden_states_index=1, + return_encoder_hidden_states_index=0) registered by the hook. + """ + + def __init__(self, delta: float): + super().__init__() + self.delta = delta + self.forward_calls = 0 + + def forward(self, hidden_states, encoder_hidden_states): + self.forward_calls += 1 + return encoder_hidden_states, hidden_states + self.delta + + +class _ToyTransformer(nn.Module): + def __init__(self, n_blocks: int, delta: float = 1.0): + super().__init__() + self.blocks = nn.ModuleList(_CountingBlock(delta) for _ in range(n_blocks)) + + def forward(self, hidden_states, encoder_hidden_states): + for block in self.blocks: + encoder_hidden_states, hidden_states = block(hidden_states, encoder_hidden_states) + return hidden_states + + @property + def total_forward_calls(self) -> int: + return sum(block.forward_calls for block in self.blocks) + + +def _make_config(num_inference_steps: int = 2) -> MagCacheConfig: + # retention_ratio=0.5 with 2 steps -> retention_step = 1: step 0 always + # computes, step 1 with mag_ratio 1.0 accumulates zero error, which is a + # guaranteed cache hit. + return MagCacheConfig( + transformer_type="Toy", + threshold=0.24, + max_skip_steps=5, + retention_ratio=0.5, + num_inference_steps=num_inference_steps, + mag_ratios=torch.ones(num_inference_steps), + ) + + +def test_skipped_step_applies_full_stack_residual_once(): + n_blocks, delta = 3, 1.0 + model = _ToyTransformer(n_blocks=n_blocks, delta=delta) + apply_mag_cache_hook(model, _make_config()) + + x = torch.zeros(1, 4, 8) + e = torch.zeros(1, 2, 8) + + # Step 0 computes and caches residual = tail_output - head_input. + out_computed = model(x, e) + assert torch.allclose(out_computed, x + n_blocks * delta) + assert model.total_forward_calls == n_blocks + + # Step 1 is a cache hit: no block computes, and the cached residual is + # applied exactly once, so the output matches the computed step. + out_skipped = model(x, e) + assert model.total_forward_calls == n_blocks, "step 1 was not skipped" + assert torch.allclose(out_skipped, out_computed), ( + f"cached residual over-applied on skip: expected " + f"{out_computed.flatten()[0].item()}, got {out_skipped.flatten()[0].item()}" + ) + + +def test_single_block_advances_schedule_on_skip(): + model = _ToyTransformer(n_blocks=1) + apply_mag_cache_hook(model, _make_config(num_inference_steps=2)) + + x = torch.zeros(1, 4, 8) + e = torch.zeros(1, 2, 8) + + model(x, e) # step 0: computed + assert model.total_forward_calls == 1 + model(x, e) # step 1: cache hit + assert model.total_forward_calls == 1, "step 1 was not skipped" + + # The head hook is also the tail here, so a skipped step must still + # advance the schedule; step 1 was the last step, which resets the + # state for the next inference run. + registry = HookRegistry.check_if_exists_or_initialize(model.blocks[0]) + hook = registry.get_hook(_MAG_CACHE_LEADER_BLOCK_HOOK) + state = hook.state_manager.get_state() + assert state.step_index == 0 + assert state._is_first_step + assert state.previous_residual is None + + # A new run therefore starts by computing, not by replaying the stale + # residual from the previous run. + model(x, e) + assert model.total_forward_calls == 2, "new run reused stale skip state" diff --git a/vllm_omni/diffusion/cache/magcache/hook.py b/vllm_omni/diffusion/cache/magcache/hook.py index 6a440e020..4b2b4ea06 100644 --- a/vllm_omni/diffusion/cache/magcache/hook.py +++ b/vllm_omni/diffusion/cache/magcache/hook.py @@ -143,6 +143,11 @@ class MagCacheHeadHook(ModelHook): if not should_compute: res = state.previous_residual + if self._is_tail: + # Single-block configuration: this hook is also the tail, so + # the step schedule must advance on skipped steps too. + self._advance_step_head(state) + if isinstance(res, tuple): res = tuple(r.to(hidden_states.device) for r in res) @@ -363,10 +368,11 @@ class MagCacheBlockHook(ModelHook): state: MagCacheState = self.state_manager.get_state() if not state.should_compute: - res = state.previous_residual - if res is None: - res = torch.zeros_like(args[0]) - + # The head hook already applied the cached residual, which is the + # residual of the whole block stack (tail output - head input). + # Non-head blocks must pass their inputs through unchanged on a + # skipped step; adding the residual here again would apply it once + # per hooked block. if hasattr(self._metadata, "hidden_states_argument_name"): arg_name = self._metadata.hidden_states_argument_name else: @@ -380,27 +386,15 @@ class MagCacheBlockHook(ModelHook): encoder_hidden_states = self._metadata._get_parameter_from_args_kwargs( "encoder_hidden_states", args, kwargs ) - - if self._strategy is not None: - out_hidden, enc_out = self._strategy.apply_residual_tuple(hidden_states, encoder_hidden_states, res) - else: - out_hidden = hidden_states + res - enc_out = encoder_hidden_states - max_idx = max( self._metadata.return_hidden_states_index, self._metadata.return_encoder_hidden_states_index ) ret_list = [None] * (max_idx + 1) - ret_list[self._metadata.return_hidden_states_index] = out_hidden - ret_list[self._metadata.return_encoder_hidden_states_index] = enc_out + ret_list[self._metadata.return_hidden_states_index] = hidden_states + ret_list[self._metadata.return_encoder_hidden_states_index] = encoder_hidden_states return tuple(ret_list) - if self._strategy is not None: - output = self._strategy.apply_residual(hidden_states, res) - else: - output = hidden_states + res - - return output + return hidden_states output = self.fn_ref.original_forward(*args, **kwargs)