[Perf] Support TeaCache for MammothModa2

docs/user_guide/diffusion_features.md

@@ -154,7 +154,7 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **MammothModa2(T2I)**    |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ✅     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

examples/offline_inference/text_to_image/text_to_image.py

@@ -527,14 +527,8 @@ def main():
             "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
         }
     elif cache_backend == "tea_cache":
-        # TeaCache configuration
-        # All parameters marked with [tea_cache only] in DiffusionCacheConfig
-        cache_config = {
-            # TeaCache parameters [tea_cache only]
-            "rel_l1_thresh": 0.2,  # Threshold for accumulated relative L1 distance
-            # Note: coefficients will use model-specific defaults based on model_type
-            #       (e.g., QwenImagePipeline or FluxPipeline)
-        }
+        # Let TeaCache select the model-specific threshold and coefficients.
+        cache_config = {}
 
     profiler_enabled = args.profiler_config is not None
 

recipes/MammothModa2/MammothModa2.md

@@ -21,9 +21,9 @@ pipeline-declared `extra_body` contract.
 MammothModa2's DiT stage runs in the shared diffusion runtime in request mode.
 The first integration intentionally supports one request and one image per
 forward only (`max_num_seqs: 1`, `num_outputs_per_prompt: 1`). Request-level
-batching, step execution, continuous batching, cache acceleration,
-compilation, quantization, parallelism, and offload are not enabled by this
-recipe.
+batching, step execution, continuous batching, compilation, quantization,
+parallelism, and offload are not enabled by this recipe. TeaCache acceleration
+is supported for the DiT stage.
 
 Image size, seed, guidance, and denoising steps use the standard diffusion
 request fields. `cfg_range` remains a MammothModa2-specific `extra_body`
@@ -176,7 +176,30 @@ sets the relative step range `[start, end]` over which CFG is applied (default
 `[0.0, 1.0]`). For compatibility, `text_guidance_scale` and
 `num_inference_steps` remain accepted `extra_body` aliases and, when non-null,
 take precedence over the standard request fields. Model extras are filtered
