[Qwen-Image 2.1] TaylorSeer cache crashes with the default KV cache (feature shape changes after the prefill step) #14829

Describe the bug

QwenImage21Pipeline crashes on the second denoising step when a TaylorSeer cache is enabled on the transformer and the KV cache is left at its default (use_kv_cache=True).

With the KV cache, the first transformer call (kv_cache_mode="extract", prefill) returns the whole joint sequence (prefix + target tokens), while the following calls (kv_cache_mode="cached") return target tokens only. The pipeline handles this by slicing noise_pred[:, -latents.size(1):], but the TaylorSeer hook on proj_out sees both shapes and fails when it takes the divided difference between step 0 and step 1:

RuntimeError: The size of tensor a (1024) must match the size of tensor b (1038) at non-singleton dimension 1

(512×512 → 1024 target tokens; 1038 = 14 text tokens + 1024. At 832×1216 it is 3952 vs 4075.)

Passing use_kv_cache=False avoids the crash, but for image-conditioned edits that means the reference-image tokens are recomputed every step, so most of the speed-up is lost (numbers below).

The mismatch comes from the sequence length, not from the device, so I expect the same on CUDA — but I have only tested on Apple Silicon (MPS). I have also only tried TaylorSeerCacheConfig; other cache hooks may or may not be affected.

Reproduction

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 -> crashes at step 2

Workaround that worked for me (not proposing this as the fix, just sharing what made it run): reset the Taylor factors when the feature shape changes, so the prefill step is simply not used for the finite difference.

import diffusers.hooks.taylorseer_cache as ts

_orig_update = ts.TaylorSeerState.update

def _update(self, outputs):
    if not self.is_inactive and self.last_update_step is not None:
        for i, feat in enumerate(outputs):
            prev = self.taylor_factors.get(i, {}).get(0)
            if prev is not None and prev.shape != feat.shape:
                self.taylor_factors = {}
                self.last_update_step = None
                break
    return _orig_update(self, outputs)

ts.TaylorSeerState.update = _update

With this patch the run completes, and the output is pixel-identical to the use_kv_cache=False + TaylorSeer output (mean abs diff 0.0 on both cases below).

Timings on an M1 Max 64GB (MPS, bf16, 832×1216, 40 steps, seed 42, true_cfg_scale=1.0, 15 full transformer passes out of 40 with the config above; "generation" = pipe(...) call only, single run each):

For the edit case, disabling the KV cache raises a full step from ~17.4 s to ~31.6 s, which is why keeping the KV cache matters.

A cleaner fix may be on the model side (e.g. only passing target tokens through norm_out/proj_out in the prefill step, since the pipeline discards the prefix part anyway), but I'll leave that to you.

Logs

  File ".../diffusers/hooks/hooks.py", line 217, in new_forward
    output = function_reference.forward(*args, **kwargs)
  File ".../diffusers/hooks/taylorseer_cache.py", line 240, in new_forward
    state.update(wrapped_outputs)
  File ".../diffusers/hooks/taylorseer_cache.py", line 154, in update
    new_factors[j + 1] = (new_factors[j] - prev.to(features.dtype)) / delta_step
                          ~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~
RuntimeError: The size of tensor a (1024) must match the size of tensor b (1038) at non-singleton dimension 1

System Info

• 🤗 Diffusers version: 0.41.0.dev0 (installed from commit 80c7ed2; pipeline_qwenimage21.py, transformer_qwenimage21.py and hooks/taylorseer_cache.py have no newer commits on main as of 2026-09-21)
• Platform: macOS-27.0-arm64-arm-64bit
• Running on Google Colab?: No
• Python version: 3.12.12
• PyTorch version (GPU?): 2.14.0 (False)
• Huggingface_hub version: 1.32.0
• Transformers version: 5.17.0
• Accelerate version: 1.15.0
• PEFT version: not installed
• Safetensors version: 0.8.0
• xFormers version: not installed
• Accelerator: Apple M1 Max (64GB unified memory)
• Using GPU in script?: Yes (MPS)
• Using distributed or parallel set-up in script?: No

Who can help?

@sayakpaul @naykun