[Bugfix][Examples][MammothModa2] Serve T2I via LLM-typed DiT stage (#7199) + benchmark harness and profiler wiring (#7293)

Signed-off-by: quan-yi-ai <82917002+quan-yi-ai@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

benchmarks/diffusion/performance_dashboard/mammoth_moda2_serving_performance.md

@@ -0,0 +1,209 @@
+# MammothModa2-Preview Serving Performance Dashboard
+
+This document describes how to deploy and benchmark **bytedance-research/MammothModa2-Preview** (T2I) using vLLM-Omni. It includes service startup configuration, benchmark methodology, dataset settings, and performance results.
+
+MammothModa2 is a two-stage AR→DiT text-to-image pipeline: stage 0 is an autoregressive LLM that encodes the prompt into a visual-token grid, and stage 1 is a generation-LLM DiT that renders the final image (`final_output_type: image`). Unlike single-stage diffusion pipelines, `/v1/images/generations` requests traverse both stages, and the DiT stage ran with `max_num_seqs: 1` in the then-default deploy config, so concurrent requests serialized on the DiT stage while the AR stage batched (request-level DiT batching landed later in #7476).
+
+> **Legacy-topology baseline**: all numbers in this dashboard were measured on the legacy topology, where the MammothModa2 DiT stage runs as an LLM-generation stage (`stage_type: llm`) built from `VllmConfig`, served through the LLM-typed image-stage compatibility path. #7134 migrates this pipeline to the shared diffusion runtime (DiT registered as a standard `DIFFUSION` stage built from `OmniDiffusionConfig`). These results should be treated as the pre-migration baseline and are expected to be rerun after #7134 lands; they are not directly comparable with post-migration numbers.
+
+---
+
+# 1. Overview
+
+This document covers:
+
+* Service launch configuration
+* Benchmark scripts and usage
+* Dataset and workload settings
+* Performance measurement results
+* Reproducibility guidelines
+
+---
+
+# 2. Test Environment
+
+| Component | Specification |
+|------------|----------------|
+| GPU | NVIDIA A800-SXM4-80GB (1x) |
+| CPU | 18 cores |
+| Model weights | `bytedance-research/MammothModa2-Preview` |
+| Deploy config | `vllm_omni/deploy/mammoth_moda2.yaml` (bundled) |
+| Stage 0 (AR) | `gpu_memory_utilization: 0.5`, `enforce_eager: true`, `max_num_seqs: 100` |
+| Stage 1 (DiT) | `gpu_memory_utilization: 0.3`, `enforce_eager: true`, `max_num_seqs: 1` |
+| Sampling params (per request) | `text_guidance_scale: 9.0`, `cfg_range: [0.0, 1.0]`, `num_inference_steps: 50` sent via `--extra-body` / `--num-inference-steps` |
+
+---
+
+# 3. Service Launch Configuration
+
+## 3.1 Basic Serving Command
+
+```bash
+vllm-omni serve bytedance-research/MammothModa2-Preview --omni \
+    --port 8091 \
+    --trust-remote-code
+```
+
+When serving from a local weights directory (no `model_index.json`), pass the bundled deploy config explicitly so the `mammoth_moda2` pipeline is selected:
+
+```bash
+vllm-omni serve /path/to/MammothModa2-Preview --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --deploy-config vllm_omni/deploy/mammoth_moda2.yaml
+```
+
+## 3.2 Key Parameters
+
+| Parameter | Description |
+|-----------|-------------|
+| `--deploy-config` | Path to the deploy YAML; defaults to the bundled `mammoth_moda2.yaml` for the HF checkpoint |
+| `size` (request) | `WxH`, e.g. `1024x1024`; routed to the AR grid and the DiT `target_w/target_h` |
+| `num_inference_steps` (request) | DiT sampling steps (default 50) |
+| `guidance_scale` (request) | Text guidance scale (default 9.0) |
+
+Note: the benchmark script's `--warmup-num-inference-steps` shortcut is not usable for this pipeline — a reduced step count truncates the AR stage's visual-token output and crashes the DiT stage. Use `--warmup-requests 0` and warm up with a real request instead.
+
+---
+
+# 4. Benchmark Script
+
+## 4.1 Benchmark Entry
+
+```bash
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --base-url http://localhost:8091 \
+    --model bytedance-research/MammothModa2-Preview \
+    --endpoint /v1/images/generations \
+    --task t2i \
+    --dataset random \
+    --num-prompts 8 \
+    --width <W> --height <H> \
+    --num-inference-steps <S> \
+    --extra-body '{"text_guidance_scale": 9.0, "cfg_range": [0.0, 1.0]}' \
+    --max-concurrency <C> \
+    --warmup-requests 0 \
+    --seed 142 \
+    --output-file <OUT>.json
+```
+
+## 4.2 Dataset
+
+`random` synthetic T2I prompts, fixed seed 142, 8 prompts per configuration (4 for 2048x2048 due to per-request latency).
+
+---
+
+# 5. Performance Results
+
+All requests succeeded (0 failures across the sweep). Latency is end-to-end per request (AR + DiT + VAE decode + base64 response).
+
+> Note: these are legacy-topology baseline results (see the callout in the overview); rerun after the #7134 diffusion-runtime migration before comparing across topologies.
+
+## 5.1 Latency
+
+| Resolution | Steps | Concurrency | Mean (s) | P50 (s) | P95 (s) | P99 (s) |
+|------------|-------|-------------|----------|---------|---------|---------|
+| 512x512    | 20    | 1           | 21.8     | 21.7    | 22.1    | 22.2    |
+| 512x512    | 20    | 2           | 24.8     | 25.2    | 25.6    | 25.8    |
+| 512x512    | 20    | 4           | 28.5     | 29.2    | 31.1    | 31.8    |
+| 1024x1024  | 50    | 1           | 96.3     | 95.8    | 99.8    | 100.8   |
+| 1024x1024  | 50    | 2           | 121.5    | 125.4   | 127.9   | 128.5   |
+| 1024x1024  | 50    | 4           | 150.9    | 157.6   | 178.6   | 186.3   |
+| 2048x2048  | 50    | 1           | 271.9    | 271.6   | 273.6   | 273.8   |
+
+## 5.2 Throughput
+
+| Resolution | Steps | Concurrency | Throughput (img/s) |
+|------------|-------|-------------|--------------------|
+| 512x512    | 20    | 1           | 0.0459             |
+| 512x512    | 20    | 2           | 0.0798             |
+| 512x512    | 20    | 4           | 0.1329             |
+| 1024x1024  | 50    | 1           | 0.0104             |
+| 1024x1024  | 50    | 2           | 0.0162             |
+| 1024x1024  | 50    | 4           | 0.0238             |
+| 2048x2048  | 50    | 1           | 0.0037             |
+
+## 5.3 Stage Breakdown (concurrency 1, from server logs)
+
+Per-stage wall time from the server's `[OmniTiming]` records across the §5.1/§5.4/§5.5 sweeps (n = 18/9/5 per row; at concurrency 1 stage wall time equals `stage_gen_time_ms` in the engine stats table):
+
+| Resolution | Steps | AR stage (s) | DiT stage (s) | E2E (s) | AR share | AR output tokens | AR ms/token |
+|------------|-------|--------------|---------------|---------|----------|------------------|-------------|
+| 512x512    | 20    | 20.3         | 1.7           | 22.0    | 92%      | 1,057            | 19.2        |
+| 1024x1024  | 50    | 78.9         | 17.2          | 96.1    | 82%      | 4,161            | 19.0        |
+| 2048x2048  | 50    | 156.3        | 114.8         | 271.2   | 58%      | 8,154            | 19.2        |
+
+Stage 0 is the AR prompt encoder (TTFT 82 ms at 1024x1024); stage 1 is the DiT renderer. Three observations:
+
+1. **AR per-token cost is constant (~19 ms/token) across resolutions**, so AR time is fully predictable: the AR stage decodes a fixed visual grid whose length scales with resolution (1,057 → 4,161 → 8,154 tokens), giving `AR ≈ tokens × 19 ms`. This is also why §5.4's step sweep is sub-linear — the AR token count is step-invariant.
+2. **DiT time scales super-linearly with pixels**: a 16x pixel increase (512 → 2048) multiplies DiT time by 67 (1.7 → 114.8 s).
+3. **The dominant stage flips with resolution**: AR is 92% of E2E at 512x512 (DiT optimization is nearly pointless there) but 58% at 2048x2048, where both stages matter.
+
+At concurrency > 1, per-stage wall times additionally include intra-stage queueing (each stage serializes at `max_num_seqs: 1`), so the c2/c4 rows in §5.1 cannot be decomposed additively from this table.
+
+## 5.4 Inference-Step Sensitivity (1024x1024 / 512x512, concurrency 1)
+
+Steps 25 at 1024x1024 (half the default 50) and steps 25 at 512x512, same protocol as §5.1:
+
+| Resolution | Steps | Concurrency | Mean (s) | P50 (s) | P95 (s) | P99 (s) |
+|------------|-------|-------------|----------|---------|---------|---------|
+| 1024x1024  | 25    | 1           | 85.2     | 85.3    | 87.8    | 88.7    |
+| 1024x1024  | 25    | 4           | 113.7    | 115.8   | 126.1   | 129.4   |
+| 512x512    | 25    | 1           | 21.9     | 21.8    | 22.2    | 22.3    |
+
+Halving DiT steps does not halve end-to-end latency (96.3 → 85.2 s at 1024): the AR stage emits a fixed count of visual tokens (4,161 at 1024x1024) regardless of step count, so its ~79 s contribution (§5.3) is step-invariant; only the DiT sampling time scales with steps. At 512x512, steps 20 → 25 adds ~0.1 s per step (21.8 → 21.9 s).
+
+## 5.5 Prompt-Length Sensitivity (long vs short prompts)
+
+A ~430-word long prompt (repeated via `--random-request-config '[{"weight":1,"prompt":"..."}]'`) against the default short synthetic prompt, same seed and settings:
+
+| Resolution | Steps | Concurrency | Prompt | Mean (s) | P99 (s) | Throughput (img/s) |
+|------------|-------|-------------|--------|----------|---------|---------------------|
+| 1024x1024  | 50    | 1           | short  | 96.3     | 100.8   | 0.0104              |
+| 1024x1024  | 50    | 1           | long   | 93.6     | 94.4    | 0.0107              |
+| 1024x1024  | 50    | 4           | short  | 150.9    | 186.3   | 0.0238              |
+| 1024x1024  | 50    | 4           | long   | 141.7    | 174.2   | 0.0255              |
+| 512x512    | 20    | 1           | short  | 21.8     | 22.2    | 0.0459              |
+| 512x512    | 20    | 1           | long   | 22.2     | 22.5    | 0.0450              |
+
+Prompt length has no material effect on latency (the long-prompt rows are within run-to-run variance and even slightly faster with tighter P99 tails). This confirms the AR stage's cost is dominated by decoding the fixed 4,161-token visual grid, not by prompt encoding: the KV prefill of a few hundred extra text tokens is negligible against 4,161 decode steps.
+
+## 5.6 Peak GPU Memory (per-stage process, nvidia-smi polling at 1 Hz)
+
+Sampled during the §5.4/§5.5 sweeps (stage processes identified by PID; both stages share the single A800):
+
+| Configuration (worst case) | stage0_ar (MB) | stage1_dit (MB) |
+|----------------------------|----------------|-----------------|
+| 1024x1024, 50 steps, c4    | 39,290         | 11,104          |
+| all other configs          | 39,194         | 11,102–11,104   |
+
+Memory is essentially flat across the sweep: the AR stage grows by only ~96 MB at concurrency 4 and the DiT stage by 2 MB, because both stages pre-allocate from `gpu_memory_utilization` budgets (0.5 / 0.3) at startup. Combined steady-state usage ≈ 50.4 GiB of the 80 GiB device.
+
+## 5.7 DiT-Stage Component Attribution (pipeline profiler)
+
+With `enable_diffusion_pipeline_profiler: true` in the stage-1 deploy config (`MammothModa2DiTPipeline` targets: `gen_transformer.forward`, `gen_vae.decode`, `gen_image_condition_refiner.forward`), three warm 1024x1024 / 50-step requests (profiler adds a `synchronize` per call, so timings are indicative, not throughput-grade):
+
+| Component | Per request (s) | Share of DiT pipeline |
+|-----------|-----------------|-----------------------|
+| `gen_transformer.forward` (100 calls: 50 steps x 2 passes) | 14.01–14.07 | 98.6% |
+| `gen_vae.decode` (1 call) | 0.148 | 1.0% |
+| `gen_image_condition_refiner.forward` (1 call) | 0.005 | 0.03% |
+| whole `MammothModa2DiTPipeline.forward` | 14.18–14.27 | — |
+
+The DiT latency is ~99% transformer sampling; VAE decode costs only ~0.15 s per image at 1024x1024 and is not a bottleneck. Per-step transformer time is ~140 ms (two forward passes per step: ~2.8 s nominal x CFG), consistent with §5.3's 17.2 s image-TTFT figure once base64 encoding and stage handoff are included.
+
+## 5.8 Endpoint Fix Verification
+
+Before the serving fix, `/v1/images/generations` returned HTTP 503 ("No diffusion stage found in multi-stage pipeline") because the pipeline's image-output stage is a generation-LLM DiT rather than a `diffusion`-typed stage. After the fix all requests return HTTP 200 with a valid base64 image (verified across every sweep configuration above).
+
+---
+
+# 6. Reproducibility
+
+1. Start the server as in §3.1 and wait for `Application startup complete`.
+2. Warm up with one real request (any size) — do not use reduced-step warmups.
+3. Run the sweep commands from §4.1 for each row in §5.
+4. Server-side per-stage timings are printed in the engine stats table (`stage_gen_time_ms`, `output_unit_count`) at `--log-stats` level info.
+5. New rows use the same seed (`--seed 142`), `--num-prompts 8`, and `--warmup-requests 0` protocol as §4.1; when saving results with `--output-file`, name them `<res>_s<steps>_c<conc>[_long].json` to keep runs comparable (steps-25, long-prompt, peak-memory and profiler runs).
+6. To reproduce §5.6 peak-memory numbers, poll `nvidia-smi --query-gpu=memory.used --format=csv -l 1` per stage PID while a sweep is in flight (e.g. `nvidia-smi pmon` / `--query-compute-apps`).
+7. To reproduce §5.7, copy the bundled deploy config, set `enable_diffusion_pipeline_profiler: true` on the stage-1 entry, restart, and grep `[DiffusionPipelineProfiler]` lines from the server log; timings include a `torch.cuda.synchronize` per call.

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,7 +14,7 @@ from diffusers.utils import numpy_to_pil
 
 from vllm_omni.diffusion.data import logger
 from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args, ar_grid_max_tokens
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -72,6 +72,24 @@ def build_text_to_image_prompt(prompt: str, negative_prompt: str | None) -> dict
     return result
 
 
+def build_parallel_knob_kwargs(args: argparse.Namespace) -> dict[str, Any]:
+    """Forward explicitly-set parallel knobs so they can override a deploy YAML.
+
+    The argparse defaults are ``None``; only explicitly supplied values enter
+    the returned dict, letting a deploy YAML's parallel_config stay in effect
+    when the flags are omitted (and letting explicit values equal to the old
+    argparse defaults, e.g. ``--ulysses-degree 1``, still override the YAML).
+    """
+    knob_names = (
+        "ulysses_degree",
+        "ring_degree",
+        "ulysses_mode",
+        "cfg_parallel_size",
+        "vae_patch_parallel_size",
+    )
+    return {name: getattr(args, name) for name in knob_names if getattr(args, name) is not None}
+
+
 def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     """Convert NumPy diffusion outputs to PIL images before saving."""
     normalized = []
@@ -83,7 +101,8 @@ def _normalize_images_for_save(images: list[Any]) -> list[Any]:
     return normalized
 
 
-def parse_args() -> argparse.Namespace:
+def build_parser() -> argparse.ArgumentParser:
+    """Build the example's CLI parser (kept separate for testability)."""
     parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
     parser.add_argument(
         "--model",
@@ -166,28 +185,31 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--ulysses-degree",
         type=int,
-        default=1,
-        help="Number of GPUs used for ulysses sequence parallelism.",
+        default=None,
+        help="Number of GPUs used for ulysses sequence parallelism. "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--ulysses-mode",
         type=str,
-        default="strict",
+        default=None,
         choices=["strict", "advanced_uaa"],
-        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
+        help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA). "
+        "Explicitly passing 'strict' overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--ring-degree",
         type=int,
-        default=1,
-        help="Number of GPUs used for ring sequence parallelism.",
+        default=None,
+        help="Number of GPUs used for ring sequence parallelism. Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--cfg-parallel-size",
         type=int,
-        default=1,
+        default=None,
         choices=[1, 2],
-        help="Number of GPUs used for classifier free guidance parallel size.",
+        help="Number of GPUs used for classifier free guidance parallel size. "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     parser.add_argument(
         "--enforce-eager",
@@ -296,8 +318,9 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--vae-patch-parallel-size",
         type=int,
-        default=1,
-        help="Number of ranks used for VAE patch/tile parallelism (decode/encode).",
+        default=None,
+        help="Number of ranks used for VAE patch/tile parallelism (decode/encode). "
+        "Explicitly passing 1 overrides a deploy YAML value.",
     )
     # NextStep-1.1 specific arguments
     parser.add_argument(
@@ -413,7 +436,11 @@ def parse_args() -> argparse.Namespace:
         ),
     )
     current_omni_platform.pre_register_and_update(parser)
-    return parser.parse_args()
+    return parser
+
+
+def parse_args() -> argparse.Namespace:
+    return build_parser().parse_args()
 
 
 def _apply_ar_stage_inputs(
@@ -556,21 +583,11 @@ def main():
 
     omni_kwargs = {
         "model": args.model,
-        "enable_layerwise_offload": args.enable_layerwise_offload,
-        "vae_use_slicing": args.vae_use_slicing,
-        "vae_use_tiling": args.vae_use_tiling,
         "cache_backend": args.cache_backend,
         "cache_config": cache_config,
-        "enable_cache_dit_summary": args.enable_cache_dit_summary,
-        "ulysses_degree": args.ulysses_degree,
-        "ring_degree": args.ring_degree,
-        "ulysses_mode": args.ulysses_mode,
-        "cfg_parallel_size": args.cfg_parallel_size,
-        "vae_patch_parallel_size": args.vae_patch_parallel_size,
         "enable_expert_parallel": args.enable_expert_parallel,
         "enable_cpu_offload": args.enable_cpu_offload,
         "log_stats": args.log_stats,
-        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         "profiler_config": args.profiler_config,
         "init_timeout": args.init_timeout,
         "stage_init_timeout": args.stage_init_timeout,
@@ -578,6 +595,8 @@ def main():
         **lora_args,
         **quant_kwargs,
     }
+    # Diffusion-only engine args: pass only when explicitly set, so pipelines
+    # without a DIFFUSION stage pass ownership validation.
     if args.tensor_parallel_size is not None:
         omni_kwargs["tensor_parallel_size"] = args.tensor_parallel_size
     if args.enforce_eager is not None:
@@ -597,6 +616,18 @@ def main():
     # gate is an engine-level config (offline analog of the server's --no-guardrails).
     if args.extra_body and "guardrails" in args.extra_body:
         omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
+    if args.enable_layerwise_offload:
+        omni_kwargs["enable_layerwise_offload"] = True
+    if args.vae_use_slicing:
+        omni_kwargs["vae_use_slicing"] = True
+    if args.vae_use_tiling:
+        omni_kwargs["vae_use_tiling"] = True
+    if args.enable_cache_dit_summary:
+        omni_kwargs["enable_cache_dit_summary"] = True
+    if args.enable_diffusion_pipeline_profiler:
+        omni_kwargs["enable_diffusion_pipeline_profiler"] = True
+    omni_kwargs.update(build_parallel_knob_kwargs(args))
+
     omni = Omni(**omni_kwargs)
     model_class_name = get_model_class_name(omni)
     declared_extra_body_params = get_extra_body_params(model_class_name)
@@ -615,11 +646,17 @@ def main():
     if ignored_layers:
         print(f"  Ignored layers: {ignored_layers}")
     tp_display = args.tensor_parallel_size if args.tensor_parallel_size is not None else "deploy/default"
+
+    def _knob_display(value: Any, default: Any) -> Any:
+        return value if value is not None else f"deploy/default ({default})"
+
     print(
         f"  Parallel configuration: tensor_parallel_size={tp_display}, "
-        f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
-        f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
-        f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
+        f"ulysses_degree={_knob_display(args.ulysses_degree, 1)}, "
+        f"ulysses_mode={_knob_display(args.ulysses_mode, 'strict')}, "
+        f"ring_degree={_knob_display(args.ring_degree, 1)}, "
+        f"cfg_parallel_size={_knob_display(args.cfg_parallel_size, 1)}, "
+        f"vae_patch_parallel_size={_knob_display(args.vae_patch_parallel_size, 1)}, "
         f"enable_expert_parallel={args.enable_expert_parallel}."
     )
     print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
@@ -730,8 +767,9 @@ def main():
             if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
                 ar_width = int(prompt_info.get("ar_width", [0])[0])
                 ar_height = int(prompt_info.get("ar_height", [0])[0])
-                if ar_width > 0 and ar_height > 0:
-                    params.max_tokens = ar_height * (ar_width + 1) + 1
+                ar_grid_budget = ar_grid_max_tokens(ar_width, ar_height)
+                if ar_grid_budget is not None:
+                    params.max_tokens = ar_grid_budget
 
     if not diffusion_replaced and len(sampling_params_list) == 1:
         sampling_params_list = [diffusion_params]

recipes/MammothModa2/MammothModa2.md

@@ -322,6 +322,44 @@ python3 examples/offline_inference/text_to_image/text_to_image.py \
 
 The first request took 85.224 seconds. The AR stage generated 4,161 visual tokens in 72.996 seconds, and the DiT stage took 12.163 seconds. AR weight loading used 21.4 GiB and took 8.250 seconds. DiT weight loading used 5.49 GiB and took 1.824 seconds. The largest one second whole device memory sample was 106.57 GiB, including the AR KV cache reserved by the 0.5 memory setting.
 
+### 1x A800 80GB, MammothModa2 Preview (serving)
+
+#### Environment
+
+- OS: Linux, x86_64
+- CPU: 18 cores
+- GPU: one NVIDIA A800-SXM4-80GB
+- Deploy config: `vllm_omni/deploy/mammoth_moda2.yaml` (default split: AR `gpu_memory_utilization` 0.5, DiT 0.3)
+- vLLM-Omni version or commit: use the commit you are deploying from
+
+#### Serving Commands
+
+```bash
+vllm-omni serve bytedance-research/MammothModa2-Preview --omni \
+    --port 8091 \
+    --trust-remote-code
+
+python benchmarks/diffusion/diffusion_benchmark_serving.py \
+    --model bytedance-research/MammothModa2-Preview \
+    --endpoint /v1/images/generations \
+    --host 127.0.0.1 --port 8091 \
+    --height 1024 --width 1024 --num-inference-steps 50 \
+    --extra-body '{"text_guidance_scale": 9.0, "cfg_range": [0.0, 1.0]}' \
+    --num-prompts 8 --seed 142 --warmup-requests 0 \
+    --output-file mm2_1024_s50_c1.json
+```
+
+#### Verification
+
+On 1024x1024 / 50 steps, single-concurrency end-to-end mean latency is ~96 s
+(P99 ~101 s) per image and concurrency-4 mean is ~151 s (~0.024 img/s);
+steady-state combined GPU memory is ~50.4 GiB (AR ~39.2 GiB, DiT ~11.1 GiB).
+The AR stage dominates (~77 s of a ~96 s request) because it decodes a fixed
+4,161-token visual grid per image, so latency is insensitive to prompt length
+and scales only partially with DiT step count. See the
+[serving performance dashboard](../../benchmarks/diffusion/performance_dashboard/mammoth_moda2_serving_performance.md)
+for the full sweep, peak-memory, and component-attribution data.
+
 The output was a valid 1024 by 1024 RGB PNG.
 
 ### Startup benchmark

tests/entrypoints/openai_api/test_image_server.py

@@ -438,6 +438,90 @@ def async_omni_stage_configs_only_client():
     return TestClient(app)
 
 
+@pytest.fixture
+def mammoth_moda2_test_client():
+    """MammothModa2-style pipeline: AR (llm) + generation-LLM DiT stage.
+
+    The DiT stage is ``stage_type == "llm"`` with ``final_output_type ==
+    "image"`` (no classical diffusion stage) — the legacy topology
+    MammothModa2 served before #7134 migrated it to the shared diffusion
+    runtime. Regression guard for the 503 (stage discovery) and the T2I
+    envelope / max_tokens wiring on that compat path.
+    """
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.async_omni import AsyncOmni
+    from vllm_omni.entrypoints.openai.api_server import router
+    from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
+
+    class FakeAsyncOmniClass(AsyncOmni):
+        def __init__(self):
+            stage_configs = [
+                SimpleNamespace(
+                    stage_type="llm",
+                    is_comprehension=False,
+                    engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+                ),
+                SimpleNamespace(
+                    stage_type="llm",
+                    is_comprehension=False,
+                    final_output=True,
+                    final_output_type="image",
+                    engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+                ),
+            ]
+            default_sampling_params_list = [
+                SamplingParams(temperature=1.0, max_tokens=16),
+                OmniDiffusionSamplingParams(),
+            ]
+            self.engine = SimpleNamespace(
+                stage_configs=stage_configs,
+                default_sampling_params_list=default_sampling_params_list,
+                # Real AsyncOmni engines always carry stage_clients (read by
+                # get_diffusion_od_config); edits-endpoint tests walk through
+                # helpers that touch it before any stage-type gate.
+                stage_clients=[],
+            )
+            self.default_sampling_params_list = default_sampling_params_list
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self._images = [Image.new("RGB", (64, 64), color="green")]
+
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
+            if sampling_params_list is not None:
+                self.captured_sampling_params_list = sampling_params_list
+            else:
+                self.captured_sampling_params_list = [sampling_params]
+            self.captured_prompt = prompt
+            images = [img.copy() for img in self._images]
+            yield MockGenerationResult(images)
+
+        def __class_getitem__(cls, item):
+            return cls
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmniClass()
+    chat_handler = object.__new__(OmniOpenAIServingChat)
+    chat_handler.engine_client = engine
+    chat_handler._diffusion_engine = None
+    app.state.openai_serving_chat = chat_handler
+    app.state.engine_client = engine
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="llm", final_output=True, final_output_type="image"),
+    ]
+    app.state.openai_serving_models = _DiffusionServingModels(
+        [BaseModelPath(name="Mammoth/MammothModa2-Preview", model_path="Mammoth/MammothModa2-Preview")]
+    )
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":50, "generator_device":"cpu"}}',
+        max_generated_image_size=1024 * 1024,
+    )
+    return TestClient(app)
+
+
 @pytest.fixture
 def streaming_image_edit_client():
     """Create a multi-stage client whose engine yields AR text before image output."""
