Fix TaylorSeer cache crash when the feature shape changes between steps

[hide3tu]

What does this PR do?

Fixes #14829

QwenImage21Pipeline with a TaylorSeer cache crashes on the second step when the KV cache is on (the default). With the KV cache, the transformer returns prefix + target tokens on the prefill step and target tokens only afterwards, so TaylorSeerState.update tries to difference two tensors of different sequence lengths.

This PR makes update restart the Taylor expansion from order 0 when the feature shape changes, instead of differencing against stale factors. Nothing changes when the shape is stable.

Tested on an M1 Max (MPS, bf16) only, on top of main at cc8644b. With this change the run completes, and the output is pixel-identical to use_kv_cache=False + TaylorSeer. For an edit with one reference image at 832×1216 / 40 steps, generation goes from 484 s (use_kv_cache=False + TaylorSeer) to 278 s (KV cache + TaylorSeer); the uncached baseline is 721 s.

If you would rather fix this on the model side (e.g. only passing target tokens through norm_out / proj_out on the prefill step), feel free to close this or push over it.

Self-review notes

I used Claude Code to help with this PR and ran the self-review skill on the diff.

• Blocking issues: none found.
• Non-blocking / left for review:
  1. The new test exercises TaylorSeerState directly. A pipeline-level regression test (tiny Qwen-Image 2.1 + TaylorSeer + KV cache) is not included, since testing.md says cache tests are added case by case.
  2. With the KV cache, the first predicted step still has to come after at least one full step in cached mode, so disable_cache_before_step needs to be >= 2. Lower values are not handled here.
  3. Only TaylorSeer was checked. Other cache hooks may hit the same shape change.
• Dead code: none.
• Verdict: READY

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue or the forum? [Qwen-Image 2.1] TaylorSeer cache crashes with the default KV cache (feature shape changes after the prefill step) #14829
• Did you make sure to update the documentation with your changes? (no public API change)
• Did you write any new necessary tests?

Who can review?

@sayakpaul

[sayakpaul]

I don't understand the test.

[sayakpaul]

And this works as expected?

[sayakpaul]

Thanks for the PR. Could you also provide a minimal working code snippet and expected results (along with the time savings)?

[sayakpaul]

/diffusers-bot pytest tests/models -k "taylorseer"

[github-actions]

✅ pytest tests/models -k "taylorseer" passed on GPU — view logs.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[DN6]

Can we structure the test in a similar manner to

diffusers/tests/models/transformers/test_models_transformer_flux.py

Line 618 in 8b3c707

class TestFluxTransformerTaylorSeerCache(FluxTransformerTesterConfig, TaylorSeerCacheTesterMixin):

[sayakpaul]

I think there's some confusion here. The current test follows what is here:
https://github.com/huggingface/diffusers/blob/main/tests/hooks/test_mag_cache.py

I guess what you're suggesting is to add a TestQwenImage21TaylorSeerCache class and add the respective QwenImage2.1 specific tests in it?

[hide3tu]

Thanks for running the tests! Here is a minimal snippet (the same one as in #14829):

import torch
from diffusers import QwenImage21Pipeline, TaylorSeerCacheConfig

device = "cuda" if torch.cuda.is_available() else "mps"
pipe = QwenImage21Pipeline.from_pretrained("Qwen/Qwen-Image-2.1", torch_dtype=torch.bfloat16).to(device)
pipe.transformer.enable_cache(
    TaylorSeerCacheConfig(cache_interval=3, disable_cache_before_step=3, max_order=1, use_lite_mode=True)
)
image = pipe(
    prompt="a cat sitting on a wall",
    width=512,
    height=512,
    num_inference_steps=8,
    true_cfg_scale=1.0,
    generator=torch.Generator("cpu").manual_seed(0),
).images[0]  # use_kv_cache defaults to True
image.save("out.png")

Expected results (M1 Max, MPS, bf16):

• On main (cc8644b): crashes at step 2 with RuntimeError: The size of tensor a (1024) must match the size of tensor b (1038) at non-singleton dimension 1.
• On this branch: the run completes. The output is byte-identical (same MD5) to the workaround from [Qwen-Image 2.1] TaylorSeer cache crashes with the default KV cache (feature shape changes after the prefill step) #14829, and that workaround's output is pixel-identical to use_kv_cache=False + TaylorSeer.

Time savings. These numbers were measured with the same fix applied as a monkeypatch (the workaround in #14829), not with this branch installed. It gives identical outputs, as above. Edit with one reference image, 832×1216, 40 steps, seed 42, cache_interval=3, one run each:

setting	generation time
no cache	721 s
TaylorSeer, use_kv_cache=False	484 s
TaylorSeer + KV cache (crashes on main, runs with this fix)	278 s

For text-to-image at the same size, it goes from 615 s (no cache) to 219 s (TaylorSeer + KV cache).

The scripts I used are here: https://github.com/hide3tu/LiltingChannelLabo/tree/main/2026/09/22/qwen-image-2-1-taylorseer-cache-m1-max (repro_issue_14829.py is the snippet above, generate.py is the timing script).

[sayakpaul]

Thanks! Can you also show the actual outputs?