[Perf] Support TeaCache for MammothModa2 (#5357)

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

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"