@@ -2418,3 +2502,190 @@ def test_image_edits_omitted_bot_task_stop_tokens_match_prompt_default(
         f"{ar_params.stop_token_ids} -- this is the full ratio range, meaning "
         "resolve_stop_token_ids disagreed with build_prompt_tokens's default."
     )
+
+
+def test_generate_images_mammoth_moda2_llm_dit_pipeline_accepted(mammoth_moda2_test_client):
+    """Regression: /v1/images/generations must not 503 on generation-LLM DiT stages.
+
+    A generation-LLM DiT stage (``stage_type == "llm"`` with
+    ``final_output_type == "image"``, MammothModa2's topology before #7134):
+    stage discovery previously required a classical diffusion stage and
+    returned 503 (issue #7199).
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+    data = response.json()
+    assert len(data["data"]) == 1
+
+
+def test_generate_images_mammoth_moda2_llm_dit_rejects_multiple_outputs(mammoth_moda2_test_client):
+    """Regression: n > 1 must be rejected on the LLM-typed DiT path.
+
+    The LLM engine path for generation-LLM DiT stages produces exactly one
+    image per request; accepting ``n > 1`` used to silently drop the extra
+    outputs (review on PR #7293). The request must fail fast with 400 instead.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 4,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 400, response.text
+    assert "exactly one image" in response.text
+
+
+def test_generate_images_mammoth_moda2_t2i_envelope_and_max_tokens(mammoth_moda2_test_client):
+    """Regression: the registry T2I builder drives the AR prompt + max_tokens.
+
+    The AR stage must receive the <|image start|>W*H<|image token|> envelope,
+    the structural additional_information, and max_tokens sized from the AR
+    grid (ar_height * (ar_width + 1) + 1) instead of SamplingParams' default.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    captured_prompt = engine.captured_prompt
+    assert captured_prompt["prompt"].endswith("<|image start|>64*64<|image token|>")
+    addi = captured_prompt["additional_information"]
+    assert addi["omni_task"] == ["t2i"]
+    assert addi["ar_width"] == [64]
+    assert addi["ar_height"] == [64]
+
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    assert captured[0].max_tokens == 64 * 65 + 1
+    assert captured[0].seed == 7
+
+
+def test_generate_images_mammoth_moda2_extra_args_routed_to_both_stages(mammoth_moda2_test_client):
+    """Regression: declared extra_body params reach AR and DiT extra_args.
+
+    text_guidance_scale / cfg_range / num_inference_steps are declared by the
+    MammothModa2 registry spec; the DiT stage (llm-typed) reads them from
+    sampling_params.extra_args, so the serving layer must route them there.
+    """
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+            "text_guidance_scale": 4.5,
+            "cfg_range": [0.0, 0.8],
+            "num_inference_steps": 30,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    for stage_params in captured:
+        extra_args = getattr(stage_params, "extra_args", None) or {}
+        assert extra_args.get("text_guidance_scale") == 4.5
+        assert extra_args.get("cfg_range") == [0.0, 0.8]
+        assert extra_args.get("num_inference_steps") == 30
+
+
+def test_generate_images_mammoth_moda2_migrated_topology_keeps_t2i_builder(mammoth_moda2_test_client):
+    """Regression: the T2I builder must stay active after #7134's migration.
+
+    Once MammothModa2's DiT stage is migrated to the shared diffusion runtime
+    (``stage_type == "diffusion"``, see #7134), the pipeline contains a
+    classical diffusion stage while stage 0 still needs the AR envelope/grid
+    metadata. An earlier guard skipped the registry T2I builder whenever any
+    diffusion stage was present, which would silently drop the
+    <|image start|>W*H<|image token|> envelope and additional_information
+    after the migration. Dispatch must be by resolved model class, not by
+    pipeline topology.
+    """
+    engine = mammoth_moda2_test_client.app.state.engine_client
+    # Simulate the post-#7134 topology in place: stage 1 becomes a classical
+    # diffusion stage (final_output metadata moves onto it unchanged).
+    engine.engine.stage_configs[1] = SimpleNamespace(
+        stage_type="diffusion",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type="image",
+        engine_args=SimpleNamespace(model_arch="MammothModa2ForConditionalGeneration"),
+    )
+    mammoth_moda2_test_client.app.state.stage_configs = list(engine.engine.stage_configs)
+
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a mammoth in the snow",
+            "n": 1,
+            "size": "1024x1024",
+            "seed": 7,
+        },
+    )
+    assert response.status_code == 200, response.text
+
+    captured_prompt = engine.captured_prompt
+    assert captured_prompt["prompt"].endswith("<|image start|>64*64<|image token|>"), (
+        "the registry T2I builder must still wrap the AR prompt when the "
+        "pipeline contains a classical diffusion stage (post-#7134 topology)"
+    )
+    addi = captured_prompt["additional_information"]
+    assert addi["omni_task"] == ["t2i"]
+    assert addi["ar_width"] == [64]
+    assert addi["ar_height"] == [64]
+
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert captured[0].max_tokens == 64 * 65 + 1
+
+
+def test_image_edits_mammoth_moda2_llm_dit_pipeline_rejected(mammoth_moda2_test_client):
+    """Regression: /v1/images/edits stays diffusion-only for now.
+
+    MammothModa2 is admitted to /v1/images/generations via the LLM-typed DiT
+    stage, but image edits (img2img) were never validated for it. Keep the
+    pre-existing diffusion-stage gate on the edits endpoint so the widened
+    generations acceptance does not silently admit it here.
+    """
+    img = make_test_image_bytes((64, 64))
+    response = mammoth_moda2_test_client.post(
+        "/v1/images/edits",
+        files=[("image", img)],
+        data={"prompt": "make it neon", "size": "512x512"},
+    )
+    assert response.status_code == 503, response.text
+    assert "No diffusion stage found" in response.json()["detail"]
+
+
+def test_is_image_generation_stage_matrix():
+    """Unit-test the stage classifier used by the image serving endpoints."""
+    from vllm_omni.entrypoints.openai.utils import is_image_generation_stage
+
+    assert is_image_generation_stage(SimpleNamespace(stage_type="diffusion"))
+    assert is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type="image"))
+    assert is_image_generation_stage({"stage_type": "llm", "final_output": True, "final_output_type": "images"})
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type="text"))
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=False))
+    assert not is_image_generation_stage(SimpleNamespace(stage_type="llm", final_output=True, final_output_type=None))

tests/examples/offline_inference/test_text_to_image_cli_forwarding.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""
+Offline inference regression tests: explicit CLI parallel-knob forwarding.
+
+Covers the argparse defaults -> ``build_parallel_knob_kwargs`` -> ``main()``'s
+``omni_kwargs`` assembly in examples/offline_inference/text_to_image/text_to_image.py.
+Three invariants:
+
+1. The shipped argparse must default the parallel knobs to ``None`` so unset
+   flags do not enter ``omni_kwargs`` and a deploy YAML value stays in effect.
+2. Explicitly supplied values equal to the old argparse defaults (e.g.
+   ``--ulysses-degree 1``, ``--ulysses-mode strict``) must still be forwarded
+   so they can override the deploy YAML, instead of being silently dropped
+   (review feedback on PR #7293).
+3. The AR grid ``max_tokens`` budget must stay owned by the shared
+   ``ar_grid_max_tokens`` helper used by both this example and the online
+   serving path (review feedback on PR #7293).
+
+Both invariants are asserted against the real script (its ``parse_args()`` and
+its ``build_parallel_knob_kwargs``), not a re-implemented copy. Invariant 3 is
+asserted directly against ``ar_grid_max_tokens`` itself.
+"""
+
+import sys
+from pathlib import Path
+
+import pytest
+
+_script_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "offline_inference" / "text_to_image")
+sys.path.insert(0, _script_dir)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_KNOBS = (
+    "ulysses_degree",
+    "ulysses_mode",
+    "ring_degree",
+    "cfg_parallel_size",
+    "vae_patch_parallel_size",
+)
+
+
+def _parse_real_args(cli_flags: list[str]):
+    """Run the example's real argparse against a synthetic argv."""
+    import text_to_image as t2i
+
+    saved = sys.argv
+    sys.argv = ["text_to_image.py", *cli_flags]
+    try:
+        return t2i.parse_args()
+    finally:
+        sys.argv = saved
+
+
+class TestParallelKnobForwarding:
+    """Drive the example's real argparse + real forwarding helper."""
+
+    def test_unset_knobs_forward_nothing(self):
+        """No CLI flags -> no keys, so the deploy YAML values stay in effect."""
+        import text_to_image as t2i
+
+        args = _parse_real_args([])
+        assert t2i.build_parallel_knob_kwargs(args) == {}
+
+    def test_explicit_default_equivalent_values_are_forwarded(self):
+        """--ulysses-degree 1 / --ulysses-mode strict must override the YAML."""
+        import text_to_image as t2i
+
+        args = _parse_real_args(["--ulysses-degree", "1", "--ulysses-mode", "strict"])
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ulysses_degree": 1,
+            "ulysses_mode": "strict",
+        }
+
+        args = _parse_real_args(["--ring-degree", "1", "--cfg-parallel-size", "1", "--vae-patch-parallel-size", "1"])
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ring_degree": 1,
+            "cfg_parallel_size": 1,
+            "vae_patch_parallel_size": 1,
+        }
+
+    def test_non_default_values_are_forwarded(self):
+        import text_to_image as t2i
+
+        args = _parse_real_args(
+            [
+                "--ulysses-degree",
+                "2",
+                "--ulysses-mode",
+                "advanced_uaa",
+                "--ring-degree",
+                "2",
+                "--cfg-parallel-size",
+                "2",
+                "--vae-patch-parallel-size",
+                "2",
+            ]
+        )
+        assert t2i.build_parallel_knob_kwargs(args) == {
+            "ulysses_degree": 2,
+            "ulysses_mode": "advanced_uaa",
+            "ring_degree": 2,
+            "cfg_parallel_size": 2,
+            "vae_patch_parallel_size": 2,
+        }
+
+    def test_script_defaults_are_none(self):
+        """The shipped argparse must default to None so unset == defer to YAML."""
+        import text_to_image as t2i
+
+        parser = t2i.build_parser()
+        defaults = {key: parser.get_default(key) for key in _KNOBS}
+        assert defaults == {key: None for key in _KNOBS}
+
+
+class TestARGridMaxTokens:
+    """The AR grid budget formula must stay shared between serving and offline.
+
+    ``ar_grid_max_tokens`` (vllm_omni.diffusion.utils.param_utils) is the
+    single owner of the ``ar_height * (ar_width + 1) + 1`` contract used by
+    both the online serving path (serving_chat.py) and this example script,
+    so a grid-contract change cannot desynchronize one copy (review feedback
+    on PR #7293).
+    """
+
+    def test_grid_budget_matches_formula(self):
+        from vllm_omni.diffusion.utils.param_utils import ar_grid_max_tokens
+
+        assert ar_grid_max_tokens(64, 64) == 64 * 65 + 1
+        assert ar_grid_max_tokens(32, 48) == 48 * 33 + 1
+
+    def test_grid_budget_none_for_absent_or_invalid_grid(self):
+        from vllm_omni.diffusion.utils.param_utils import ar_grid_max_tokens
+
+        # Callers treat None as "keep the existing max_tokens default".
+        assert ar_grid_max_tokens(0, 0) is None
+        assert ar_grid_max_tokens(0, 64) is None
+        assert ar_grid_max_tokens(64, 0) is None

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -337,7 +338,7 @@ class _RequestConditioning:
     image_mask: torch.Tensor
 
 