-against the declared `extra_body_params` (see
+against the declared `extra_body_params`.
+
+TeaCache can be enabled for the DiT stage with the same user-facing sampling
+parameters:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model ./MammothModa2-Preview \
+  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --guidance-scale 4.0 \
+  --num-inference-steps 50 \
+  --cache-backend tea_cache \
+  --extra-body '{"cfg_range": [0.0, 1.0]}' \
+  --output mammoth_t2i_teacache.png
+```
+
+The bundled TeaCache coefficients were fitted from MammothModa2 full-compute
+traces. MammothModa2 uses the model-specific default `rel_l1_thresh=0.075`,
+selected for the evaluated 1024x1024, 50-step configuration.
+
+The model-specific keys are declared in
 [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
 so unknown MammothModa2 extras may be dropped.
 

tests/diffusion/cache/test_cache_backends.py

@@ -571,6 +571,57 @@ def test_refresh(self, mock_apply_hook):
         backend.refresh(mock_pipeline, num_inference_steps=50)
         mock_registry.reset_hook.assert_called_once()
 
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_enable_mammoth_moda2_uses_gen_transformer(self, mock_apply_hook):
+        """MammothModa2 exposes its DiT as gen_transformer, not transformer."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "MammothModa2DiTPipeline"
+        mock_pipeline.gen_transformer = Mock()
+        del mock_pipeline.transformer
+
+        backend = TeaCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        mock_apply_hook.assert_called_once()
+        transformer_arg, config_arg = mock_apply_hook.call_args.args
+        assert transformer_arg is mock_pipeline.gen_transformer
+        assert config_arg.transformer_type == "MammothModa2Transformer2DModel"
+        assert config_arg.rel_l1_thresh == 0.075
+        assert mock_pipeline.transformer is mock_pipeline.gen_transformer
+
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_enable_mammoth_moda2_preserves_explicit_threshold(self, mock_apply_hook):
+        pipeline = Mock()
+        pipeline.__class__.__name__ = "MammothModa2DiTPipeline"
+        pipeline.gen_transformer = Mock()
+        del pipeline.transformer
+
+        TeaCacheBackend(DiffusionCacheConfig(rel_l1_thresh=0.05)).enable(pipeline)
+
+        assert mock_apply_hook.call_args.args[1].rel_l1_thresh == 0.05
+
+    @patch("vllm_omni.diffusion.cache.teacache.backend.apply_teacache_hook")
+    def test_refresh_mammoth_moda2_resets_gen_transformer_hook(self, mock_apply_hook):
+        """Refresh should find the hook through the transformer alias set by the enabler."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "MammothModa2DiTPipeline"
+        mock_pipeline.gen_transformer = Mock()
+        del mock_pipeline.transformer
+
+        backend = TeaCacheBackend(DiffusionCacheConfig())
+        backend.enable(mock_pipeline)
+
+        mock_hook = Mock()
+        mock_registry = Mock()
+        mock_registry.get_hook = Mock(return_value=mock_hook)
+        mock_registry.reset_hook = Mock()
+        mock_pipeline.gen_transformer._hook_registry = mock_registry
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+
+        mock_registry.reset_hook.assert_called_once_with("teacache")
+
 
 class TestCacheSelector:
     """Test cache backend selector function."""

tests/diffusion/cache/test_teacache_extractors.py

@@ -28,17 +28,28 @@
 import torch.nn as nn
 
 from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
 from vllm_omni.diffusion.cache.teacache.extractors import (
     extract_flux2_context,
     extract_flux2_klein_context,
     extract_flux_context,
+    extract_mammoth_moda2_context,
     extract_minimax_h3_context,
     extract_zimage_context,
 )
+from vllm_omni.diffusion.cache.teacache.hook import apply_teacache_hook
+from vllm_omni.diffusion.config import set_current_diffusion_config
+from vllm_omni.diffusion.data import AttentionConfig, DiffusionCacheConfig
 from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
+from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+    Transformer2DModel as MammothModa2Transformer2DModel,
+)
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import MammothModa2DiTPipeline
+from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal
 
 pytestmark = [pytest.mark.core_model]
 
@@ -135,6 +146,164 @@ def get_sample_inputs(self, sample_inputs, /):
         pass
 
 
+@pytest.mark.cpu
+class TestMammothModa2Extractor(BaseExtractorTest):
+    """Test extract_mammoth_moda2_context function."""
+
+    def get_extractor(self):
+        return extract_mammoth_moda2_context
+
+    @pytest.fixture
+    def mammoth_module(self):
+        # These tests run on CPU even when the host has CUDA available.
+        od_config = SimpleNamespace(
+            diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"),
+            parallel_config=SimpleNamespace(ring_degree=1),
+        )
+        with set_current_diffusion_config(od_config):
+            yield MammothModa2Transformer2DModel(
+                patch_size=2,
+                in_channels=4,
+                hidden_size=16,
+                num_layers=2,
+                num_refiner_layers=1,
+                num_attention_heads=2,
+                num_kv_heads=1,
+                axes_dim_rope=(2, 2, 4),
+                axes_lens=(16, 16, 16),
+                text_feat_dim=8,
+            )
+
+    def get_module(self, mammoth_module):
+        return mammoth_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        return {
+            "hidden_states": torch.randn(1, 4, 4, 4),
+            "timestep": torch.tensor([500.0]),
+            "text_hidden_states": torch.randn(1, 3, 8),
+            "text_attention_mask": torch.ones(1, 3, dtype=torch.bool),
+            "freqs_cis": RotaryPosEmbedReal.get_freqs_real((2, 2, 4), (16, 16, 16), theta=10000),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    def test_context_shapes_and_postprocess(self, mammoth_module, sample_inputs):
+        context = extract_mammoth_moda2_context(mammoth_module, **sample_inputs)
+
+        assert context.modulated_input.shape == context.hidden_states.shape
+        assert context.encoder_hidden_states is None
+        assert context.temb.shape[0] == 1
+
+        outputs = context.run_transformer_blocks()
+        assert len(outputs) == 1
+        assert outputs[0].shape == context.hidden_states.shape
+
+        output = context.postprocess(context.hidden_states)
+        assert output.shape == sample_inputs["hidden_states"].shape
+
+    def test_branch_hint_is_preserved(self, mammoth_module, sample_inputs):
+        context = extract_mammoth_moda2_context(
+            mammoth_module,
+            **sample_inputs,
+            teacache_branch="positive",
+        )
+
+        assert context.extra_states is not None
+        assert context.extra_states["teacache_branch"] == "positive"
+
+    def test_hook_keeps_positive_and_negative_branch_state_separate(self, mammoth_module, sample_inputs, monkeypatch):
+        config = TeaCacheConfig(
+            transformer_type="MammothModa2Transformer2DModel",
+            coefficients=[0, 0, 0, 0, 0],
+            rel_l1_thresh=999.0,
+        )
+        apply_teacache_hook(mammoth_module, config)
+
+        original_apply_layers = mammoth_module._apply_transformer_layers
+        apply_layers_mock = Mock(wraps=original_apply_layers)
+        monkeypatch.setattr(mammoth_module, "_apply_transformer_layers", apply_layers_mock)
+
+        mammoth_module(**sample_inputs, teacache_branch="positive")
+        assert apply_layers_mock.call_count == 1
+
+        mammoth_module(**sample_inputs, teacache_branch="positive")
+        assert apply_layers_mock.call_count == 1
+
+        mammoth_module(**sample_inputs, teacache_branch="negative")
+        assert apply_layers_mock.call_count == 2
+
+    def test_hook_rejects_invalid_branch_hint(self, mammoth_module, sample_inputs):
+        config = TeaCacheConfig(
+            transformer_type="MammothModa2Transformer2DModel",
+            coefficients=[0, 0, 0, 0, 0],
+            rel_l1_thresh=999.0,
+        )
+        apply_teacache_hook(mammoth_module, config)
+
+        with pytest.raises(ValueError, match="teacache_branch"):
+            mammoth_module(**sample_inputs, teacache_branch="bad")
+
+    @pytest.mark.parametrize("change_shape", [False, True], ids=["same_shape", "different_shape"])
+    @torch.no_grad()
+    def test_backend_refresh_clears_both_branches(self, mammoth_module, sample_inputs, monkeypatch, change_shape):
+        pipeline = MammothModa2DiTPipeline.__new__(MammothModa2DiTPipeline)
+        nn.Module.__init__(pipeline)
+        pipeline.gen_transformer = mammoth_module
+        # Zero coefficients guarantee cache hits after each branch's first call.
+        backend = TeaCacheBackend(DiffusionCacheConfig(coefficients=[0, 0, 0, 0, 0]))
+        backend.enable(pipeline)
+        apply_layers = Mock(wraps=mammoth_module._apply_transformer_layers)
+        monkeypatch.setattr(mammoth_module, "_apply_transformer_layers", apply_layers)
+
+        for branch in ("positive", "negative"):
+            mammoth_module(**sample_inputs, teacache_branch=branch)
+            mammoth_module(**sample_inputs, teacache_branch=branch)
+        assert apply_layers.call_count == 2
+
+        next_inputs = dict(sample_inputs)
+        next_inputs["hidden_states"] = torch.randn(1, 4, 4, 6 if change_shape else 4)
+        text_length = 5 if change_shape else 3
+        next_inputs["text_hidden_states"] = torch.randn(1, text_length, 8)
+        next_inputs["text_attention_mask"] = torch.ones(1, text_length, dtype=torch.bool)
+        # Compute an uncached reference through the original model forward.
+        expected = MammothModa2Transformer2DModel.forward(mammoth_module, **next_inputs)
+        apply_layers.reset_mock()
+
+        backend.refresh(pipeline, num_inference_steps=4)
+
+        for count, branch in enumerate(("positive", "negative"), start=1):
+            actual = mammoth_module(**next_inputs, teacache_branch=branch)
+            assert apply_layers.call_count == count
+            torch.testing.assert_close(actual, expected)
+            # Refresh must preserve caching for subsequent steps of request B.
+            cached = mammoth_module(**next_inputs, teacache_branch=branch)
+            assert apply_layers.call_count == count
+            torch.testing.assert_close(cached, expected)
+
+    def test_forced_full_compute_matches_original_forward(self, mammoth_module, sample_inputs, monkeypatch):
+        with torch.no_grad():
+            expected = mammoth_module(**sample_inputs, teacache_branch="positive")
+
+        config = TeaCacheConfig(
+            transformer_type="MammothModa2Transformer2DModel",
+            coefficients=[0, 0, 0, 0, 0],
+            rel_l1_thresh=0.1,
+        )
+        apply_teacache_hook(mammoth_module, config)
+        hook = mammoth_module._hook_registry.get_hook("teacache")
+        should_compute_full = Mock(return_value=True)
+        monkeypatch.setattr(hook, "_should_compute_full_transformer", should_compute_full)
+
+        with torch.no_grad():
+            actual = mammoth_module(**sample_inputs, teacache_branch="positive")
+
+        should_compute_full.assert_called_once()
+        torch.testing.assert_close(actual, expected)
+
+
 class TestFlux2KleinExtractor(BaseExtractorTest):
     """Test extract_flux2_klein_context function."""