Repository navigation
[hooks] Support Flux2 in MagCache - #14756
Merged
Merged
Conversation
github-actions
Bot
added
documentation
rootonchair
force-pushed
the
fix-magcache-flux2
branch
2 times, most recently
from
September 11, 2026 16:37
14c38d9 to
9d4b362
Compare
MagCache could not be enabled on `Flux2Transformer2DModel` (Flux2 and Flux2 Klein): 1. `Flux2TransformerBlock` and `Flux2SingleTransformerBlock` were not registered in `TransformerBlockRegistry`, so `apply_mag_cache` raised `ValueError: Model class ... not registered`. 2. Once registered, the tail hook raised a shape error. The head hook records the image-only input of `transformer_blocks[0]`, but the Flux2 model concatenates the text and image streams before the single-stream blocks and calls them with `encoder_hidden_states=None`, so the last single block returns the fused `[text; image]` sequence. The residual `tail_output - head_input` could not be formed. Register the two Flux2 blocks (the single block with `return_encoder_hidden_states_index=None`, matching how the model calls it) and add `_align_to_hidden_states`, which keeps only the trailing image tokens of the fused tensor when computing the residual. The cached residual is therefore image-only, which is what the reference MagCache4FLUX implementation does, and it matches the head input on skipped steps. This also replaces a dead "fallback" branch in the tail hook that subtracted mismatched shapes anyway. Tests: `MagCacheTesterMixin` for `Flux2Transformer2DModel`, `Flux2Pipeline` and `Flux2KleinPipeline`. FirstBlockCache has the same head/tail structure and still fails on Flux2; left for a follow-up.
rootonchair
force-pushed
the
fix-magcache-flux2
branch
2 times, most recently
from
September 11, 2026 16:47
25586ef to
4658765
Compare
github-actions
Bot
added
size/S
Implement the Flux2 residual fix inside the tail hook's sequence-length mismatch branch instead of a separate `_align_to_hidden_states` helper. The surrounding structure stays as it was: the `diff == 0` case still covers a batch-size mismatch via broadcasting, and only the mismatched-length case changes to subtract the trailing image tokens.
rootonchair
force-pushed
the
fix-magcache-flux2
branch
from
September 11, 2026 16:49
4658765 to
97696e9
Compare
rootonchair
marked this pull request as ready for review
September 13, 2026 06:06
Contributor
|
Hi @rootonchair, thanks for the PR! It does not appear to link an issue it fixes. If this PR addresses an existing issue, please add a closing keyword (e.g. Please note that PRs without a linked issue are likely to be automatically closed 10 days after this notice. Once the PR links an issue (or gets the |
Contributor
Author
|
Hi @DN6, can you help review this PR? |
DN6
approved these changes
Oct 14, 2026
|
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. |
rootonchair
added a commit
to rootonchair/diffusers
that referenced
this pull request
Oct 14, 2026
`QwenImage21TransformerBlock` was not registered in `TransformerBlockRegistry`, so enabling MagCache on `QwenImage21Transformer2DModel` raised `ValueError: Model class ... not registered`. Register it as a single-stream block (`return_hidden_states_index=0`, no encoder hidden states), matching how the model calls it. With the pipeline's KV cache, the hooked blocks see prefix + target tokens on the prefill step and only the target tokens afterwards. A step skipped right after the prefill therefore received a residual longer than its input and fell through to the "dimension mismatch" branch, returning the input without any residual; calibration crashed on the same shape change. Apply / compare the trailing (target) tokens of the residual instead, following the trailing-token alignment used for Flux2 in huggingface#14756. Tests: `MagCacheTesterMixin` for `QwenImage21Transformer2DModel` and `QwenImage21Pipeline`, plus KV-cache skip and calibration tests on the transformer. Refs: huggingface#14756
rootonchair
added a commit
to rootonchair/diffusers
that referenced
this pull request
Oct 14, 2026
`QwenImage21TransformerBlock` was not registered in `TransformerBlockRegistry`, so enabling MagCache on `QwenImage21Transformer2DModel` raised `ValueError: Model class ... not registered`. Register it as a single-stream block (`return_hidden_states_index=0`, no encoder hidden states), matching how the model calls it. With the pipeline's KV cache, the hooked blocks see prefix + target tokens on the prefill step and only the target tokens afterwards. A step skipped right after the prefill therefore received a residual longer than its input and fell through to the "dimension mismatch" branch, returning the input without any residual; calibration crashed on the same shape change. Apply / compare the trailing (target) tokens of the residual instead, following the trailing-token alignment used for Flux2 in huggingface#14756. Tests: `MagCacheTesterMixin` for `QwenImage21Transformer2DModel` and `QwenImage21Pipeline`, plus KV-cache skip and calibration tests on the transformer. Refs: huggingface#14756
rootonchair
added a commit
to rootonchair/diffusers
that referenced
this pull request
Oct 14, 2026
Match the Flux2 MagCache tests (huggingface#14756): drop the bespoke KV-cache test and keep the shared denoising helpers in the TaylorSeer test as on main.
5 of 10 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
MagCache could not be enabled on
Flux2Transformer2DModel(Flux2 and Flux2 Klein):Flux2TransformerBlockandFlux2SingleTransformerBlockwere not registered inTransformerBlockRegistry, soapply_mag_cacheraisedValueError: Model class ... not registered.transformer_blocks[0], but the Flux2 model concatenates the text and image streams before the single-stream blocks and calls them withencoder_hidden_states=None, so the last single block returns the fused[text; image]sequence. The residualtail_output - head_inputcould not be formed.Register the two Flux2 blocks (the single block with
return_encoder_hidden_states_index=None, matching how the model calls it) and add_align_to_hidden_states, which keeps only the trailing image tokens of the fused tensor when computing the residual and when adding it back on skipped steps. The cached residual is therefore image-only, which is what the reference MagCache4FLUX implementation does. This also replaces a dead "fallback" branch in the tail hook that subtracted mismatched shapes anyway.Tests: hook unit tests with a Flux2-shaped dummy (fused tail, skip and calibration), and
MagCacheTesterMixinforFlux2Transformer2DModel,Flux2PipelineandFlux2KleinPipeline.FirstBlockCache has the same head/tail structure and still fails on Flux2; left for a follow-up.
What does this PR do?
Fixes # (issue)
Before submitting
self-reviewskill on the diff?documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.