-class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
+class MammothModa2DiTPipeline(nn.Module, DiffusionPipelineProfilerMixin, SupportsComponentDiscovery):
     """
     MammothModa2 DiT + VAE generation stage (non-autoregressive).
 
@@ -424,12 +425,6 @@ 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":
@@ -440,6 +435,14 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 else DiffusionCacheConfig.from_dict(cache_config or {})
             )
 
+        profiler_targets = ["gen_transformer.forward", "gen_vae.decode"]
+        if self.gen_image_condition_refiner is not None:
+            profiler_targets.append("gen_image_condition_refiner.forward")
+        self.setup_diffusion_pipeline_profiler(
+            profiler_targets=profiler_targets,
+            enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
+        )
+
     def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
         """Adopt a runner-installed Cache-DiT backend (request-scoped protocol)."""
 

vllm_omni/diffusion/utils/param_utils.py

@@ -25,3 +25,18 @@ def apply_declared_extra_args(
     if not declared:
         return
     sampling_params.extra_args = {**(sampling_params.extra_args or {}), **declared}
+
+
+def ar_grid_max_tokens(ar_width: int, ar_height: int) -> int | None:
+    """Return the AR-stage ``max_tokens`` for MammothModa2-style visual-token grids.
+
+    The AR stage emits one visual token per grid cell, one EOL token per row,
+    and one final look-ahead token, so the generation budget is
+    ``ar_height * (ar_width + 1) + 1``. Shared by the online serving path and
+    the offline example so a grid-contract change cannot desynchronize one
+    copy. Returns ``None`` when the grid is absent or non-positive, letting
+    callers keep their existing default.
+    """
+    if ar_width <= 0 or ar_height <= 0:
+        return None
+    return ar_height * (ar_width + 1) + 1

vllm_omni/entrypoints/openai/api_server.py

@@ -2167,6 +2167,8 @@ async def generate_images(
                 extra_body["system_prompt"] = request.system_prompt
             if request.return_stage_metrics is not None:
                 extra_body["return_stage_metrics"] = request.return_stage_metrics
+            for extra_key, extra_value in (request.model_extra or {}).items():
+                extra_body.setdefault(extra_key, extra_value)
 
             generation_result = await chat_handler.generate_diffusion_images(
                 prompt=request.prompt,
@@ -2426,18 +2428,20 @@ async def edit_images(
         # 3.0 Init with system default values
         app_state_args = getattr(raw_request.app.state, "args", None)
         default_sample_param = getattr(app_state_args, "default_sampling_params", None)
-        # Currently only have one diffusion stage.
-        diffusion_stage_ids = [i for i, cfg in enumerate(stage_configs) if get_stage_type(cfg) == "diffusion"]
-        if not diffusion_stage_ids:
+        # Image edits (img2img) remain gated to classical diffusion stages;
+        # MammothModa2 is rejected at route level by its pipeline's endpoint
+        # restrictions (test_mammoth_moda2_shared_runtime.py).
+        image_stage_ids = [i for i, cfg in enumerate(stage_configs) if get_stage_type(cfg) == "diffusion"]
+        if not image_stage_ids:
             raise HTTPException(
                 status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
                 detail="No diffusion stage found in multi-stage pipeline.",
             )
-        diffusion_stage_id = diffusion_stage_ids[0]
+        image_stage_id = image_stage_ids[0]
         apply_stage_default_sampling_params(
             default_sample_param,
             gen_params,
-            str(diffusion_stage_id),
+            str(image_stage_id),
         )
         _update_if_not_none(gen_params, "num_outputs_per_prompt", n)
         # 3.1 Parse per-request LoRA (compatible with chat's extra_body.lora shape).

vllm_omni/entrypoints/openai/app_state.py

@@ -17,7 +17,7 @@ from vllm_omni.entrypoints.openai.serving_audio_generate import OmniOpenAIServin
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
 from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
-from vllm_omni.entrypoints.openai.utils import get_stage_type
+from vllm_omni.entrypoints.openai.utils import is_image_generation_stage
 
 ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format"
 
@@ -57,7 +57,8 @@ def _get_engine_and_model(raw_request: Request):
             detail="Multi-stage engine not initialized. Start server with a multi-stage omni model.",
         )
 
-    # Check if there's a diffusion stage.
+    # Check if there's an image generation stage (classical diffusion, or a
+    # legacy generation-LLM DiT stage that emits images).
     # Prefer app state (compat layer populated at startup), then fall back to
     # the engine client's stage configs for refactored AsyncOmni paths.
     stage_configs = getattr(raw_request.app.state, "stage_configs", None)
@@ -70,12 +71,12 @@ def _get_engine_and_model(raw_request: Request):
         )
 
     normalized_stage_configs = list(stage_configs)
-    has_diffusion_stage = any(get_stage_type(stage_cfg) == "diffusion" for stage_cfg in normalized_stage_configs)
+    has_image_stage = any(is_image_generation_stage(stage_cfg) for stage_cfg in normalized_stage_configs)
 
-    if not has_diffusion_stage:
+    if not has_image_stage:
         raise HTTPException(
             status_code=HTTPStatus.SERVICE_UNAVAILABLE.value,
-            detail="No diffusion stage found in multi-stage pipeline.",
+            detail="No image generation stage found in multi-stage pipeline.",
         )
 
     # Get server's loaded model name

vllm_omni/entrypoints/openai/protocol/images.py

@@ -18,7 +18,7 @@ from typing import Any, Literal
 
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, Field, field_validator
+from pydantic import BaseModel, ConfigDict, Field, field_validator
 
 from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers
 
@@ -59,6 +59,8 @@ class ImageGenerationRequest(BaseModel):
     for advanced diffusion parameters.
     """
 
