[Bugfix][Diffusion] Apply MagCache residual once per skipped step (#5561)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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"

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)