[nagisa-kunhah]

Purpose

part of #1217
Related: #7106

Add TeaCache support for the MammothModa2 DiT stage (text-to-image).

Implementation Approach

The MammothModa2 DiT stage already runs in the shared diffusion runtime, so cache lifecycle is handled by diffusion_model_runner.py on main: the backend is enabled once at startup and refreshed per batch through _refresh_cache_for_requests. This PR does not add any stage- or model-specific cache plumbing; it only adds what TeaCache needs to understand MammothModa2's transformer.

• Extractor (teacache/extractors.py): extract_mammoth_moda2_context, registered for MammothModa2Transformer2DModel. It reuses the transformer's own _validate_inputs / _prepare_embeddings / _apply_refiners helpers, builds the joint text+image sequence, and uses layers[0].norm1(...) as the modulated input. The transformer layers and norm_out + unpatchify are exposed as run_transformer_blocks / postprocess.
• Calibrated defaults (teacache/config.py): fitted MammothModa2 coefficients and a model-specific default rel_l1_thresh=0.075 (see Coefficient Fitting).
• Enabler (teacache/backend.py): enable_mammoth_moda2_teacache, registered in CUSTOM_TEACACHE_ENABLERS["MammothModa2DiTPipeline"]. It hooks gen_transformer and aliases it as pipeline.transformer, following the existing Bagel pattern. The transformer then appears under two names in state_dict(); the parameters are shared, so no extra memory is used.
• Explicit CFG branch hints: the batched _denoise_group forward passes teacache_branch="positive" on the conditional gen_transformer(...) call and teacache_branch="negative" on the unconditional call. MammothModa2's transformer does not set do_true_cfg, so without these hints both CFG branches would share the same TeaCache state. The transformer forward accepts and drops the kwarg.
• Hook (teacache/hook.py): TeaCacheHook.new_forward prefers ctx.extra_states["teacache_branch"] when present (validated to be positive / negative), and falls back to the existing do_true_cfg / CFG-parallel detection otherwise. Other models are unaffected.
• Example: text_to_image.py now passes cache_config = {} for tea_cache, so TeaCache selects each model's own default threshold and coefficients. This changes the effective threshold for every model using the example, from the previously hardcoded 0.2 to the per-model default (e.g. 0.075 for MammothModa2, 0.17 for MiniMax-H3, 0.2 for models without a specific default).
• Docs: MammothModa2 TeaCache marked as supported in docs/user_guide/diffusion_features.md; recipe documents the TeaCache command and the default threshold.

