From c11596c63711326ba1af4527f949f1558c63a8b4 Mon Sep 17 00:00:00 2001 From: shihongzhi Date: Tue, 22 Sep 2026 09:34:12 +0800 Subject: [PATCH] [Model] MammothModa2: enable Cache-DiT acceleration for the DiT stage (#7291) Signed-off-by: shihongzhi --- docs/design/feature/cache_dit.md | 48 ++++- .../diffusion/cache_acceleration/cache_dit.md | 33 ++++ docs/user_guide/diffusion_features.md | 3 +- .../text_to_image/text_to_image.py | 6 +- tests/diffusion/cache/test_cache_dit.py | 108 ++++++++++ .../test_pipeline_mammothmoda2_dit.py | 186 +++++++++++++++++- vllm_omni/deploy/mammoth_moda2.yaml | 10 + .../cache/cachedit/model_specific.py | 42 ++++ .../mammoth_moda2/mammothmoda2_dit_model.py | 10 +- .../pipeline_mammothmoda2_dit.py | 75 ++++++- 10 files changed, 507 insertions(+), 14 deletions(-) diff --git a/docs/design/feature/cache_dit.md b/docs/design/feature/cache_dit.md index 38b597ca3..69a8de613 100644 --- a/docs/design/feature/cache_dit.md +++ b/docs/design/feature/cache_dit.md @@ -180,17 +180,52 @@ cache_dit.enable_cache( > **Note:** For single transformer with multiple block lists, `refresh_context` works the same as standard models. +### Example 3: Sequential-CFG Diffusion Pipeline (MammothModa2) + +MammothModa2's DiT stage runs on the diffusion runner, which enables the +configured backend at startup and transfers ownership to the pipeline through +the request-scoped protocol (`adopt_cache_dit_backend` / `is_cache_dit_enabled`); +`forward()` then reconciles per-request state through +`RequestScopedCacheDiTRuntime` (MiniMax H3 precedent). + +**Key difference:** The enabler caches only the repeated main-layer stack and +marks the pipeline as requiring paired CFG forwards, because cache-dit's +separate-CFG accounting distinguishes cond/uncond passes purely by forward +parity: + +```python +pipeline._cache_dit_requires_paired_cfg = True +block_adapter = BlockAdapter( + transformer=pipeline.gen_transformer, + blocks=[pipeline.gen_transformer.layers], # Refiners stay outside + forward_pattern=[ForwardPattern.Pattern_3], + has_separate_cfg=True, + check_forward_pattern=True, +) +``` + +Per request, the pipeline calls `runtime.prepare(spec)` for CFG requests +(refreshing the step context) and `runtime.prepare(None)` for `guidance=1.0` +requests (disabling hooks — their single forward per step cannot be expressed +by the parity accounting). + ### Registering Custom Implementations -After writing your custom enabler, register it in `CUSTOM_DIT_ENABLERS` in +After writing your custom enabler, register it in +`CUSTOM_DIT_ENABLERS` via `register_custom_dit_enablers()` in `vllm_omni/diffusion/cache/cachedit/model_specific.py`: ```python -CUSTOM_DIT_ENABLERS = { - "Wan22Pipeline": enable_cache_for_wan22, - "LongCatImagePipeline": enable_cache_for_longcat_image, - "YourCustomPipeline": enable_cache_for_your_model, # Add here -} +CUSTOM_DIT_ENABLERS.update( + { + "Wan22Pipeline": enable_cache_for_wan22, + "Cosmos3OmniDiffusersPipeline": enable_cache_for_cosmos3, + "Krea2Pipeline": enable_cache_for_krea2, + "Magi2Pipeline": enable_cache_for_magi2, + "MammothModa2DiTPipeline": enable_cache_for_mammothmoda2, + "YourCustomPipeline": enable_cache_for_your_model, # Add here + } +) ``` --- @@ -274,6 +309,7 @@ Complete examples in the codebase: | **Wan2.2** | `cachedit.model_specific::enable_cache_for_wan22` | Single or dual-transformer | Auto-detects mode based on transformer_2 presence | | **LongCat** | `cachedit.config::CacheDiTAdapterConfig` | Declarative block adapter | Two block lists in one transformer | | **BAGEL** | `cachedit.model_specific::BagelCachedAdapter` | Custom cached adapter | Complex architecture | +| **MammothModa2** | `cachedit.model_specific::enable_cache_for_mammothmoda2` | Sub-stack + paired CFG | Generation-runner pipeline; lifecycle owned by `pipeline_mammothmoda2_dit.py` | --- diff --git a/docs/user_guide/diffusion/cache_acceleration/cache_dit.md b/docs/user_guide/diffusion/cache_acceleration/cache_dit.md index 4d1f69068..3b686a32e 100644 --- a/docs/user_guide/diffusion/cache_acceleration/cache_dit.md +++ b/docs/user_guide/diffusion/cache_acceleration/cache_dit.md @@ -72,6 +72,39 @@ omni = Omni( ## Example Script +### Multi-Stage Models (MammothModa2) + +MammothModa2 is a multi-stage AR→DiT pipeline; the diffusion runner owns the +Cache-DiT lifecycle of its DiT stage. Configure it on the **DiT stage entry** +of the deploy YAML so the acceleration only applies to the denoising stage: + +```yaml +# deploy YAML (see vllm_omni/deploy/mammoth_moda2.yaml) +stages: + - stage_id: 1 + # ... other DiT stage settings ... + cache_backend: cache_dit + cache_config: + Fn_compute_blocks: 1 + Bn_compute_blocks: 0 + max_warmup_steps: 4 + residual_diff_threshold: 0.24 + max_continuous_cached_steps: 3 + enable_cache_dit_summary: true # log skip-ratio stats per request +``` + +Behavior notes: + +- Only the repeated main-layer stack is cached; the Q-Former refiners always run. +- Sequential-CFG parity requires the unconditional pass on every denoise step, + so the `cfg_range` skip optimization is not used: outside the interval CFG is + neutralized with `scale=1.0` instead of skipping the uncond forward. +- Requests with `text_guidance_scale = 1.0` run with cache hooks disabled. +- `residual_diff_threshold` defaults to `0.24` (measured ~1.7-2.2x at 50 steps, + PSNR 27-34 dB vs uncached). For 20-step generation — especially with a + partial `cfg_range` — use `0.12` for noticeably closer parity at a modest + speed cost. + ### Offline Inference Use the example script under `examples/offline_inference/text_to_image`: diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index a8fb6d972..703a83f38 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -130,7 +130,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** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | @@ -154,6 +154,7 @@ The following tables show which models support each feature: > 4. Krea 2 currently supports single-GPU inference plus LoRA, Cache-DiT, HSDP, CPU/layerwise offload, and VAE-patch-parallel (decode). TP/SP/CFG-Parallel are not yet wired. The few-step distilled (Turbo) checkpoint uses `is_distilled=true` (fixed timestep shift `mu=1.15`); generate at 2048x2048 by default with `num_inference_steps≈8` and `guidance_scale=0`. The Raw checkpoint uses 1024x1024, `num_inference_steps=28`, and `guidance_scale=4.5`. > 5. HunyuanImage3 supports step execution. Multi-request step execution requires `TORCH_SDPA`; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). > 6. BAGEL step execution supports image generation with `bagel.yaml`, `bagel_think.yaml`, and `bagel_single_stage.yaml`; two-stage Thinker execution and explicit single-stage text output remain on their existing complete-request paths. Image requests require `num_inference_steps >= 2`. BAGEL step execution cannot currently be combined with sequence parallelism or a diffusion cache backend; see [Diffusion Execution Modes](diffusion/execution_modes.md#step-execution). +> 7. MammothModa2 runs its DiT stage on the diffusion runner (`StageExecutionType.DIFFUSION`); Cache-DiT is enabled through the standard diffusion-stage knobs on the stage entry of the deploy YAML (`cache_backend: cache_dit`, plus optional `cache_config` / `enable_cache_dit_summary`). The runner installs the backend at startup and the pipeline adopts it per request. Only the repeated main-layer stack is cached; requests with `text_guidance_scale = 1.0` bypass the cache hooks. ### VideoGen 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 c74236316..a8624b9ae 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -157,7 +157,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--enable-cache-dit-summary", action="store_true", - help="Enable cache-dit summary logging after diffusion forward passes.", + default=None, + help=( + "Enable cache-dit summary logging after diffusion forward passes. " + "Default: unset (defer to the deploy YAML's enable_cache_dit_summary)." + ), ) parser.add_argument( "--ulysses-degree", diff --git a/tests/diffusion/cache/test_cache_dit.py b/tests/diffusion/cache/test_cache_dit.py index c1d061ef2..7de2aaa38 100644 --- a/tests/diffusion/cache/test_cache_dit.py +++ b/tests/diffusion/cache/test_cache_dit.py @@ -60,6 +60,7 @@ def test_custom_cache_dit_enablers_are_registered_explicitly(): "Cosmos3OmniPipeline": cd_model_specific.enable_cache_for_cosmos3, "Krea2Pipeline": cd_model_specific.enable_cache_for_krea2, "Magi2Pipeline": cd_model_specific.enable_cache_for_magi2, + "MammothModa2DiTPipeline": cd_model_specific.enable_cache_for_mammothmoda2, } with patch.dict(cd_backend.CUSTOM_DIT_ENABLERS, {}, clear=True): @@ -116,6 +117,33 @@ def test_magi2_cache_dit_targets_only_nested_repeated_layers(mock_block_adapter, assert get_transformer(pipeline) is transformer_block +@patch("vllm_omni.diffusion.cache.cachedit.model_specific.enable_cache_for_dit") +@patch("vllm_omni.diffusion.cache.cachedit.model_specific.BlockAdapter") +def test_mammothmoda2_cache_dit_targets_only_main_layers(mock_block_adapter, mock_enable_cache): + pipeline = Mock() + transformer = pipeline.gen_transformer + layers = torch.nn.ModuleList([torch.nn.Identity()]) + transformer.layers = layers + adapter = mock_block_adapter.return_value + refresh = Mock() + mock_enable_cache.return_value = refresh + + result = cd_model_specific.enable_cache_for_mammothmoda2(pipeline, SAMPLE_CACHE_CONFIG) + + mock_block_adapter.assert_called_once() + adapter_kwargs = mock_block_adapter.call_args.kwargs + assert adapter_kwargs["transformer"] is transformer + assert adapter_kwargs["blocks"] == [layers] + assert adapter_kwargs["has_separate_cfg"] is True + assert adapter_kwargs["check_forward_pattern"] is True + assert result.refresh is refresh + assert result.targets == (adapter,) + get_transformer = mock_enable_cache.call_args.kwargs["get_pipeline_transformer"] + assert get_transformer(pipeline) is transformer + # Sequential-CFG parity is mandatory while cache-dit is installed. + assert pipeline._cache_dit_requires_paired_cfg is True + + @patch("vllm_omni.diffusion.cache.cachedit.backend.cache_dit.summary") @patch("vllm_omni.diffusion.cache.cachedit.backend.BlockAdapter.is_cached", return_value=True) def test_cache_summary_uses_custom_nested_targets(mock_is_cached, mock_summary): @@ -338,3 +366,83 @@ def test_summary_with_no_transformer_is_nonfatal(): pass cache_summary(pipeline=FakePipeline()) + + +# This test is skipped on ROCm since rocm_unquantized_gemm doesn't support CPU backend +@pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="vLLM ROCm custom ops lack CPU fallback", +) +def test_mammothmoda2_cache_dit_runs_end_to_end_on_tiny_model(request: pytest.FixtureRequest): + """A tiny MammothModa2 DiT runs through Cache-DiT with paired CFG forwards. + + Exercises the Pattern_3 contract end to end: after enabling the custom + enabler, blocks must be callable as ``block(hidden_states, **kwargs)`` + and the pipeline's two-forwards-per-step CFG cadence must drive the + cache context without shape or parity errors. + """ + from cache_dit import BlockAdapter + + from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import Transformer2DModel + from vllm_omni.diffusion.models.mammoth_moda2.rope_real import RotaryPosEmbedReal + + # MammothModa2 attention resolves through the shared Omni attention layer, + # which picks the platform-default backend at construction time. Pin + # TORCH_SDPA so a CUDA-visible host does not select FA3 for CPU tensors. + with _force_torch_sdpa(): + model = Transformer2DModel( + patch_size=2, + in_channels=4, + hidden_size=96, + num_layers=4, + num_refiner_layers=1, + num_attention_heads=2, + num_kv_heads=2, + multiple_of=8, + axes_dim_rope=(16, 16, 16), + axes_lens=(300, 512, 512), + text_feat_dim=16, + ) + model.eval() + + MammothModa2Pipeline = type("MammothModa2DiTPipeline", (), {}) + pipeline = MammothModa2Pipeline() + pipeline.gen_transformer = model + backend = CacheDiTBackend(DiffusionCacheConfig()) + backend.enable(pipeline) + request.addfinalizer(lambda: backend.disable(pipeline)) + assert BlockAdapter.is_cached(model) + + num_inference_steps = 8 + backend.refresh(pipeline, num_inference_steps=num_inference_steps) + + freqs_cis = RotaryPosEmbedReal.get_freqs_real((16, 16, 16), (300, 512, 512), theta=10000) + latents = torch.randn(1, 4, 8, 8) + text = torch.randn(1, 5, 16) + negative = torch.randn(1, 3, 16) + text_mask = torch.ones(1, 5, dtype=torch.bool) + negative_mask = torch.ones(1, 3, dtype=torch.bool) + + with torch.no_grad(): + for step in range(num_inference_steps): + timestep = torch.full((1,), 1.0 - step / num_inference_steps) + # Conditional pass ... + cond = model( + hidden_states=latents, + timestep=timestep, + text_hidden_states=text, + freqs_cis=freqs_cis, + text_attention_mask=text_mask, + ) + # ... always followed by the unconditional pass (parity). + uncond = model( + hidden_states=latents, + timestep=timestep, + text_hidden_states=negative, + freqs_cis=freqs_cis, + text_attention_mask=negative_mask, + ) + latents = uncond + 4.0 * (cond - uncond) + + assert latents.shape == (1, 4, 8, 8) + assert torch.isfinite(latents).all() diff --git a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py index 776ca5bc6..b49ada83e 100644 --- a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py +++ b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py @@ -1,13 +1,27 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from contextlib import contextmanager from dataclasses import dataclass +from types import SimpleNamespace +import cache_dit import pytest import torch +from cache_dit import BlockAdapter from torch import nn -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig, TransformerConfig +from vllm_omni.diffusion.cache.cachedit import CacheDiTBackend, RequestScopedCacheDiTRuntime +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.config import set_current_diffusion_config +from vllm_omni.diffusion.data import ( + AttentionConfig, + DiffusionCacheConfig, + DiffusionOutput, + OmniDiffusionConfig, + TransformerConfig, +) +from vllm_omni.diffusion.models.interface import adopt_request_scoped_cache_dit from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( MammothModa2DiTPipeline, @@ -17,6 +31,7 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import ( from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -110,6 +125,10 @@ def _pipeline_shell() -> MammothModa2DiTPipeline: pipeline.device = torch.device("cpu") pipeline.config = _build_mammoth_config(_od_config()) pipeline._llm_hidden_size = 8 + # Cache-DiT runtime state that __init__ normally creates (disabled here, + # matching a pipeline built without a cache_backend). + pipeline._cache_dit_runtime = RequestScopedCacheDiTRuntime(pipeline) + pipeline._cache_dit_config = None return pipeline @@ -118,12 +137,14 @@ def _batch( request_id: str = "req-7", prompt: object | None = None, sampling: OmniDiffusionSamplingParams | None = None, + height: int = 32, + width: int = 48, ) -> DiffusionRequestBatch: if prompt is None: prompt = { "prompt": "", - "height": 32, - "width": 48, + "height": height, + "width": width, "additional_information": { "full_hidden_states": torch.arange(32, dtype=torch.float32).reshape(4, 8), "full_token_ids": [10, 11, 100, 101], @@ -390,3 +411,162 @@ def test_forward_rejects_missing_visual_tokens_before_model_access() -> None: } with pytest.raises(ValueError, match="no visual-token hidden states.*req-empty"): _pipeline_shell().forward(_batch(request_id="req-empty", prompt=prompt)) + + +@contextmanager +def _force_torch_sdpa(): + """Pin TORCH_SDPA so CPU shape tests do not pick CUDA-only backends (FA3).""" + od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default="TORCH_SDPA"), + parallel_config=SimpleNamespace(ring_degree=1), + ) + with set_current_diffusion_config(od_config): + yield + + +def _constructible_raw_config() -> dict: + """Checkpoint config whose DiT/VAE are small enough to actually run on CPU.""" + return { + "model_type": "mammothmoda2", + "llm_config": { + "model_type": "mammothmoda2_qwen2_5_vl", + "text_config": { + "model_type": "mammothmoda2_qwen2_5_vl_text", + "hidden_size": 8, + "gen_vocab_start_index": 100, + }, + }, + "gen_vae_config": {"block_out_channels": [8, 8], "norm_num_groups": 8}, + "gen_dit_config": { + "patch_size": 2, + "in_channels": 4, + "hidden_size": 96, + "num_layers": 4, + "num_refiner_layers": 1, + "num_attention_heads": 2, + "num_kv_heads": 2, + "multiple_of": 8, + "axes_dim_rope": [16, 16, 16], + "axes_lens": [300, 512, 512], + "text_feat_dim": 16, + }, + # Top-level axes feed the pipeline-owned rotary table (see pipeline __init__). + "gen_axes_dim_rope": [16, 16, 16], + "gen_axes_lens": [300, 512, 512], + } + + +def _cache_dit_od_config() -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model="/models/MammothModa2-Preview", + model_class_name="MammothModa2DiTPipeline", + tf_model_config=TransformerConfig.from_dict(_constructible_raw_config()), + cache_backend="cache_dit", + cache_config=DiffusionCacheConfig(), + ) + + +@pytest.mark.skipif( + current_omni_platform.is_rocm(), + reason="vLLM ROCm custom ops lack CPU fallback", +) +def test_forward_transitions_request_scoped_cache_dit_across_requests(request: pytest.FixtureRequest) -> None: + """Cache-DiT request lifecycle on one production-built pipeline instance. + + Mirrors the diffusion runner startup (``get_cache_backend`` → ``enable`` → + ``adopt_request_scoped_cache_dit``) on a tiny pipeline built through the + real ``__init__``, then drives three requests through the real ``forward()``: + + 1. a partial-CFG request installs the request-scoped cache context: + cond/uncond forwards stay paired on every denoise step (parity + accounting), even outside ``cfg_range``, and blocks record cached steps; + 2. a guidance=1.0 request runs with hooks disabled (``prepare(None)``) and + its output is bit-identical to a reference captured before the backend + was ever enabled, proving a clean uninstall with no stale module state; + 3. a CFG request with a different step count and image size re-installs the + hooks through a fresh context, so no stale shape or step count survives. + """ + torch.manual_seed(0) # Deterministic random init → deterministic cache hits. + with _force_torch_sdpa(): + pipeline = MammothModa2DiTPipeline(od_config=_cache_dit_od_config()) + pipeline.eval() + assert pipeline._cache_dit_config is not None + request.addfinalizer(pipeline._cache_dit_runtime.disable) + + # The default prompt carries 2 text + 2 visual condition tokens, so the + # conditional pass sees 4 tokens and the uncond pass (empty prompt) sees 0. + cond_tokens, uncond_tokens = 4, 0 + seq_lens: list[int] = [] + + def record_call(_module, _args, kwargs): + text = kwargs.get("text_hidden_states") + seq_lens.append(int(text.shape[1]) if text is not None else -1) + + hook = pipeline.gen_transformer.register_forward_pre_hook(record_call, with_kwargs=True) + + unguided_sampling = OmniDiffusionSamplingParams( + height=32, width=48, seed=4321, guidance_scale=1.0, num_inference_steps=8 + ) + partial_cfg_sampling = OmniDiffusionSamplingParams( + height=32, + width=48, + seed=1234, + guidance_scale=4.0, + num_inference_steps=8, + extra_args={"cfg_range": [0.2, 0.8]}, + ) + changed_cfg_sampling = OmniDiffusionSamplingParams( + height=48, + width=64, + seed=777, + guidance_scale=6.0, + num_inference_steps=10, + extra_args={"cfg_range": [0.1, 0.9]}, + ) + + # Reference run for phase 2: identical request while hooks were never + # installed on this pipeline. + assert not pipeline.is_cache_dit_enabled() + reference = pipeline.forward(_batch(sampling=unguided_sampling)).output + assert seq_lens == [cond_tokens] * 8 + assert torch.isfinite(reference.float()).all() + + # Runner startup: enable once, then the pipeline owns all transitions. + backend = get_cache_backend("cache_dit", DiffusionCacheConfig()) + assert isinstance(backend, CacheDiTBackend) + backend.enable(pipeline) + assert adopt_request_scoped_cache_dit(pipeline, backend) + assert pipeline.is_cache_dit_enabled() + + # Phase 1: partial CFG keeps cond/uncond paired on every step, including + # the steps whose normalized index falls outside cfg_range [0.2, 0.8]. + seq_lens.clear() + pipeline.forward(_batch(sampling=partial_cfg_sampling)) + assert seq_lens == [cond_tokens, uncond_tokens] * 8 + assert pipeline.is_cache_dit_enabled() + assert BlockAdapter.is_cached(pipeline.gen_transformer) + hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0] + assert hits.cached_steps or hits.cfg_cached_steps + + # Phase 2: no-CFG request disables the hooks and must reproduce the + # never-enabled reference bit-for-bit. + seq_lens.clear() + unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output + assert seq_lens == [cond_tokens] * 8 + assert not pipeline.is_cache_dit_enabled() + assert not BlockAdapter.is_cached(pipeline.gen_transformer) + assert torch.equal(unguided, reference) + + # Phase 3: CFG with a new step count and image size re-installs the hooks + # on a fresh context; output shape follows the new request geometry. + seq_lens.clear() + regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output + assert seq_lens == [cond_tokens, uncond_tokens] * 10 + assert pipeline.is_cache_dit_enabled() + assert BlockAdapter.is_cached(pipeline.gen_transformer) + hits = cache_dit.summary(pipeline.gen_transformer, logging=False)[0] + assert hits.cached_steps or hits.cfg_cached_steps + assert regenerated.shape == (1, 3, 12, 16) + assert torch.isfinite(regenerated.float()).all() + + hook.remove() diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml index 76027424c..e0b652915 100644 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ b/vllm_omni/deploy/mammoth_moda2.yaml @@ -20,3 +20,13 @@ stages: enforce_eager: true trust_remote_code: true enable_prefix_caching: false + # Optional Cache-DiT acceleration for the DiT stage (~2x faster denoise, + # training-free). Uncomment to enable; guidance=1.0 requests bypass it. + # cache_backend: cache_dit + # cache_config: + # Fn_compute_blocks: 1 + # Bn_compute_blocks: 0 + # max_warmup_steps: 4 + # residual_diff_threshold: 0.24 + # max_continuous_cached_steps: 3 + # enable_cache_dit_summary: true diff --git a/vllm_omni/diffusion/cache/cachedit/model_specific.py b/vllm_omni/diffusion/cache/cachedit/model_specific.py index 0a25abd3c..8bc36a8b0 100644 --- a/vllm_omni/diffusion/cache/cachedit/model_specific.py +++ b/vllm_omni/diffusion/cache/cachedit/model_specific.py @@ -867,6 +867,47 @@ def enable_cache_for_magi2(pipeline: Any, cache_config: Any) -> CacheDiTEnableRe return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,)) +def _get_mammothmoda2_transformer(pipeline: Any) -> torch.nn.Module: + return pipeline.gen_transformer + + +def enable_cache_for_mammothmoda2(pipeline: Any, cache_config: Any) -> CacheDiTEnableResult: + """Cache only MammothModa2's repeated main DiT stack. + + ``Transformer2DModel`` runs three Q-Former refiners (noise / ref-image / + context) whose inputs change every denoise step, plus the ``layers`` + stack that dominates per-step compute. Only ``layers`` is a repeated + residual stack, so the ``BlockAdapter`` wraps it and the refiners stay + outside the cached region. Blocks take ``hidden_states`` plus keyword + step context and return only hidden states (``Pattern_3``). + + The pipeline runs sequential CFG: each conditional forward is followed by + an unconditional forward. cache-dit tells cond/uncond apart purely by + transformer-forward parity (``has_separate_cfg=True``), so the + ``cfg_range`` optimization that skips the unconditional pass outside the + interval would desync that accounting. Like Cosmos3, we keep the passes + paired and neutralize CFG via scale=1.0 outside the interval; the + pipeline also disables hooks for no-CFG requests (single forward per + step), whose parity the accounting cannot represent. + """ + pipeline._cache_dit_requires_paired_cfg = True + transformer = _get_mammothmoda2_transformer(pipeline) + block_adapter = BlockAdapter( + transformer=transformer, + blocks=[transformer.layers], + forward_pattern=[ForwardPattern.Pattern_3], + has_separate_cfg=True, + check_forward_pattern=True, + ) + refresh = enable_cache_for_dit( + pipeline, + cache_config, + block_adapter, + get_pipeline_transformer=_get_mammothmoda2_transformer, + ) + return CacheDiTEnableResult(refresh=refresh, targets=(block_adapter,)) + + def register_custom_dit_enablers() -> None: """Register model-specific Cache-DiT enablers. @@ -885,6 +926,7 @@ def register_custom_dit_enablers() -> None: "Cosmos3OmniPipeline": enable_cache_for_cosmos3, "Krea2Pipeline": enable_cache_for_krea2, "Magi2Pipeline": enable_cache_for_magi2, + "MammothModa2DiTPipeline": enable_cache_for_mammothmoda2, } ) 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 e2c086d4a..51ea8f431 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py @@ -695,8 +695,16 @@ class Transformer2DModel(ModelMixin, ConfigMixin): return text_hidden_states, img_tokens def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb): + # Call blocks with keyword arguments: Cache-DiT's Pattern_3 wrapper + # re-invokes blocks as block(hidden_states, **kwargs), so the cached + # and uncached paths share this exact calling convention. for layer in self.layers: - hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb) + hidden_states = layer( + hidden_states, + attention_mask=attention_mask, + image_rotary_emb=rotary_emb, + temb=temb, + ) return hidden_states def forward( 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 21977f488..2cc2cdcab 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py @@ -15,7 +15,12 @@ from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.cache.cachedit import ( + CacheDiTBackend, + CacheDiTRequestSpec, + RequestScopedCacheDiTRuntime, +) +from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery @@ -28,6 +33,9 @@ from .schedulers import FlowMatchEulerDiscreteScheduler logger = init_logger(__name__) +# Identifies the pipeline-owned Cache-DiT installation across requests. +_MAMMOTHMODA2_CACHE_DIT_KEY = "mammothmoda2:cache_dit" + def _first_request_value(value: object) -> object: if isinstance(value, (list, tuple)): @@ -169,6 +177,54 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): self._llm_hidden_size = llm_hidden_size + # Cache-DiT lifecycle: the diffusion runner enables the configured + # backend (``cache_backend`` on the deploy YAML stage entry) at startup + # and transfers ownership here via the request-scoped protocol; + # forward() then reconciles per-request state (step count, CFG parity). + # The runner also emits the cache summary when + # ``enable_cache_dit_summary`` is set. + self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self) + self._cache_dit_config: DiffusionCacheConfig | None = None + if str(getattr(od_config, "cache_backend", "") or "").lower() == "cache_dit": + cache_config = od_config.cache_config + self._cache_dit_config = ( + cache_config + if isinstance(cache_config, DiffusionCacheConfig) + else DiffusionCacheConfig.from_dict(cache_config or {}) + ) + + def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None: + """Adopt a runner-installed Cache-DiT backend (request-scoped protocol).""" + + self._cache_dit_runtime.adopt(backend, installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY) + + def is_cache_dit_enabled(self) -> bool: + """Return the request-scoped Cache-DiT installation state.""" + + return self._cache_dit_runtime.is_enabled + + def _prepare_cache_dit_for_request(self, *, num_inference_steps: int, cfg_active: bool) -> None: + """Reconcile Cache-DiT hooks with this request before the denoise loop. + + cache-dit's separate-CFG accounting assumes exactly two transformer + forwards per denoise step (conditional then unconditional), so CFG + requests refresh the context (per-request step counts vary) while + no-CFG requests run with hooks disabled: their single forward per + step cannot be represented by the parity accounting. + """ + if self._cache_dit_config is None: + return + if not cfg_active: + self._cache_dit_runtime.prepare(None) + return + self._cache_dit_runtime.prepare( + CacheDiTRequestSpec( + installation_key=_MAMMOTHMODA2_CACHE_DIT_KEY, + cache_config=self._cache_dit_config, + num_inference_steps=num_inference_steps, + ) + ) + def _reinit_caption_embedder(self, in_features: int) -> None: # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`: # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features). @@ -432,6 +488,18 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): # Run diffusion loop (CFG supported when text_guidance_scale > 1.0) total_steps = max(1, len(scheduler.timesteps)) + self._prepare_cache_dit_for_request( + num_inference_steps=total_steps, + cfg_active=request.text_guidance_scale > 1.0, + ) + # cache-dit distinguishes the conditional vs unconditional passes + # purely by transformer-forward parity, so when it is active the + # uncond pass must run on every step. Outside cfg_range CFG is + # neutralized via scale=1.0 (uncond + 1.0 * (cond - uncond) == cond), + # matching the skip optimization's cond-only update. + requires_paired_cfg = self._cache_dit_runtime.is_enabled and getattr( + self, "_cache_dit_requires_paired_cfg", False + ) for i, t in enumerate(scheduler.timesteps): timestep = t.expand(latents.shape[0]).to(latents.dtype) model_pred = self.gen_transformer( @@ -447,7 +515,10 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery): guidance_scale = ( request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.cfg_range[1] else 1.0 ) - if guidance_scale > 1.0 and negative_prompt_embeds is not None: + run_uncond = guidance_scale > 1.0 and negative_prompt_embeds is not None + if requires_paired_cfg and negative_prompt_embeds is not None: + run_uncond = True + if run_uncond: model_pred_uncond = self.gen_transformer( hidden_states=latents, timestep=timestep,