[Xenoryn]

Purpose

Enable regional torch.compile on the MammothModa2 DiT stage (Transformer2DModel) with fullgraph=True, and validate parity plus no shape-driven recompile across resolutions and CFG text lengths.

Sub-issue of #7075: closes #7139.

#7134 has landed, so stage 1 now runs on the shared DiffusionModelRunner. Its _compile_transformer already reads the pipeline's _dit_modules = ["gen_transformer"] and calls regionally_compile; all this PR has to do is give it a compile target. The earlier [TEMPORARY] bridge commit has been dropped — the branch is now two commits, both permanent.

AR compile / DiT CUDA graph capture / TeaCache / TP / SP are tracked in their own PRs and out of scope here.

What changed in the DiT forward

Block topology is untouched. One class attribute wires the compile target, and after the rebase that is the entire source change.

Declare the compile target. Transformer2DModel gets _repeated_blocks = ["TransformerBlock"]. Without it regionally_compile has nothing to walk into and logs "Regional compilation skipped" — which is what happens on main today.

The two block-body edits this PR originally carried have since landed on main independently, and were dropped on rebase:

• The empty-sequence skip now lives one level up, in _apply_refiners (if text_hidden_states.shape[1] > 0), added upstream for the Online-FP8 M=0 constraint. It is the same hoist this PR originally made, for a different reason, and it solves the same compile problem: the block-level if sequence_length == 0 short-circuit guards on a backed SymInt (sequence_length comes from hidden_states.shape), so Dynamo traces it fine but specializes on it and compiles a second graph the first time CFG's zero-token unconditional stream arrives. With the skip in _apply_refiners, that stream never reaches a block at all. _apply_refiners is the only in-tree caller of context_refiner (the extractor at teacache/extractors.py that also iterates a context_refiner is extract_zimage_context, a different model), and noise_refiner / ref_image_refiner / layers always carry img_len >= 1.
• sdpa.py's GQA-expansion path already guards its logger.debug behind torch.compiler.is_compiling(), so Dynamo no longer raises Unsupported: logging.Logger method not supported for non-export cases under fullgraph=True.

Both fixes are on main, so this PR is now a single class attribute plus a test. The AttnProcessor block-level short-circuit stays as a backstop; it is unreachable on the CFG path for the reason above.

How to enable

Set Stage 1 enforce_eager: false in vllm_omni/deploy/mammoth_moda2.yaml. This PR adds that as a commented line next to the Cache-DiT block, and documents it in the recipe (per review feedback). It combines with Cache-DiT.

No other deploy-YAML changes are needed. Earlier revisions of this PR described cudagraph_mode: NONE and custom_ops: [all] workarounds; those were artifacts of stage 1 running as LLM_GENERATION before #7134. DiffusionModelRunner builds its own CompilationConfig() (vllm_omni/diffusion/vllm_config.py:200) and never reads the stage's compilation_config, so both settings are inert on a diffusion stage.

Test Plan

• Compile parity unit tests — tests/diffusion/models/mammoth_moda2/test_dit_compile_parity.py (new) drives the same entry point the runner uses, regionally_compile(model, dynamic=True, fullgraph=True), so a graph break fails the test instead of silently falling back to eager. Asserts eager/compiled parity across text_len ∈ {0, 4, 77} (the 0 case exercises the empty-text skip in _apply_refiners), and asserts no additional frame compilations across a resolution and text-length sweep by patching torch._dynamo.convert_frame._compile.
• Empty-sequence coverage — unchanged from main: test_dit_attention.py covers the _apply_refiners skip, test_dit_attention_cuda.py covers the block-level short-circuit on the default backend. This PR no longer modifies either file.
• Real-model end-to-end on 4× H800 with MammothModa2-Dev: enforce_eager: true versus false, image parity by PSNR, plus TORCH_LOGS=graph_breaks,recompiles.

vLLM Version: 0.31.0
vLLM-Omni Commit: rebased onto bea9b9e9e

Test Result

