Add First Block Cache support for Flux 2

[guptaishaan]

Fixes #14280

Flux2TransformerBlock and Flux2SingleTransformerBlock were never registered in
TransformerBlockRegistry, so enable_cache(FirstBlockCacheConfig(...)) raised
ValueError: Model class ... not registered. Registering them is not enough on its own: the model
forward did the txt/img cat once between the two block loops, so the FBCache head block (double,
2-tuple of separate streams) and tail block (single, bare concatenated tensor) had incompatible
outputs and the residual bridge failed with TypeError: unsupported operand type(s) for -: 'Tensor' and 'list'.

Changes:

• Register both Flux 2 block classes with the same indices as Flux v1
  (return_hidden_states_index=1, return_encoder_hidden_states_index=0).
• Move the cat/split into Flux2SingleTransformerBlock.forward, matching
  FluxSingleTransformerBlock. Every Flux 2 block now takes (hidden_states, encoder_hidden_states) and returns (encoder_hidden_states, hidden_states). The block already
  carried split_hidden_states / text_seq_len parameters for this; nothing ever passed them, so
  they are removed.
• Add TestFlux2TransformerFBCCache and FirstBlockCacheTesterMixin on
  Flux2KleinPipelineFastTests.

This is direction (A) from the issue. Direction (B) would put one model's double/single boundary
into the generic apply_first_block_cache.

Verified on 8x A40 (sm_86), torch 2.13.0+cu126, bf16, with dummy models:

• The 7 new tests fail before the patch and pass after.
• With only the registration half applied, the failure moves to the Tensor - list bridge error, so
  the model change is load-bearing.
• The model refactor is bit-for-bit identical on all three forward paths: plain, kv_cache_mode="extract",
  kv_cache_mode="cached". The Klein KV-cache path including ref-token modulation blending is
  unchanged.
• Cost of the extra per-block cat: 288.0 / 288.7 ms per forward with the patch vs 287.6 / 287.7 ms
  without, on an idle A40 at 8 double + 48 single blocks, dim 3072, 1024 image + 512 text tokens.
  Noise floor.
• tests/models/transformers/test_models_transformer_flux2.py has an identical failure set before
  and after (10 torch.compile tests that die in the inductor C++ backend on this box, unrelated).
  tests/hooks/ and the Flux v1 cache tests pass.

Not verified: no pretrained FLUX.2 checkpoint was loaded, so there is no image-quality check, and no
FBCache speedup numbers. All torch.compile tests fail on this machine for an unrelated toolchain
reason, before and after, so the compiled path is untested. Flux2Pipeline (non-Klein) still does
not open a cache_context, so FBCache raises No context is set there; that is a separate gap and
I left it alone.

Reported by @sqhuang, whose issue also worked out the root cause and both candidate directions.

[guptaishaan]

Rebased onto current main (6f2010e).

The only conflict was in tests/pipelines/flux2/test_pipeline_flux2_klein.py, which #14336 refactored to the new config + mixin structure while this PR was open. I dropped my old edit (adding FirstBlockCacheTesterMixin to Flux2KleinPipelineFastTests) and re-expressed it in the new layout as a separate test class, matching how Flux v1 does it in tests/pipelines/flux/test_pipeline_flux.py:

class TestFlux2KleinPipelineFirstBlockCache(Flux2KleinPipelineTesterConfig, FirstBlockCacheTesterMixin):
    """First Block Cache tests for the Flux2 Klein pipeline."""

The import moved from ..test_pipelines_common to ..testing_utils accordingly. No other file conflicted, and the src/ change is unchanged from before.

Re-ran on the rebased tree (8x A40, sm_86):

• tests/models/transformers/test_models_transformer_flux2.py -k FBCCache -> 6 passed
• tests/pipelines/flux2/test_pipeline_flux2_klein.py -k FirstBlockCache -> 1 passed
• tests/hooks/ -> 53 passed, 1 failed

The one tests/hooks/ failure is test_mag_cache_calibration dying on OSError: No space left on device, which is this box running out of disk, not the patch. The same disk problem also makes a few save_pretrained-based tests in the Klein pipeline file error out on both main and this branch.