huggingface / huggingface/diffusers

[Cache] Add First Block Cache (FBCache) support for Flux 2 (Flux2Transformer2DModel)

オープン
#14,280 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

**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).

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

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.

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning, performance
issue の種類
機能追加
難易度
5/5
見積もり時間
1週間以上
活発さ
静か
明瞭さ
おおむね明確
初心者へのやさしさ
42/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。