Tech7 min read

Qwen-Image 2.1 + TaylorSeer fix merged in Diffusers: 4-day turnaround

IkesanContents

In a previous post on speeding up Qwen-Image 2.1 using TaylorSeer caching, running TaylorSeer alongside the default KV cache triggered a tensor shape mismatch at step 2. We initially worked around this by monkey-patching TaylorSeerState.update.

After reporting the issue and opening a pull request with the same patch upstream to Diffusers, the fix was merged into the main branch on September 25. The entire turnaround took under four days from the initial issue on the night of September 21.

The Original Issue

The Qwen-Image 2.1 pipeline enables KV caching (use_kv_cache=True) by default.
The prompt and reference image portions are computed and cached during step 1; from step 2 onward, only the target image tokens are computed.
Because of this design, the transformer output length is longer only in the first step. For a 512×512 image generation run, step 1 outputs 1,038 tokens (14 text tokens plus 1,024 image tokens), while step 2 and subsequent steps output 1,024 tokens.

TaylorSeer predicts the next step value by computing the difference between consecutive steps. When it attempted to subtract two tensors of different sequence lengths, it threw an error:

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

Setting use_kv_cache=False worked around the crash, but during image-to-image editing, recalculating the reference image tokens at every single step caused editing time on an M1 Max (832×1216 resolution, 40 steps) to jump from 277.6 seconds to 484.4 seconds.

From Issue to Pull Request

On September 21 at 22:17 JST, I opened issue #14829.
The issue included reproduction code, full error logs, the monkey-patch workaround, and benchmark numbers on an M1 Max. I noted that verification was limited to Apple Silicon (MPS) and TaylorSeer.
I also noted that fixing this on the model side might be cleaner: because the Qwen-Image 2.1 pipeline discards prompt tokens from the step 1 output anyway, passing only image tokens through the final projection layer would keep tensor lengths uniform.

Just 38 minutes later, maintainer Sayak Paul replied asking if I could submit a pull request.
I opened PR #14831 at 23:45 JST that same evening.

The fix and PR description were drafted using Claude Code.
The Diffusers contribution guide includes a dedicated section on contributing with AI agents. The workflow requirements applicable to this PR are summarized below:

RequirementWhat We Did
Discuss the plan in an issue first and wait for explicit maintainer consent before opening a PROpened the PR only after receiving explicit maintainer encouragement in the issue
Run the repository’s self-review skill to inspect diffsExecuted the self-review skill via Claude Code
Include self-review findings in the PR description or commentStated open questions and validation status directly in the PR description

self-review is an agent skill located in .ai/skills/ within the repository. It audits diffs against the same review standards used by the Claude-powered automated review running in Diffusers CI.
Our run reported no blockers and passed with status READY.
The open points raised for maintainer review were that end-to-end pipeline tests were not included, combining with KV cache requires setting disable_cache_before_step to at least 2 before predictions can begin, and caching methods beyond TaylorSeer were unverified.

The Upstream Fix

The upstream change modified a single location in src/diffusers/hooks/taylorseer_cache.py, adding one condition and three lines of comments:

# 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

TaylorSeer tracks zeroth-order terms (raw output) and first-order or higher derivative terms computed from differences between consecutive steps.
Before the fix, it attempted to difference against the previous state whenever one was present. With this fix, if tensor shapes mismatch, it halts expansion and stores the current output as a fresh zeroth-order baseline. Unchanged shapes behave exactly as before.

The monkey patch from the earlier article discarded stored states and the last updated step count whenever shapes changed before invoking the original logic.
While the implementation approaches differ, both treat shape-changed steps as fresh initial steps. Images generated under identical conditions showed matching MD5 hashes down to the byte.

Review Requirements

Maintainers raised few concerns about the core fix itself; the focus was on verification and regression testing. Timestamps below are in JST.

Date / TimeMaintainer RequestWhat Was Provided
Sep 22, ~18:00Unit test clarification (“I do not understand this test”) and confirmation on the modified line. Requested minimal repro code, expected behavior, and timing improvementsProvided identical repro code from the issue, unpatched main error log, and benchmark table on Sep 23 at 00:11
Sep 23, 00:14Visual proof requested with sample output imagesProvided comparison grid on Sep 25 at 00:03 comparing un-cached vs. TaylorSeer + KV cache outputs
Sep 25, 12:42Requested a dedicated Qwen-Image 2.1 test class mirroring the existing Flux TaylorSeer test suiteAdded dedicated test class in Qwen-Image 2.1 transformer test file on Sep 25 at 15:01

I clarified that benchmark numbers came from the identical monkey-patch script rather than fresh runs on the PR branch, as output bit-identity was already confirmed.

The initial test simply passed mismatched tensor shapes directly into TaylorSeer’s internal state to ensure it handled differences without crashing.
Another maintainer suggested patterning tests after Flux, which Sayak Paul clarified as adding a dedicated test class for Qwen-Image 2.1.

The resulting test class inherits the shared TaylorSeer test suite (matching Flux’s structure) and adds a KV cache integration test.
This test runs a lightweight Qwen-Image 2.1 transformer through an 8-step denoising loop, caching prefix tokens at step 1 and enabling KV cache from step 2 onward. It verifies that outputs match between KV-cached and un-cached runs across all steps. On unpatched main, this reproduces the exact issue crash; with the fix, it passes cleanly.

Because KV-cached and un-cached runs will trivially match if TaylorSeer never engages, the test also asserts that step 4 output differs from an un-cached run.
Sayak Paul noted on this check: “This is a very meaningful assertion!”

The PR received approval at 15:31 JST. After maintainers verified tests on GPU runners via PR comments, it was merged at 16:04 JST.

Using the Fix

As of September 26, the latest Diffusers release is v0.40.0 (released August 20), which does not yet include this patch.
Until the next release, install Diffusers directly from GitHub main or use the monkey patch from our previous post:

pip install git+https://github.com/huggingface/diffusers

The fix landed in TaylorSeer’s hook implementation rather than the Qwen-Image 2.1 model code. The alternative model-side approach mentioned in the original issue was not discussed during review.

References