Note on request batching: the coefficients were fitted on single-request traces. With request batching (#7476) the whole DiT batch shares one skip decision per step. The batched e2e run below shows quality comparable to the single-request run at the default threshold, but quality under batching can still vary with the batch composition.

Test Plan

Hardware and runtime:

• NVIDIA RTX PRO 6000 Blackwell Server Edition, sm_120, 97887 MiB, driver 595.58.03
• vLLM 0.30.0, PyTorch 2.13.0+cu130
• BF16, eager (enforce_eager), single GPU, diffusion attention CUDNN_ATTN
• Model: bytedance-research/MammothModa2-Preview, snapshot ef5a5e41dbf0de1ef6275586b7580f0d4248b4c6
• Deploy config: vllm_omni/deploy/mammoth_moda2.yaml (DiT stage max_num_seqs: 8)

Sources:

Arm	vLLM-Omni commit	Cache
main baseline	ee8fdab1dc26ea8f09d7910cdc98a17c4b0ea02a (current main)	disabled
this PR	c237ea745678de5f3e2c3463492c16343526a337 (rebased onto ee8fdab1d)	disabled
this PR	same	TeaCache, rel_l1_thresh = 0.075 / 0.10 / 0.15 / 0.20

Request settings for all runs:

height=1024
width=1024
seed=142
text_guidance_scale=4.0   (CFG enabled)
cfg_range=[0.0, 1.0]
num_inference_steps=50

Prompts:

Slug	Prompt
nyc_motorcycle	A stylish woman riding a motorcycle through rainy New York City at night, cinematic movie poster style, neon reflections, sharp focus.
cat_laptop	A fluffy orange cat sitting on a laptop keyboard in a cozy home office, warm desk lamp, shallow depth of field, detailed fur.
futuristic_city	A futuristic coastal city at sunset with glass skyscrapers, flying transit vehicles, sailboats on calm water, detailed architecture, soft pink sky.
hf_city_zh	这张图片展示了一座现代化城市的美丽景象，画面中显眼的是一座高耸入云的摩天大楼，其外立面在夕阳余晖的映照下显得格外醒目。周围环绕着多栋风格各异的高楼大厦，这些大楼的窗户透出点点灯光，显示出城市的繁华。

Method:

• Each arm runs in a separate engine process through Omni(...) with cache_backend="tea_cache", cache_config={"rel_l1_thresh": <thresh>} (or no cache). Model loading is excluded from all timings, and one full warmup request is discarded.
• Sequential: each prompt is submitted alone, 4 prompts x 2 repeats.
• Batched: all 4 prompts are submitted together in one generate() call, 2 repeats. The DiT stage actually formed a batch of 1 followed by a batch of 3 (the first AR request finishes slightly earlier), so the batched run covers several requests sharing one TeaCache skip decision.
• Stage timings come from stage_durations (stage_0_gen_ms = AR, stage_1_gen_ms = DiT).
• Peak memory: the engine-reported peak_memory_mb, plus nvidia-smi sampled every second for the whole process.
• Quality: TeaCache outputs are compared with the main baseline for the same prompt, seed, and mode (sequential vs sequential, batched vs batched). Metrics are PSNR, grayscale SSIM, MAE / RMSE (0-255 scale), and LPIPS (SqueezeNet).

Equivalent CLI (TeaCache, default threshold):

python examples/offline_inference/text_to_image/text_to_image.py \
  --model /path/to/MammothModa2-Preview \
  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
  --prompt "${PROMPT}" \
  --height 1024 --width 1024 --seed 142 \
  --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
  --cache-backend tea_cache \
  --output "${SLUG}_teacache.png"

Test Result

Correctness without cache

With the cache disabled, this PR produced output pixel-identical to main for all 4 prompts in both sequential and batched mode. Latency also matched main (E2E 1.002x, DiT 1.001x). Every arm was deterministic: the two repeats of each prompt were pixel-identical, in both modes.

Sequential (single request) — speed and quality vs main

Mean of 4 prompts x 2 repeats.

Arm	E2E (s)	AR (s)	DiT (s)	E2E speedup	DiT speedup	PSNR	SSIM gray	MAE	RMSE	LPIPS
main, no cache	69.17	60.98	8.134	1.000x	1.000x	–	–	–	–	–
PR, no cache	69.06	60.89	8.127	1.002x	1.001x	identical	1.0000	0.00	0.00	0.000
PR, TeaCache 0.075 (default)	67.54	60.66	6.830	1.024x	1.191x	33.77	0.9731	2.47	5.71	0.021
PR, TeaCache 0.10	66.79	60.58	6.161	1.036x	1.320x	32.54	0.9695	2.90	6.33	0.025
PR, TeaCache 0.15	65.44	60.34	5.055	1.057x	1.609x	27.99	0.9406	4.52	10.33	0.041
PR, TeaCache 0.20	65.60	61.21	4.344	1.054x	1.873x	25.42	0.9131	6.70	13.90	0.067

E2E standard deviation was about 0.2-0.5 s per arm, and DiT standard deviation was at most 0.1 s. The AR stage (~61 s) dominates end-to-end latency on this GPU and is not affected by TeaCache, so the useful signal is the DiT-stage speedup.

Per-prompt results at the default threshold (0.075):

Prompt	main E2E (s)	TeaCache E2E (s)	E2E speedup	main DiT (s)	TeaCache DiT (s)	DiT speedup	PSNR	SSIM gray	MAE	LPIPS
nyc_motorcycle	69.58	67.13	1.036x	8.127	6.905	1.177x	35.33	0.9872	1.75	0.010
cat_laptop	69.15	67.61	1.023x	8.137	6.903	1.179x	28.60	0.9656	4.23	0.044
futuristic_city	69.05	67.56	1.022x	8.092	6.790	1.192x	38.57	0.9845	1.37	0.009
hf_city_zh	68.88	67.84	1.015x	8.183	6.721	1.217x	32.58	0.9551	2.53	0.021
Average	69.17	67.54	1.024x	8.134	6.830	1.191x	33.77	0.9731	2.47	0.021

Per-prompt threshold sweep (sequential):

Prompt	rel_l1_thresh	E2E speedup	DiT speedup	PSNR	SSIM gray	MAE
nyc_motorcycle	0.075	1.036x	1.177x	35.33	0.9872	1.75
nyc_motorcycle	0.10	1.043x	1.325x	32.41	0.9791	2.64
nyc_motorcycle	0.15	1.065x	1.610x	28.42	0.9596	3.59
nyc_motorcycle	0.20	1.062x	1.880x	26.41	0.9404	6.07
cat_laptop	0.075	1.023x	1.179x	28.60	0.9656	4.23
cat_laptop	0.10	1.035x	1.308x	28.37	0.9606	4.86
cat_laptop	0.15	1.063x	1.590x	26.06	0.9421	6.41
cat_laptop	0.20	1.045x	1.868x	23.34	0.9319	8.72
futuristic_city	0.075	1.022x	1.192x	38.57	0.9845	1.37
futuristic_city	0.10	1.032x	1.314x	35.96	0.9788	1.76
futuristic_city	0.15	1.046x	1.615x	30.38	0.9551	3.24
futuristic_city	0.20	1.058x	1.873x	27.46	0.9171	4.80
hf_city_zh	0.075	1.015x	1.217x	32.58	0.9551	2.53
hf_city_zh	0.10	1.033x	1.334x	33.42	0.9594	2.35
hf_city_zh	0.15	1.053x	1.622x	27.11	0.9057	4.82
hf_city_zh	0.20	1.052x	1.869x	24.46	0.8629	7.21

Batched (4 concurrent requests, DiT max_num_seqs: 8) — speed and quality vs main

Each repeat submits all 4 prompts at once. Results are the mean of 2 repeats. Quality is compared with main's batched output for the same prompt.

Arm	Batch wall (s)	Throughput (img/min)	Wall speedup	Mean per-request DiT (s)	DiT speedup	PSNR	SSIM gray	MAE	RMSE	LPIPS
main, no cache	99.12	2.42	1.000x	26.81	1.000x	–	–	–	–	–
PR, no cache	98.64	2.43	1.005x	26.83	0.999x	identical	1.0000	0.00	0.00	0.000
PR, TeaCache 0.075 (default)	94.16	2.55	1.053x	22.15	1.211x	34.70	0.9767	1.94	4.80	0.014
PR, TeaCache 0.10	90.39	2.66	1.097x	19.76	1.357x	31.33	0.9620	2.94	7.07	0.025
PR, TeaCache 0.15	85.60	2.80	1.158x	15.97	1.679x	28.85	0.9437	4.47	9.37	0.040
PR, TeaCache 0.20	82.87	2.90	1.196x	13.65	1.964x	28.14	0.9347	5.10	10.05	0.046

At the default threshold, batched quality (SSIM 0.977, LPIPS 0.014) is in line with the sequential run (SSIM 0.973, LPIPS 0.021). The shared per-batch skip decision did not degrade quality in this sample. The CFG positive/negative branches stayed separated: outputs were deterministic across repeats and close to the no-cache baseline.

Note: on main itself, batched output differs from sequential output for the same prompt and seed (mean SSIM 0.87 on this set). The difference already appears for the request that ran its DiT step alone, so it comes from AR-stage batching, not from the DiT or this PR. For that reason, batched TeaCache output is compared with batched main output.

GPU memory

Arm	Engine-reported peak (MiB)	nvidia-smi process peak (MiB)
main, no cache	10450	69670
PR, no cache	10450	69670
PR, TeaCache 0.075	10450	69670
PR, TeaCache 0.10	10618	69838
PR, TeaCache 0.15	10618	69838
PR, TeaCache 0.20	10450	69670

The engine-reported value is the sequential DiT peak. The nvidia-smi peak covers the whole process, including AR KV-cache pre-allocation and batched DiT. TeaCache adds at most 168 MiB (+1.6% of the reported peak), and adds nothing at the default threshold.

Coefficient Fitting

The bundled MammothModa2 TeaCache coefficients were fitted from full-compute traces at 1024x1024 with 50 denoising steps. A 70-prompt calibration set was collected with the TeaCache hook installed but forced to full compute, recording each CFG branch separately. A degree-4 polynomial was then fitted, mapping the adjacent-step relative L1 distance of the modulated input to the relative L1 distance of the transformer output.

Fitted coefficients (default in vllm_omni/diffusion/cache/teacache/config.py):

[-1761.242764481119, 859.398730831851, -126.78044436937441, 9.16173963457422, -0.1102826051200547]

Fitting data:

70 train prompts x 2 CFG branches x 49 transitions = 6860 valid points
input_rel_l1 range: 0.02318 - 0.27042
output_rel_l1 range: 0.03098 - 1.12319
degree-4 ordinary least-squares fit: MAE=0.04749, RMSE=0.07383, R^2=0.6663

The default rel_l1_thresh=0.075 gives about 1.19-1.21x DiT speedup with SSIM about 0.97-0.98 and LPIPS about 0.01-0.02. Higher thresholds trade quality for speed, as shown in the sweeps above.

Regression Tests

pytest -q \
  tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_teacache.py \
  tests/diffusion/cache/test_cache_backends.py \
  tests/diffusion/cache/test_teacache_extractors.py
# 87 passed

pytest -q tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py
# 55 passed

test_mammoth_moda2_teacache.py drives the batched forward(req) -> list[DiffusionOutput] path (_parse_request -> _denoise_group) with 1 and 2 requests. It checks:

• the branch order for non-CFG, CFG, and partial cfg_range cases;
• that AR image conditioning is forwarded to the positive branch.

BEFORE SUBMITTING: read CONTRIBUTING.md and run the precheck-pr skill with the code agent for a self-check against project conventions.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[hsliuustc0106]

Thanks for the rework — the extractor, the hook branch separation, and the calibration write-up are solid, and all five issues from the July 30 round are resolved (the string cache_config path is now covered by main's _normalize_cache_config, and Cache-DiT reaches Mammoth via the request-scoped path). Two blockers remain before this can merge.

1. Rebase required — the branch is still based on pre-#7476 main (blocking)

The head's pipeline_mammothmoda2_dit.py still has supports_request_batch = False and the old single-request forward, while main (since #7476, merged Sep 25) ships the batched _denoise_group forward. GitHub currently reports MERGEABLE, but a textual merge would effectively revert #7476. Please rebase onto current main — ideally after #7954 lands, since it refactors the same file — and port the branch hints into the batched forward (see the inline comment on the pipeline hunk for the exact placement and failure mode). The new test file also targets the old forward and needs rework along the lines of test_pipeline_mammothmoda2_dit.py::_pipeline_shell (second inline comment).

After the rebase, please re-run the e2e suite on the new base SHA and add at least one batched run (max_num_seqs: 8, 2+ concurrent requests) with --cache-backend tea_cache — the current numbers are single-request-era, and with batching the whole batch shares one skip decision.

2. The PR description no longer matches the diff (blocking, quick fix)

The Implementation Approach describes plumbing that was dropped during the September rebase: the stage engine_args forwarding, the MammothModa2ForConditionalGeneration enable path, and the pipeline-stored backend + in-forward refresh. Those responsibilities now live in diffusion_model_runner.py (enable at startup, _refresh_cache_for_requests per batch — this path is complete on main, and TeaCacheBackend.refresh ignores num_inference_steps, so the extra-body steps resolution is safe). Please rewrite the section to match what the PR actually contains today: the MammothModa2 extractor, the calibrated coefficients + rel_l1_thresh=0.075 default, the MammothModa2DiTPipeline enabler (gen_transformer + transformer alias), the pipeline branch hints, hook teacache_branch support, tests, and docs.

Non-blocking notes

• The example change (cache_config = {} for tea_cache) shifts the effective threshold for every model using the example (hardcoded 0.2 → per-model defaults, e.g. 0.17 for MiniMax-H3). Reasonable, but worth a sentence in the description.
• The coefficients were fitted single-request; with request batching the whole batch shares one skip decision. A note in the recipe would help set expectations.
• pipeline.transformer = transformer in the enabler follows the existing Bagel pattern — fine; just be aware the transformer then appears under two names in state_dict() (shared parameters, no extra memory).

[hsliuustc0106]

This hunk is against the pre-#7476 single-request forward. After the rebase, the hints must move into the batched _denoise_group forward: teacache_branch="positive" on the conditional gen_transformer(...) call and teacache_branch="negative" on the unconditional one (main: pipeline_mammothmoda2_dit.py:801 / :815, shifted by #7954's refactor).

Without these, TeaCacheHook.new_forward falls back to do_true_cfg detection — MammothModa2's transformer doesn't set it — so both CFG branches land in the same teacache_positive state and CFG outputs are silently corrupted. That's exactly the failure mode the branch hints exist to prevent, so please make sure the rebased e2e run exercises CFG (text_guidance_scale > 1.0) with the cache enabled.

[nagisa-kunhah]

Thanks, addressed after rebasing onto current main.

The branch hints now live in the batched _denoise_group forward:

• pipeline_mammothmoda2_dit.py:810: teacache_branch="positive" on the conditional gen_transformer(...) call (main :801)
• pipeline_mammothmoda2_dit.py:823: teacache_branch="negative" on the unconditional call (main :815)

With the hints present, the hook uses them instead of the do_true_cfg fallback, so the two branches keep separate teacache_positive / teacache_negative state. test_mammoth_moda2_teacache.py now drives _denoise_group with 1 and 2 requests and asserts the branch order for the non-CFG, CFG and partial cfg_range cases.

The rebased e2e runs all use CFG (text_guidance_scale=4.0, cfg_range=[0,1]) with TeaCache enabled, both sequential and batched (4 concurrent requests, max_num_seqs: 8). At the default rel_l1_thresh=0.075, compared with main without cache:

• sequential: SSIM 0.973, LPIPS 0.021, DiT 1.19x faster
• batched: SSIM 0.977, LPIPS 0.014, DiT 1.21x faster

See "Test Result → Sequential" and "Test Result → Batched" in the PR description.

[hsliuustc0106]

This fake targets the old single-request forward(req) -> DiffusionOutput. After the rebase the pipeline exposes forward(req) -> list[DiffusionOutput] via _parse_request → _denoise_group; please rework the shell the same way tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py::_pipeline_shell does (patched FlowMatchEulerDiscreteScheduler / randn_tensor, etc.). Keep _cache_dit_runtime on the fake — _denoise_group calls _prepare_cache_dit_for_request.

[nagisa-kunhah]

Thanks for your carefully review. The test now drives the real MammothModa2DiTPipeline.forward(req) -> list[DiffusionOutput] path (_parse_request → _denoise_group) instead of the old single-request fake, following test_pipeline_mammothmoda2_dit.py::_pipeline_shell:

• _build_pipeline patches FlowMatchEulerDiscreteScheduler and randn_tensor, and keeps _cache_dit_runtime = RequestScopedCacheDiTRuntime(pipe) for _prepare_cache_dit_for_request. The transformer and VAE are replaced with fakes; the fake transformer records the teacache_branch of each call.
• _run_pipeline calls pipe(DiffusionRequestBatch(requests)), asserts that it returns a list of DiffusionOutput with one entry per request, and that every transformer call gets timestep.shape == (num_requests,).
• The branch-order tests (non-CFG, CFG, partial cfg_range) are parametrized over num_requests=[1, 2].

See "Regression Tests" in the PR description: this file plus test_cache_backends.py and test_teacache_extractors.py give 87 passed, and test_pipeline_mammothmoda2_dit.py gives 55 passed.

[nagisa-kunhah]

@hsliuustc0106 Thanks for the careful review! Sorry about these issues: I didn't anticipate how #7476 changed the MammothModa2 DiT interfaces, so the previous revision still targeted the old single-request forward.

The branch is now rebased onto the latest main (ee8fdab1d), and the interface problems are fixed:

• The teacache_branch hints are now on the conditional and unconditional calls in the batched _denoise_group forward.
• test_mammoth_moda2_teacache.py now drives the real _parse_request → _denoise_group path, with batches of 1 and 2 requests.

I also re-ran the comparison between this PR and current main and updated the PR description.

Could you please take another look at the changes? Thanks!

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on cursor (cursor-grok-4.6-high) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.