huggingface / huggingface/diffusers
[Cache] Add First Block Cache (FBCache) support for Flux 2 (Flux2Transformer2DModel)
- Vorherrschende Sprache
- Python
- Sterne
- 34.5k
- Forks
- 7.3k
- Ø Merge
- 3 T. 3 Std.
- Gemergte PRs (30 T.)
- 91
Beschreibung
**Is your feature request related to a problem? Please describe.**
`FirstBlockCache` works for other DiT image models (e.g. Qwen-Image, Flux, CogView4) via the generic `enable_cache` API, but it fails for Flux 2 (the Klein pipelines that use `Flux2Transformer2DModel`):
```python
import torch
from diffusers import Flux2KleinPipeline, FirstBlockCacheConfig
pipe = Flux2KleinPipeline.from_pretrained("black-forest-labs/FLUX.2-klein-base-9B", torch_dtype=torch.bfloat16)
pipe.transformer.enable_cache(FirstBlockCacheConfig(threshold=0.2))
# raises: ValueError: Model class not registered.
```
`Flux2Transformer2DModel` already inherits `CacheMixin`, so `enable_cache` is exposed and dispatches to `apply_first_block_cache`. The failure happens immediately at `enable_cache` time (not during a forward pass): `apply_first_block_cache` registers `FBCHeadBlockHook`/`FBCBlockHook` on the blocks, and each hook's `initialize_hook` calls `TransformerBlockRegistry.get(block.__class__)` (`hooks/first_block_cache.py`). `Flux2TransformerBlock` and `Flux2SingleTransformerBlock` are not present in `_register_transformer_blocks_metadata()` (`hooks/_helpers.py`), so the lookup raises.
**Describe the solution you'd like.**
Add FBCache support for Flux 2. Concretely this needs:
1. Register the Flux 2 block classes in `TransformerBlockRegistry` (`hooks/_helpers.py`). The double-stream block `Flux2TransformerBlock` returns `(encoder_hidden_states, hidden_states)`, matching Flux v1's `return_hidden_states_index=1, return_encoder_hidden_states_index=0`.
2. Handle the double/single block-shape asymmetry (see below), which registration alone does not resolve.
3. Add the standard cache tests (`FirstBlockCacheTesterMixin`) for the Flux 2 transformer/pipelines, and list Flux 2 in the caching docs.
**The non-trivial part — Flux 2's block I/O is heterogeneous, unlike every currently supported model.**
`apply_first_block_cache` collects all block `ModuleList`s (`transformer_blocks` + `single_transformer_blocks`, both already valid identifiers in `_ALL_TRANSFORMER_BLOCK_IDENTIFIERS`) into one list, then bridges a residual from the head block (first double block) to the tail block (last single block). This assumes a homogeneous per-block I/O contract, which holds for all currently registered models (Flux v1, Qwen-Image, HunyuanVideo, … all keep `hidden_states`/`encoder_hidden_states` as a consistent tuple across every block). Flux 2 breaks that assumption:
- **Flux v1** does the txt/img `cat` and `split` *inside* each single block, so `FluxSingleTransformerBlock.forward` still returns `(encoder_hidden_states, hidden_states)` — uniform with the double blocks (`transformer_flux.py:405-406`).
- **Flux 2** does the `cat` *once in the model forward* between the two block loops (`transformer_flux2.py:1330`); the single blocks run on the concatenated `[txt, img]` sequence and return a **bare tensor** (`transformer_flux2.py:849-853`), with the final `split` also done in the model forward (`transformer_flux2.py:1372`).
So the FBCache head block (double → 2-tuple over separate streams, seq = `num_img`) and tail block (single → bare tensor over the concatenated stream, seq = `num_txt + num_img`) disagree in both structure (list vs tensor) and sequence length, and the `cat`/`split` boundary lives outside the hooked blocks where FBCache cannot see it. Registering the classes alone would move the failure from `enable_cache` to a shape/type error (or silently wrong output) inside the head/tail residual bridge.
Two directions, and I'd like maintainer guidance on the preferred one:
- **(A)** Bring Flux 2 single blocks in line with the Flux v1 pattern (return `(encoder_hidden_states, hidden_states)` per block), so the whole stack is homogeneous and FBCache "just works" after registration. This also matches the design intent stated in #11180, where the FBCache author notes that if all blocks take `hidden_states`/`encoder_hidden_states` and always return `(hidden_states, encoder_hidden_states)`, a lot of the hook-based design is simplified. Downside: this touches the model forward and must not regress the Klein KV-cache path (`kv_cache_mode` extract/cached, ref-token modulation blending), which relies on the single-stream concatenated layout.
- **(B)** Keep the model as-is and make FBCache aware of the double→single boundary (e.g. apply the cache over a homogeneous sub-list, or special-case the head/tail selection for this architecture).
**Describe alternatives you've considered.**
- Using another supported cache method (FasterCache/PAB/TaylorSeer) instead — but those have the same per-block registration/assumption story and don't sidestep the underlying heterogeneity.
- Applying FBCache only over `single_transformer_blocks` (homogeneous bare-tensor blocks) and leaving the double blocks uncached — simpler, but likely leaves speedup on the table; worth measuring.
**Additional context.**
- FBCache was introduced in #11180.
- Flux 2 model/pipeline review meta-issue: #13579 (does not currently track FBCache/`enable_cache`).
- Happy to open the PR once we agree on direction (A) vs (B).
Beitragsleitfaden
Rechercherichtung
Start with hooks/_helpers.py, hooks/first_block_cache.py, transformer_flux2.py, and the existing FirstBlockCacheTesterMixin tests to trace block registration and the double-to-single boundary. Compare the proposed homogeneous-block and boundary-aware approaches, then verify Flux 2 transformer and pipeline cache tests, caching documentation, and the Klein KV-cache path without regressions.
Vom Indexierungsmodell aus dem Issue-Text verfasst.
Bewertung
- Tech-Stack
- python, pytorch
- Bereich
- machine-learning, performance
- Issue-Typ
- Feature
- Schwierigkeit
- 5/5
- Geschätzter Aufwand
- Über eine Woche
- Aktivitätsstatus
- Ruhig
- Klarheit
- Größtenteils klar
- Anfängerfreundlichkeit
- 42/100