[Bugfix][Examples][MammothModa2] Serve T2I via LLM-typed DiT stage (#7199) + benchmark harness and profiler wiring by quan-yi-ai · Pull Request #7293 · vllm-project/vllm-omni

quan-yi-ai · 2026-09-08T17:55:21Z

Purpose

Fixes #7199 — serve MammothModa2-Preview (T2I) through /v1/images/generations, plus the MammothModa2 text-to-image benchmark harness and profiler wiring for #7215 (Refs #7075).

MammothModa2 is a two-stage AR→DiT pipeline whose DiT stage runs as a generation-LLM stage (stage_type: llm, final_output_type: image) rather than a diffusion-typed stage. The image endpoints only looked for diffusion-typed stages, so every /v1/images/generations request returned HTTP 503 "No diffusion stage found" (#7199). This PR makes the serving path accept LLM-typed image stages end to end, and adds the profiler/benchmark tooling to measure it.

Changes

A. Online serving fix (#7199) — the largest part of the diff:

entrypoints/openai/utils.py: new is_image_generation_stage() helper — accepts both classical diffusion stages and LLM-typed stages whose final_output_type == "image". Used by endpoint admission (app_state.py::_get_engine_and_model), stage lookup (api_server.py::edit_images), model-class resolution, and DiT param injection, so all four call sites share one semantic.
serving_chat.py: multi-stage generation-LLM image serving — _build_multistage_generation_inputs routes image requests through the T2I prompt builder (injects the |image start|>W*H|image token|> envelope), _resolve_multistage_model_class_name() (od_config first, LLM-wrapper fallback, cached) resolves the DiT model class when the pipeline has no diffusion stage, and the AR stage gets extra_args/max_tokens wiring for the visual-token grid budget.
api_server.py: /v1/images/generations route forwards request.model_extra into extra_body and hands multi-stage image requests to the chat handler; edit_images uses the shared stage-lookup semantic.
protocol/images.py: ImageGenerationRequest is extra="allow" so per-model knobs reach the pipeline.

B. Benchmark harness + profiler wiring (#7215, component-level baseline):

pipeline_mammothmoda2_dit.py: DiffusionPipelineProfilerMixin on MammothModa2DiTPipeline (same pattern as flux/bagel). This stage runs on the legacy LLM-generation path (built from VllmConfig, no OmniDiffusionConfig), so the switch travels through additional_config, preserved for LLM stages per [Bugfix] Preserve additional_config for LLM stages #7272.
mammoth_moda2.yaml: expose additional_config.enable_diffusion_pipeline_profiler (default false — zero overhead when off).
text_to_image.py example: forward diffusion-only engine args only when explicitly set (_validate_global_stage_cli_ownership rejects them for pipelines without a DIFFUSION stage); explicit CLI parallel-knob overrides equal to defaults are preserved (review feedback), and the AR grid max_tokens budget comes from the shared ar_grid_max_tokens() helper.
diffusion/utils/param_utils.py: ar_grid_max_tokens(ar_width, ar_height) — single owner of the ar_height * (ar_width + 1) + 1 grid contract, shared by the online serving path and the offline example (review feedback; third copy lives in [Benchmark] Add MammothModa2 startup/loading breakdown #7312's bench_startup.py).
Performance dashboard (benchmarks/diffusion/performance_dashboard/mammoth_moda2_serving_performance.md + raw JSON): serving latency/throughput sweep with stage breakdown. All numbers are legacy-topology baseline (DiT on the LLM path) and are expected to be rerun after [Model] Migrate MammothModa2 DiT to the shared diffusion runtime #7134 lands — see the callout in the dashboard.

Coordination with #7134: #7134 owns the migration of this pipeline to the shared diffusion runtime (DiT as a standard DIFFUSION stage built from OmniDiffusionConfig). Until it lands, the LLM-typed compatibility path added here is the only way to serve this model, so the serving fix lives in this PR and was verified against current main (7e520f9c). Once #7134 lands, the profiler switch moves to the first-class od_config.enable_diffusion_pipeline_profiler field and the dashboard gets rerun — happy to follow up in a small PR as promised below.

Test Plan

vLLM Version: 0.29.0

vLLM-Omni Commit: rebased on main 7e520f9c

Serving smoke (single A800-80G):

vllm-omni serve /path/to/MammothModa2-Preview --omni \
  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml --port 8000
curl /v1/images/generations  # size 512x512

Unit: pytest tests/entrypoints/openai_api/test_image_server.py tests/examples/offline_inference/test_text_to_image_cli_forwarding.py -q

Test Result

[Bug] MammothModa2 online serving: /v1/images/generations returns 503 (no diffusion stage found) #7199 fixed: /v1/images/generations (512x512) returns HTTP 200 with a valid base64 image (28.5s end-to-end: AR + DiT + VAE). Before the fix: 503 on every request.
Full sweep after rebase to 0.29.0/main 7e520f9c: all configs 8/8 succeeded, no perf regression vs baseline — c1 mean 22.19s (baseline 21.8s), c2 24.72s (24.8s), c4 27.57s (28.5s). Raw JSON in mammoth_moda2_serving_results/.
Unit tests: 93 passed, 1 xfailed — incl. 4 mammoth serving tests, TestARGridMaxTokens (grid formula + None-edge), and CLI knob-forwarding regressions.
Determinism (offline, profiler path): anchor config (1024x1024 / 50 steps / guidance 4.0 / seed 42) — identical md5 across repeated runs, profiler on and off.
Zero profiler overhead when off: no profiler records, identical output.
No behavior change for pure-diffusion pipelines (flux/bagel): engine-side defaults for the now-conditionally-forwarded args equal the previously hard-coded values.

cc @Levius-Fubuki (coordination on #7134 overlap), @hsliuustc0106 (entrypoints review + the vllm bench serve --omni entry for PR2).

chatgpt-codex-connector · 2026-09-08T17:55:27Z

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 · 2026-09-08T17:56:34Z

This PR appears to belong to: docs/design/module/vllm_omni_config.md, docs/design/module/entrypoints.md, docs/design/module/diffusion/offloader.md, docs/design/module/diffusion/index.md, docs/design/module/profiling.md.

Module owners: @david6666666 @wtomin @Isotr0py

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

@quan-yi-ai, 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.

Levius-Fubuki · 2026-09-09T00:48:19Z

Thanks for coordinating the overlap.

#7134 is now Ready for review at head 68c7dfdf. After that migration, the MammothModa2 DiT stage runs as a DIFFUSION stage and MammothModa2DiTPipeline is constructed from OmniDiffusionConfig.

When rebasing or stacking #7293, the profiler switch can therefore use the existing first-class od_config.enable_diffusion_pipeline_profiler field instead of routing it through VllmConfig.additional_config. The offline diffusion-argument forwarding change should also be rechecked against the migrated stage topology.

#7134 will continue to keep the profiler mixin, reusable benchmark harness, and online sweeps out of scope, so ownership remains with #7293/#7215.

hsliuustc0106 · 2026-09-10T00:16:27Z

+        omni_kwargs["enable_cache_dit_summary"] = True
+    if args.enable_diffusion_pipeline_profiler:
+        omni_kwargs["enable_diffusion_pipeline_profiler"] = True
+    if args.ulysses_degree and args.ulysses_degree > 1:


[P2] Preserve explicit CLI overrides equal to defaults

These conditions silently discard explicit values such as --ulysses-degree 1 and --ulysses-mode strict. With a deploy YAML specifying different values, users can no longer override them, potentially retaining unwanted multi-GPU execution. Use None as the argparse default and forward every explicitly supplied value; add a regression test covering YAML overrides.

Addressed in 1f574db2 ("Preserve explicit CLI parallel-knob overrides equal to defaults"): argparse defaults for the five knobs are now None, every explicitly supplied value is forwarded (so --ulysses-degree 1 / --ulysses-mode strict can override the deploy YAML again), and TestParallelKnobForwarding::test_script_defaults_are_none + override-forwarding regressions pin the behavior.

…aults Address review feedback on vllm-project#7293: parallel knobs (ulysses_degree, ulysses_mode, ring_degree, cfg_parallel_size, vae_patch_parallel_size) now default to None and every explicitly supplied value is forwarded, so flags like --ulysses-degree 1 can override a deploy YAML instead of being silently dropped. Adds regression tests against the real argparse and forwarding helper.

…aults Address review feedback on vllm-project#7293: parallel knobs (ulysses_degree, ulysses_mode, ring_degree, cfg_parallel_size, vae_patch_parallel_size) now default to None and every explicitly supplied value is forwarded, so flags like --ulysses-degree 1 can override a deploy YAML instead of being silently dropped. Adds regression tests against the real argparse and forwarding helper. Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com>

Levius-Fubuki · 2026-09-11T12:37:38Z

I noticed that the latest commit 2846dfe1 expands #7293 from profiler/benchmark wiring into the #7199 online-serving fix for LLM-typed image stages.

For MammothModa2, that compatibility path becomes obsolete once #7134 lands: stage 1 is registered as a standard DIFFUSION stage and should use the existing diffusion endpoint discovery and request path.

Could we keep the scopes separated by either:

dropping or splitting the LLM-typed serving workaround from [Bugfix][Examples][MammothModa2] Serve T2I via LLM-typed DiT stage (#7199) + benchmark harness and profiler wiring #7293, or
rebasing onto [Model] Migrate MammothModa2 DiT to the shared diffusion runtime #7134 and validating the online endpoint through the migrated DIFFUSION topology?

If generic support for other LLM-typed image stages such as Dynin is still needed, that would be better handled as a separate serving PR coordinated with #7199's owner.

The current seven-point serving results should also be identified as legacy-topology baseline results until they are rerun after the #7134 migration.

hsliuustc0106 · 2026-09-12T01:50:46Z

+                ar_width = int((ar_prompt_info.get("ar_width") or [0])[0])
+                ar_height = int((ar_prompt_info.get("ar_height") or [0])[0])
+                if ar_width > 0 and ar_height > 0 and hasattr(default_stage_params, "max_tokens"):
+                    default_stage_params.max_tokens = ar_height * (ar_width + 1) + 1


This AR grid sizing now exists in three places: here, the offline example, and bench_startup.py in #7312 (ar_height * (ar_width + 1) + 1). Worth lifting into the model_extras spec (e.g. have the T2I builder or a registered helper own it) so a grid-contract change can't silently desynchronize one copy.

Addressed: the formula now has a single owner — ar_grid_max_tokens(ar_width, ar_height) in vllm_omni/diffusion/utils/param_utils.py (next to apply_declared_extra_args, already shared by examples and entrypoints). Both this call site and the offline example consume it, and TestARGridMaxTokens pins the contract (formula + None-for-absent-grid semantics), so a grid-contract change can't desynchronize a copy. The third copy you flagged in #7312's bench_startup.py is outside this PR's tree — worth switching it over in #7312 if that harness sticks around post-#7134.

hsliuustc0106 · 2026-09-12T01:50:47Z

The body's Changes section doesn't cover the largest part of this diff: the serving-path work (serving_chat.py +126/-24 — multistage generation-LLM image serving, _resolve_multistage_model_class_name, AR-stage extra_args/max_tokens wiring — plus api_server.py model_extra forwarding, extra="allow" on ImageGenerationRequest, and the new is_image_generation_stage helper). That's a production entrypoints feature riding in a PR titled as benchmark/profiler wiring — please update the body (or split it out) so entrypoints reviewers see it. I checked the scoping: bagel's DiT stage declares execution_type=DIFFUSION, so the new no-diffusion-stage prompt-rewrite path is mammoth-only, good. Also note the PR currently conflicts with main and overlaps #7134's rewrite of the same pipeline — the rebase will be non-trivial.

quan-yi-ai · 2026-09-17T15:37:35Z

Thank you @zhongc1211 for the independent verification — the config sweep results and the latency comparison against the legacy path are exactly the evidence this PR needed. Glad #7199 is confirmed fixed end-to-end. (For the record: the review comments on the benchmark command and dashboard link were addressed in 7257574, along with rejecting n > 1 on the LLM-path DiT stage.)

quan-yi-ai · 2026-09-20T14:27:49Z

Rebased onto the latest main (head 55f8b12f6), conflicts resolved. @linyueqian could you cycle the ready label and take a review pass when you have a chance?

quan-yi-ai · 2026-09-21T11:54:05Z

Found the failing test — it was ours: the deploy yaml's default_sampling_params collided with the shared-runtime contract test from #7134. Removed the block (the anchor sampling params now go per request via --extra-body in the benchmark commands), fixed in 9a17400eb. @hsliuustc0106 could you cycle the labels for a re-run when convenient?

quan-yi-ai · 2026-09-22T13:40:07Z

Sorry for the extra ping @hsliuustc0106 — gentle reminder: the fix for the failing lane test is up (9a17400eb, GitHub checks are green). Could you cycle the ready/cuda-test labels for a re-run when you get a chance? Thanks!

quan-yi-ai · 2026-09-24T15:05:21Z

The GPU lane and all GitHub checks are green on 22bb5d8e2 (build #15886). Anything else needed from our side, or is this ready to merge? Thanks!

MrlixiangWE · 2026-09-24T16:35:05Z

@quan-yi-ai No worries. Since this is a collaborative project, the PRs need to be merged in order. Once it passes Buildkite, you can just wait for the merge.

quan-yi-ai · 2026-09-28T12:38:52Z

Rebased onto v0.30.0 / post-#7476 main — head 1e2cb35b9, GitHub checks all green. @hsliuustc0106 could you cycle the ready/cuda-test labels for a lane re-run when convenient? Thanks!

vllm-omni-review-bot · 2026-09-29T00:20:21Z

Omni ReviewBot: superseded

The CI failure noted on 62f82a0c8c3e refers to an earlier head; the pull request now points at 09095de2536e.

…T stage - Admit generation-LLM DiT stages (stage_type="llm" + final_output_type="image") to image serving via a shared is_image_generation_stage classifier; /v1/images/edits keeps the pre-existing diffusion-only gate (fixes vllm-project#7199 503). - Route the registry T2I prompt builder by resolved model class so the AR envelope (<|image start|>W*H<|image token|>), additional_information and grid-sized max_tokens survive the vllm-project#7134 topology migration. - Forward undeclared request extras into extra_body and apply declared extra_args (text_guidance_scale / cfg_range / num_inference_steps) to both AR and DiT stages; fall back to extract_images_from_outputs for LLM-DiT multimodal_output payloads. - Extract ar_grid_max_tokens into param_utils so online serving and the offline example share one grid contract; make parallel knobs (ulysses/ring/cfg/vae-patch) default to None so deploy YAML values stay in effect unless explicitly overridden. - Add MammothModa2 serving performance dashboard, deploy yaml knobs, recipe docs, and regression tests for the serving surface. Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com>

…time contract The migrated pipeline's config test requires mammoth_moda2.yaml to keep stage-level default_sampling_params unset so request-level sampling fields are never shadowed. Remove the block (plus the no-op profiler flag line) and pass the benchmark anchor sampling params (text_guidance_scale 9.0, cfg_range [0.0, 1.0], num_inference_steps 50) per request via --extra-body in the serving benchmark commands so the recorded numbers stay reproducible. Fixes the 'Simple - Other Test' failure on build #15723. Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com>

quan-yi-ai · 2026-09-29T12:24:24Z

The previous lane failure looks infra-side: the only failed step was Diffusion · MiniMax-H3 DLO DP2 T2VA Test on the H100 pool — unrelated to this diff — and recent main builds were red on H100-pool steps around the same time. We've rebased onto the latest main (09095de25, now includes #8253's MiniMax test fix); GitHub checks are all green. @hsliuustc0106 could you cycle the labels for a lane re-run when convenient? Thanks!

hsliuustc0106

Both 09-12 comments are resolved and the body now covers the serving-path scope: the CLI-override P2 is fixed properly (argparse defaults moved to None so only explicit values forward, with --ulysses-degree 1 overriding the YAML as the docstring itself notes) and pinned by the new 140-line test_text_to_image_cli_forwarding.py; the AR grid-sizing triplication is lifted into the shared ar_grid_max_tokens() helper used by both the online path and the offline example. Mergeable and checks green. Good to merge.

vllm-omni-review-bot · 2026-10-04T02:45:09Z

Omni ReviewBot routing record

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

quan-yi-ai requested review from Gaohan123, alex-jw-brooks, hsliuustc0106, lishunyang12, wtomin and ywang96 as code owners September 8, 2026 17:55

quan-yi-ai mentioned this pull request Sep 8, 2026

[Performance]: MammothModa2-Preview text-to-image benchmark #7215

Open

hsliuustc0106 added diffusion codes related to diffusion models benchmark/profiler/metrics/logger Codes related to benchmarks, profiler, metrics and logger system labels Sep 8, 2026

quan-yi-ai force-pushed the mammoth-moda2-t2i-benchmark branch from 55abe2e to 5e52105 Compare September 9, 2026 13:21

hsliuustc0106 added the high priority high priority issue, needs to be done asap label Sep 10, 2026

hsliuustc0106 reviewed Sep 10, 2026

View reviewed changes

quan-yi-ai requested review from NickCao and yenuo26 as code owners September 10, 2026 20:11

quan-yi-ai force-pushed the mammoth-moda2-t2i-benchmark branch from 7eeb651 to ca920b7 Compare September 10, 2026 20:25

quan-yi-ai requested review from Isotr0py, RuixiangMa, SamitHuang, ZJY0516, linyueqian, princepride and xuechendi as code owners September 11, 2026 01:47

hsliuustc0106 reviewed Sep 12, 2026

View reviewed changes

quan-yi-ai force-pushed the mammoth-moda2-t2i-benchmark branch from 7257574 to 55f8b12 Compare September 20, 2026 14:22

hsliuustc0106 added bug Something isn't working ready label to trigger buildkite CI cuda-test Used to trigger vllm-omni cuda CI separately. and removed ready label to trigger buildkite CI labels Sep 20, 2026

quan-yi-ai force-pushed the mammoth-moda2-t2i-benchmark branch from 9a17400 to e9a1122 Compare September 22, 2026 13:52

hsliuustc0106 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Sep 22, 2026

quan-yi-ai force-pushed the mammoth-moda2-t2i-benchmark branch from 22bb5d8 to 1e2cb35 Compare September 28, 2026 11:31

hsliuustc0106 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Sep 28, 2026

quan-yi-ai added 2 commits September 29, 2026 09:18

quan-yi-ai force-pushed the mammoth-moda2-t2i-benchmark branch from 62f82a0 to 09095de Compare September 29, 2026 01:18

Merge branch 'main' into mammoth-moda2-t2i-benchmark

2701a19

hsliuustc0106 added ready label to trigger buildkite CI and removed ready label to trigger buildkite CI labels Sep 29, 2026

hsliuustc0106 approved these changes Oct 4, 2026

View reviewed changes

hsliuustc0106 merged commit 65f5a1d into vllm-project:main Oct 4, 2026
6 checks passed

1fanwang mentioned this pull request Oct 8, 2026

[Bugfix] Reject unsupported quality before multi-stage image generation #8622

Draft

Conversation

quan-yi-ai commented Sep 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

chatgpt-codex-connector Bot commented Sep 8, 2026

Uh oh!

vllm-omni-review-bot commented Sep 8, 2026

Uh oh!

Levius-Fubuki commented Sep 9, 2026

Uh oh!

hsliuustc0106 Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

quan-yi-ai Sep 12, 2026

Choose a reason for hiding this comment

Uh oh!

Levius-Fubuki commented Sep 11, 2026

Uh oh!

hsliuustc0106 Sep 12, 2026

Choose a reason for hiding this comment

Uh oh!

quan-yi-ai Sep 12, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 commented Sep 12, 2026

Uh oh!

quan-yi-ai commented Sep 17, 2026

Uh oh!

quan-yi-ai commented Sep 20, 2026

Uh oh!

quan-yi-ai commented Sep 21, 2026

Uh oh!

quan-yi-ai commented Sep 22, 2026

Uh oh!

quan-yi-ai commented Sep 24, 2026

Uh oh!

MrlixiangWE commented Sep 24, 2026

Uh oh!

quan-yi-ai commented Sep 28, 2026

Uh oh!

vllm-omni-review-bot commented Sep 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Omni ReviewBot: superseded

Uh oh!

quan-yi-ai commented Sep 29, 2026

Uh oh!

hsliuustc0106 left a comment

Choose a reason for hiding this comment

Uh oh!

vllm-omni-review-bot commented Oct 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

quan-yi-ai commented Sep 8, 2026 •

edited

Loading

vllm-omni-review-bot commented Sep 29, 2026 •

edited

Loading