+    model_config = ConfigDict(extra="allow")
+
     # Required fields
     prompt: str = Field(..., description="Text description of the desired image(s)")
     bot_task: str | None = Field(

vllm_omni/entrypoints/openai/serving_chat.py

@@ -28,7 +28,8 @@ from vllm.parser.utils import (
     count_chat_history_tool_calls as get_history_tool_calls_cnt,
 )
 
-from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
+from vllm_omni.diffusion.utils.image_output import extract_images_from_outputs
+from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args, ar_grid_max_tokens
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.diffusion_request_utils import (
     apply_normalized_diffusion_request_extra_args,
@@ -48,8 +49,10 @@ from vllm_omni.metrics.modality import (
 )
 from vllm_omni.model_executor.models.minicpmo_4_5.pipeline import MINICPMO45_REFERENCE_AUDIO_KEY
 from vllm_omni.model_extras import (
+    build_text_to_image_prompt,
     get_extra_body_params,
     get_extra_output_params,
+    should_init_extra_args_for_non_diffusion_stages,
 )
 
 try:
@@ -142,6 +145,7 @@ from vllm_omni.entrypoints.openai.stage_params import (
 from vllm_omni.entrypoints.openai.utils import (
     get_stage_type,
     get_supported_speakers_from_hf_config,
+    is_image_generation_stage,
     is_single_stage_diffusion,
     parse_lora_request,
     resolve_diffusion_od_config,
@@ -389,15 +393,74 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
 
         params: frozenset[str] = frozenset()
         try:
-            od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine)
-            if od_config is not None and getattr(od_config, "model_class_name", None):
-                params = get_extra_body_params(od_config.model_class_name)
+            try:
+                od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine)
+            except Exception:
+                od_config = None
+            model_class_name = getattr(od_config, "model_class_name", None) if od_config is not None else None
+            if model_class_name is None:
+                model_class_name = self._resolve_multistage_model_class_name(self.engine_client)
+            if model_class_name:
+                params = get_extra_body_params(model_class_name)
         except Exception as e:
             logger.warning("Failed to read model extra_body params: %s", e)
 
         self._diffusion_extra_body_params = params
         return params
 
+    def _resolve_multistage_model_class_name(self, engine: Any) -> str | None:
+        """Resolve the multistage pipeline's diffusion model class name.
+
+        Prefers the diffusion od_config (single source of truth for the
+        registry spec); falls back to the wrapper model arch declared on the
+        final image-output stage so multistage wrappers without a diffusion
+        od_config (e.g. the legacy generation-LLM DiT topology MammothModa2
+        used before #7134) still resolve.
+        """
+        cached = getattr(self, "_multistage_model_class_name", None)
+        if cached is not None:
+            return cached
+
+        model_class_name: str | None = None
+        try:
+            od_config = resolve_diffusion_od_config(self.engine_client, self._diffusion_engine)
+            if od_config is not None:
+                model_class_name = getattr(od_config, "model_class_name", None)
+        except Exception as e:
+            logger.warning("Failed to read diffusion od_config model class: %s", e)
+
+        if model_class_name is None:
+            for stage_cfg in getattr(engine, "stage_configs", None) or []:
+                if not is_image_generation_stage(stage_cfg):
+                    continue
+                stage_engine_args = getattr(stage_cfg, "engine_args", None)
+                if stage_engine_args is None and hasattr(stage_cfg, "get"):
+                    try:
+                        stage_engine_args = stage_cfg.get("engine_args")
+                    except Exception:
+                        stage_engine_args = None
+                for attr in ("model_class_name", "model_arch"):
+                    model_class_name = None
+                    if isinstance(stage_engine_args, dict):
+                        model_class_name = stage_engine_args.get(attr)
+                    elif stage_engine_args is not None and hasattr(stage_engine_args, "get"):
+                        try:
+                            model_class_name = stage_engine_args.get(attr)
+                        except Exception:
+                            model_class_name = None
+                    elif stage_engine_args is not None:
+                        model_class_name = getattr(stage_engine_args, attr, None)
+                    if model_class_name is None:
+                        model_class_name = getattr(stage_cfg, attr, None)
+                    if model_class_name is not None:
+                        break
+                if model_class_name is not None:
+                    break
+
+        if model_class_name is not None:
+            self._multistage_model_class_name = model_class_name
+        return model_class_name
+
     def _normalize_diffusion_request_args(
         self,
         request: ChatCompletionRequest,
@@ -3129,6 +3192,31 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             else:
                 engine_prompt_data = {"image": reference_images}
 
+        ar_prompt_info: dict[str, Any] | None = None
+        model_class_name = self._resolve_multistage_model_class_name(engine)
+        # Topology-independent: dispatch is by resolved model class, not by
+        # stage_type. The registry returns the prompt untouched for models
+        # without a text_to_image_prompt_builder, and #7134 migrates
+        # MammothModa2's DiT to a classical "diffusion" stage while still
+        # needing this AR envelope, so gating on "no diffusion stage present"
+        # would silently skip the builder after that migration.
+        if (
+            model_class_name is not None
+            and not reference_images
+            and bot_task is None
+            and use_system_prompt is None
+            and custom_system_prompt is None
+        ):
+            model_t2i_prompt = build_text_to_image_prompt(
+                model_class_name,
+                {"prompt": prompt, "negative_prompt": negative_prompt},
+                height=height,
+                width=width,
+            )
+            if model_t2i_prompt is not None and model_t2i_prompt.get("prompt") != prompt:
+                prompt = model_t2i_prompt["prompt"]
+                ar_prompt_info = model_t2i_prompt.get("additional_information")
+
         prompt_token_ids: list[int] | None = None
         system_prompt_type: str | None = None
         build_kwargs: dict[str, Any] = {}
@@ -3201,6 +3289,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         engine_prompt["modalities"] = modalities
         if negative_prompt is not None:
             engine_prompt["negative_prompt"] = negative_prompt
+        if ar_prompt_info is not None:
+            engine_prompt["additional_information"] = ar_prompt_info
 
         mm_processor_kwargs: dict[str, Any] = {}
         if height is not None:
@@ -3234,6 +3324,26 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             if stage_type == "llm" and ar_stop_token_ids is not None:
                 default_stage_params.stop_token_ids = ar_stop_token_ids
 
+            if stage_type == "llm" and idx == 0 and ar_prompt_info is not None:
+                if should_init_extra_args_for_non_diffusion_stages(model_class_name):
+                    extra_args = getattr(default_stage_params, "extra_args", None)
+                    if extra_args is None:
+                        extra_args = {}
+                        default_stage_params.extra_args = extra_args
+                    apply_declared_extra_args(
+                        default_stage_params,
+                        self._get_diffusion_extra_body_params(),
+                        extra_body,
+                    )
+                    seed_for_ar = extra_body.get("seed")
+                    if seed_for_ar is not None and hasattr(default_stage_params, "seed"):
+                        default_stage_params.seed = seed_for_ar
+                ar_width = int((ar_prompt_info.get("ar_width") or [0])[0])
+                ar_height = int((ar_prompt_info.get("ar_height") or [0])[0])
+                ar_grid_budget = ar_grid_max_tokens(ar_width, ar_height)
+                if ar_grid_budget is not None and hasattr(default_stage_params, "max_tokens"):
+                    default_stage_params.max_tokens = ar_grid_budget
+
             if (
                 comprehension_idx is not None
                 and idx == comprehension_idx
@@ -3283,6 +3393,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                                 default_stage_params.lora_scale = lora_scale
                     except Exception as e:  # pragma: no cover - safeguard
                         logger.warning("Failed to parse LoRA request: %s", e)
+            elif stage_type == "llm" and idx > 0 and is_image_generation_stage(stage_cfg):
+                if num_outputs_per_prompt is not None and num_outputs_per_prompt > 1:
+                    raise ValueError(
+                        "n > 1 is not supported for image generation stages "
+                        "served on the LLM path; each request produces "
+                        "exactly one image"
+                    )
+                self._set_if_supported(
+                    default_stage_params,
+                    height=height,
+                    width=width,
+                    seed=seed,
+                    generator_device=generator_device,
+                )
+                apply_declared_extra_args(
+                    default_stage_params,
+                    self._get_diffusion_extra_body_params(),
+                    extra_body,
+                )
 
         return engine_prompt, sampling_params_list
 
@@ -3473,32 +3602,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 request_id=request_id,
             )
 
-        images = getattr(result, "images", [])
+        # Generation-LLM DiT stages (e.g. MammothModa2) attach the decoded
+        # image to ``multimodal_output`` on the completion instead of the
+        # ``images`` field, so fall back to the shared extractor used by the
+        # offline example path before giving up.
+        images = getattr(result, "images", None) or extract_images_from_outputs(result)
         stage_durations = result.stage_durations
         peak_memory_mb = result.peak_memory_mb
         response_metrics = getattr(result, "metrics", None) if return_stage_metrics else None
         cot_output = None
 
-        req_out = result
-        if req_out:
-            prompt_obj = getattr(req_out, "prompt", None)
-            if isinstance(prompt_obj, dict):
-                extra = prompt_obj.get("extra", {})
-                if isinstance(extra, dict):
-                    ar_text = extra.get("ar_generated_text")
-                    if isinstance(ar_text, str) and ar_text.strip():
-                        cot_output = ar_text
-
-        req_out = result
-        if req_out:
-            prompt_obj = getattr(req_out, "prompt", None)
-            if isinstance(prompt_obj, dict):
-                extra = prompt_obj.get("extra", {})
-                if isinstance(extra, dict):
-                    ar_text = extra.get("ar_generated_text")
-                    if isinstance(ar_text, str) and ar_text.strip():
-                        cot_output = ar_text
-
         req_out = result
         if req_out:
             prompt_obj = getattr(req_out, "prompt", None)

vllm_omni/entrypoints/openai/utils.py

@@ -38,6 +38,34 @@ def get_stage_type(stage_cfg: Any) -> str:
     return getattr(stage_cfg, "stage_type", "llm")
 
 
+def is_image_generation_stage(stage_cfg: Any) -> bool:
+    """Return whether a stage produces the pipeline's final image output.
+
+    Covers both classical diffusion stages (``stage_type == "diffusion"``) and
+    generation-LLM stages that emit images (``stage_type == "llm"`` with
+    ``final_output_type == "image"`` — the legacy topology MammothModa2 used
+    before #7134 migrated it to the shared diffusion runtime). Serving
+    endpoints use this so image routes accept either topology.
+    """
+    if get_stage_type(stage_cfg) == "diffusion":
+        return True
+    final_output_type = None
+    if isinstance(stage_cfg, dict):
+        final_output = stage_cfg.get("final_output", False)
+        final_output_type = stage_cfg.get("final_output_type")
+    elif hasattr(stage_cfg, "get"):
+        try:
+            final_output = stage_cfg.get("final_output", False)
+            final_output_type = stage_cfg.get("final_output_type")
+        except Exception:
+            final_output = getattr(stage_cfg, "final_output", False)
+            final_output_type = getattr(stage_cfg, "final_output_type", None)
+    else:
+        final_output = getattr(stage_cfg, "final_output", False)
+        final_output_type = getattr(stage_cfg, "final_output_type", None)
+    return bool(final_output) and final_output_type in {"image", "images"}
+
+
 def is_video_generation_pipeline(stage_configs: list[Any] | None) -> bool:
     """Return whether a pipeline declares a final video output stage."""
     for stage in stage_configs or ():