diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 7a80ac781..7bbb08951 100644 --- a/docs/user_guide/diffusion_features.md +++ b/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** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index a8624b9ae..7e1b393e6 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/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 diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 9b6fab76c..7c1e8bfd8 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/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. diff --git a/tests/diffusion/cache/test_cache_backends.py b/tests/diffusion/cache/test_cache_backends.py index 29cda91a3..b9bc5c6a2 100644 --- a/tests/diffusion/cache/test_cache_backends.py +++ b/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.""" diff --git a/tests/diffusion/cache/test_teacache_extractors.py b/tests/diffusion/cache/test_teacache_extractors.py index cdaef2c7c..5ce2cfb4e 100644 --- a/tests/diffusion/cache/test_teacache_extractors.py +++ b/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.""" diff --git a/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_teacache.py b/tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_teacache.py new file mode 100644 index 000000000..9d68a2a64 --- /dev/null +++ b/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" diff --git a/vllm_omni/diffusion/cache/teacache/backend.py b/vllm_omni/diffusion/cache/teacache/backend.py index 508d63eec..6da9d953b 100644 --- a/vllm_omni/diffusion/cache/teacache/backend.py +++ b/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, } diff --git a/vllm_omni/diffusion/cache/teacache/config.py b/vllm_omni/diffusion/cache/teacache/config.py index b83f600d6..832e13ca0 100644 --- a/vllm_omni/diffusion/cache/teacache/config.py +++ b/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, } diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index d1b04151d..2ca3e22bd 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/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, diff --git a/vllm_omni/diffusion/cache/teacache/hook.py b/vllm_omni/diffusion/cache/teacache/hook.py index 664de792f..730d4936e 100644 --- a/vllm_omni/diffusion/cache/teacache/hook.py +++ b/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() diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 3a0291398..4b5edf717 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/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 ) diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index fd716aa9a..6a35955ae 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/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)