Provenance note (post-rebase): every number below was measured on the pre-rebase revision of this branch. The compiled region has changed since — the AttnProcessor short-circuit is back in place, unreachable on this path as explained above — so these figures are being re-confirmed on the current revision and will be updated if they move.

Unit tests

Scope	Result
tests/diffusion/models/mammoth_moda2 + tests/diffusion/layers/test_norm.py	186 passed, 1 skipped
same, with -m "core_model and cpu" (the "Simple · Diffusion Test" lane marker, one process)	143 passed

The second row is the case @MrlixiangWE reported failing. Reproduced it exactly against the pre-fix file (1 failed, 142 passed, assert 0 >= 1) and confirmed the fix clears it. Details in the review reply.

End-to-end T2I

4× H800, MammothModa2-Dev, 1024×1024, 50 steps, guidance 4.0, seed 42, FLASH_ATTN, two prompts (studio cat, Samoyed) × (1 warmup + 2 timed). Eager and compiled run as separate processes on separate GPUs.

Stage 1	DiT stage time	Startup
enforce_eager: true	4,681 ms	62.3 s
enforce_eager: false	4,024 ms (−14.03%)	67.3 s (+5 s compile)

Per-run spread is tight — eager [4664, 4673, 4697, 4689], compiled [4017, 4018, 4028, 4034] — so the 657 ms gap is well outside noise.

Image parity against eager: PSNR 50.63 dB (cat), 48.04 dB (Samoyed); max abs diff 31/26 on uint8. Visually identical.

On the gap versus @MrlixiangWE's −29.87%: that measurement was 1× A800 with MammothModa2-Preview, where the eager DiT stage takes 15,640 ms. Here the eager baseline is 4,681 ms — faster hardware, a different checkpoint, and a DiT stage time that also covers non-compiled work (VAE decode). Same direction, different denominator; the compiled region itself is not regressing. This is also why the earlier revision of this description claimed a 2.40× DiT speedup: those numbers were taken pre-rebase on vLLM 0.29 with a per-forward CUDA-event probe injected into the LLM_GENERATION engine core, before main landed the fused QK-norm/RoPE work that shrank the eager baseline. They are not comparable to the current code and have been removed rather than restated.

E2E wall-clock is not a useful figure here: the AR stage is ~96% of it (≈100 s of ≈105 s), and it drifted 4% between the two processes despite being identical code, which swamps the DiT delta. AR compile is tracked separately under #7075.

Graph breaks and recompiles

TORCH_LOGS=graph_breaks,recompiles on a compiled run:

• No "Regional compilation skipped" — compile is engaged.
• Graph breaks only at flash_attn.py:734 — the if use_masked_path: branch, on the torch.any(~attention_mask) result computed at :730. Nine occurrences, no other break site. Pre-existing and backend-side, matching @MrlixiangWE's finding at line 463 of cb5f508; not introduced or worsened here. This box resolves FA3, so the FA2 dense-fallback break they also saw does not appear.
• Exactly one recompile, from a norm1 type guard — RMSNorm in the context refiner versus LuminaRMSNormZero elsewhere. That is the two-block-kind startup compile @MrlixiangWE called expected.

Known limitation, not addressed here

FlashAttentionImpl._flash_varlen_wrapper builds torch.arange(0, (batch_size + 1) * q_len, q_len) (flash_attn.py:407), which raises RuntimeError: step must be nonzero when q_len == 0. Reaching it needs a varlen-only backend and no mask metadata, so no MammothModa2 path hits it — AttnProcessor always supplies AttentionMetadata when a mask is present, and the masked path unpads fine at length 0 (verified on GPU). Worth hardening at the backend layer, but that is independent of this PR.

[MrlixiangWE]

Measured with this PR merged onto main (cb5f508) on 1x A800 with vLLM 0.30.0, MammothModa2-Preview text-to-image at 1024x1024, 50 steps, guidance 4.0, seed 42, studio-cat and Samoyed prompts, FLASH_ATTN, mean of 4 timed requests (2 per prompt after one warmup), the two runs back to back on one GPU:

