From 0833470d4172652b9c6a61de851837c144a6fa23 Mon Sep 17 00:00:00 2001 From: Lancer Date: Mon, 23 Mar 2026 10:19:25 +0800 Subject: [PATCH] [Feat] support TeaCache for Flux2 klein (#1234) Signed-off-by: Lancer --- docs/user_guide/diffusion_acceleration.md | 2 +- .../cache/test_teacache_extractors.py | 170 ++++++++++++++++++ vllm_omni/diffusion/cache/teacache/backend.py | 24 ++- vllm_omni/diffusion/cache/teacache/config.py | 9 + .../diffusion/cache/teacache/extractors.py | 144 +++++++++++++++ vllm_omni/diffusion/cache/teacache/hook.py | 30 ++-- 6 files changed, 365 insertions(+), 14 deletions(-) create mode 100644 tests/diffusion/cache/test_teacache_extractors.py diff --git a/docs/user_guide/diffusion_acceleration.md b/docs/user_guide/diffusion_acceleration.md index db05d5f3c..0140819f0 100644 --- a/docs/user_guide/diffusion_acceleration.md +++ b/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 diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py new file mode 100644 index 000000000..5ba52ddfe --- /dev/null +++ b/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)), + ) diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index bf328d43d..132dcce0d 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/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): diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index 30b1745f4..9fc7dc5b1 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/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 diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 780297919..c1475fc1a 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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, diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index 65f764c43..d38148366 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/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()