[Feat] support TeaCache for Flux2 klein (#1234)

Signed-off-by: Lancer <maruixiang6688@gmail.com>

docs/user_guide/diffusion_acceleration.md

@@ -69,7 +69,7 @@ The following table shows which models are currently supported by each accelerat
 | **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
-| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 ### VideoGen

tests/diffusion/cache/test_teacache_extractors.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Unit tests for TeaCache extractor functions.
+
+This module provides a generic testing framework for model-specific extractor functions
+used by TeaCache. Each model's extractor can be tested by:
+1. Creating a fixture that returns model module
+2. Creating a fixture that returns sample inputs for that model
+3. Creating a test class that inherits from BaseExtractorTest
+4. Implementing any model-specific test methods
+
+Currently implemented:
+- TestFlux2KleinExtractor: Flux2Klein model extractor
+"""
+
+from abc import ABC, abstractmethod
+from unittest.mock import MagicMock, Mock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_klein_context
+from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
+    Flux2Transformer2DModel,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_tp_group():
+    """Set up TP group for each test function"""
+    with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1):
+        with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group:
+            mock_tp_group = MagicMock()
+            mock_tp_group.world_size = 1
+            mock_get_tp_group.return_value = mock_tp_group
+            yield
+
+
+class BaseExtractorTest(ABC):
+    """Base class for testing TeaCache extractors.
+
+    Subclasses should implement:
+    - get_extractor(): Return extractor function
+    - get_module(): Return model module
+    - get_sample_inputs(): Return sample inputs for model
+    """
+
+    @abstractmethod
+    def get_extractor(self):
+        """Return extractor function to test."""
+        pass
+
+    @abstractmethod
+    def get_module(self):
+        """Return model module instance."""
+        pass
+
+    @abstractmethod
+    def get_sample_inputs(self):
+        """Return sample inputs for model."""
+        pass
+
+
+class TestFlux2KleinExtractor(BaseExtractorTest):
+    """Test extract_flux2_klein_context function."""
+
+    def get_extractor(self):
+        return extract_flux2_klein_context
+
+    @pytest.fixture
+    def flux2_klein_module(self):
+        """Create a minimal Flux2Transformer2DModel for testing."""
+        model = Flux2Transformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=48,
+            attention_head_dim=128,
+            joint_attention_dim=15360,
+        )
+        return model
+
+    def get_module(self, flux2_klein_module):
+        return flux2_klein_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        """Create sample input tensors for Flux2Klein.
+
+        Note: hidden_states uses in_channels=128 (default for Flux2Klein),
+        not inner_dim=6144. The x_embedder projects from 128 -> 6144.
+        encoder_hidden_states uses joint_attention_dim=15360 (model default),
+        which then gets projected to inner_dim=6144 by context_embedder.
+        """
+        batch_size = 1
+        img_seq_len = 1024
+        txt_seq_len = 512
+        in_channels = 128  # Model default in_channels
+        txt_dim = 15360  # Model default joint_attention_dim
+
+        return {
+            "hidden_states": torch.randn(batch_size, img_seq_len, in_channels),
+            "encoder_hidden_states": torch.randn(batch_size, txt_seq_len, txt_dim),
+            "timestep": torch.tensor([500]),
+            "img_ids": torch.randint(0, 64, (batch_size, img_seq_len, 4)),
+            "txt_ids": torch.randint(0, 64, (batch_size, txt_seq_len, 4)),
+            "guidance": torch.tensor([3.5]),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    def test_modulated_input_shape(self, flux2_klein_module, sample_inputs):
+        """Test that modulated_input has correct shape matching the model's inner_dim.
+
+        Note: After x_embedder projection, hidden_states are projected from
+        in_channels (128) to inner_dim (6144), so modulated_input should match
+        the projected shape, not the input shape.
+        """
+        context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
+
+        batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
+        inner_dim = flux2_klein_module.inner_dim
+        assert context.modulated_input.shape == (batch_size, img_seq_len, inner_dim)
+
+    def test_run_transformer_blocks_callable(self, flux2_klein_module, sample_inputs):
+        """Test that run_transformer_blocks is callable."""
+        context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
+        assert callable(context.run_transformer_blocks)
+
+    def test_postprocess_callable(self, flux2_klein_module, sample_inputs):
+        """Test that postprocess is callable."""
+        context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
+        assert callable(context.postprocess)
+
+    def test_extra_states_contains_full_transformer(self, flux2_klein_module, sample_inputs):
+        """Test that extra_states contains run_flux2_full_transformer_with_single."""
+        context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
+
+        assert context.extra_states is not None
+        assert "run_flux2_full_transformer_with_single" in context.extra_states
+        assert callable(context.extra_states["run_flux2_full_transformer_with_single"])
+
+    def test_without_guidance(self, flux2_klein_module, sample_inputs):
+        """Test context extraction works without guidance (no CFG)."""
+        inputs = sample_inputs.copy()
+        inputs["guidance"] = None
+
+        context = extract_flux2_klein_context(flux2_klein_module, **inputs)
+
+        assert context is not None
+        assert context.temb is not None
+
+    def test_invalid_module_raises_error(self):
+        """Test that invalid module without transformer_blocks raises ValueError."""
+        invalid_module = Mock()
+        invalid_module.transformer_blocks = []
+
+        with pytest.raises(ValueError, match="Module must have transformer_blocks"):
+            extract_flux2_klein_context(
+                invalid_module,
+                hidden_states=torch.randn(1, 1024, 6144),
+                encoder_hidden_states=torch.randn(1, 512, 15360),
+                timestep=torch.tensor([500]),
+                img_ids=torch.randint(0, 64, (1, 1024, 4)),
+                txt_ids=torch.randint(0, 64, (1, 512, 4)),
+            )

vllm_omni/diffusion/cache/teacache/backend.py

@@ -48,7 +48,29 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     )
 
 
-CUSTOM_TEACACHE_ENABLERS = {"BagelPipeline": enable_bagel_teacache}
+def enable_flux2_klein_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
+    """
+    Enable TeaCache for Flux2 Klein model.
+    """
+    teacache_config = TeaCacheConfig(
+        transformer_type="Flux2Klein",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    transformer = pipeline.transformer
+
+    apply_teacache_hook(transformer, teacache_config)
+
+    logger.info(
+        f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
+        f"transformer_class={teacache_config.transformer_type}"
+    )
+
+
+CUSTOM_TEACACHE_ENABLERS = {
+    "BagelPipeline": enable_bagel_teacache,
+    "Flux2KleinPipeline": enable_flux2_klein_teacache,
+}
 
 
 class TeaCacheBackend(CacheBackend):

vllm_omni/diffusion/cache/teacache/config.py

@@ -15,6 +15,15 @@ _MODEL_COEFFICIENTS = {
         -3.82021401e00,
         2.64230861e-01,
     ],
+    # Flux2 Klein transformer coefficients
+    # Same as FLUX.1 (similar dual-stream architecture)
+    "Flux2Klein": [
+        4.98651651e02,
+        -2.83781631e02,
+        5.58554382e01,
+        -3.82021401e00,
+        2.64230861e-01,
+    ],
     # Qwen-Image transformer coefficients from ComfyUI-TeaCache
     # Tuned specifically for Qwen's dual-stream transformer architecture
     # Used for all Qwen-Image Family pipelines, in general

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -566,6 +566,149 @@ def extract_zimage_context(
     )
 
 
+def extract_flux2_klein_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor | None = None,
+    timestep: torch.LongTensor = None,
+    img_ids: torch.Tensor = None,
+    txt_ids: torch.Tensor = None,
+    guidance: torch.Tensor | None = None,
+    joint_attention_kwargs: dict[str, Any] | None = None,
+    **kwargs: Any,
+) -> CacheContext:
+    """
+    Extract cache context for Flux2Klein model.
+
+    Caches the full transformer output (including single_transformer_blocks).
+    When cache is reused, single_transformer_blocks is skipped to achieve maximum speedup.
+
+    Args:
+        module: Flux2Transformer2DModel instance
+        hidden_states: Input image hidden states tensor
+        encoder_hidden_states: Input text hidden states tensor
+        timestep: Current diffusion timestep
+        img_ids: Image position IDs for RoPE
+        txt_ids: Text position IDs for RoPE
+        guidance: Optional guidance scale for CFG
+        joint_attention_kwargs: Additional attention kwargs
+
+    Returns:
+        CacheContext with all information needed for generic caching
+    """
+    from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+    if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
+        raise ValueError("Module must have transformer_blocks")
+
+    # ============================================================================
+    # PREPROCESSING (Flux2-specific)
+    # ============================================================================
+    dtype = hidden_states.dtype
+
+    num_txt_tokens = encoder_hidden_states.shape[1]
+
+    timestep = timestep.to(dtype=dtype) * 1000
+    if guidance is not None:
+        guidance = guidance.to(dtype=dtype) * 1000
+
+    temb = module.time_guidance_embed(timestep, guidance)
+
+    double_stream_mod_img = module.double_stream_modulation_img(temb)
+    double_stream_mod_txt = module.double_stream_modulation_txt(temb)
+    single_stream_mod = module.single_stream_modulation(temb)[0]
+
+    hidden_states = module.x_embedder(hidden_states)
+    encoder_hidden_states = module.context_embedder(encoder_hidden_states)
+
+    if img_ids.ndim == 3:
+        img_ids = img_ids[0]
+    if txt_ids.ndim == 3:
+        txt_ids = txt_ids[0]
+
+    image_rotary_emb = module.pos_embed(img_ids)
+    text_rotary_emb = module.pos_embed(txt_ids)
+    concat_rotary_emb = (
+        torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
+        torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
+    )
+
+    # ============================================================================
+    # EXTRACT MODULATED INPUT (for cache decision)
+    # ============================================================================
+    block = module.transformer_blocks[0]
+
+    norm_hidden_states = block.norm1(hidden_states)
+    norm_hidden_states = (1 + double_stream_mod_img[0][1]) * norm_hidden_states + double_stream_mod_img[0][0]
+
+    modulated_input = norm_hidden_states
+
+    # ============================================================================
+    # DEFINE TRANSFORMER EXECUTION (Flux2-specific)
+    # ============================================================================
+    def run_flux2_transformer_blocks():
+        h = hidden_states
+        c = encoder_hidden_states
+        for block in module.transformer_blocks:
+            c, h = block(
+                hidden_states=h,
+                encoder_hidden_states=c,
+                temb_mod_params_img=double_stream_mod_img,
+                temb_mod_params_txt=double_stream_mod_txt,
+                image_rotary_emb=concat_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+        return (h, c)
+
+    def run_flux2_full_transformer_with_single(ori_h, ori_c):
+        h = ori_h
+        c = ori_c
+        for block in module.transformer_blocks:
+            c, h = block(
+                hidden_states=h,
+                encoder_hidden_states=c,
+                temb_mod_params_img=double_stream_mod_img,
+                temb_mod_params_txt=double_stream_mod_txt,
+                image_rotary_emb=concat_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+        h_concat = torch.cat([c, h], dim=1)
+        for block in module.single_transformer_blocks:
+            h_concat = block(
+                hidden_states=h_concat,
+                encoder_hidden_states=None,
+                temb_mod_params=single_stream_mod,
+                image_rotary_emb=concat_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+        final_hidden_states = h_concat[:, num_txt_tokens:, ...]
+        return final_hidden_states, c
+
+    # ============================================================================
+    # DEFINE POSTPROCESSING (Flux2-specific)
+    # ============================================================================
+    return_dict = kwargs.get("return_dict", True)
+
+    def postprocess(h):
+        h = module.norm_out(h, temb)
+        h = module.proj_out(h)
+        if not return_dict:
+            return (h,)
+        return Transformer2DModelOutput(sample=h)
+
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        temb=temb,
+        run_transformer_blocks=run_flux2_transformer_blocks,
+        postprocess=postprocess,
+        extra_states={
+            "run_flux2_full_transformer_with_single": run_flux2_full_transformer_with_single,
+        },
+    )
+
+
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -576,6 +719,7 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "QwenImageTransformer2DModel": extract_qwen_context,
     "Bagel": extract_bagel_context,
     "ZImageTransformer2DModel": extract_zimage_context,
+    "Flux2Klein": extract_flux2_klein_context,
     # Future models:
     # "FluxTransformer2DModel": extract_flux_context,
     # "CogVideoXTransformer3DModel": extract_cogvideox_context,

vllm_omni/diffusion/cache/teacache/hook.py

@@ -157,20 +157,26 @@ class TeaCacheHook(ModelHook):
                 ctx.encoder_hidden_states.clone() if ctx.encoder_hidden_states is not None else None
             )
 
-            # Run transformer blocks using model-specific callable
-            outputs = ctx.run_transformer_blocks()
+            # Handle models with additional blocks (e.g., Flux2 single_transformer_blocks)
+            if getattr(ctx, "extra_states", None) and "run_flux2_full_transformer_with_single" in ctx.extra_states:
+                run_full = ctx.extra_states["run_flux2_full_transformer_with_single"]
+                ctx.hidden_states, ctx.encoder_hidden_states = run_full(ori_hidden_states, ori_encoder_hidden_states)
+                output = ctx.hidden_states
+                state.previous_residual = (ctx.hidden_states - ori_hidden_states).detach()
+            else:
+                # Run transformer blocks using model-specific callable
+                outputs = ctx.run_transformer_blocks()
+                # Update context with outputs
+                ctx.hidden_states = outputs[0]
+                if len(outputs) > 1 and ctx.encoder_hidden_states is not None:
+                    ctx.encoder_hidden_states = outputs[1]
 
-            # Update context with outputs
-            ctx.hidden_states = outputs[0]
-            if len(outputs) > 1 and ctx.encoder_hidden_states is not None:
-                ctx.encoder_hidden_states = outputs[1]
+                output = ctx.hidden_states
 
-            # Cache residuals for next timestep
-            state.previous_residual = (ctx.hidden_states - ori_hidden_states).detach()
-            if ori_encoder_hidden_states is not None:
-                state.previous_residual_encoder = (ctx.encoder_hidden_states - ori_encoder_hidden_states).detach()
-
-            output = ctx.hidden_states
+                # Cache residuals for next timestep
+                state.previous_residual = (ctx.hidden_states - ori_hidden_states).detach()
+                if ori_encoder_hidden_states is not None:
+                    state.previous_residual_encoder = (ctx.encoder_hidden_states - ori_encoder_hidden_states).detach()
 
         # Update state
         state.previous_modulated_input = ctx.modulated_input.detach()