Stage 1	DiT stage time
enforce_eager: true	15,640 ms
enforce_eager: false (regional compile)	10,968 ms (-29.87%)

Compiled images against eager (PSNR 50.05 dB for the cat, 25.16 dB for the Samoyed):

Compiled, it also runs with Cache-DiT, TORCH_SDPA, #7969 and #7489. On main, enforce_eager: false only logs "Regional compilation skipped" and the images match eager bit for bit.

On a compiled FLASH_ATTN run of the cat prompt with TORCH_LOGS=graph_breaks,recompiles (line numbers at main cb5f508), the compiled blocks break at vllm_omni/diffusion/attention/backends/flash_attn.py:463 (if attention_mask is not None and torch.any(~attention_mask), data-dependent branching; the same branch is at line 450 of this PR's base). On this model's path the check sits in front of the kernel call whichever FA version is loaded. With only vLLM's bundled FA2 available, flash_attn_func is None, so each block breaks again in the dense fallback (flash_attn.py:483), at _vllm_fa2_C.varlen_fwd, which has no fake-tensor kernel. The two block kinds (norm1 is Qwen2RMSNorm in the context refiner and LuminaRMSNormZero elsewhere) compile once each during startup, which is expected.

Stage 1 in vllm_omni/deploy/mammoth_moda2.yaml keeps enforce_eager: true, and the recipe does not mention compile; a commented enforce_eager: false next to the Cache-DiT block would show how to turn it on.

[MrlixiangWE]

Run with the rest of this file, this test fails with regional compile did not trigger on first shape (assert 0 >= 1; 1 failed, 3 passed) and passes on its own. The parity cases above compile the same TransformerBlock.forward code object with dynamic=True, so the first shape here hits dynamo's cache and the counter sees no compile. The file is core_model and cpu, which puts it in the ready and merge "Simple · Diffusion Test" lane that runs tests/diffusion in one process. With torch._dynamo.reset() as the first line of this test, the file passes (4 passed), and so do tests/diffusion/models/mammoth_moda2 and tests/diffusion/layers/test_norm.py (99 passed).

[Xenoryn]

Confirmed, reproduces exactly as you describe — thanks. Fixed in b28b1d1 with an autouse fixture rather than a single reset() here, so the parity cases above are isolated from each other too (they compile the same code object, so whichever runs first owns the cache).

Your command, single process, one H800 box:

test_dit_compile_parity.py	result
pre-fix	1 failed, 142 passed — assert 0 >= 1
post-fix	143 passed

Without the marker filter: 186 passed, 1 skipped.

[MrlixiangWE]

Hi @Xenoryn, gentle ping: could you update this PR and merge the latest main when you have a chance? Thanks!

[Xenoryn]

@MrlixiangWE Thanks for measuring this end to end — it caught a real bug. All three points addressed; rebased onto f4108d77e, [TEMPORARY] bridge dropped now that #7134 has landed. PR description rewritten.

1. Dynamo cache bleed — fixed, see inline reply.
2. YAML + recipe — commented enforce_eager: false next to the Cache-DiT block, recipe notes the two combine. No percentage in the comment; see below.
3. Validation — 4× H800, MammothModa2-Dev, 1024², 50 steps, guidance 4.0, seed 42, FLASH_ATTN, 2 prompts × (1 warmup + 2 timed):

Stage 1	DiT stage time
enforce_eager: true	4,681 ms
enforce_eager: false	4,024 ms (−14.03%)

PSNR 50.63 dB (cat) / 48.04 dB (Samoyed). Your three observations all reproduce: compile engaged, one recompile from the norm1 type guard, graph breaks only at flash_attn.py:734 (FA3 here, so no FA2 dense-fallback break).

[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.

[vllm-omni-review-bot]

This PR appears to belong to: docs/design/module/diffusion/index.md, docs/design/module/diffusion/diffusion_model_integration.md.

Module owners: @wtomin @RuixiangMa @david6666666

Routing: @wtomin via module of the changed files, module named in the PR description, semantic router, CODEOWNERS; @RuixiangMa via module of the changed files, module named in the PR description, semantic router; @david6666666 via module of the changed files, module named in the PR description, CODEOWNERS

@Xenoryn, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Strict on zcode (GLM-5.3-Flash) under experiment fleet-strict-cursor-grok46-zcode-glm53flash-5050-c5-z10-20261002.

[MrlixiangWE]

@hsliuustc0106 The review points are addressed at b28b1d1, and its two diffusion lanes from test-ready.yml pass on our 2x A800 with 0 failures. Could you help trigger CI and merge it when you have a chance?

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 7 days

@Xenoryn this pull request has had no human commit, comment or review since 2026-10-12. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[MrlixiangWE]

@Xenoryn The branch now conflicts with main in vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py, vllm_omni/diffusion/attention/backends/sdpa.py, vllm_omni/deploy/mammoth_moda2.yaml and tests/diffusion/models/mammoth_moda2/test_dit_attention_cuda.py. Could you merge main? I'll recheck the compile path on the merged head.

[Xenoryn]

@MrlixiangWE Rebased onto bea9b9e9e (head fa83b66c3); all four conflicts resolved.

Heads-up on what the rebase dropped: two of the three edits this PR carried are now on main independently — the empty-sequence skip lives in _apply_refiners (the upstream Online-FP8 M=0 guard), and sdpa.py already guards its GQA-expansion logger.debug behind torch.compiler.is_compiling(). The AttnProcessor block-level short-circuit is back in place as a backstop. Net, the branch is now _repeated_blocks + test_dit_compile_parity.py.

[MrlixiangWE]

The Test Result numbers predate #7221, which moved the TransformerBlock attention and feed-forward linears to vLLM layers; please re-measure them at the new head.

[MrlixiangWE]

Since #7221, TransformerBlock builds vLLM ReplicatedLinear layers, and at fa83b66 all four tests in this file fail, alone and in the Simple · Diffusion Models Test lane command (pytest tests/diffusion/models -m 'core_model and cpu'), with AssertionError: tensor model parallel group is not initialized. Adding mock_tp1 is not enough on its own: copy.deepcopy(eager) then raises TypeError: BasevLLMParameter.__new__() takes 2 positional arguments but 3 were given, and the vLLM linears are left uninitialized (torch.empty), so in our run the eager output came out NaN for text_len 4 and 77. Please add pytest.mark.usefixtures("mock_tp1") to pytestmark, call initialize_block_weights(model) from .helpers at the end of _build_model(), and build compiled with a second _build_model() instead of copy.deepcopy. With those three changes the file passes (4 passed).

[MrlixiangWE]

With mock_tp1, initialize_block_weights and a second _build_model() in place, setting _repeated_blocks = [] still passes all four tests. vLLM compiles SiluAndMul.forward_native (vllm/model_executor/layers/activation.py:140) by itself, so the first shape counts one frame without any regional compile, and the parity cases then compare eager with eager. Please count only TransformerBlock.forward here, for example def _tracked(code, *args, **kwargs): with if code is TransformerBlock.forward.__code__: outer.count += 1 and TransformerBlock imported from the DiT model module. With that, emptying _repeated_blocks turns test_forward_no_shape_driven_recompile_across_resolutions_and_text_len red.

[Xenoryn]

@MrlixiangWE Both applied at 333774728 — thanks for the precise repros, the second one especially.

1. mock_tp1 in pytestmark, initialize_block_weights at the end of _build_model(), and compiled from a second _build_model().
2. _CompileCounter now counts only TransformerBlock.forward.__code__.

Confirmed on our machine: 4 passed standalone and under the lane marker, and emptying _repeated_blocks turns the recompile test red as you described.

[MrlixiangWE]

Both inline points are fixed at 3337747. The PR description is unchanged since fa83b66: the unit-test table (186 passed, 1 skipped), the end-to-end numbers and vLLM Version: 0.31.0 predate #7221, and the base pins vLLM 0.32.0 since #9066. Please re-measure at 3337747 and update the Test Result section.