Fix TaylorSeer cache crash when the feature shape changes between steps (#14831)

* Fix TaylorSeer cache crash when the feature shape changes between steps

* Add QwenImage 2.1 x TaylorSeer tests covering the KV-cache prefill shape change

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

src/diffusers/hooks/taylorseer_cache.py

@@ -149,7 +149,10 @@ class TaylorSeerState:
                     prev_factors = self.taylor_factors.get(i, {})
                     for j in range(self.max_order):
                         prev = prev_factors.get(j)
-                        if prev is None:
+                        # The feature shape can change between steps, e.g. Qwen-Image 2.1 returns prefix + target
+                        # tokens on the KV-cache prefill step and target tokens only afterwards. Stale factors
+                        # cannot be differenced against the new features, so restart the expansion from order 0.
+                        if prev is None or prev.shape != new_factors[j].shape:
                             break
                         new_factors[j + 1] = (new_factors[j] - prev.to(features.dtype)) / delta_step
                 self.taylor_factors[i] = {

tests/hooks/test_taylorseer_cache.py

@@ -0,0 +1,52 @@
+# Copyright 2026 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.taylorseer_cache import TaylorSeerState
+
+
+def _step(state: TaylorSeerState, features: torch.Tensor) -> None:
+    state.current_step += 1
+    state.update((features,))
+
+
+def test_update_restarts_when_feature_shape_changes():
+    # Qwen-Image 2.1 with the KV cache returns prefix + target tokens on the prefill step and target tokens only
+    # afterwards, so the hooked module sees two different sequence lengths within one denoising loop.
+    state = TaylorSeerState(taylor_factors_dtype=torch.float32, max_order=1)
+
+    _step(state, torch.ones(1, 10, 4))
+    _step(state, torch.ones(1, 6, 4))
+
+    assert set(state.taylor_factors[0]) == {0}
+    assert state.taylor_factors[0][0].shape == (1, 6, 4)
+
+    _step(state, torch.full((1, 6, 4), 3.0))
+
+    assert set(state.taylor_factors[0]) == {0, 1}
+    torch.testing.assert_close(state.taylor_factors[0][1], torch.full((1, 6, 4), 2.0))
+
+    state.current_step += 1
+    (prediction,) = state.predict()
+    torch.testing.assert_close(prediction, torch.full((1, 6, 4), 5.0))
+
+
+def test_update_keeps_finite_difference_when_shape_is_stable():
+    state = TaylorSeerState(taylor_factors_dtype=torch.float32, max_order=1)
+
+    _step(state, torch.zeros(1, 6, 4))
+    _step(state, torch.full((1, 6, 4), 2.0))
+
+    torch.testing.assert_close(state.taylor_factors[0][1], torch.full((1, 6, 4), 2.0))

tests/models/transformers/test_models_transformer_qwenimage21.py

@@ -28,6 +28,7 @@ from ..testing_utils import (
     MemoryTesterMixin,
     ModelTesterMixin,
     SingleFileTesterMixin,
+    TaylorSeerCacheTesterMixin,
     TrainingTesterMixin,
 )
 
@@ -297,6 +298,72 @@ class TestQwenImage21TransformerAttention(QwenImage21TransformerTesterConfig, At
     pass
 
 
+class TestQwenImage21TransformerTaylorSeerCache(QwenImage21TransformerTesterConfig, TaylorSeerCacheTesterMixin):
+    """TaylorSeerCache tests for QwenImage 2.1 Transformer."""
+
+    def _run_denoising_loop(self, model, step_inputs, use_kv_cache):
+        from diffusers.models.transformers.transformer_qwenimage21 import QwenImage21KVCache
+
+        # Mirrors `QwenImage21Pipeline`: the first step prefills the cache, later steps only recompute the target.
+        kv_cache = QwenImage21KVCache(self.get_init_dict()["num_layers"]) if use_kv_cache else None
+        outputs = []
+        with torch.no_grad(), model.cache_context("cond"):
+            for index, inputs in enumerate(step_inputs):
+                kv_cache_mode = ("extract" if index == 0 else "cached") if use_kv_cache else None
+                output = model(**inputs, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, return_dict=False)[0]
+                outputs.append(output[:, -inputs["hidden_states"].shape[1] :])
+        return outputs
+
+    @pytest.mark.parametrize("use_lite_mode", [False, True])
+    def test_taylorseer_cache_with_kv_cache(self, use_lite_mode):
+        """
+        With the KV cache, the hooked modules see prefix + target tokens on the prefill step and target tokens only
+        afterwards. TaylorSeer must restart its expansion on the shape change instead of differencing the two, and
+        still predict the same target tokens as a loop that recomputes the full sequence every step.
+        """
+        from diffusers import TaylorSeerCacheConfig
+
+        num_steps = 8
+        base_inputs = self.get_dummy_inputs()
+        generator = torch.Generator("cpu").manual_seed(0)
+        step_inputs = [
+            dict(
+                base_inputs,
+                hidden_states=randn_tensor(
+                    base_inputs["hidden_states"].shape, generator=generator, device=torch_device
+                ),
+                timestep=torch.tensor([1.0 - index / num_steps], device=torch_device),
+            )
+            for index in range(num_steps)
+        ]
+
+        torch.manual_seed(0)
+        model = self.model_class(**self.get_init_dict()).to(torch_device).eval()
+        uncached = self._run_denoising_loop(model, step_inputs, use_kv_cache=False)
+
+        # Steps 0-2 compute fully, so the first cached steps refresh the factors before step 3 is predicted.
+        model.enable_cache(
+            TaylorSeerCacheConfig(
+                cache_interval=3,
+                disable_cache_before_step=3,
+                max_order=1,
+                taylor_factors_dtype=torch.float32,
+                use_lite_mode=use_lite_mode,
+            )
+        )
+        with_kv_cache = self._run_denoising_loop(model, step_inputs, use_kv_cache=True)
+        model._reset_stateful_cache()
+        without_kv_cache = self._run_denoising_loop(model, step_inputs, use_kv_cache=False)
+        model.disable_cache()
+
+        for step, (cached, reference) in enumerate(zip(with_kv_cache, without_kv_cache)):
+            assert cached.shape == reference.shape
+            torch.testing.assert_close(cached, reference, atol=2e-5, rtol=2e-5, msg=f"step {step} diverged")
+
+        # The prediction steps must actually be approximated, otherwise the comparison above proves nothing.
+        assert not torch.allclose(with_kv_cache[3], uncached[3], atol=1e-5)
+
+
 class TestQwenImage21TransformerSingleFile(QwenImage21TransformerTesterConfig, SingleFileTesterMixin):
     @property
     def ckpt_path(self):