[Perf] Support TeaCache for MammothModa2 (#5357)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>

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
+        # Let TeaCache select the model-specific threshold and coefficients.
-        # All parameters marked with [tea_cache only] in DiffusionCacheConfig
+        cache_config = {}
-        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)
-        }
 
     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,
+batching, step execution, continuous batching, compilation, quantization,
-compilation, quantization, parallelism, and offload are not enabled by this
+parallelism, and offload are not enabled by this recipe. TeaCache acceleration
-recipe.
+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 @@ class TestTeaCacheBackend:
         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
 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 @@ class BaseExtractorTest(ABC):
         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."""
 

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_teacache.py

@@ -0,0 +1,274 @@
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.cache.cachedit import RequestScopedCacheDiTRuntime
+from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
+from vllm_omni.diffusion.cache.teacache.extractors import extract_mammoth_moda2_context
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit as mammoth_pipeline_module
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import MammothModa2DiTPipeline
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.cpu]
+
+
+def test_mammoth_moda2_calibrated_teacache_defaults():
+    config = TeaCacheConfig(transformer_type="MammothModa2Transformer2DModel")
+
+    assert config.rel_l1_thresh == 0.075
+    assert config.coefficients == [
+        -1761.242764481119,
+        859.398730831851,
+        -126.78044436937441,
+        9.16173963457422,
+        -0.1102826051200547,
+    ]
+
+
+class _FakeScheduler:
+    def set_timesteps(self, num_inference_steps, device, num_tokens):  # noqa: ARG002
+        self.timesteps = torch.arange(num_inference_steps, 0, -1, device=device, dtype=torch.float32)
+
+    def step(self, model_pred, t, latents, return_dict=False):  # noqa: ARG002
+        return (latents,)
+
+
+class _FakeTransformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(in_channels=4)
+        self.time_caption_embed = SimpleNamespace(image_embedder=None)
+        self.param = nn.Parameter(torch.zeros(()))
+        self.branches: list[str | None] = []
+        self.calls: list[dict] = []
+
+    def forward(self, hidden_states, **kwargs):
+        self.branches.append(kwargs.get("teacache_branch"))
+        self.calls.append(kwargs)
+        return torch.zeros_like(hidden_states)
+
+
+class _FakeVAE(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.param = nn.Parameter(torch.zeros(()))
+        self.config = SimpleNamespace(scaling_factor=None, shift_factor=None)
+
+    def decode(self, latents, return_dict=False):  # noqa: ARG002
+        return (latents,)
+
+
+def _build_pipeline(monkeypatch):
+    monkeypatch.setattr(mammoth_pipeline_module, "FlowMatchEulerDiscreteScheduler", _FakeScheduler)
+    monkeypatch.setattr(
+        mammoth_pipeline_module,
+        "randn_tensor",
+        lambda shape, *, device, dtype, generator=None: torch.zeros(shape, device=device, dtype=dtype),
+    )
+
+    pipe = MammothModa2DiTPipeline.__new__(MammothModa2DiTPipeline)
+    nn.Module.__init__(pipe)
+    pipe.gen_transformer = _FakeTransformer()
+    pipe.gen_vae = _FakeVAE()
+    pipe.gen_image_condition_refiner = None
+    pipe.gen_freqs_cis = []
+    pipe._llm_hidden_size = 8
+    pipe._cache_dit_config = None
+    pipe._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipe)
+    pipe.config = SimpleNamespace(
+        llm_config=SimpleNamespace(gen_vocab_start_index=100),
+        image_token_id=900,
+        video_token_id=901,
+        vision_start_token_id=902,
+        vision_end_token_id=903,
+    )
+    return pipe
+
+
+def _run_pipeline(pipe, *, text_guidance_scale, cfg_range, num_inference_steps=4, num_requests=1):
+    request = OmniDiffusionRequest(
+        request_id="req-teacache",
+        prompt={
+            "prompt": "",
+            "height": 32,
+            "width": 32,
+            "additional_information": {
+                "full_hidden_states": torch.randn(3, 8),
+                "full_token_ids": [10, 11, 100],
+                "answer_start_index": 2,
+            },
+        },
+        sampling_params=OmniDiffusionSamplingParams(
+            height=32,
+            width=32,
+            guidance_scale=text_guidance_scale,
+            num_inference_steps=num_inference_steps,
+            extra_args={"cfg_range": cfg_range},
+        ),
+    )
+    requests = [
+        OmniDiffusionRequest(
+            request_id=f"req-teacache-{i}",
+            prompt=request.prompt,
+            sampling_params=request.sampling_params,
+        )
+        for i in range(num_requests)
+    ]
+    outputs = pipe(DiffusionRequestBatch(requests))
+    assert isinstance(outputs, list)
+    assert len(outputs) == num_requests
+    assert all(isinstance(output, DiffusionOutput) and output.output.shape == (1, 4, 4, 4) for output in outputs)
+    assert all(call["timestep"].shape == (num_requests,) for call in pipe.gen_transformer.calls)
+    return pipe.gen_transformer.branches
+
+
+@pytest.mark.parametrize("num_requests", [1, 2])
+def test_mammoth_moda2_non_cfg_passes_positive_teacache_branch(monkeypatch, num_requests):
+    pipe = _build_pipeline(monkeypatch)
+
+    branches = _run_pipeline(pipe, text_guidance_scale=1.0, cfg_range=[0.0, 1.0], num_requests=num_requests)
+
+    assert branches == ["positive", "positive", "positive", "positive"]
+
+
+def test_mammoth_moda2_dev_forwards_ar_image_conditioning_to_positive_branch(monkeypatch):
+    """Dev keeps AR image tokens separate and forwards them to the positive DiT call."""
+    pipe = _build_pipeline(monkeypatch)
+    pipe.gen_transformer.time_caption_embed.image_embedder = object()
+
+    _run_pipeline(
+        pipe,
+        text_guidance_scale=1.0,
+        cfg_range=[0.0, 1.0],
+        num_inference_steps=1,
+    )
+
+    assert len(pipe.gen_transformer.calls) == 1
+    call = pipe.gen_transformer.calls[0]
+    assert call["teacache_branch"] == "positive"
+    assert call["text_hidden_states"].shape == (1, 2, 8)
+    assert call["ar_image_hidden_states"].shape == (1, 1, 8)
+    assert torch.equal(call["ar_image_attention_mask"], torch.ones(1, 1, dtype=torch.bool))
+
+
+@pytest.mark.parametrize("num_requests", [1, 2])
+def test_mammoth_moda2_cfg_passes_positive_then_negative_teacache_branch(monkeypatch, num_requests):
+    pipe = _build_pipeline(monkeypatch)
+
+    branches = _run_pipeline(pipe, text_guidance_scale=4.0, cfg_range=[0.0, 1.0], num_requests=num_requests)
+
+    assert branches == [
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+    ]
+
+
+@pytest.mark.parametrize("num_requests", [1, 2])
+def test_mammoth_moda2_cfg_range_only_uses_negative_inside_range(monkeypatch, num_requests):
+    pipe = _build_pipeline(monkeypatch)
+
+    branches = _run_pipeline(pipe, text_guidance_scale=4.0, cfg_range=[0.5, 1.0], num_requests=num_requests)
+
+    assert branches == [
+        "positive",
+        "positive",
+        "positive",
+        "negative",
+        "positive",
+        "negative",
+    ]
+
+
+def test_mammoth_moda2_teacache_forwards_ar_image_conditioning():
+    """The TeaCache extractor follows the full embedding contract, including the expanded text mask."""
+    hidden_size = 4
+    hidden_states = torch.zeros(1, hidden_size, 1, 1)
+    timestep = torch.zeros(1)
+    text_hidden_states = torch.zeros(1, 1, hidden_size)
+    text_attention_mask = torch.ones(1, 1, dtype=torch.bool)
+    freqs_cis = object()
+    ar_image_hidden_states = torch.randn(1, 2, hidden_size)
+    ar_image_attention_mask = torch.ones(1, 2, dtype=torch.bool)
+
+    temb = torch.zeros(1, hidden_size)
+    prepared_text_hidden_states = torch.randn(1, 3, hidden_size)
+    prepared_text_attention_mask = torch.ones(1, 3, dtype=torch.bool)
+    img_tokens = torch.randn(1, 1, hidden_size)
+    img_mask = torch.ones(1, 1, dtype=torch.bool)
+    context_rotary_emb = object()
+    noise_rotary_emb = object()
+    rotary_emb = object()
+
+    first_layer = SimpleNamespace(
+        norm1=Mock(side_effect=lambda joint_hidden_states, _temb: (joint_hidden_states,)),
+    )
+    module = SimpleNamespace(
+        layers=[first_layer],
+        config=SimpleNamespace(hidden_size=hidden_size, patch_size=1),
+        _validate_inputs=Mock(return_value=(1, 1, 1)),
+        _prepare_embeddings=Mock(
+            return_value=(
+                temb,
+                prepared_text_hidden_states,
+                prepared_text_attention_mask,
+                img_tokens,
+                img_mask,
+                1,
+                context_rotary_emb,
+                noise_rotary_emb,
+                rotary_emb,
+                [3],
+                [4],
+            )
+        ),
+        _apply_refiners=Mock(return_value=(prepared_text_hidden_states, img_tokens)),
+    )
+
+    ctx = extract_mammoth_moda2_context(
+        module,
+        hidden_states=hidden_states,
+        timestep=timestep,
+        text_hidden_states=text_hidden_states,
+        text_attention_mask=text_attention_mask,
+        ref_image_hidden_states=None,
+        ar_image_hidden_states=ar_image_hidden_states,
+        ar_image_attention_mask=ar_image_attention_mask,
+        freqs_cis=freqs_cis,
+        teacache_branch="positive",
+    )
+
+    module._prepare_embeddings.assert_called_once_with(
+        hidden_states,
+        timestep,
+        text_hidden_states,
+        text_attention_mask,
+        freqs_cis,
+        1,
+        1,
+        1,
+        ar_image_hidden_states,
+        ar_image_attention_mask,
+    )
+    module._apply_refiners.assert_called_once_with(
+        prepared_text_hidden_states,
+        prepared_text_attention_mask,
+        context_rotary_emb,
+        img_tokens,
+        img_mask,
+        noise_rotary_emb,
+        temb,
+    )
+    assert ctx.hidden_states.shape == (1, 4, hidden_size)
+    assert ctx.extra_states["teacache_branch"] == "positive"

vllm_omni/diffusion/cache/teacache/backend.py

@@ -8,8 +8,9 @@ This module provides the TeaCache backend that implements the CacheBackend
 interface using the hooks-based TeaCache system.
 """
 
-from typing import Any
+from typing import Any, Protocol
 
+import torch.nn as nn
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.cache.base import CacheBackend
@@ -20,6 +21,11 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig
 logger = init_logger(__name__)
 
 
+class _MammothModa2TeaCachePipeline(Protocol):
+    gen_transformer: nn.Module
+    transformer: nn.Module
+
+
 def enable_hunyuan_image3_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """
     Enable TeaCache for HunyuanImage3 model.
@@ -57,6 +63,23 @@ def enable_bagel_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     )
 
 
+def enable_mammoth_moda2_teacache(pipeline: _MammothModa2TeaCachePipeline, config: DiffusionCacheConfig) -> None:
+    """Enable TeaCache for MammothModa2's DiT stage."""
+    teacache_config = TeaCacheConfig(
+        transformer_type="MammothModa2Transformer2DModel",
+        rel_l1_thresh=config.rel_l1_thresh,
+        coefficients=config.coefficients,
+    )
+    transformer = pipeline.gen_transformer
+    apply_teacache_hook(transformer, teacache_config)
+    pipeline.transformer = transformer
+
+    logger.info(
+        f"TeaCache applied with rel_l1_thresh={teacache_config.rel_l1_thresh}, "
+        f"transformer_class={teacache_config.transformer_type}"
+    )
+
+
 def enable_sensenova_u1_teacache(pipeline: Any, config: DiffusionCacheConfig) -> None:
     """Enable TeaCache for SenseNova-U1 denoising forwards."""
     teacache_config = TeaCacheConfig(
@@ -124,6 +147,7 @@ CUSTOM_TEACACHE_ENABLERS = {
     "BagelPipeline": enable_bagel_teacache,
     "Flux2KleinPipeline": enable_flux2_klein_teacache,
     "HunyuanImage3Pipeline": enable_hunyuan_image3_teacache,
+    "MammothModa2DiTPipeline": enable_mammoth_moda2_teacache,
     "MiniMaxH3Pipeline": enable_minimax_h3_teacache,
     "SenseNovaU1Pipeline": enable_sensenova_u1_teacache,
 }

vllm_omni/diffusion/cache/teacache/config.py

@@ -83,6 +83,15 @@ _MODEL_COEFFICIENTS = {
     ],
     # LongCat Image transformer coefficients
     "LongCatImageTransformer2DModel": [652.5980, -424.1615, 84.5526, -4.5923, 0.1694],
+    # MammothModa2 DiT coefficients fitted from full-compute traces
+    # at 1024x1024 with 50 denoising steps.
+    "MammothModa2Transformer2DModel": [
+        -1761.242764481119,
+        859.398730831851,
+        -126.78044436937441,
+        9.16173963457422,
+        -0.1102826051200547,
+    ],
     # MiniMax-H3 FL2VA coefficients.
     "MiniMaxH3DiTModel": [
         2.283704065852778e03,
@@ -95,6 +104,7 @@ _MODEL_COEFFICIENTS = {
 
 _DEFAULT_REL_L1_THRESH = 0.2
 _MODEL_DEFAULT_REL_L1_THRESH = {
+    "MammothModa2Transformer2DModel": 0.075,
     "MiniMaxH3DiTModel": 0.17,
 }
 

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -20,6 +20,7 @@ from typing import Any
 
 import torch
 import torch.nn as nn
+from einops import rearrange
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.forward_context import get_forward_context
@@ -823,6 +824,114 @@ def extract_longcat_context(
     )
 
 
+def extract_mammoth_moda2_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    text_hidden_states: torch.Tensor,
+    freqs_cis: torch.Tensor,
+    text_attention_mask: torch.Tensor,
+    ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
+    ar_image_hidden_states: torch.Tensor | None = None,
+    ar_image_attention_mask: torch.Tensor | None = None,
+    return_dict: bool = False,
+    teacache_branch: str | None = None,
+) -> CacheContext:
+    """Extract cache context for MammothModa2's Lumina-style DiT."""
+    if not hasattr(module, "layers") or len(module.layers) == 0:
+        raise ValueError("Module must have layers")
+
+    batch_size, height, width = module._validate_inputs(
+        hidden_states,
+        text_hidden_states,
+        text_attention_mask,
+        ref_image_hidden_states,
+        return_dict,
+    )
+
+    (
+        temb,
+        text_hidden_states,
+        text_attention_mask,
+        img_tokens,
+        img_mask,
+        img_len,
+        context_rotary_emb,
+        noise_rotary_emb,
+        rotary_emb,
+        encoder_seq_lengths,
+        seq_lengths,
+    ) = module._prepare_embeddings(
+        hidden_states,
+        timestep,
+        text_hidden_states,
+        text_attention_mask,
+        freqs_cis,
+        batch_size,
+        height,
+        width,
+        ar_image_hidden_states,
+        ar_image_attention_mask,
+    )
+
+    text_hidden_states, img_tokens = module._apply_refiners(
+        text_hidden_states,
+        text_attention_mask,
+        context_rotary_emb,
+        img_tokens,
+        img_mask,
+        noise_rotary_emb,
+        temb,
+    )
+
+    max_seq_len = max(seq_lengths)
+    attention_mask = img_tokens.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
+    joint_hidden_states = img_tokens.new_zeros(batch_size, max_seq_len, module.config.hidden_size)
+    for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+        attention_mask[i, :seq_len] = True
+        joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len]
+        joint_hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] = img_tokens[i, :img_len]
+
+    modulated_input = module.layers[0].norm1(joint_hidden_states, temb)[0]
+
+    def run_transformer_blocks() -> tuple[torch.Tensor]:
+        h = module._apply_transformer_layers(joint_hidden_states, attention_mask, rotary_emb, temb)
+        return (h,)
+
+    def postprocess(h: torch.Tensor) -> torch.Tensor:
+        h = module.norm_out(h, temb)
+        p = module.config.patch_size
+        img_hidden_states = torch.stack(
+            [
+                h[i, encoder_seq_len : encoder_seq_len + img_len]
+                for i, encoder_seq_len in enumerate(encoder_seq_lengths)
+            ],
+            dim=0,
+        )
+        return rearrange(
+            img_hidden_states,
+            "b (h w) (p1 p2 c) -> b c (h p1) (w p2)",
+            h=height // p,
+            w=width // p,
+            p1=p,
+            p2=p,
+        )
+
+    extra_states = {}
+    if teacache_branch is not None:
+        extra_states["teacache_branch"] = teacache_branch
+
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=joint_hidden_states,
+        encoder_hidden_states=None,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+        extra_states=extra_states,
+    )
+
+
 def extract_stable_audio_context(
     module: nn.Module,
     hidden_states: torch.Tensor,
@@ -1556,6 +1665,7 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "StableAudioDiTModel": extract_stable_audio_context,
     "Flux2Transformer2DModel": extract_flux2_context,
     "LongCatImageTransformer2DModel": extract_longcat_context,
+    "MammothModa2Transformer2DModel": extract_mammoth_moda2_context,
     "FluxTransformer2DModel": extract_flux_context,
     "SenseNovaU1ForCausalLM": extract_sensenova_u1_context,
     "MiniMaxH3DiTModel": extract_minimax_h3_context,

vllm_omni/diffusion/cache/teacache/hook.py

@@ -145,7 +145,15 @@ class TeaCacheHook(ModelHook):
         #   - cfg_rank 0: positive branch
         #   - cfg_rank > 0: negative branch
         # Without CFG-parallel, branches alternate within a single rank
-        if getattr(module, "do_true_cfg", False):
+        branch_hint = None
+        if getattr(ctx, "extra_states", None):
+            branch_hint = ctx.extra_states.get("teacache_branch")
+
+        if branch_hint is not None:
+            if branch_hint not in ("positive", "negative"):
+                raise ValueError(f"Invalid teacache_branch={branch_hint!r}; expected 'positive' or 'negative'.")
+            cache_branch = branch_hint
+        elif getattr(module, "do_true_cfg", False):
             cfg_parallel_size = get_classifier_free_guidance_world_size()
             if cfg_parallel_size > 1:
                 cfg_rank = get_classifier_free_guidance_rank()

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -743,7 +743,11 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         ar_image_hidden_states: torch.Tensor | None = None,
         ar_image_attention_mask: torch.Tensor | None = None,
         return_dict: bool = False,
+        teacache_branch: str | None = None,
     ) -> torch.Tensor:
+        # Consumed by TeaCacheHook/extractor before this plain forward runs.
+        del teacache_branch
+
         batch_size, height, width = self._validate_inputs(
             hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict
         )

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -807,6 +807,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 ar_image_hidden_states=ar_image_embeds,
                 ar_image_attention_mask=ar_image_attention_mask,
                 freqs_cis=self.gen_freqs_cis,
+                teacache_branch="positive",
             )
             run_uncond = (any_active_per_step is not None and any_active_per_step[i]) if needs_uncond else False
             if requires_paired_cfg and negative_prompt_embeds is not None:
@@ -819,6 +820,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                     text_attention_mask=negative_prompt_attention_mask,
                     ref_image_hidden_states=None,
                     freqs_cis=self.gen_freqs_cis,
+                    teacache_branch="negative",
                 )
                 # Fused CFG blend: torch.lerp fuses (uncond + scale * (cond - uncond)) into 1 kernel
                 blended = torch.lerp(model_pred_uncond, model_pred, scale_vec)