[Fix] Restore MammothModa2 support

Signed-off-by: hsliuustc0106 <hsliuustc0106@users.noreply.github.com>

docs/api/README.md

@@ -114,6 +114,10 @@ Configuration classes.
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
 - [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
 - [vllm_omni.transformers_utils.configs.cosyvoice3.CosyVoice3Config][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
+- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
 - [vllm_omni.transformers_utils.configs.omnivoice.OmniVoiceConfig][]
 - [vllm_omni.transformers_utils.configs.voxcpm2.VoxCPM2Config][]
 

docs/configuration/stage_configs.md

@@ -59,7 +59,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
 | `edges` | list | optional | `null` | Explicit edge list for the KV transfer graph. Auto-derived from stage inputs if omitted. |
 | `stages` | list | optional | `[]` | Per-stage runtime overrides matched by `stage_id`. Pipeline stages are still created from `PipelineConfig` when this list is empty. |
 | `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
-| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
+| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only` / `qwen3_omni_moe_thinker_only`). |
 | `trust_remote_code` | bool \| null | optional | `null` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage when specified. |
 | `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
 | `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |

docs/models/supported_models.md

@@ -69,6 +69,8 @@ th {
 | `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | | — |
 | `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | | — |
 | `DotsTTSForConditionalGeneration` | dots.tts | `rednote-hilab/dots.tts-soar` | ✅︎ | | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | | — |
+| `MammothModa2ForConditionalGeneration` | MammothModa2-Dev (AR-only image understanding) | `bytedance-research/MammothModa2-Dev` | ✅︎ | | | | — |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | — |
 | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | — |

docs/user_guide/diffusion_features.md

@@ -129,6 +129,7 @@ The following tables show which models support each feature:
 | **Krea 2**               |     ❌     |     ✅      |           ❌           |       ❌        |         ❌         |          ❌          |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LongCat-Image**        |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **LongCat-Image-Edit**   |     ✅     |     ✅      |           ✅           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
+| **MammothModa2(T2I)**    |     ❌     |     ❌      |           ❌           |       ❌        |         ❌         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Nextstep_1(T2I)**      |     ❓     |     ❓      |           ❌           |       ✅        |         ✅         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **OmniGen2**             |     ❌     |     ✅      |           ✅           |       ❌        |         ✅         |          ❌          |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
 | **Ovis-Image**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -28,6 +28,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --deploy-config /path/to/deploy_config_file
 ```
 
+To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
+weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
+`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+```
+
+Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
+the same single-stage pipeline without `--deploy-config`.
+
 ### Launch individual stages (stage-based CLI)
 
 Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -0,0 +1,117 @@
+"""
+Offline inference example: MammothModa2 image summarization (single AR stage).
+
+Example:
+  uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+    --model path/to/MammothModa2-Preview \
+    --deploy-config vllm_omni/deploy/mammoth_moda2_ar.yaml \
+    --image /path/to/input.jpg \
+    --question "Please summarize the content of this image."
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+
+from PIL import Image
+from vllm import SamplingParams
+from vllm.multimodal.image import convert_image_mode
+
+from vllm_omni import Omni
+
+DEFAULT_SYSTEM = "You are a helpful assistant."
+DEFAULT_QUESTION = "Please summarize the content of this image."
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
+    parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
+    parser.add_argument("--deploy-config", type=str, required=True, help="Path to deploy config yaml.")
+    parser.add_argument("--image", type=str, required=True, help="Path to input image.")
+    parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
+    parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=512,
+        help="Max new tokens to generate.",
+    )
+    parser.add_argument("--temperature", type=float, default=0.2)
+    parser.add_argument("--top-p", type=float, default=0.9)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
+    return parser.parse_args()
+
+
+def build_prompt(system: str, question: str) -> str:
+    return (
+        f"<|im_start|>system\n{system}<|im_end|>\n"
+        "<|im_start|>user\n"
+        "<|vision_start|><|image_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+    )
+
+
+def main() -> None:
+    args = parse_args()
+
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Image file not found: {args.image}")
+
+    pil_image = Image.open(args.image)
+    image_data = convert_image_mode(pil_image, "RGB")
+    prompt = build_prompt(args.system, args.question)
+
+    omni = Omni(
+        model=args.model,
+        deploy_config=args.deploy_config,
+        trust_remote_code=args.trust_remote_code,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
+    try:
+        sp = SamplingParams(
+            temperature=float(args.temperature),
+            top_p=float(args.top_p),
+            top_k=-1,
+            max_tokens=int(args.max_tokens),
+            seed=int(args.seed),
+            detokenize=True,
+        )
+        # NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
+        # Consume it inside the try block so the worker isn't closed early.
+        outputs = list(
+            omni.generate(
+                [
+                    {
+                        "prompt": prompt,
+                        "multi_modal_data": {"image": image_data},
+                        "additional_information": {"omni_task": ["chat"]},
+                    }
+                ],
+                [sp],
+            )
+        )
+    finally:
+        omni.close()
+
+    lines: list[str] = []
+    for stage_outputs in outputs:
+        ro = stage_outputs
+        text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+        lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+        lines.append("answer:\n")
+        lines.append(text.strip() + "\n")
+        lines.append("\n")
+
+    print("\n".join(lines))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/text_to_image/text_to_image.py

@@ -599,6 +599,17 @@ def main():
             if args.seed is not None and hasattr(params, "seed"):
                 params.seed = args.seed
 
+            # MammothModa2's AR stage emits one visual token per grid cell,
+            # one EOL token per row, and one final look-ahead token whose hidden
+            # state is unavailable. Size the first stage from the prompt metadata
+            # instead of SamplingParams' default of 16 tokens.
+            prompt_info = prompt_dict.get("additional_information", {})
+            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
+
     if not diffusion_replaced and len(sampling_params_list) == 1:
         sampling_params_list = [diffusion_params]
 

examples/offline_inference/x_to_text/README.md

@@ -10,10 +10,12 @@ Generate text from text or image inputs with vLLM-Omni's shared offline entrypoi
 | --- | --- | --- | --- |
 | `ByteDance-Seed/BAGEL-7B-MoT` | Yes | Yes | `vllm_omni/deploy/bagel.yaml` |
 | `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | `vllm_omni/deploy/hunyuan_image3_ar.yaml` |
+| `bytedance-research/MammothModa2-Preview` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
+| `bytedance-research/MammothModa2-Dev` | Yes | Yes | `vllm_omni/deploy/mammoth_moda2_ar.yaml` |
 
 The script recognizes these model families from `config.json`, applies the
-model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3.
-BAGEL uses its registered default deploy.
+model-specific prompt format, and selects an AR-only deploy for HunyuanImage-3
+and MammothModa2. BAGEL uses its registered default deploy.
 
 ## Text-To-Text
 
@@ -33,6 +35,22 @@ python examples/offline_inference/x_to_text/x_to_text.py \
   --prompt "Explain multimodal inference in three concise sentences."
 ```
 
+### MammothModa2
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model bytedance-research/MammothModa2-Preview \
+  --prompt "Explain multimodal inference in three concise sentences."
+```
+
+The same command supports the Dev checkpoint:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model bytedance-research/MammothModa2-Dev \
+  --prompt "Explain multimodal inference in three concise sentences."
+```
+
 ## Image-To-Text
 
 Pass one image with `--image`. The default question can be replaced with any
@@ -54,6 +72,16 @@ python examples/offline_inference/x_to_text/x_to_text.py \
   --prompt "Describe this image in detail."
 ```
 
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model bytedance-research/MammothModa2-Preview \
+  --image image.png \
+  --prompt "Describe this image in detail."
+```
+
+Use `bytedance-research/MammothModa2-Dev` in the same command to run I2T with
+the Dev checkpoint.
+
 ## Key Arguments
 
 | Argument | Default | Description |
@@ -74,4 +102,4 @@ python examples/offline_inference/x_to_text/x_to_text.py \
 
 The committed HunyuanImage-3 AR-only default uses four GPUs with TP=4. To use a
 different topology, pass a custom AR-only YAML through `--deploy-config`.
-BAGEL uses its repository deploy default unless overridden.
+BAGEL and MammothModa2 use their repository deploy defaults unless overridden.

examples/offline_inference/x_to_text/x_to_text.py

@@ -15,6 +15,7 @@ from vllm_omni.model_extras import build_x_to_text_prompt, get_x_to_text_model_f
 
 _REPO_ROOT = Path(__file__).resolve().parents[3]
 _HUNYUAN_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml"
+_MAMMOTH_AR_DEPLOY_CONFIG = _REPO_ROOT / "vllm_omni" / "deploy" / "mammoth_moda2_ar.yaml"
 
 
 def parse_args() -> argparse.Namespace:
@@ -61,7 +62,7 @@ def main() -> None:
     omni_kwargs: dict[str, Any] = {
         "model": args.model,
         "mode": "image-to-text" if image is not None else "text-to-text",
-        "trust_remote_code": args.trust_remote_code or family == "hunyuan_image3",
+        "trust_remote_code": args.trust_remote_code or family in {"hunyuan_image3", "mammoth_moda2"},
         "log_stats": args.log_stats,
     }
     if args.enforce_eager is not None:
@@ -70,6 +71,8 @@ def main() -> None:
         omni_kwargs["deploy_config"] = args.deploy_config
     elif family == "hunyuan_image3":
         omni_kwargs["deploy_config"] = str(_HUNYUAN_AR_DEPLOY_CONFIG)
+    elif family == "mammoth_moda2":
+        omni_kwargs["deploy_config"] = str(_MAMMOTH_AR_DEPLOY_CONFIG)
 
     omni = Omni(**omni_kwargs)
     try:

examples/online_serving/minicpmo/README.md

@@ -192,6 +192,64 @@ stronger `response-required` mode is diagnostic: it requires a purpose-built
 two-response WAV, its `--input-sha256`, and an
 `--expect-second-response-substring` value.
 
+## Run Omni-DuplexEval
+
+Omni-DuplexEval generation uses the vLLM-Omni native MiniCPM-o duplex
+endpoint. Evaluation uses a separate OpenAI-compatible multimodal judge
+served from the same vLLM-Omni environment. This validates the vLLM-Omni
+duplex client and local judge integration; it does not reproduce the paper's
+original MiniCPM-o inference implementation.
+
+Start the duplex generation server:
+
+```bash
+vllm serve openbmb/MiniCPM-o-4_5 --omni --trust-remote-code \
+    --deploy-config vllm_omni/deploy/minicpmo_4_5_duplex.yaml \
+    --served-model-name openbmb/MiniCPM-o-4_5 \
+    --host 0.0.0.0 --port 8099
+```
+
+Start a separate multimodal judge. The allowed path must contain any local
+videos passed with `--judge-video-mode video_url`:
+
+```bash
+vllm serve Qwen/Qwen2.5-VL-7B-Instruct \
+    --served-model-name Qwen/Qwen2.5-VL-7B-Instruct \
+    --allowed-local-media-path /data/omni-duplex-eval \
+    --host 0.0.0.0 --port 8000
+```
+
+Generate, evaluate, and summarize:
+
+```bash
+vllm bench omni-duplex-eval --omni generate \
+    --url ws://127.0.0.1:8099/v1/realtime?duplex=1 \
+    --model openbmb/MiniCPM-o-4_5 \
+    --ref-audio /data/ref.wav \
+    --dataset Hothan/Omni-DuplexEval \
+    --concurrency 2 \
+    --response-root /data/omni-duplex-eval/responses
+
+vllm bench omni-duplex-eval --omni evaluate \
+    --dataset Hothan/Omni-DuplexEval \
+    --response-root /data/omni-duplex-eval/responses \
+    --score-root /data/omni-duplex-eval/scores \
+    --judge-base-url http://127.0.0.1:8000 \
+    --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+    --judge-video-mode video_url \
+    --eval-workers 4
+
+vllm bench omni-duplex-eval --omni summarize \
+    --score-root /data/omni-duplex-eval/scores
+```
+
+The dataset names such as `RTD_OCR` and `PR_correction` are Hugging Face
+splits, not dataset configurations. Pass one with `--split`, or omit it to
+run all nine splits. `--limit 1` is useful for a smoke test. Realtime
+generation records `clock=media`; artifacts generated with
+`--pace as-fast-as-possible` record `clock=invalid` and evaluation rejects
+them unless `--allow-invalid-clock` is explicit.
+
 ## Related examples
 
 - [Offline MiniCPM-o inference](../../offline_inference/minicpmo/)

examples/online_serving/qwen3_omni/README.md

@@ -24,6 +24,18 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
     --deploy-config /path/to/your_deploy_config.yaml
 ```
 
+To serve **thinker-only** (text output, no talker / code2wav loaded) with Instruct
+weights, pass the bundled thinker-only deploy YAML. The `pipeline:` key
+`qwen3_omni_moe_thinker_only` overrides the HF `enable_audio_output` resolver:
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --deploy-config vllm_omni/deploy/qwen3_omni_moe_thinking.yaml
+```
+
+Captioner / Thinking checkpoints (`enable_audio_output=false`) still auto-select
+the same single-stage pipeline without `--deploy-config`.
+
 For a 3x-GPU multi-replica layout (talker/code2wav scale-out on cuda:1,2),
 use `--stage-overrides` on top of the default config:
 

recipes/MammothModa2/MammothModa2.md

@@ -0,0 +1,212 @@
+# MammothModa2
+
+> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
+
+## Summary
+
+- Vendor: ByteDance Research
+- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
+- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
+- Mode: Offline inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to run MammothModa2-Preview through shared task-oriented
+offline examples. Text-to-image uses the shared image example (`text_to_image.py`) instead of a model-specific script.
+The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
+and forwards MammothModa2-specific generation parameters through the
+pipeline-declared `extra_body` contract.
+
+MammothModa2's DiT stage consumes its inputs through the multi-stage kwargs
+interface (not `OmniDiffusionRequest`), so its generation knobs
+(`text_guidance_scale`, `cfg_range`, `num_inference_steps`) are passed via
+`--extra-body` rather than the standard `--num-inference-steps` / `--cfg-scale`
+flags. Image size uses the standard `--height` / `--width` flags.
+
+## References
+
+- Upstream model:
+  [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
+- Dev model:
+  [`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
+- Related offline example:
+  [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
+- Related T2T/I2T example:
+  [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
+- Declared parameters:
+  [`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)
+- Deploy config:
+  [`vllm_omni/deploy/mammoth_moda2.yaml`](../../vllm_omni/deploy/mammoth_moda2.yaml)
+
+## Hardware Support
+
+The default deploy config runs both the AR and DiT stages on a single GPU
+(`devices: "0"`). The committed `gpu_memory_utilization` split (stage-0 AR `0.5`,
+stage-1 DiT `0.3`) is sized for an ~80 GB GPU. The model also fits on a 48 GB GPU
+after rebalancing the split so the AR weights (~23 GB) leave room for the KV
+cache — see the note under *1x L40S 48GB*.
+
+## GPU
+
+### 1x L40S 48GB
+
+> **48 GB config adjustment:** the committed
+> `vllm_omni/deploy/mammoth_moda2.yaml` uses
+> `gpu_memory_utilization` 0.5 / 0.3 (sized for ~80 GB). To fit on a 48 GB L40S,
+> set the stage-0 (AR) value to `0.8` and the stage-1 (DiT) value to `0.16`
+> before running. (On an ~80 GB GPU, leave the defaults unchanged.)
+
+#### Environment
+
+- OS: Linux
+- Python: Match the repository requirements for your checkout
+- Driver / runtime: NVIDIA CUDA environment with one L40S 48 GB (verified) or an
+  ~80 GB GPU for the default config
+- vLLM version: Match the repository requirements for your checkout
+- vLLM-Omni version or commit: Use the commit you are deploying from
+
+#### Offline Commands
+
+Download the model:
+
+```bash
+hf download bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
+```
+
+Run text-to-image with the shared offline example from the repository root. The
+deploy config sets `trust_remote_code`, so no extra flag is needed. Forward the
+MammothModa2 generation parameters as a JSON object through `--extra-body`:
+
+```bash
+python examples/offline_inference/text_to_image/text_to_image.py \
+  --model ./MammothModa2-Preview \
+  --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
+  --output mammoth_t2i.png
+```
+
+The `--extra-body` JSON forwards MammothModa2-specific parameters into
+`OmniDiffusionSamplingParams.extra_args`. Keys are filtered against the model's
+declared `extra_body_params` (see
+[`vllm_omni/model_extras/mammothmodal2_preview.py`](../../vllm_omni/model_extras/mammothmodal2_preview.py)),
+so unknown keys for MammothModa2 are silently dropped:
+
+- `text_guidance_scale` — classifier-free guidance scale for the DiT stage
+  (default `9.0`; CFG is active only when `> 1.0`).
+- `cfg_range` — relative step range `[start, end]` over which CFG is applied
+  (default `[0.0, 1.0]`).
+- `num_inference_steps` — number of DiT denoising steps (default `50`).
+
+`--height` and `--width` must be multiples of 16.
+
+Run text-to-text through the shared understanding example. It recognizes the
+MammothModa2 checkpoint and automatically selects `mammoth_moda2_ar.yaml`:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Preview \
+  --prompt "Explain multimodal generation in three sentences."
+```
+
+Add an image for image-to-text or image summarization. The shared example
+uses MammothModa2's chat and vision-token template:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Preview \
+  --image /path/to/input.jpg \
+  --prompt "Please summarize the content of this image."
+```
+
+#### Verification
+
+The example writes the generated image to the `--output` path. Confirm the file
+exists and is a valid image:
+
+```bash
+ls -lh mammoth_t2i.png
+python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
+```
+
+### 1x AMD MI300X, MammothModa2 Preview
+
+#### Environment
+
+- OS: Linux 6.8.0-134-generic, x86_64
+- Container: official ROCm image built from `docker/Dockerfile.rocm`
+- Python: 3.12.13
+- PyTorch: 2.11.0+gitd0c8b1f
+- Driver / runtime: AMD 6.19.14.31400000 / ROCm 7.2.53211
+- GPU: one AMD Instinct MI300X, `gfx942:sramecc+:xnack-`, 191.69 GiB visible HBM
+- vLLM version: 0.27.0+rocm723
+- vLLM Omni version or commit: `73e1368c7bb940efe1a025859c9d6c8eeeb2e3f0`
+- Installed vLLM Omni package metadata: `0.27.0rc2.dev44+g55abdade9.rocm`
+
+#### Offline Commands
+
+The checked run used the committed stage split, with `gpu_memory_utilization` set to 0.5 for AR and 0.3 for DiT:
+
+```bash
+python3 examples/offline_inference/text_to_image/text_to_image.py \
+    --model bytedance-research/MammothModa2-Preview \
+    --deploy-config vllm_omni/deploy/mammoth_moda2.yaml \
+    --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+    --height 1024 \
+    --width 1024 \
+    --seed 42 \
+    --extra-body '{"text_guidance_scale": 4.0, "cfg_range": [0.0, 1.0], "num_inference_steps": 50}' \
+    --enable-diffusion-pipeline-profiler \
+    --log-stats \
+    --output mammoth_t2i.png
+```
+
+#### Verification
+
+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.
+
+The output was a valid 1024 by 1024 RGB PNG.
+
+## MammothModa2-Dev unified inference
+
+MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
+Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
+`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
+required.
+
+Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
+Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
+then sends the generated visual tokens and hidden states to the DiT stage.
+
+Download the checkpoint:
+
+```bash
+hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
+```
+
+Run text-to-text through the shared understanding example. It recognizes the
+Dev checkpoint as MammothModa2 and automatically selects
+`mammoth_moda2_ar.yaml`:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Dev \
+  --prompt "Explain multimodal generation in three sentences."
+```
+
+Add an image for image-to-text or image summarization:
+
+```bash
+python examples/offline_inference/x_to_text/x_to_text.py \
+  --model ./MammothModa2-Dev \
+  --image ./image.png \
+  --prompt "Please summarize the content of this image."
+```
+
+The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
+run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
+and encoder caches. Allow additional GPU memory for those caches and the input
+image.

tests/benchmarks/duplex/test_omni_duplex_eval.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import hashlib
+import json
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.benchmarks.duplex import omni_duplex_eval_eval as eval_module
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_clock import (
+    extract_timed_sentences,
+    normalize_response_items,
+    split_text,
+    validate_clock,
+)
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import (
+    DuplexSample,
+    canonical_task_type,
+    family_for_split,
+    load_samples,
+    task_type_for_split,
+)
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_metrics import (
+    PROTOCOL_PIN,
+    build_content_prompt,
+    build_reminder_prompt,
+    build_temporal_prompt,
+    parse_judge_json,
+    reminder_window,
+    summarize_pr_results,
+    temporal_window,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_dataset_mapping_and_manifest(tmp_path):
+    manifest = tmp_path / "samples.json"
+    manifest.write_text(json.dumps([{"id": "a", "split": "RTD_OCR", "video": "clip.mp4"}]), encoding="utf-8")
+    sample = load_samples(manifest, media_root=tmp_path)[0]
+    assert sample.family == "rtd"
+    assert sample.video == str(tmp_path / "clip.mp4")
+    assert family_for_split("PR_event_reminder") == "pr"
+    assert canonical_task_type("post-event-reminder") == "post_event_reminder"
+    assert task_type_for_split("PR_correction") == "correction"
+    assert load_samples([{"id": "pr", "split": "PR_correction"}])[0].task_type == "correction"
+
+
+def test_hugging_face_subset_names_are_splits(monkeypatch):
+    calls = []
+
+    def fake_load_dataset(name, **kwargs):
+        calls.append((name, kwargs))
+        return [{"id": "pr", "question_text": "Correct this."}]
+
+    monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
+    sample = load_samples("Hothan/Omni-DuplexEval", split="PR_correction")[0]
+    assert calls == [("Hothan/Omni-DuplexEval", {"split": "PR_correction"})]
+    assert sample.split == "PR_correction"
+    assert sample.task_type == "correction"
+
+
+def test_response_aliases_and_clock_guard():
+    assert split_text("One. Two!") == ["One.", "Two!"]
+    assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
+        {"sentence": "x", "start": 800.0, "end": 800.0}
+    ]
+    with pytest.raises(ValueError, match="clock=invalid"):
+        validate_clock({"clock": "invalid"})
+    validate_clock({"clock": "invalid"}, allow_invalid=True)
+    timed = extract_timed_sentences([{"type": "response.output_text.delta", "delta": "Done.", "_media_clock_ms": 800}])
+    assert [item.as_dict() for item in timed] == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
+    with pytest.raises(ValueError, match="only clock=media"):
+        extract_timed_sentences([], clock="wall")
+
+
+def test_protocol_windows_and_parsing():
+    assert PROTOCOL_PIN == "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
+    assert temporal_window(4, 5, 10) == (2.0, 3.0)
+    assert temporal_window(1, 1.2, 10) == (0.0, 0.5)
+    assert reminder_window(5) == (5.0, 15.0)
+    assert parse_judge_json('noise {"content_score": 2.345, "is_relevant": 1}') == {
+        "content_score": 2.35,
+        "is_relevant": 1,
+    }
+    assert parse_judge_json('noise "success_score": 1') == {"success_score": 1}
+    assert summarize_pr_results([{"task_type": "correction", "all_success": 1}])["mean_all_success"] == 1.0
+
+
+def test_protocol_prompt_wording_is_pinned():
+    prompts = [
+        build_temporal_prompt(1, 2, "response", "question"),
+        build_content_prompt("response", "question", ["one", "two"]),
+        build_reminder_prompt("instruction", "response", "proactive_reminder"),
+        build_reminder_prompt("instruction", "response", "post_event_reminder"),
+        build_reminder_prompt("instruction", "response", "correction", "answer"),
+    ]
+    assert [hashlib.sha256(prompt.encode()).hexdigest() for prompt in prompts] == [
+        "5f7f5fcb4a37e4d2eea8ba23beabf74630e6fc2fb5b2661ceb94403254bb04d5",
+        "1fccbed6c1ed8ea1c76d5564d7289ed6adca3678f6a0da2cc735942123d64457",
+        "d0a0e3a9d873d2f368c5c99b9fe46681f683cc36a9fe0e09b99d6cbca1d6dd0c",
+        "08f9951243c859d271f82ee53da22cd0bdab1fe93c6ca04f4870d3b959c0803a",
+        "3de131a7e36cb4b8ac369bb05e87ede518be20221e214df3644211c179131a73",
+    ]
+
+
+def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
+    response = tmp_path / "response.json"
+    response.write_text(json.dumps([{"sentence": "A person moves.", "start": 1, "end": 2}]), encoding="utf-8")
+    response.with_name("response.meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
+    video = tmp_path / "video.mp4"
+    video.write_bytes(b"video")
+    monkeypatch.setattr(eval_module, "extract_jpeg", lambda *args, **kwargs: b"\xff\xd8jpeg")
+
+    class Judge:
+        def temporal(self, *args, **kwargs):
+            return '{"temporal_score": 3, "is_relevant": 1}'
+
+        def content(self, prompt, video, frames=None, *, mode="video_url"):
+            assert mode == "frame-sample"
+            assert frames == [b"\xff\xd8jpeg", b"\xff\xd8jpeg"]
+            return '{"content_score": 3}'
+
+    sample = DuplexSample(
+        id="rtd",
+        split="RTD_OCR",
+        family="rtd",
+        task_type=None,
+        video=video,
+        video_duration=4,
+    )
+    score = eval_module.evaluate_sample(
+        sample,
+        response,
+        tmp_path / "score.json",
+        Judge(),
+        judge_video_mode="frame-sample",
+    )
+    assert score["content"]["frame_count"] == 2

tests/config/test_config_factory.py

@@ -534,6 +534,7 @@ class TestPipelineDiscovery:
         """Check that specific models are in OMNI_PIPELINES."""
         assert "qwen2_5_omni" in OMNI_PIPELINES
         assert "qwen3_omni_moe" in OMNI_PIPELINES
+        assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
     def test_registry_resolver_qwen3_omni_all_stages(self):
@@ -558,6 +559,18 @@ class TestPipelineDiscovery:
         assert pipeline.default_deploy_config_name is None
         assert len(pipeline.stages) == 1  # thinker only
 
+    def test_registry_qwen3_omni_thinker_only_key_is_static(self):
+        """Explicit registry key must not re-enter the HF-config resolver."""
+        pipeline = resolve_pipeline_config(
+            "qwen3_omni_moe_thinker_only",
+            Q3_OMNI_ALL_STAGES_HF_CONFIG,
+        )
+        assert pipeline is OMNI_PIPELINES["qwen3_omni_moe_thinker_only"]
+        assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
+        assert len(pipeline.stages) == 1
+        assert pipeline.stages[0].engine_output_type == "text"
+        assert pipeline.stages[0].final_output_type == "text"
+
     @pytest.mark.parametrize(
         "pipeline",
         [registered for registered in OMNI_PIPELINES.values() if isinstance(registered, PipelineConfig)],
@@ -1225,14 +1238,22 @@ stages:
         assert deploy.connectors is not None
         assert deploy.platforms is not None
 
-    def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self):
+    @pytest.mark.parametrize(
+        ("hf_config", "model"),
+        [
+            (Q3_OMNI_THINKER_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Thinking"),
+            (Q3_OMNI_ALL_STAGES_HF_CONFIG, "Qwen/Qwen3-Omni-30B-A3B-Instruct"),
+        ],
+        ids=["thinker_checkpoint", "instruct_forced_thinker_only"],
+    )
+    def test_qwen3_omni_thinker_only_deploy_config_uses_resolved_pipeline(self, hf_config, model):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe_thinking.yaml"))
         with patch(
             "vllm_omni.config.config_factory.get_config",
-            return_value=Q3_OMNI_THINKER_HF_CONFIG,
+            return_value=hf_config,
         ):
             pipeline = StageConfigFactory.get_pipeline_config(
-                "Qwen/Qwen3-Omni-30B-A3B-Thinking",
+                model,
                 trust_remote_code=False,
                 deploy_config_path=str(deploy_path),
             )
@@ -1241,6 +1262,7 @@ stages:
         assert pipeline.model_type == "qwen3_omni_moe_thinker_only"
 
         deploy = load_deploy_config(deploy_path)
+        assert deploy.pipeline == "qwen3_omni_moe_thinker_only"
         stages = merge_pipeline_deploy(pipeline, deploy)
 
         assert len(stages) == 1
@@ -1292,6 +1314,8 @@ stages:
     @pytest.mark.parametrize(
         ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
         [
+            ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
+            ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
             ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
             ("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
             ("step_audio_2.yaml", "step_audio_2", 2, "audio", None),

tests/e2e/features/fullduplex/test_omni_duplex_eval_cli.py

@@ -0,0 +1,194 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import argparse
+import base64
+import json
+import threading
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+import pytest
+
+from vllm_omni.benchmarks.duplex import omni_duplex_eval_runner as runner
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DuplexSample
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
+from vllm_omni.entrypoints.cli.benchmark import omni_duplex_eval as cli
+from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.benchmark]
+
+
+def test_cli_generate_evaluate_summarize_flow(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys):
+    manifest = tmp_path / "manifest.json"
+    manifest.write_text(
+        json.dumps(
+            [
+                {
+                    "id": "sample-1",
+                    "split": "PR_correction",
+                    "question_text": "What changed?",
+                    "answer1": "The object moved.",
+                }
+            ]
+        ),
+        encoding="utf-8",
+    )
+    response_root = tmp_path / "responses"
+    score_root = tmp_path / "scores"
+
+    async def fake_generate(sample, *, output_root, **kwargs):
+        output = Path(output_root) / sample.split / f"{sample.id}.json"
+        output.parent.mkdir(parents=True, exist_ok=True)
+        output.write_text(json.dumps([{"sentence": "It moved.", "start": 0, "end": 1}]), encoding="utf-8")
+        output.with_name(output.stem + ".meta.json").write_text(json.dumps({"clock": "media"}), encoding="utf-8")
+        return output
+
+    class FakeJudge:
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def chat(self, *args, **kwargs):
+            return '{"success_score": 1, "is_relevant": 1}'
+
+    monkeypatch.setattr(cli, "generate_sample", fake_generate)
+    monkeypatch.setattr(cli, "DuplexJudge", FakeJudge)
+
+    parser = argparse.ArgumentParser()
+    OmniDuplexEvalSubcommand.add_cli_args(parser)
+
+    common = ["--dataset", str(manifest), "--family", "pr"]
+    generate = parser.parse_args(
+        [
+            "generate",
+            *common,
+            "--model",
+            "mock",
+            "--ref-audio",
+            str(manifest),
+            "--response-root",
+            str(response_root),
+            "--concurrency",
+            "2",
+        ]
+    )
+    OmniDuplexEvalSubcommand.cmd(generate)
+    evaluate = parser.parse_args(
+        [
+            "evaluate",
+            *common,
+            "--response-root",
+            str(response_root),
+            "--score-root",
+            str(score_root),
+            "--judge-model",
+            "mock-judge",
+            "--eval-workers",
+            "2",
+        ]
+    )
+    OmniDuplexEvalSubcommand.cmd(evaluate)
+    OmniDuplexEvalSubcommand.cmd(parser.parse_args(["summarize", "--score-root", str(score_root)]))
+
+    summary = json.loads(capsys.readouterr().out)
+    assert summary["samples"] == 1
+    assert summary["pr"]["mean_all_success"] == 1.0
+
+
+@pytest.mark.asyncio
+async def test_generate_exercises_realtime_socket_and_media_clock(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    import websockets
+
+    received = []
+
+    async def handler(websocket):
+        response_started = False
+        async for raw in websocket:
+            event = json.loads(raw)
+            received.append(event)
+            if event["type"] == "session.update":
+                await websocket.send(json.dumps({"type": "session.created"}))
+            elif event["type"] == "input_audio_buffer.append" and not response_started:
+                response_started = True
+                await websocket.send(json.dumps({"type": "response.created", "response": {"id": "r1"}}))
+                await websocket.send(
+                    json.dumps({"type": "response.output_text.delta", "response_id": "r1", "delta": "Done."})
+                )
+                await websocket.send(
+                    json.dumps(
+                        {
+                            "type": "response.audio.delta",
+                            "response_id": "r1",
+                            "delta": base64.b64encode(b"\0\0" * 240).decode(),
+                            "sample_rate_hz": 24_000,
+                        }
+                    )
+                )
+            elif event["type"] == "input_audio_buffer.commit":
+                await websocket.send(json.dumps({"type": "response.done", "response": {"id": "r1"}}))
+            elif event["type"] == "session.close":
+                await websocket.send(json.dumps({"type": "session.closed"}))
+                return
+
+    monkeypatch.setattr(runner, "read_audio_pcm16", lambda path: b"\0\0" * (16_000 * 8 // 10))
+    monkeypatch.setattr(runner, "video_duration", lambda path: 0.8)
+    monkeypatch.setattr(runner, "iter_jpegs", lambda *args, **kwargs: iter([(0.0, b"jpeg")]))
+    audio = tmp_path / "question.wav"
+    video = tmp_path / "video.mp4"
+    ref = tmp_path / "ref.wav"
+    for path in (audio, video, ref):
+        path.write_bytes(b"data")
+    sample = DuplexSample("sample", "PR_correction", "pr", "correction", video, audio)
+
+    async with websockets.serve(handler, "127.0.0.1", 0) as server:
+        port = server.sockets[0].getsockname()[1]
+        output = await runner.generate_sample(
+            sample,
+            url=f"ws://127.0.0.1:{port}/v1/realtime?duplex=1",
+            model="mock",
+            ref_audio=ref,
+            output_root=tmp_path / "responses",
+        )
+
+    assert json.loads(output.read_text(encoding="utf-8")) == [{"sentence": "Done.", "start": 0.8, "end": 0.8}]
+    meta = json.loads(output.with_name("sample.meta.json").read_text(encoding="utf-8"))
+    assert meta["response_done"] is True
+    assert meta["drain_timeout"] is None
+    assert any(event["type"] == "playback.ack" for event in received)
+
+
+def test_judge_exercises_openai_http_schema():
+    requests = []
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_POST(self):
+            length = int(self.headers["Content-Length"])
+            requests.append((self.path, self.headers["Authorization"], json.loads(self.rfile.read(length))))
+            payload = json.dumps({"choices": [{"message": {"content": '{"success_score": 1}'}}]}).encode()
+            self.send_response(200)
+            self.send_header("Content-Type", "application/json")
+            self.send_header("Content-Length", str(len(payload)))
+            self.end_headers()
+            self.wfile.write(payload)
+
+        def log_message(self, *args):
+            pass
+
+    server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    try:
+        judge = DuplexJudge(f"http://127.0.0.1:{server.server_port}", "judge", api_key="token")
+        assert judge.chat("prompt") == '{"success_score": 1}'
+    finally:
+        server.shutdown()
+        thread.join()
+        server.server_close()
+
+    path, authorization, payload = requests[0]
+    assert path == "/v1/chat/completions"
+    assert authorization == "Bearer token"
+    assert payload["model"] == "judge"
+    assert payload["messages"] == [{"role": "user", "content": "prompt"}]

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

@@ -0,0 +1,204 @@
+"""
+End-to-end test for MammothModa2 text-to-image generation.
+
+Verifies that the AR->DiT pipeline produces an image tensor whose pixel values
+match a golden reference.
+
+Model Hub repo id: ``bytedance-research/MammothModa2-Preview``.
+Deploy config: ``get_deploy_config_path("mammoth_moda2.yaml")`` -> ``vllm_omni/deploy/mammoth_moda2.yaml``
+
+Golden pixel file: ``tests/e2e/offline_inference/fixtures/mammoth_moda2_t2i_golden.json``
+  Regenerate with: ``UPDATE_GOLDEN=1 pytest tests/e2e/offline_inference/test_mammoth_moda2_expansion.py``
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from huggingface_hub import snapshot_download
+from vllm.sampling_params import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+_IMAGE_TOKEN_ID = 151655  # "<|image_pad|>"
+_VIDEO_TOKEN_ID = 151656  # "<|video_pad|>"
+_VISION_START_TOKEN_ID = 151652  # "<|vision_start|>"
+_VISION_END_TOKEN_ID = 151653  # "<|vision_end|>"
+_AR_PATCH_SIZE = 16
+
+MODEL_PATH = "bytedance-research/MammothModa2-Preview"
+T2I_DEPLOY_CONFIG = get_deploy_config_path("mammoth_moda2.yaml")
+
+_OMNI_RUNNER_PARAM = (MODEL_PATH, T2I_DEPLOY_CONFIG)
+
+# Golden pixel reference file.  Set UPDATE_GOLDEN=1 to regenerate.
+_GOLDEN_T2I_PATH = Path(__file__).parent / "fixtures" / "mammoth_moda2_t2i_golden.json"
+# Fixed sampling coordinates: (channel, row_fraction, col_fraction)
+# Covers corners, centre, and mid-edges across all 3 channels.
+_PIXEL_SAMPLE_COORDS = [
+    (0, 0.0, 0.0),
+    (0, 0.5, 0.5),
+    (0, 1.0, 1.0),
+    (0, 0.25, 0.75),
+    (1, 0.0, 1.0),
+    (1, 0.5, 0.0),
+    (1, 0.75, 0.25),
+    (1, 1.0, 0.5),
+    (2, 0.0, 0.5),
+    (2, 0.5, 1.0),
+    (2, 0.75, 0.75),
+    (2, 1.0, 0.0),
+]
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+def _load_t2i_gen_config(repo_id: str) -> dict:
+    weights_dir = Path(snapshot_download(repo_id))
+    cfg_path = weights_dir / "t2i_generation_config.json"
+    if not cfg_path.exists():
+        pytest.skip(f"t2i_generation_config.json not found at {cfg_path}")
+    with cfg_path.open() as f:
+        return json.load(f)
+
+
+def _format_t2i_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
+    return (
+        "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+        f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
+        "<|im_start|>assistant\n"
+        f"<|image start|>{ar_width}*{ar_height}<|image token|>"
+    )
+
+
+def _sample_pixels(img_tensor: torch.Tensor) -> list[float]:
+    """Sample pixel values at fixed fractional coordinates from a (C, H, W) tensor."""
+    t = img_tensor.float().clamp(0.0, 1.0)
+    if t.ndim == 4:
+        t = t[0]  # unbatch
+    C, H, W = t.shape
+    values = []
+    for c, rh, rw in _PIXEL_SAMPLE_COORDS:
+        ri = min(int(rh * (H - 1)), H - 1)
+        ci = min(int(rw * (W - 1)), W - 1)
+        values.append(round(float(t[c, ri, ci]), 6))
+    return values
+
+
+# ---------------------------------------------------------------------------
+# End-to-end test
+# ---------------------------------------------------------------------------
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.diffusion,
+    pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
+]
+
+
+@pytest.mark.skip(reason="https://github.com/vllm-project/vllm-omni/issues/3201")
+@hardware_test(res={"cuda": "H100"})
+def test_mammothmoda2_t2i_e2e(omni_runner: OmniRunner):
+    """
+    End-to-end text-to-image generation with MammothModa2 (AR -> DiT).
+
+    Verifies:
+      - Omni pipeline initialises with the two-stage YAML config.
+      - DiT stage outputs an image tensor with the correct shape.
+      - A fixed set of pixel values matches a golden reference
+        (regenerate with ``UPDATE_GOLDEN=1``).
+    """
+    gen_cfg = _load_t2i_gen_config(MODEL_PATH)
+    eol_token_id = int(gen_cfg["eol_token_id"])
+    visual_start = int(gen_cfg["visual_token_start_id"])
+    visual_end = int(gen_cfg["visual_token_end_id"])
+
+    height, width = 256, 256  # small for CI speed
+    ar_height, ar_width = height // _AR_PATCH_SIZE, width // _AR_PATCH_SIZE
+    expected_grid_tokens = ar_height * (ar_width + 1)
+
+    prompt_text = "A cat sitting on a laptop keyboard"
+    formatted_prompt = _format_t2i_prompt(prompt_text, ar_width, ar_height)
+
+    omni = omni_runner.omni
+    ar_sampling = SamplingParams(
+        temperature=0.0,
+        top_k=1,
+        max_tokens=max(1, expected_grid_tokens + 1),
+        detokenize=False,
+    )
+    dit_sampling = SamplingParams(temperature=0.0, max_tokens=1, detokenize=False)
+
+    outputs = list(
+        omni.generate(
+            [
+                {
+                    "prompt": formatted_prompt,
+                    "additional_information": {
+                        "omni_task": ["t2i"],
+                        "ar_width": [ar_width],
+                        "ar_height": [ar_height],
+                        "eol_token_id": [eol_token_id],
+                        "visual_token_start_id": [visual_start],
+                        "visual_token_end_id": [visual_end],
+                        "image_height": [height],
+                        "image_width": [width],
+                        "num_inference_steps": [2],
+                        "text_guidance_scale": [1.0],
+                        "cfg_range": [0.0, 1.0],
+                        "visual_ids": [
+                            _IMAGE_TOKEN_ID,
+                            _VIDEO_TOKEN_ID,
+                            _VISION_START_TOKEN_ID,
+                            _VISION_END_TOKEN_ID,
+                        ],
+                    },
+                }
+            ],
+            [ar_sampling, dit_sampling],
+        )
+    )
+
+    assert len(outputs) > 0, "Pipeline produced no outputs"
+
+    found_image = False
+    for out in outputs:
+        ro_list = out
+        if not isinstance(ro_list, list):
+            ro_list = [ro_list]
+        for ro in ro_list:
+            completion_outputs = getattr(ro, "outputs", None)
+            if not isinstance(completion_outputs, list):
+                continue
+            for completion in completion_outputs:
+                mm = getattr(completion, "multimodal_output", None)
+                if not (isinstance(mm, dict) and "image" in mm):
+                    continue
+                img_list = mm["image"] if isinstance(mm["image"], list) else [mm["image"]]
+                for img_tensor in img_list:
+                    assert isinstance(img_tensor, torch.Tensor), f"Expected image tensor, got {type(img_tensor)}"
+                    assert img_tensor.ndim in (3, 4), f"Expected 3D or 4D image tensor, got {img_tensor.ndim}D"
+
+                    sampled = _sample_pixels(img_tensor)
+
+                    if os.environ.get("UPDATE_GOLDEN"):
+                        _GOLDEN_T2I_PATH.parent.mkdir(parents=True, exist_ok=True)
+                        _GOLDEN_T2I_PATH.write_text(json.dumps({"pixels": sampled}, indent=2))
+                        print(f"\nGolden file written to {_GOLDEN_T2I_PATH}")
+                    elif _GOLDEN_T2I_PATH.exists():
+                        golden = json.loads(_GOLDEN_T2I_PATH.read_text())["pixels"]
+                        for i, (got, exp) in enumerate(zip(sampled, golden)):
+                            assert abs(got - exp) < 1e-4, f"Pixel {i} mismatch: got {got}, expected {exp}"
+
+                    found_image = True
+
+    assert found_image, "No image tensor found in pipeline output"

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Unit tests for MammothModa2 config classes.
+
+These tests verify the correct behavior of the Mammothmoda2Config class,
+particularly the get_text_config() method which is called during parent
+class initialization (PretrainedConfig.__init__ -> validation -> get_text_config).
+
+Regression test for: AttributeError: 'Mammothmoda2Config' object has no attribute 'llm_config'
+"""
+
+import pytest
+
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import (
+    Mammothmoda2Config,
+    Mammothmoda2Qwen2_5_VLConfig,
+    Mammothmoda2Qwen2_5_VLTextConfig,
+    Mammothmoda2Qwen3VLConfig,
+    Mammothmoda2Qwen3VLTextConfig,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestMammothmoda2Config:
+    """Tests for Mammothmoda2Config class."""
+
+    def test_get_text_config_returns_nested_text_config(self):
+        """
+        Verify that get_text_config() returns the nested text_config from llm_config.
+
+        This is a regression test for the AttributeError that occurred when
+        get_text_config() was called during parent class initialization before
+        llm_config was set.
+        """
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen2_5_vl",
+            }
+        )
+
+        assert config.llm_config is not None
+        assert isinstance(config.llm_config, Mammothmoda2Qwen2_5_VLConfig)
+
+        text_config = config.get_text_config()
+        assert text_config is not None
+        assert isinstance(text_config, Mammothmoda2Qwen2_5_VLTextConfig)
+
+        # The text_config should be the same object as llm_config.text_config
+        assert text_config is config.llm_config.text_config
+
+    def test_get_text_config_returns_none_when_llm_config_is_none(self):
+        """Verify that get_text_config() returns None when llm_config is None."""
+        config = Mammothmoda2Config(llm_config=None)
+
+        assert config.llm_config is None
+        assert config.get_text_config() is None
+
+    def test_init_sets_llm_config_before_parent_init(self):
+        """
+        Verify that llm_config is set before parent __init__ is called.
+
+        This is critical because the parent PretrainedConfig.__init__ calls
+        validation which in turn calls get_text_config(). If llm_config is not
+        set before super().__init__(), an AttributeError will be raised.
+        """
+        # This should not raise an AttributeError
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen2_5_vl",
+            }
+        )
+        assert config.llm_config is not None
+
+
+class TestMammothmoda2Qwen2_5_VLConfig:
+    """Tests for Mammothmoda2Qwen2_5_VLConfig class."""
+
+    def test_has_text_config_attribute(self):
+        """Verify that Mammothmoda2Qwen2_5_VLConfig has text_config attribute."""
+        config = Mammothmoda2Qwen2_5_VLConfig()
+        assert hasattr(config, "text_config")
+        assert isinstance(config.text_config, Mammothmoda2Qwen2_5_VLTextConfig)
+
+    def test_text_config_has_expected_attributes(self):
+        """Verify that the nested text_config has expected attributes for token validation."""
+        config = Mammothmoda2Qwen2_5_VLConfig()
+        text_config = config.text_config
+
+        assert hasattr(text_config, "vocab_size")
+        assert hasattr(text_config, "image_token_id")
+        assert hasattr(text_config, "video_token_id")
+
+
+class TestMammothmoda2Qwen3VLConfig:
+    def test_dev_config_uses_qwen3_vl_subconfigs(self):
+        config = Mammothmoda2Config(
+            llm_config={
+                "model_type": "mammothmoda2_qwen3_vl",
+                "text_config": {
+                    "model_type": "mammothmoda2_qwen3_vl_text",
+                    "vocab_size": 151936,
+                    "gen_vocab_size": 32800,
+                    "gen_vocab_start_index": 152064,
+                },
+                "vision_config": {
+                    "model_type": "mammothmoda2_qwen3_vl_vision",
+                    "deepstack_visual_indexes": [8, 16, 24],
+                },
+            }
+        )
+
+        assert isinstance(config.llm_config, Mammothmoda2Qwen3VLConfig)
+        assert isinstance(config.get_text_config(), Mammothmoda2Qwen3VLTextConfig)
+        # vLLM validates sampling against the complete base + visual vocabulary.
+        assert config.get_text_config().vocab_size == 152064 + 32800
+        assert config.get_text_config().gen_vocab_start_index == 152064
+        assert config.llm_config.vision_config.deepstack_visual_indexes == [8, 16, 24]

tests/model_executor/models/test_kv_cache_scale_mapper.py

@@ -14,6 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 _REPO_ROOT = Path(__file__).resolve().parents[3]
 _STALE_API_FILES = [
     "vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py",
+    "vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py",
     "vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py",
     "vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py",
 ]
@@ -117,6 +118,10 @@ def test_auto_weights_loader_delegated_loaders_accept_mapped_cache_scale(
 @pytest.mark.parametrize(
     ("module_name", "class_name"),
     [
+        (
+            "vllm_omni.model_executor.models.mammoth_moda2.mammoth_moda2",
+            "MammothModa2Qwen2ForCausalLM",
+        ),
         (
             "vllm_omni.model_executor.models.mimo_audio.mimo_audio_llm",
             "MiMoAudioLLMForConditionalGeneration",

tests/model_extras/test_model_extras.py

@@ -31,6 +31,7 @@ from vllm_omni.model_extras import (
     [
         ({"model_type": "bagel"}, "bagel"),
         ({"architectures": ["HunyuanImage3ForConditionalGeneration"]}, "hunyuan_image3"),
+        ({"model_type": "mammothmoda2"}, "mammoth_moda2"),
         ({"model_type": "qwen2_vl"}, "generic"),
     ],
 )
@@ -50,6 +51,22 @@ def test_bagel_x_to_text_prompt_builder() -> None:
     assert stop_token_ids is None
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_mammoth_x_to_text_prompt_builder() -> None:
+    prompt, stop_token_ids = build_x_to_text_prompt("mammoth_moda2", "unused", "Hello.", has_image=False)
+    assert prompt == {
+        "prompt": (
+            "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+            "<|im_start|>user\nHello.<|im_end|>\n"
+            "<|im_start|>assistant\n"
+        ),
+        "modalities": ["text"],
+        "additional_information": {"omni_task": ["chat"]},
+    }
+    assert stop_token_ids is None
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_generic_x_to_text_prompt_builder() -> None:
@@ -555,3 +572,64 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
         "cfg_text_scale": 4.0,
         "think": False,
     }
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None:
+    for model_class_name in (
+        "MammothModa2DiTPipeline",
+        "MammothModa2ForConditionalGeneration",
+        "Mammothmoda2Model",
+    ):
+        assert get_extra_body_params(model_class_name) == frozenset(
+            {
+                "text_guidance_scale",
+                "cfg_range",
+                "num_inference_steps",
+            }
+        )
+        assert get_extra_output_params(model_class_name) == frozenset()
+        assert should_init_extra_args_for_non_diffusion_stages(model_class_name) is True
+
+    wrapper_prompt = build_text_to_image_prompt(
+        "MammothModa2ForConditionalGeneration",
+        {"prompt": "a cat", "modalities": ["image"]},
+        height=256,
+        width=256,
+    )
+    assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_mammothmoda2_text_to_image_prompt_builder() -> None:
+    # Image dims are converted to the AR grid (width/16 x height/16); the negative
+    # prompt is ignored (MammothModa2 t2i uses CFG, not an explicit negative path).
+    assert build_text_to_image_prompt(
+        "MammothModa2DiTPipeline",
+        {
+            "prompt": "a cat",
+            "modalities": ["image"],
+            "negative_prompt": "blurry",
+        },
+        height=512,
+        width=768,
+    ) == {
+        "prompt": (
+            "<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
+            "<|im_start|>user\na cat<|im_end|>\n"
+            "<|im_start|>assistant\n"
+            "<|image start|>48*32<|image token|>"
+        ),
+        "additional_information": {
+            "omni_task": ["t2i"],
+            "ar_width": [48],
+            "ar_height": [32],
+            "image_height": [512],
+            "image_width": [768],
+            "eol_token_id": [152064],
+            "visual_token_start_id": [152072],
+            "visual_token_end_id": [168456],
+        },
+    }

tests/platforms/npu/test_diffusion_attn_backend_selector.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the NPU diffusion attention backend selector.
+
+Regression coverage for the eager ``mindiesd`` import in
+``NPUOmniPlatform.get_diffusion_attn_backend_cls``: the import must run only
+for backends that reach mindiesd kernels (FLASH_ATTN, and RAINFUSION_ATTN via
+its dense FlashAttention fallback), so that a broken optional mindiesd
+install cannot block backends that never use it (e.g. TORCH_SDPA).
+
+The platform module is loaded from source with fake vllm/vllm-ascend
+dependencies, so the tests run on CPU without NPU or mindiesd installed.
+"""
+
+from __future__ import annotations
+
+import builtins
+import importlib.util
+import sys
+import types
+from enum import Enum
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_BACKEND_NAMES = (
+    "FLASH_ATTN",
+    "FLASH_ATTN_HUB",
+    "FLASH_ATTN_3_HUB",
+    "TORCH_SDPA",
+    "RAINFUSION_ATTN",
+)
+
+
+def _repo_root() -> Path:
+    marker = Path("vllm_omni") / "platforms" / "npu" / "platform.py"
+    for parent in Path(__file__).resolve().parents:
+        if (parent / marker).is_file():
+            return parent
+    raise FileNotFoundError(f"could not locate repo root containing {marker}")
+
+
+def _install_fake_module(monkeypatch: pytest.MonkeyPatch, name: str, **attrs):
+    module = types.ModuleType(name)
+    for key, value in attrs.items():
+        setattr(module, key, value)
+    monkeypatch.setitem(sys.modules, name, module)
+    return module
+
+
+def _make_fake_backend_enum():
+    members = {
+        name: SimpleNamespace(get_path=lambda n=name: f"fake.backends.{n.lower()}.Backend") for name in _BACKEND_NAMES
+    }
+
+    class _FakeEnumMeta(type):
+        def __getitem__(cls, name):
+            return members[name]
+
+        def __getattr__(cls, name):
+            try:
+                return members[name]
+            except KeyError:
+                raise AttributeError(name) from None
+
+    class FakeDiffusionAttentionBackendEnum(metaclass=_FakeEnumMeta):
+        pass
+
+    return FakeDiffusionAttentionBackendEnum
+
+
+class _FakeOmniPlatformEnum(Enum):
+    CUDA = "cuda"
+    NPU = "npu"
+
+
+class _FakeNPUPlatform:
+    pass
+
+
+class _FakeOmniPlatform:
+    _omni_enum = _FakeOmniPlatformEnum.NPU
+
+    @classmethod
+    def validate_diffusion_attn_backend(cls, selected_backend: str) -> None:
+        pass
+
+
+def _load_platform_module(monkeypatch: pytest.MonkeyPatch):
+    """Load vllm_omni/platforms/npu/platform.py with fake dependencies."""
+    _install_fake_module(monkeypatch, "vllm")
+    _install_fake_module(
+        monkeypatch,
+        "vllm.logger",
+        init_logger=lambda name: SimpleNamespace(
+            debug=lambda *a, **k: None,
+            info=lambda *a, **k: None,
+            warning=lambda *a, **k: None,
+            error=lambda *a, **k: None,
+        ),
+    )
+    _install_fake_module(monkeypatch, "vllm_ascend")
+    _install_fake_module(monkeypatch, "vllm_ascend.platform", NPUPlatform=_FakeNPUPlatform)
+    for pkg in (
+        "vllm_omni",
+        "vllm_omni.diffusion",
+        "vllm_omni.diffusion.attention",
+        "vllm_omni.diffusion.attention.backends",
+        "vllm_omni.platforms",
+    ):
+        _install_fake_module(monkeypatch, pkg)
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.diffusion.attention.backends.registry",
+        DiffusionAttentionBackendEnum=_make_fake_backend_enum(),
+    )
+    _install_fake_module(
+        monkeypatch,
+        "vllm_omni.platforms.interface",
+        OmniPlatform=_FakeOmniPlatform,
+        OmniPlatformEnum=_FakeOmniPlatformEnum,
+    )
+
+    path = _repo_root() / "vllm_omni" / "platforms" / "npu" / "platform.py"
+    spec = importlib.util.spec_from_file_location("npu_platform_under_test", path)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"cannot load import spec for {path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.NPUOmniPlatform
+
+
+@pytest.fixture
+def mindiesd_tracker(monkeypatch: pytest.MonkeyPatch):
+    """Make mindiesd importable and record eager import attempts."""
+    monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
+    attempts = []
+    real_import = builtins.__import__
+
+    def fake_import(name, *args, **kwargs):
+        if name == "mindiesd":
+            attempts.append(name)
+            return types.ModuleType("mindiesd")
+        return real_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", fake_import)
+    monkeypatch.setattr(
+        "importlib.util.find_spec",
+        lambda name, *a, **k: object() if name == "mindiesd" else None,
+    )
+    return attempts
+
+
+@pytest.mark.parametrize(
+    ("selected_backend", "expect_mindiesd_import"),
+    [
+        ("FLASH_ATTN", True),
+        # HF-hub variants fall back to local FLASH_ATTN on NPU.
+        ("FLASH_ATTN_HUB", True),
+        ("FLASH_ATTN_3_HUB", True),
+        ("RAINFUSION_ATTN", True),
+        # Backends that never reach mindiesd must not trigger the import.
+        ("TORCH_SDPA", False),
+    ],
+)
+def test_eager_mindiesd_import_scoped_to_mindiesd_backends(
+    monkeypatch: pytest.MonkeyPatch,
+    mindiesd_tracker: list[str],
+    selected_backend: str,
+    expect_mindiesd_import: bool,
+):
+    platform_cls = _load_platform_module(monkeypatch)
+
+    path = platform_cls.get_diffusion_attn_backend_cls(selected_backend, head_size=64)
+
+    assert (len(mindiesd_tracker) > 0) is expect_mindiesd_import
+    expected = "FLASH_ATTN" if selected_backend in ("FLASH_ATTN_HUB", "FLASH_ATTN_3_HUB") else selected_backend
+    assert path == f"fake.backends.{expected.lower()}.Backend"
+
+
+def test_explicit_flash_attn_without_mindiesd_still_resolves(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    platform_cls = _load_platform_module(monkeypatch)
+    monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
+    monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
+
+    path = platform_cls.get_diffusion_attn_backend_cls("FLASH_ATTN", head_size=64)
+
+    assert path == "fake.backends.flash_attn.Backend"
+
+
+def test_default_backend_prefers_flash_attn_when_mindiesd_available(
+    monkeypatch: pytest.MonkeyPatch,
+    mindiesd_tracker: list[str],
+):
+    platform_cls = _load_platform_module(monkeypatch)
+
+    path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
+
+    assert path == "fake.backends.flash_attn.Backend"
+    assert len(mindiesd_tracker) > 0
+
+
+def test_default_backend_falls_back_to_sdpa_without_mindiesd(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    platform_cls = _load_platform_module(monkeypatch)
+    monkeypatch.delitem(sys.modules, "mindiesd", raising=False)
+    monkeypatch.setattr("importlib.util.find_spec", lambda name, *a, **k: None)
+
+    path = platform_cls.get_diffusion_attn_backend_cls(None, head_size=64)
+
+    assert path == "fake.backends.torch_sdpa.Backend"

tests/worker/test_omni_connector_mixin.py

@@ -379,6 +379,7 @@ class TestLoadCustomFuncSelection(unittest.TestCase):
     def test_skips_non_payload_stage_input_processors_for_full_payload_mode(self):
         incompatible_paths = [
             "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav",
+            "vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit",
             "vllm_omni.model_executor.stage_input_processors.cosyvoice3.text2flow",
             "vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion",
         ]

vllm_omni/benchmarks/duplex/__init__.py

@@ -0,0 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CPU-friendly Omni-DuplexEval generation and scoring helpers."""
+
+from .omni_duplex_eval_dataset import DuplexSample, load_samples
+from .omni_duplex_eval_metrics import PROTOCOL_PIN
+
+__all__ = ["DuplexSample", "PROTOCOL_PIN", "load_samples"]

vllm_omni/benchmarks/duplex/omni_duplex_eval_clock.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Media-clock sentence extraction and response JSON normalization."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any
+
+import regex as re
+
+_TERMINAL = re.compile(r"(.+?[.!?。！？]+)(?:\s+|$)", re.S)
+
+
+@dataclass
+class TimedSentence:
+    sentence: str
+    start: float
+    end: float
+
+    def as_dict(self) -> dict[str, Any]:
+        return {"sentence": self.sentence, "start": round(self.start, 3), "end": round(self.end, 3)}
+
+
+def split_text(text: str) -> list[str]:
+    text = str(text or "").strip()
+    if not text:
+        return []
+    parts: list[str] = []
+    position = 0
+    for match in _TERMINAL.finditer(text):
+        parts.append(match.group(1).strip())
+        position = match.end()
+    tail = text[position:].strip()
+    if tail:
+        parts.append(tail)
+    return [part for part in parts if part]
+
+
+def extract_timed_sentences(events: Iterable[dict[str, Any]], *, clock: str = "media") -> list[TimedSentence]:
+    """Split output deltas while anchoring text to input ``audio_end_ms``."""
+    if clock != "media":
+        raise ValueError("only clock=media is supported")
+    sentences: list[TimedSentence] = []
+    pending = ""
+    start: float | None = None
+    last: float | None = None
+    for event in events:
+        event_type = event.get("type")
+        if event_type not in {"response.output_text.delta", "response.audio_transcript.delta", "response.text.delta"}:
+            continue
+        delta = event.get("delta")
+        if not isinstance(delta, str) or not delta.strip():
+            continue
+        timestamp = event.get("audio_end_ms")
+        if timestamp is None:
+            timestamp = event.get("_media_clock_ms", event.get("current_time_ms", 0.0))
+        try:
+            # Every accepted event field is explicitly millisecond-valued.
+            now = float(timestamp) / 1000.0
+        except (TypeError, ValueError):
+            now = last or 0.0
+        pending += delta
+        start = now if start is None else start
+        last = now
+        pieces = split_text(pending)
+        complete = pieces if pieces and re.search(r"[.!?。！？]$", pending.strip()) else pieces[:-1]
+        if complete:
+            for piece in complete:
+                sentences.append(TimedSentence(piece, start, max(start, now)))
+            consumed = " ".join(complete)
+            pending = pending[len(consumed) :].strip()
+            start = now if pending else None
+    if pending.strip():
+        sentences.append(TimedSentence(pending.strip(), start or last or 0.0, last or start or 0.0))
+    return sentences
+
+
+def normalize_response_items(raw: Any) -> list[dict[str, Any]]:
+    if isinstance(raw, dict):
+        items = raw.get("sentences", raw.get("chunks"))
+    else:
+        items = raw
+    if not isinstance(items, list):
+        raise ValueError("response JSON must be a list or contain sentences/chunks")
+    normalized = []
+    for item in items:
+        if not isinstance(item, dict):
+            continue
+        text = item.get("sentence", item.get("text", item.get("content", "")))
+        if not isinstance(text, str) or not text.strip():
+            continue
+        try:
+            start = float(item.get("start", item.get("current_time", item.get("time", 0.0))) or 0.0)
+            end = float(item.get("end", item.get("current_time", item.get("time", start))) or start)
+        except (TypeError, ValueError):
+            continue
+        normalized.append({"sentence": text.strip(), "start": start, "end": max(start, end)})
+    return normalized
+
+
+def validate_clock(meta: dict[str, Any], *, allow_invalid: bool = False) -> None:
+    if meta.get("clock") == "invalid" and not allow_invalid:
+        raise ValueError("generate artifact has clock=invalid; pass --allow-invalid-clock to score it")

vllm_omni/benchmarks/duplex/omni_duplex_eval_dataset.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Dataset and split normalization for Omni-DuplexEval.
+
+The loader accepts a Hugging Face dataset id, a JSON/JSONL manifest, or an
+already materialized iterable.  Media is deliberately resolved at use time so
+the benchmark remains usable in air-gapped environments.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+DEFAULT_DATASET = "Hothan/Omni-DuplexEval"
+RTD_SPLITS = frozenset(
+    {
+        "RTD_world_knowledge",
+        "RTD_counting",
+        "RTD_fine_grained_movement",
+        "RTD_interaction_relation",
+        "RTD_OCR",
+        "RTD_Omni",
+    }
+)
+PR_SPLITS = frozenset({"PR_correction", "PR_event_reminder", "PR_post_event_reminder"})
+_TASK_ALIASES = {
+    "correction": "correction",
+    "pr_correction": "correction",
+    "event_reminder": "proactive_reminder",
+    "pr_event_reminder": "proactive_reminder",
+    "proactive_reminder": "proactive_reminder",
+    "post_event_reminder": "post_event_reminder",
+    "pr_post_event_reminder": "post_event_reminder",
+}
+
+
+def canonical_task_type(value: str) -> str:
+    key = str(value).strip().lower().replace("-", "_").replace(" ", "_")
+    if key not in _TASK_ALIASES:
+        raise ValueError(f"unsupported Omni-DuplexEval task type: {value!r}")
+    return _TASK_ALIASES[key]
+
+
+def task_type_for_split(split: str) -> str:
+    """Match the task routing used by the official HF batch evaluator."""
+    key = str(split).strip().lower()
+    if "correction" in key:
+        return "correction"
+    if "post_event" in key:
+        return "post_event_reminder"
+    if key.startswith("pr"):
+        return "proactive_reminder"
+    raise ValueError(f"cannot infer proactive-reminder task from split {split!r}")
+
+
+def family_for_split(split: str, task_type: str | None = None) -> str:
+    if split in RTD_SPLITS or str(split).upper().startswith("RTD"):
+        return "rtd"
+    if split in PR_SPLITS or str(split).upper().startswith("PR"):
+        return "pr"
+    if task_type:
+        return "rtd" if str(task_type).lower().startswith("rtd") else "pr"
+    raise ValueError(f"cannot infer benchmark family from split {split!r}")
+
+
+def _value(row: dict[str, Any], *keys: str, default: Any = None) -> Any:
+    for key in keys:
+        if key in row and row[key] is not None:
+            return row[key]
+    return default
+
+
+def _float(value: Any, default: float | None = None) -> float | None:
+    try:
+        return None if value is None else float(value)
+    except (TypeError, ValueError):
+        return default
+
+
+@dataclass(frozen=True)
+class DuplexSample:
+    id: str
+    split: str
+    family: str
+    task_type: str | None
+    video: Any
+    question_audio: Any = None
+    question_text: str = ""
+    answer1: str = ""
+    answer2: str = ""
+    reminder1: Any = None
+    reminder2: Any = None
+    video_duration: float | None = None
+    video_type: str = ""
+    raw: dict[str, Any] | None = None
+
+    @classmethod
+    def from_row(cls, row: dict[str, Any], *, split: str | None = None, media_root: Path | None = None) -> DuplexSample:
+        chosen_split = str(split or _value(row, "split", "subset", "config", default=""))
+        task_value = _value(row, "task_type", "task", "type")
+        family = str(_value(row, "family", default="") or family_for_split(chosen_split, task_value))
+        task = None
+        if family == "pr":
+            task = canonical_task_type(str(task_value)) if task_value else task_type_for_split(chosen_split)
+        sample_id = str(_value(row, "id", "sample_id", "uid", "name", default=""))
+        video = _value(row, "video", "video_path", "video_file")
+        audio = _value(row, "question_audio", "audio", "question_wav")
+        if media_root:
+
+            def resolve(value: Any) -> Any:
+                if not isinstance(value, str):
+                    return value
+                candidate = Path(value).expanduser()
+                return str(candidate if candidate.is_absolute() else media_root / candidate)
+
+            video, audio = resolve(video), resolve(audio)
+        return cls(
+            id=sample_id,
+            split=chosen_split,
+            family=family,
+            task_type=task,
+            video=video,
+            question_audio=audio,
+            question_text=str(_value(row, "question_text", "question", "instruction", default="") or ""),
+            answer1=str(_value(row, "answer1", "answer", "ground_answer", default="") or ""),
+            answer2=str(_value(row, "answer2", default="") or ""),
+            reminder1=_value(row, "reminder1", "reminder_1"),
+            reminder2=_value(row, "reminder2", "reminder_2"),
+            video_duration=_float(_value(row, "video_duration", "duration")),
+            video_type=str(_value(row, "video_type", default="") or ""),
+            raw=dict(row),
+        )
+
+
+def _read_manifest(path: Path) -> list[dict[str, Any]]:
+    if path.suffix.lower() == ".jsonl":
+        return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
+    payload = json.loads(path.read_text(encoding="utf-8"))
+    if isinstance(payload, dict):
+        payload = payload.get("data", payload.get("samples", [payload]))
+    if not isinstance(payload, list):
+        raise ValueError("manifest must contain a JSON list")
+    return payload
+
+
+def load_samples(
+    dataset: str | Path | Iterable[dict[str, Any]] = DEFAULT_DATASET,
+    *,
+    split: str | None = None,
+    family: str = "all",
+    media_root: str | Path | None = None,
+    limit: int | None = None,
+    ids: Iterable[str] | None = None,
+) -> list[DuplexSample]:
+    if isinstance(dataset, (str, Path)) and Path(str(dataset)).exists():
+        rows = _read_manifest(Path(str(dataset)))
+    elif not isinstance(dataset, (str, Path)):
+        rows = list(dataset)
+    else:
+        try:
+            from datasets import load_dataset
+        except ImportError as exc:
+            raise RuntimeError("datasets is required for Hugging Face loading; use a local manifest instead") from exc
+        kwargs: dict[str, Any] = {}
+        if split and split != "all":
+            kwargs["split"] = split
+        loaded = load_dataset(str(dataset), **kwargs)
+        if isinstance(loaded, dict):
+            rows = []
+            for name, table in loaded.items():
+                rows.extend({**dict(row), "split": name} for row in table)
+        else:
+            rows = [dict(row) for row in loaded]
+    wanted = set(str(item) for item in ids) if ids else None
+    root = Path(media_root).expanduser() if media_root else None
+    result = []
+    for row in rows:
+        sample = DuplexSample.from_row(row, split=None if split in (None, "all") else split, media_root=root)
+        if split and split != "all" and sample.split != split:
+            continue
+        if family != "all" and sample.family != family:
+            continue
+        if wanted is not None and sample.id not in wanted:
+            continue
+        result.append(sample)
+    if limit is not None:
+        result = result[: max(0, limit)]
+    return result

vllm_omni/benchmarks/duplex/omni_duplex_eval_eval.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Evaluate and summarize generated Omni-DuplexEval artifacts."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from .omni_duplex_eval_clock import normalize_response_items, validate_clock
+from .omni_duplex_eval_dataset import DuplexSample
+from .omni_duplex_eval_judge import DuplexJudge
+from .omni_duplex_eval_media import extract_jpeg, materialize_media, video_duration
+from .omni_duplex_eval_metrics import (
+    PROTOCOL_PIN,
+    build_content_prompt,
+    build_reminder_prompt,
+    build_temporal_prompt,
+    parse_judge_json,
+    reminder_window,
+    summarize_pr_results,
+    summarize_temporal_results,
+    temporal_window,
+)
+
+
+def _read(path: Path) -> Any:
+    return json.loads(path.read_text(encoding="utf-8"))
+
+
+def _write(path: Path, value: Any) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(json.dumps(value, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def _text(items: list[dict[str, Any]]) -> str:
+    return " ".join(str(item.get("sentence", item.get("text", ""))).strip() for item in items).strip()
+
+
+def evaluate_sample(
+    sample: DuplexSample,
+    response_path: Path,
+    score_path: Path,
+    judge: DuplexJudge,
+    *,
+    judge_fps: int = 2,
+    judge_video_mode: str = "video_url",
+    window_size: float = 10.0,
+    allow_invalid_clock: bool = False,
+) -> dict[str, Any]:
+    raw = _read(response_path)
+    meta = (
+        _read(response_path.with_name(response_path.stem + ".meta.json"))
+        if response_path.with_name(response_path.stem + ".meta.json").exists()
+        else {}
+    )
+    validate_clock(meta, allow_invalid=allow_invalid_clock)
+    items = normalize_response_items(raw)
+    result: dict[str, Any] = {
+        "id": sample.id,
+        "family": sample.family,
+        "task_type": sample.task_type,
+        "protocol_pin": PROTOCOL_PIN,
+        "response_meta": meta,
+        "judge_model": getattr(judge, "model", None),
+        "judge_video_mode": judge_video_mode,
+    }
+    if sample.family == "rtd":
+        video_path = materialize_media(sample.video, score_path.parent / ".media", sample.id, ".mp4")
+        duration = sample.video_duration or video_duration(video_path)
+        temporal_rows = []
+        for item in items:
+            window = temporal_window(item["start"], item["end"], duration)
+            if not window:
+                temporal_rows.append(
+                    {
+                        **item,
+                        "sentence_start": item["start"],
+                        "sentence_end": item["end"],
+                        "sentence_duration": max(0.0, item["end"] - item["start"]),
+                        "error": "No valid temporal window",
+                        "temporal_score": 0,
+                        "is_relevant": 0,
+                    }
+                )
+                continue
+            frames = _extract_frames(video_path, _times(*window, fps=judge_fps))
+            parsed = parse_judge_json(
+                judge.temporal(build_temporal_prompt(*window, item["sentence"], sample.question_text), frames)
+            )
+            temporal_rows.append(
+                {
+                    **item,
+                    "sentence_start": item["start"],
+                    "sentence_end": item["end"],
+                    "sentence_duration": max(0.0, item["end"] - item["start"]),
+                    "window_start": window[0],
+                    "window_end": window[1],
+                    "error": None,
+                    **parsed,
+                }
+            )
+        content_frames = None
+        if judge_video_mode == "frame-sample":
+            content_frames = _extract_frames(video_path, _content_frame_times(duration))
+        content = parse_judge_json(
+            judge.content(
+                build_content_prompt(_text(items), sample.question_text, [sample.answer1, sample.answer2]),
+                video_path,
+                frames=content_frames,
+                mode=judge_video_mode,
+            )
+        )
+        if content_frames is not None:
+            content["frame_count"] = len(content_frames)
+        result.update(
+            {
+                "temporal": {"sentences": temporal_rows, "summary": summarize_temporal_results(temporal_rows)},
+                "content": content,
+            }
+        )
+    else:
+        if sample.task_type is None:
+            raise ValueError(f"missing proactive-reminder task type for split {sample.split!r}")
+        task = sample.task_type
+        times = _reminder_times(sample)
+        event_rows = []
+        if task == "correction" or not times:
+            response = _text(items)
+            parsed = parse_judge_json(
+                judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
+            )
+            event_rows.append({"response_segment": response, **parsed})
+        else:
+            for start in times:
+                lo, hi = reminder_window(start, window_size)
+                response = " ".join(item["sentence"] for item in items if lo <= item["start"] <= hi)
+                if not response:
+                    event_rows.append({"start_time": start, "response_segment": "", "success_score": 0})
+                else:
+                    parsed = parse_judge_json(
+                        judge.chat(build_reminder_prompt(sample.question_text, response, task, sample.answer1))
+                    )
+                    event_rows.append({"start_time": start, "response_segment": response, **parsed})
+        scores = [int(row.get("success_score", 0)) for row in event_rows]
+        result["pr"] = {
+            "events": event_rows,
+            "all_success": bool(scores) and all(score == 1 for score in scores),
+            "total_score": int(bool(scores) and all(score == 1 for score in scores)),
+        }
+    _write(score_path, result)
+    return result
+
+
+def _times(start: float, end: float, *, fps: int) -> list[float]:
+    count = max(1, int((end - start) * fps))
+    return (
+        [start + (end - start) / 2] if count == 1 else [start + i * (end - start) / (count - 1) for i in range(count)]
+    )
+
+
+def _extract_frames(path: str | Path, timestamps: list[float]) -> list[bytes]:
+    """Extract only decodable JPEG frames; clips can end on non-keyframes."""
+    frames = []
+    for timestamp in timestamps:
+        try:
+            frame = extract_jpeg(path, timestamp=timestamp)
+        except (OSError, subprocess.CalledProcessError, ValueError):
+            continue
+        if frame.startswith(b"\xff\xd8"):
+            frames.append(frame)
+    return frames
+
+
+def _content_frame_times(duration: float) -> list[float]:
+    """Sample the full video at approximately one frame every three seconds."""
+    if duration <= 0:
+        return []
+    times = []
+    timestamp = 0.0
+    while timestamp < duration:
+        times.append(min(timestamp, max(0.0, duration - 0.01)))
+        timestamp += 3.0
+    # Qwen2.5-VL has a finite per-request image-token budget. Preserve
+    # coverage while bounding long clips to a safely portable frame count.
+    if len(times) > 16:
+        stride = (len(times) - 1) / 15
+        times = [times[round(i * stride)] for i in range(16)]
+    return times
+
+
+def _reminder_times(sample: DuplexSample) -> list[float]:
+    values = []
+    for value in (sample.reminder1, sample.reminder2):
+        if isinstance(value, dict):
+            value = value.get("start", value.get("time"))
+        try:
+            if value is not None:
+                values.append(float(value))
+        except (TypeError, ValueError):
+            pass
+    return values
+
+
+def summarize_scores(score_root: str | Path) -> dict[str, Any]:
+    rows = []
+    for path in sorted(Path(score_root).rglob("*.json")):
+        if not path.name.endswith("_summary.json"):
+            row = _read(path)
+            row.setdefault("split", path.parent.name)
+            rows.append(row)
+    rtd_temporal = [
+        row["temporal"]["summary"] | {"content_score": row.get("content", {}).get("content_score", 0.0)}
+        for row in rows
+        if "temporal" in row
+    ]
+    pr = [{"task_type": row.get("task_type"), **row.get("pr", {})} for row in rows if "pr" in row]
+    result = {"protocol_pin": PROTOCOL_PIN, "samples": len(rows)}
+    if rtd_temporal:
+        by_task: dict[str, list[dict[str, float]]] = {}
+        for row in rows:
+            if "temporal" in row:
+                by_task.setdefault(str(row.get("split", row.get("task_type", "unknown"))), []).append(
+                    {
+                        "content": float(row.get("content", {}).get("content_score", 0.0)),
+                        "temporal": float(row["temporal"]["summary"].get("avg_temporal_score", 0.0)),
+                    }
+                )
+        result["rtd"] = {
+            "mean_content_score": sum(item["content_score"] for item in rtd_temporal) / len(rtd_temporal),
+            "mean_avg_temporal_score": sum(item["avg_temporal_score"] for item in rtd_temporal) / len(rtd_temporal),
+            "by_task": {
+                task: {
+                    "mean_content_score": sum(item["content"] for item in values) / len(values),
+                    "mean_avg_temporal_score": sum(item["temporal"] for item in values) / len(values),
+                }
+                for task, values in by_task.items()
+            },
+        }
+    if pr:
+        result["pr"] = summarize_pr_results(pr)
+    return result

vllm_omni/benchmarks/duplex/omni_duplex_eval_judge.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OpenAI-compatible local judge client."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pybase64 as base64
+import requests
+
+
+def _build_openai_url(base_url: str, api_path: str) -> str:
+    base = base_url.rstrip("/")
+    normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
+    if base.endswith(normalized_path):
+        return base
+    if base.endswith("/v1"):
+        return f"{base}{normalized_path}"
+    return f"{base}/v1{normalized_path}"
+
+
+class DuplexJudge:
+    def __init__(self, base_url: str, model: str, *, api_key: str = "EMPTY", timeout: int = 600) -> None:
+        self.base_url, self.model, self.api_key, self.timeout = base_url, model, api_key, timeout
+
+    def chat(self, content: Any, *, system: str | None = None, max_tokens: int = 1200) -> str:
+        messages = ([{"role": "system", "content": system}] if system else []) + [{"role": "user", "content": content}]
+        response = requests.post(
+            _build_openai_url(self.base_url, "/chat/completions"),
+            json={
+                "model": self.model,
+                "messages": messages,
+                "temperature": 0.1,
+                "max_tokens": max_tokens,
+            },
+            headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
+            timeout=self.timeout,
+        )
+        if not response.ok:
+            raise requests.HTTPError(
+                f"{response.status_code} response from judge: {response.text[:500]}", response=response
+            )
+        return str(response.json()["choices"][0]["message"]["content"])
+
+    def temporal(self, prompt: str, frames: list[bytes]) -> str:
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        content.extend(
+            {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()}}
+            for frame in frames
+        )
+        return self.chat(content, system="You are a careful evaluator for real-time multimodal systems.")
+
+    def content(
+        self,
+        prompt: str,
+        video: str | Path | None = None,
+        frames: list[bytes] | None = None,
+        *,
+        mode: str = "video_url",
+    ) -> str:
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        if mode == "video_url" and video is not None:
+            content.append({"type": "video_url", "video_url": {"url": "file://" + str(Path(video).resolve())}})
+        else:
+            content.extend(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": "data:image/jpeg;base64," + base64.b64encode(frame).decode()},
+                }
+                for frame in (frames or [])
+            )
+        return self.chat(content)

vllm_omni/benchmarks/duplex/omni_duplex_eval_media.py

@@ -0,0 +1,113 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Small ffmpeg-backed media helpers used by the duplex runner and judge."""
+
+from __future__ import annotations
+
+import subprocess
+import wave
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+from vllm_omni.experimental.fullduplex.client import PCM16_BYTES_PER_SAMPLE, PCM16_SAMPLE_RATE
+
+
+def _run_ffmpeg(args: list[str]) -> bytes:
+    result = subprocess.run(["ffmpeg", "-hide_banner", "-loglevel", "error", *args], capture_output=True, check=True)
+    return result.stdout
+
+
+def materialize_media(value: Any, output_dir: str | Path, stem: str, suffix: str) -> Path:
+    """Resolve common Hugging Face media values to a local file."""
+    if isinstance(value, str | Path):
+        return Path(value)
+    path = value.get("path") if isinstance(value, dict) else getattr(value, "path", None)
+    if path:
+        return Path(path)
+    payload = value.get("bytes") if isinstance(value, dict) else None
+    if payload is None and isinstance(value, bytes | bytearray):
+        payload = value
+    if payload is None:
+        raise ValueError(f"cannot materialize media value for {stem!r}")
+    destination = Path(output_dir) / f"{stem}{suffix}"
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    destination.write_bytes(bytes(payload))
+    return destination
+
+
+def video_duration(path: str | Path) -> float:
+    result = subprocess.run(
+        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(path)],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    return max(0.0, float(result.stdout.strip()))
+
+
+def extract_jpeg(path: str | Path, *, timestamp: float, quality: int = 3) -> bytes:
+    return _run_ffmpeg(
+        [
+            "-ss",
+            f"{max(0.0, timestamp):.3f}",
+            "-i",
+            str(path),
+            "-frames:v",
+            "1",
+            "-q:v",
+            str(quality),
+            "-f",
+            "image2",
+            "pipe:1",
+        ]
+    )
+
+
+def iter_jpegs(
+    path: str | Path, *, fps: float = 1.0, duration: float | None = None, quality: int = 3
+) -> Iterator[tuple[float, bytes]]:
+    end = duration if duration is not None else video_duration(path)
+    step = 1.0 / fps
+    timestamp = 0.0
+    while timestamp < end:
+        try:
+            frame = extract_jpeg(path, timestamp=timestamp, quality=quality)
+        except (OSError, subprocess.CalledProcessError):
+            break
+        if frame:
+            yield timestamp, frame
+        timestamp += step
+
+
+def read_audio_pcm16(path: str | Path) -> bytes:
+    source = Path(path)
+    if source.suffix.lower() == ".wav":
+        try:
+            with wave.open(str(source), "rb") as wav_file:
+                if (
+                    wav_file.getnchannels() == 1
+                    and wav_file.getsampwidth() == 2
+                    and wav_file.getframerate() == PCM16_SAMPLE_RATE
+                    and wav_file.getcomptype() == "NONE"
+                ):
+                    return wav_file.readframes(wav_file.getnframes())
+        except wave.Error:
+            pass
+    return _run_ffmpeg(["-i", str(source), "-f", "s16le", "-ac", "1", "-ar", str(PCM16_SAMPLE_RATE), "pipe:1"])
+
+
+def iter_av_units(
+    audio_pcm16: bytes, frames: Iterator[tuple[float, bytes]], *, unit_ms: int = 1000
+) -> Iterator[tuple[bytes, bytes | None]]:
+    unit_bytes = PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * unit_ms // 1000
+    frame_iter = iter(frames)
+    next_frame = next(frame_iter, None)
+    for offset in range(0, len(audio_pcm16), unit_bytes):
+        timestamp = offset / (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
+        frame = None
+        if next_frame is not None and next_frame[0] <= timestamp + unit_ms / 1000:
+            frame = next_frame[1]
+            next_frame = next(frame_iter, None)
+        yield audio_pcm16[offset : offset + unit_bytes], frame

vllm_omni/benchmarks/duplex/omni_duplex_eval_metrics.py

@@ -0,0 +1,260 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Protocol prompts, windows, parsing, and aggregation.
+
+Prompt wording and window math are copied verbatim from the MIT-licensed
+OpenBMB/Omni-DuplexEval evaluator at :data:`PROTOCOL_PIN`.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable
+from typing import Any
+
+import regex as re
+
+# Copied from OpenBMB/Omni-DuplexEval@PROTOCOL_PIN (MIT).
+# Do not change prompt wording without bumping the pin and recording a score delta.
+PROTOCOL_PIN = "ca3c122b4d4bf67afd6b18ea5e724b4561bdde48"
+
+
+def temporal_window(sentence_start: float, sentence_end: float, video_duration: float) -> tuple[float, float] | None:
+    start = max(0.0, float(sentence_start) - 2.0)
+    end = max(start + 0.5, float(sentence_end) - 2.0)
+    return (start, end) if end <= float(video_duration) and end - start >= 0.5 else None
+
+
+def reminder_window(start_time: float, window_size: float = 10.0) -> tuple[float, float]:
+    return float(start_time), float(start_time) + float(window_size)
+
+
+def parse_judge_json(text: str) -> dict[str, Any]:
+    for match in re.findall(r"\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}", str(text), re.S):
+        try:
+            value = json.loads(re.sub(r",\s*([}\]])", r"\1", match))
+        except (TypeError, ValueError, json.JSONDecodeError):
+            continue
+        if isinstance(value, dict):
+            break
+    else:
+        value = {}
+        patterns = [
+            (r'"temporal_score"\s*:\s*(\d+)', "temporal_score", int),
+            (r'"is_relevant"\s*:\s*(\d+)', "is_relevant", int),
+            (r'"success_score"\s*:\s*([01])', "success_score", int),
+            (r'"content_score"\s*:\s*(\d+(?:\.\d+)?)', "content_score", float),
+        ]
+        for pattern, key, caster in patterns:
+            if match := re.search(pattern, str(text)):
+                parsed = caster(match.group(1))
+                value[key] = round(parsed, 2) if caster is float else parsed
+        for pattern, key in [
+            (r'"temporal_reasoning"\s*:\s*"([^"]*)"', "temporal_reasoning"),
+            (r'"content_reasoning"\s*:\s*"([^"]*)"', "content_reasoning"),
+            (r'"reasoning"\s*:\s*"([^"]*)"', "reasoning"),
+        ]:
+            if match := re.search(pattern, str(text), re.S):
+                value[key] = match.group(1).strip()
+                break
+    for key in ("temporal_score", "is_relevant", "success_score"):
+        if key in value:
+            try:
+                value[key] = int(float(value[key]))
+            except (TypeError, ValueError):
+                value[key] = 0
+    if "content_score" in value:
+        try:
+            value["content_score"] = round(float(value["content_score"]), 2)
+        except (TypeError, ValueError):
+            value["content_score"] = 0.0
+    return value
+
+
+def summarize_temporal_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]:
+    rows = list(results)
+    valid = [row for row in rows if not row.get("error")]
+    relevant = [row for row in valid if int(row.get("is_relevant", 0)) == 1]
+    irrelevant = [row for row in valid if int(row.get("is_relevant", 0)) == 0]
+    scores = [int(row.get("temporal_score", 0)) for row in relevant]
+    total_valid_duration = sum(float(row.get("sentence_duration", 0.0)) for row in valid)
+    irrelevant_duration = sum(float(row.get("sentence_duration", 0.0)) for row in irrelevant)
+    return {
+        "avg_temporal_score": round(sum(scores) / len(scores), 4) if scores else 0.0,
+        "total_sentences": len(rows),
+        "evaluated_sentences": len(valid),
+        "error_count": len(rows) - len(valid),
+        "relevant_sentences_count": len(relevant),
+        "irrelevant_sentences_count": len(irrelevant),
+        "score_distribution": {
+            "3_points": scores.count(3),
+            "2_points": scores.count(2),
+            "1_point": scores.count(1),
+            "0_points": scores.count(0),
+        },
+        "relevance_stats": {
+            "first_relevant_time": min((row.get("sentence_start") for row in relevant), default=None),
+            "total_relevant_duration": round(sum(float(row.get("sentence_duration", 0.0)) for row in relevant), 4),
+            "total_irrelevant_duration": round(irrelevant_duration, 4),
+            "irrelevant_duration_ratio": round(
+                irrelevant_duration / total_valid_duration if total_valid_duration > 0 else 0.0, 4
+            ),
+        },
+        "excluded_irrelevant_sentences": [
+            {
+                "sentence": row.get("sentence", ""),
+                "start": row.get("sentence_start"),
+                "end": row.get("sentence_end"),
+                "temporal_score": row.get("temporal_score", 0),
+            }
+            for row in irrelevant
+        ],
+    }
+
+
+def summarize_pr_results(results: Iterable[dict[str, Any]]) -> dict[str, Any]:
+    rows = list(results)
+    successes = [int(row.get("all_success", row.get("total_score", 0))) for row in rows]
+    by_task: dict[str, list[int]] = {}
+    for row, score in zip(rows, successes, strict=False):
+        by_task.setdefault(str(row.get("task_type", "unknown")), []).append(score)
+    return {
+        "samples": len(rows),
+        "mean_all_success": sum(successes) / len(successes) if successes else 0.0,
+        "by_task": {task: sum(values) / len(values) for task, values in by_task.items()},
+    }
+
+
+def build_temporal_prompt(start_time: float, end_time: float, response_text: str, question: str) -> str:
+    return f"""
+You are evaluating a real-time video description system.
+
+Basic information:
+- Video segment: {start_time:.2f}s to {end_time:.2f}s.
+- User instruction: {question}
+- Model response sentence: "{response_text}"
+
+Evaluation principles:
+- Judge only the current video segment.
+- Ignore response content that clearly belongs to earlier or later segments unless it contradicts the current segment.
+- Language does not matter; evaluate semantic correctness and temporal alignment only.
+- Decide whether the sentence is a substantive response or only a filler/polite phrase.
+
+Temporal sensitivity score:
+- 3: Excellent temporal alignment. The sentence accurately describes the current segment and follows the instruction.
+- 2: Mostly aligned. Minor inaccuracies, omissions, or harmless references to nearby segments are allowed.
+- 1: Poor alignment. Major inaccuracies, wrong timing, or mostly irrelevant content.
+- 0: No meaningful alignment with the current segment.
+
+Relevance label:
+- 1: The sentence contains substantive task- or video-related information.
+- 0: The sentence is only filler, acknowledgement, meta-commentary, or generic text without useful content.
+
+Return exactly this JSON object:
+{{
+  "temporal_score": <0, 1, 2, or 3>,
+  "temporal_reasoning": "<brief explanation>",
+  "is_relevant": <0 or 1>
+}}
+"""
+
+
+def build_content_prompt(response_text: str, question: str, references: list[str] | None = None) -> str:
+    reference_block = ""
+    if references:
+        reference_block = "Reference annotations, for reference only:\n"
+        for index, reference in enumerate(references, 1):
+            reference_block += f"{index}. {reference}\n"
+
+    return f"""
+You are a precise content-accuracy evaluator for video description.
+
+Basic information:
+- User instruction: {question}
+- Model response: "{response_text}"
+{reference_block}
+
+Evaluation goal:
+Judge whether the model response is factually accurate for the whole video and aligned with the instruction.
+Consider object/action/color/count/spatial/event errors, hallucinations, omissions, and irrelevant content.
+Reference annotations are optional guidance; the primary evidence is the video itself.
+
+Scoring:
+- Use a decimal score from 0.00 to 3.00.
+- Start from 3.00 and deduct for each error.
+- Use 0.00 only when the response is empty, completely irrelevant, or contains no correct video facts.
+- Output exactly two decimal places.
+
+Return exactly this JSON object:
+{{
+  "content_score": <decimal from 0.00 to 3.00>,
+  "content_reasoning": "<brief explanation with main errors and final score>"
+}}
+"""
+
+
+def build_reminder_prompt(instruction: str, response: str, task_type: str, ground_answer: str = "") -> str:
+    if task_type == "correction":
+        return f"""
+You are judging whether a model successfully completed a correction task.
+
+Task description:
+The user instruction contains an incorrect statement about the video.
+The model should identify the incorrect part and provide the correct information.
+
+Input:
+- User instruction: {instruction}
+- Reference correction: {ground_answer}
+- Model response: "{response}"
+
+Success criteria:
+1. Identify the error implied by the user instruction and the reference correction.
+2. Check whether the model corrected all required error points.
+3. The corrected content must preserve the correct context, including subject, object, action, and attributes.
+4. Ignore extra information that is unrelated to both the instruction and the reference, unless it contradicts them.
+
+Scoring:
+- 1 = all required error points are corrected and the context is consistent.
+- 0 = at least one required correction is missing, wrong, or contextually inconsistent.
+
+Return exactly this JSON object:
+{{
+  "success_score": <0 or 1>,
+  "reasoning": "<brief explanation covering the required correction points>"
+}}
+"""
+
+    if task_type == "post_event_reminder":
+        timing_note = "The response is evaluated immediately after the target event has occurred."
+    else:
+        timing_note = "The response is evaluated at the target event time or immediately after it."
+
+    return f"""
+You are judging whether a model successfully completed an event-reminder task.
+
+Task description:
+The user gave an instruction asking the system to remind them when a specific event happens.
+{timing_note}
+The provided model text is the response segment generated in the evaluation window.
+
+Input:
+- User instruction: {instruction}
+- Model response segment: "{response}"
+
+Success criteria:
+1. The response clearly refers to the target event in the instruction.
+2. The response communicates a reminder, notification, or confirmation that the event has happened.
+3. Vague narration, unrelated descriptions, or wrong-event reminders are failures.
+4. The response does not need to match the instruction wording exactly if the meaning is clear.
+
+Scoring:
+- 1 = successful reminder.
+- 0 = unsuccessful reminder.
+
+Return exactly this JSON object:
+{{
+  "success_score": <0 or 1>,
+  "reasoning": "<brief explanation>"
+}}
+"""

vllm_omni/benchmarks/duplex/omni_duplex_eval_runner.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Single-sample MiniCPM-o native-duplex generation."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+import pybase64 as base64
+
+from vllm_omni.experimental.fullduplex.client import RealtimeDuplexClient, build_realtime_url, wait_for
+
+from .omni_duplex_eval_clock import extract_timed_sentences
+from .omni_duplex_eval_media import iter_av_units, iter_jpegs, materialize_media, read_audio_pcm16, video_duration
+
+
+def _ref_audio(path: str | Path) -> str:
+    value = Path(path).expanduser().read_bytes()
+    return "data:audio/wav;base64," + base64.b64encode(value).decode("ascii")
+
+
+async def generate_sample(
+    sample: Any,
+    *,
+    url: str,
+    model: str,
+    ref_audio: str | Path,
+    output_root: str | Path,
+    fps: float = 1.0,
+    mix: str = "question",
+    pace: str = "realtime",
+    clock: str = "media",
+    overwrite: bool = False,
+    unit_ms: int = 1000,
+) -> Path:
+    output = Path(output_root) / sample.split / f"{sample.id}.json"
+    meta_path = output.with_name(output.stem + ".meta.json")
+    if output.exists() and not overwrite:
+        return output
+    if mix != "question":
+        raise NotImplementedError("v1 supports mix=question; soundtrack mixing is reserved for P1")
+    media_dir = output.parent / ".media"
+    audio_path = materialize_media(sample.question_audio, media_dir, f"{sample.id}_question", ".wav")
+    video_path = materialize_media(sample.video, media_dir, sample.id, ".mp4")
+    pcm = read_audio_pcm16(audio_path)
+    duration = sample.video_duration or video_duration(video_path)
+    frames = iter_jpegs(video_path, fps=fps, duration=duration)
+    realtime = pace == "realtime"
+    if pace not in {"realtime", "as-fast-as-possible"}:
+        raise ValueError("pace must be realtime or as-fast-as-possible")
+    client = RealtimeDuplexClient(build_realtime_url(url, model, autostart=False))
+    response_done = False
+    drain_timeout = None
+    close_timeout = None
+    async with client:
+        await client.configure(model, ref_audio=_ref_audio(ref_audio), instructions="Streaming Omni Conversation.")
+        ack_task = asyncio.create_task(_ack_playback(client))
+        try:
+            await client.stream_av_units(iter_av_units(pcm, frames, unit_ms=unit_ms), realtime=realtime)
+            await client.commit()
+            try:
+                await wait_for(
+                    lambda: client.events.count("response.done") > 0,
+                    timeout_s=max(20.0, duration + 20.0),
+                    label="response.done",
+                )
+                response_done = True
+            except TimeoutError as exc:
+                drain_timeout = str(exc)
+        finally:
+            ack_task.cancel()
+            try:
+                await ack_task
+            except asyncio.CancelledError:
+                pass
+        await client.acknowledge_playback()
+        try:
+            await client.close_session(timeout_s=20.0)
+        except TimeoutError as exc:
+            close_timeout = str(exc)
+        events = list(client.events.events)
+    timed = [sentence.as_dict() for sentence in extract_timed_sentences(events, clock=clock)]
+    output.parent.mkdir(parents=True, exist_ok=True)
+    output.write_text(json.dumps(timed, ensure_ascii=False, indent=2), encoding="utf-8")
+    meta = {
+        "id": sample.id,
+        "split": sample.split,
+        "clock": clock if realtime else "invalid",
+        "pace": pace,
+        "mix": mix,
+        "fps": fps,
+        "unit_ms": unit_ms,
+        "model": model,
+        "response_done": response_done,
+        "drain_timeout": drain_timeout,
+        "close_timeout": close_timeout,
+        "ref_audio_sha256": hashlib.sha256(Path(ref_audio).read_bytes()).hexdigest(),
+    }
+    meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+    return output
+
+
+async def _ack_playback(client: RealtimeDuplexClient) -> None:
+    while True:
+        await asyncio.sleep(0.25)
+        await client.acknowledge_playback()

vllm_omni/config/pipeline_registry.py

@@ -69,6 +69,10 @@ from vllm_omni.model_executor.models.indextts2.pipeline import (
     INDEXTTS25_PIPELINE,
 )
 from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE
+from vllm_omni.model_executor.models.mammoth_moda2.pipeline import (
+    MAMMOTH_MODA2_AR_PIPELINE,
+    MAMMOTH_MODA2_PIPELINE,
+)
 from vllm_omni.model_executor.models.mimo_audio.pipeline import MIMO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.ming_flash_omni.pipeline import (
     MING_FLASH_OMNI_IMAGE_PIPELINE,
@@ -97,7 +101,10 @@ from vllm_omni.model_executor.models.qwen2_5_omni.pipeline import (
     QWEN2_5_OMNI_PIPELINE,
     QWEN2_5_OMNI_THINKER_ONLY_PIPELINE,
 )
-from vllm_omni.model_executor.models.qwen3_omni.pipeline import resolve_qwen3_omni_pipeline
+from vllm_omni.model_executor.models.qwen3_omni.pipeline import (
+    QWEN3_OMNI_THINKER_ONLY_PIPELINE,
+    resolve_qwen3_omni_pipeline,
+)
 from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
 from vllm_omni.model_executor.models.step_audio2.pipeline import (
     STEP_AUDIO2_ASR_PIPELINE,
@@ -123,6 +130,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     # has no model_type key).
     "nemotron_labs_voicechat": NEMOTRON_VOICECHAT_PIPELINE,
     "qwen3_omni_moe": resolve_qwen3_omni_pipeline,
+    "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE,
     "qwen3_tts": QWEN3_TTS_PIPELINE,
     "step_audio_2": STEP_AUDIO2_PIPELINE,
     "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE,
@@ -164,6 +172,8 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "ming_flash_omni_image": MING_FLASH_OMNI_IMAGE_PIPELINE,
     "moss_tts_nano": MOSS_TTS_NANO_PIPELINE,
     "omnivoice": OMNIVOICE_PIPELINE,
+    "mammoth_moda2": MAMMOTH_MODA2_PIPELINE,
+    "mammoth_moda2_ar": MAMMOTH_MODA2_AR_PIPELINE,
     "moss_tts_delay": MOSS_TTS_PIPELINE,
     "moss_tts_realtime": MOSS_TTS_REALTIME_PIPELINE,
     "moss_tts_local": MOSS_TTS_LOCAL_PIPELINE,

vllm_omni/deploy/mammoth_moda2.yaml

@@ -0,0 +1,27 @@
+# MammothModa2 AR→DiT (text-to-image).
+
+async_chunk: false
+pipeline: mammoth_moda2
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 100
+    max_model_len: 8192
+    gpu_memory_utilization: 0.5
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+
+  - stage_id: 1
+    devices: "0"
+    max_num_seqs: 1
+    gpu_memory_utilization: 0.3
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    default_sampling_params:
+      extra_args:
+        text_guidance_scale: 9.0
+        cfg_range: [0.0, 1.0]
+        num_inference_steps: 50

vllm_omni/deploy/mammoth_moda2_ar.yaml

@@ -0,0 +1,14 @@
+# MammothModa2 AR-only (image/video understanding and summarization).
+
+async_chunk: false
+pipeline: mammoth_moda2_ar
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 16
+    max_model_len: 8192
+    gpu_memory_utilization: 0.5
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false

vllm_omni/deploy/qwen3_omni_moe_thinking.yaml

@@ -1,6 +1,8 @@
-# Deploy config for Qwen3-Omni-MoE thinker-only checkpoints.
+# Deploy config for Qwen3-Omni-MoE thinker-only (text output, no talker/code2wav).
+# Works for Captioner/Thinking checkpoints and for Instruct weights when this
+# YAML is passed via --deploy-config (pipeline key forces the single-stage topology).
 # Verified on 2x H100-80G GPUs.
-pipeline: qwen3_omni_moe
+pipeline: qwen3_omni_moe_thinker_only
 distributed_executor_backend: mp
 enable_prefix_caching: true
 

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -0,0 +1,839 @@
+import math
+
+import torch
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention_processor import Attention
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange, repeat
+from torch import nn
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+
+from .rope_real import RotaryPosEmbedReal
+
+try:
+    from transformers.modeling_flash_attention_utils import (  # type: ignore
+        flash_attn_varlen_func,  # pyright: ignore[reportAttributeAccessIssue]
+        is_flash_attn_available,
+    )
+except Exception:  # pragma: no cover - best-effort compatibility
+    flash_attn_varlen_func = None  # type: ignore[assignment]
+
+    def is_flash_attn_available() -> bool:  # type: ignore[override]
+        return False
+
+
+from .rope_real import apply_real_rotary_emb
+
+_HAS_FLASH_ATTN_VARLEN = bool(is_flash_attn_available()) and flash_attn_varlen_func is not None
+
+
+class LuminaRMSNormZero(nn.Module):
+    """
+    Norm layer adaptive RMS normalization zero.
+
+    Parameters:
+        embedding_dim (`int`): The size of each embedding vector.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        norm_eps: float,
+        norm_elementwise_affine: bool,
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(
+            min(embedding_dim, 1024),
+            4 * embedding_dim,
+            bias=True,
+        )
+
+        self.norm = Qwen2RMSNorm(embedding_dim, eps=norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        emb: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        emb = self.linear(self.silu(emb))
+        scale_msa, gate_msa, scale_mlp, gate_mlp = emb.chunk(4, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None])
+        return x, gate_msa, scale_mlp, gate_mlp
+
+
+class LuminaFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inner_dim: int,
+        multiple_of: int | None = 256,
+        ffn_dim_multiplier: float | None = None,
+    ):
+        super().__init__()
+
+        # custom hidden_size factor multiplier
+        if ffn_dim_multiplier is not None:
+            inner_dim = int(ffn_dim_multiplier * inner_dim)
+        inner_dim = multiple_of * ((inner_dim + multiple_of - 1) // multiple_of)
+
+        self.linear_1 = nn.Linear(
+            dim,
+            inner_dim,
+            bias=False,
+        )
+        self.linear_2 = nn.Linear(
+            inner_dim,
+            dim,
+            bias=False,
+        )
+        self.linear_3 = nn.Linear(
+            dim,
+            inner_dim,
+            bias=False,
+        )
+
+    def swiglu(self, x, y):
+        return F.silu(x.float(), inplace=False).to(x.dtype) * y
+
+    def forward(self, x):
+        h1, h2 = self.linear_1(x), self.linear_3(x)
+        return self.linear_2(self.swiglu(h1, h2))
+
+
+class LuminaLayerNormContinuous(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        conditioning_embedding_dim: int,
+        # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters
+        # because the output is immediately scaled and shifted by the projected conditioning embeddings.
+        # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters.
+        # However, this is how it was implemented in the original code, and it's rather likely you should
+        # set `elementwise_affine` to False.
+        elementwise_affine=True,
+        eps=1e-5,
+        bias=True,
+        norm_type="layer_norm",
+        out_dim: int | None = None,
+    ):
+        super().__init__()
+
+        self.silu = nn.SiLU()
+        self.linear_1 = nn.Linear(conditioning_embedding_dim, embedding_dim, bias=bias)
+
+        if norm_type == "layer_norm":
+            self.norm = nn.LayerNorm(embedding_dim, eps, elementwise_affine, bias)
+        elif norm_type == "rms_norm":
+            self.norm = Qwen2RMSNorm(embedding_dim, eps=eps)
+        else:
+            raise ValueError(f"unknown norm_type {norm_type}")
+
+        self.linear_2 = None
+        if out_dim is not None:
+            self.linear_2 = nn.Linear(embedding_dim, out_dim, bias=bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        conditioning_embedding: torch.Tensor,
+    ) -> torch.Tensor:
+        scale = self.linear_1(self.silu(conditioning_embedding).to(x.dtype))
+        x = self.norm(x) * (1 + scale)[:, None, :]
+
+        if self.linear_2 is not None:
+            x = self.linear_2(x)
+
+        return x
+
+
+class Lumina2CombinedTimestepCaptionEmbedding(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int = 2520,
+        text_feat_dim: int = 3584,
+        frequency_embedding_size: int = 256,
+        norm_eps: float = 1e-5,
+        timestep_scale: float = 1.0,
+    ) -> None:
+        super().__init__()
+
+        self.time_proj = Timesteps(
+            num_channels=frequency_embedding_size, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=timestep_scale
+        )
+
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=frequency_embedding_size, time_embed_dim=min(hidden_size, 1024)
+        )
+
+        self.caption_embedder = nn.Sequential(
+            Qwen2RMSNorm(text_feat_dim, eps=norm_eps),
+            nn.Linear(text_feat_dim, hidden_size, bias=True),
+        )
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        dtype: torch.dtype,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        timestep_proj = self.time_proj(timestep).to(dtype=dtype)
+        time_embed = self.timestep_embedder(timestep_proj)
+        caption_embed = self.caption_embedder(text_hidden_states)
+        return time_embed, caption_embed
+
+
+class SimpleQFormerImageRefiner(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        output_hidden_size: int | None = None,
+        num_queries: int = 128,
+        num_layers: int = 2,
+        num_heads: int | None = None,
+        dropout: float = 0.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+        input_hidden_size = hidden_size
+        hidden_size = output_hidden_size or hidden_size
+        self.hidden_size = hidden_size
+        self.num_queries = num_queries
+        # ensure num_heads divides hidden_size
+        if num_heads is None:
+            num_heads = max(1, hidden_size // 128)
+        self.num_heads = self._choose_valid_num_heads(hidden_size, num_heads)
+        self.input_proj = nn.Sequential(
+            Qwen2RMSNorm(input_hidden_size, eps=norm_eps),
+            nn.Linear(input_hidden_size, hidden_size, bias=True),
+        )
+
+        # Learnable query embeddings
+        scale = hidden_size**-0.5
+        self.query = nn.Parameter(scale * torch.randn(1, num_queries, hidden_size))
+
+        # Decoder layers
+        self.layers = nn.ModuleList()
+        for _ in range(num_layers):
+            self.layers.append(
+                nn.ModuleDict(
+                    dict(
+                        ln_q1=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        self_attn=nn.MultiheadAttention(
+                            embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True
+                        ),
+                        ln_q2=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        cross_attn=nn.MultiheadAttention(
+                            embed_dim=hidden_size, num_heads=self.num_heads, dropout=dropout, batch_first=True
+                        ),
+                        ln_ffn=Qwen2RMSNorm(hidden_size, eps=norm_eps),
+                        ffn=LuminaFeedForward(dim=hidden_size, inner_dim=4 * hidden_size),
+                    )
+                )
+            )
+
+    @staticmethod
+    def _choose_valid_num_heads(hidden_size: int, proposed_heads: int, preferred_head_dim: int = 128) -> int:
+        """Pick a number of heads that divides hidden_size, close to proposed or preferred."""
+        # If proposed is valid, use it
+        if proposed_heads > 0 and hidden_size % proposed_heads == 0:
+            return proposed_heads
+        # target based on preferred head dim
+        target = max(1, round(hidden_size / preferred_head_dim))
+        # collect divisors up to 128 heads (more than enough)
+        max_heads_cap = min(128, hidden_size)
+        divisors = [d for d in range(1, max_heads_cap + 1) if hidden_size % d == 0]
+        # choose closest to target
+        best = min(divisors, key=lambda d: (abs(d - target), -d))
+        return best
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x: Tensor of shape (batch, seq_len, input_dim)
+        Returns:
+            Tensor of shape (batch, num_queries, hidden_size)
+        """
+        batch, _, _ = x.shape
+        kv = self.input_proj(x)
+        q = self.query.repeat(batch, 1, 1).to(kv.dtype)
+
+        for layer in self.layers:
+            # Self-attention on queries
+            q_norm = layer["ln_q1"](q)
+            attn_out, _ = layer["self_attn"](q_norm, q_norm, q_norm, need_weights=False)
+            q = q + attn_out
+
+            # Cross-attention: queries attend to inputs
+            q_norm = layer["ln_q2"](q)
+            cross_out, _ = layer["cross_attn"](q_norm, kv, kv, need_weights=False, key_padding_mask=attention_mask)
+            q = q + cross_out
+
+            # Feed-forward
+            q = q + layer["ffn"](layer["ln_ffn"](q))
+
+        return q
+
+
+class AttnProcessor:
+    def __init__(self) -> None:
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError("AttnProcessor requires PyTorch 2.0+ (F.scaled_dot_product_attention).")
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: torch.Tensor | None = None,
+        base_sequence_length: int | None = None,
+    ) -> torch.Tensor:
+        batch_size, sequence_length, _ = hidden_states.shape
+
+        # Get Query-Key-Value Pair
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(encoder_hidden_states)
+        value = attn.to_v(encoder_hidden_states)
+
+        query_dim = query.shape[-1]
+        inner_dim = key.shape[-1]
+        head_dim = query_dim // attn.heads
+        dtype = query.dtype
+
+        # Get key-value heads
+        kv_heads = inner_dim // head_dim
+
+        # Reshape tensors for attention computation
+        query = query.view(batch_size, -1, attn.heads, head_dim)
+        key = key.view(batch_size, -1, kv_heads, head_dim)
+        value = value.view(batch_size, -1, kv_heads, head_dim)
+
+        # Apply Query-Key normalization
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # Apply Rotary Position Embeddings
+        if image_rotary_emb is not None:
+            query = apply_real_rotary_emb(query, image_rotary_emb[0], image_rotary_emb[1])
+            key = apply_real_rotary_emb(key, image_rotary_emb[0], image_rotary_emb[1])
+
+        query, key = query.to(dtype), key.to(dtype)
+
+        # Calculate attention scale
+        if base_sequence_length is not None:
+            softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
+        else:
+            softmax_scale = attn.scale
+
+        if _HAS_FLASH_ATTN_VARLEN and attention_mask is not None and hidden_states.is_cuda:
+            # Flash-Attn varlen expects packed tokens + cu_seqlens. Here we only need
+            # the self-attention case (q/k/v share the same padding mask).
+            attention_mask = attention_mask.to(torch.bool)
+            seqlens = attention_mask.sum(dim=-1, dtype=torch.int32)
+            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+            max_seqlen = int(seqlens.max().item())
+            cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0))
+
+            query_states = query.reshape(batch_size * sequence_length, attn.heads, head_dim)[indices]
+            key_states = key.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices]
+            value_states = value.reshape(batch_size * sequence_length, kv_heads, head_dim)[indices]
+
+            if kv_heads < attn.heads:
+                key_states = repeat(key_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
+                value_states = repeat(value_states, "l h c -> l (h k) c", k=attn.heads // kv_heads)
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                dropout_p=0.0,
+                causal=False,
+                softmax_scale=softmax_scale,
+            )
+
+            out = torch.zeros(
+                (batch_size * sequence_length, attn.heads, head_dim),
+                device=attn_output_unpad.device,
+                dtype=attn_output_unpad.dtype,
+            )
+            out[indices] = attn_output_unpad
+            hidden_states = out.view(batch_size, sequence_length, attn.heads, head_dim).flatten(-2)
+            hidden_states = hidden_states.type_as(query)
+        else:
+            # PyTorch SDPA path.
+            attn_mask = None
+            if attention_mask is not None:
+                attention_mask = attention_mask.to(torch.bool)
+                attn_mask = attention_mask.view(batch_size, 1, 1, -1)
+
+            query = query.transpose(1, 2)  # [B, H, S, D]
+            key = key.transpose(1, 2)  # [B, H_kv, S, D]
+            value = value.transpose(1, 2)
+
+            if kv_heads < attn.heads:
+                key = key.repeat_interleave(attn.heads // kv_heads, dim=1)
+                value = value.repeat_interleave(attn.heads // kv_heads, dim=1)
+
+            hidden_states = F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                attn_mask=attn_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=softmax_scale,
+            )
+
+            if attention_mask is not None:
+                # Keep padding tokens consistent with the flash-varlen path (zero output).
+                hidden_states = hidden_states * attention_mask[:, None, :, None]
+
+            hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+            hidden_states = hidden_states.type_as(query)
+
+        # Apply output projection
+        hidden_states = attn.to_out[0](hidden_states)
+        hidden_states = attn.to_out[1](hidden_states)
+
+        return hidden_states
+
+
+class TransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        num_kv_heads: int,
+        multiple_of: int,
+        ffn_dim_multiplier: float,
+        norm_eps: float,
+        modulation: bool = True,
+    ) -> None:
+        """Initialize the transformer block."""
+        super().__init__()
+        self.head_dim = dim // num_attention_heads
+        self.modulation = modulation
+
+        processor = AttnProcessor()
+
+        # Initialize attention layer
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,
+            dim_head=dim // num_attention_heads,
+            qk_norm=None,
+            heads=num_attention_heads,
+            kv_heads=num_kv_heads,
+            eps=1e-5,
+            bias=False,
+            out_bias=False,
+            processor=processor,
+        )
+        # 显式使用 transformers 的 Qwen2RMSNorm，避免依赖 diffusers 内部创建的 `RMSNorm` 再做递归替换。
+        self.attn.norm_q = Qwen2RMSNorm(self.head_dim, eps=1e-5)
+        self.attn.norm_k = Qwen2RMSNorm(self.head_dim, eps=1e-5)
+
+        # Initialize feed-forward network
+        self.feed_forward = LuminaFeedForward(
+            dim=dim, inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier
+        )
+
+        # Initialize normalization layers
+        if modulation:
+            self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True)
+        else:
+            self.norm1 = Qwen2RMSNorm(dim, eps=norm_eps)
+
+        self.ffn_norm1 = Qwen2RMSNorm(dim, eps=norm_eps)
+        self.norm2 = Qwen2RMSNorm(dim, eps=norm_eps)
+        self.ffn_norm2 = Qwen2RMSNorm(dim, eps=norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        image_rotary_emb: torch.Tensor,
+        temb: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.modulation:
+            if temb is None:
+                raise ValueError("temb must be provided when modulation is enabled")
+
+            norm_hidden_states, gate_msa, scale_mlp, gate_mlp = self.norm1(hidden_states, temb)
+            attn_output = self.attn(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_hidden_states,
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + gate_msa.unsqueeze(1).tanh() * self.norm2(attn_output)
+            mlp_output = self.feed_forward(self.ffn_norm1(hidden_states) * (1 + scale_mlp.unsqueeze(1)))
+            hidden_states = hidden_states + gate_mlp.unsqueeze(1).tanh() * self.ffn_norm2(mlp_output)
+        else:
+            norm_hidden_states = self.norm1(hidden_states)
+            attn_output = self.attn(
+                hidden_states=norm_hidden_states,
+                encoder_hidden_states=norm_hidden_states,
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + self.norm2(attn_output)
+            mlp_output = self.feed_forward(self.ffn_norm1(hidden_states))
+            hidden_states = hidden_states + self.ffn_norm2(mlp_output)
+
+        return hidden_states
+
+
+class Transformer2DModel(ModelMixin, ConfigMixin):
+    """MammothModa2 DiT transformer"""
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 16,
+        out_channels: int | None = None,
+        hidden_size: int = 2304,
+        num_layers: int = 26,
+        num_refiner_layers: int = 2,
+        num_attention_heads: int = 24,
+        num_kv_heads: int = 8,
+        multiple_of: int = 256,
+        ffn_dim_multiplier: float | None = None,
+        norm_eps: float = 1e-5,
+        axes_dim_rope: tuple[int, int, int] = (32, 32, 32),
+        axes_lens: tuple[int, int, int] = (300, 512, 512),
+        text_feat_dim: int = 1024,
+        timestep_scale: float = 1.0,
+    ) -> None:
+        """Initialize the  transformer model."""
+        super().__init__()
+        self.hidden_size = hidden_size
+
+        # Validate configuration
+        if (hidden_size // num_attention_heads) != sum(axes_dim_rope):
+            raise ValueError(
+                f"hidden_size // num_attention_heads ({hidden_size // num_attention_heads}) "
+                f"must equal sum(axes_dim_rope) ({sum(axes_dim_rope)})"
+            )
+
+        self.out_channels = out_channels or in_channels
+        # Initialize embeddings
+        self.rope_embedder = RotaryPosEmbedReal(
+            theta=10000,
+            axes_dim=axes_dim_rope,
+            axes_lens=axes_lens,
+            patch_size=patch_size,
+        )
+
+        self.x_embedder = nn.Linear(
+            in_features=patch_size * patch_size * in_channels,
+            out_features=hidden_size,
+        )
+
+        self.ref_image_patch_embedder = nn.Linear(
+            in_features=patch_size * patch_size * in_channels,
+            out_features=hidden_size,
+        )
+
+        self.time_caption_embed = Lumina2CombinedTimestepCaptionEmbedding(
+            hidden_size=hidden_size,
+            text_feat_dim=text_feat_dim,
+            norm_eps=norm_eps,
+            timestep_scale=timestep_scale,
+        )
+
+        # Initialize transformer blocks
+        self.noise_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True,
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        self.ref_image_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True,
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        self.context_refiner = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=False,
+                )
+                for _ in range(num_refiner_layers)
+            ]
+        )
+
+        # 3. Transformer blocks
+        self.layers = nn.ModuleList(
+            [
+                TransformerBlock(
+                    hidden_size,
+                    num_attention_heads,
+                    num_kv_heads,
+                    multiple_of,
+                    ffn_dim_multiplier,
+                    norm_eps,
+                    modulation=True,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        # 4. Output norm & projection
+        self.norm_out = LuminaLayerNormContinuous(
+            embedding_dim=hidden_size,
+            conditioning_embedding_dim=min(hidden_size, 1024),
+            elementwise_affine=False,
+            eps=1e-6,
+            bias=True,
+            out_dim=patch_size * patch_size * self.out_channels,
+        )
+
+        # Add learnable embeddings to distinguish different images
+        self.image_index_embedding = nn.Parameter(torch.randn(5, hidden_size))  # support max 5 ref images
+
+    def _validate_inputs(
+        self,
+        hidden_states: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        ref_image_hidden_states: list[list[torch.Tensor]] | None,
+        return_dict: bool,
+    ) -> tuple[int, int, int]:
+        if return_dict:
+            raise ValueError("return_dict=True is not supported in vLLM inference.")
+        if ref_image_hidden_states is not None:
+            raise ValueError("ref_image_hidden_states is not supported in vLLM inference.")
+        if hidden_states.ndim != 4:
+            raise ValueError(f"Expected hidden_states to be 4D [B,C,H,W], got shape={tuple(hidden_states.shape)}")
+
+        batch_size, _channels, height, width = hidden_states.shape
+        if batch_size != text_hidden_states.shape[0] or batch_size != text_attention_mask.shape[0]:
+            raise ValueError(
+                "Batch size mismatch: "
+                f"hidden_states={batch_size}, text_hidden_states={text_hidden_states.shape[0]}, "
+                f"text_attention_mask={text_attention_mask.shape[0]}"
+            )
+
+        p = self.config.patch_size
+        if height % p != 0 or width % p != 0:
+            raise ValueError(f"Input latent H/W must be divisible by patch_size={p}, got {height}x{width}")
+        return batch_size, height, width
+
+    def _prepare_embeddings(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        batch_size: int,
+        height: int,
+        width: int,
+        ar_image_hidden_states: torch.Tensor | None = None,
+        ar_image_attention_mask: torch.Tensor | None = None,
+    ):
+        device = hidden_states.device
+        p = self.config.patch_size
+
+        temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype)
+        image_embedder = getattr(self.time_caption_embed, "image_embedder", None)
+        if image_embedder is not None and ar_image_hidden_states is not None:
+            if ar_image_attention_mask is None:
+                ar_image_attention_mask = torch.ones(
+                    ar_image_hidden_states.shape[:2],
+                    dtype=torch.bool,
+                    device=ar_image_hidden_states.device,
+                )
+            image_hidden_states = image_embedder(
+                ar_image_hidden_states,
+                ~ar_image_attention_mask.bool(),
+            )
+            image_attention_mask = torch.ones(
+                image_hidden_states.shape[:2],
+                dtype=torch.bool,
+                device=image_hidden_states.device,
+            )
+            text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1)
+            text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+
+        img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
+        img_tokens = self.x_embedder(img_tokens)
+
+        img_len = (height // p) * (width // p)
+        img_mask = torch.ones((batch_size, img_len), dtype=torch.bool, device=device)
+        l_effective_img_len = [img_len for _ in range(batch_size)]
+        img_sizes = [(height, width) for _ in range(batch_size)]
+
+        l_effective_ref_img_len = [[] for _ in range(batch_size)]
+        ref_img_sizes = [None for _ in range(batch_size)]
+
+        (
+            context_rotary_emb,
+            _ref_img_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        ) = self.rope_embedder(
+            freqs_cis,
+            text_attention_mask,
+            l_effective_ref_img_len,
+            l_effective_img_len,
+            ref_img_sizes,
+            img_sizes,
+            device,
+        )
+
+        return (
+            temb,
+            text_hidden_states,
+            text_attention_mask,
+            img_tokens,
+            img_mask,
+            img_len,
+            context_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        )
+
+    def _apply_refiners(
+        self,
+        text_hidden_states: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        context_rotary_emb: torch.Tensor,
+        img_tokens: torch.Tensor,
+        img_mask: torch.Tensor,
+        noise_rotary_emb: torch.Tensor,
+        temb: torch.Tensor,
+    ):
+        for layer in self.context_refiner:
+            text_hidden_states = layer(text_hidden_states, text_attention_mask, context_rotary_emb)
+
+        for layer in self.noise_refiner:
+            img_tokens = layer(img_tokens, img_mask, noise_rotary_emb, temb)
+
+        return text_hidden_states, img_tokens
+
+    def _apply_transformer_layers(self, hidden_states, attention_mask, rotary_emb, temb):
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask, rotary_emb, temb)
+        return hidden_states
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        text_hidden_states: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        text_attention_mask: torch.Tensor,
+        ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
+        ar_image_hidden_states: torch.Tensor | None = None,
+        ar_image_attention_mask: torch.Tensor | None = None,
+        return_dict: bool = False,
+    ) -> torch.Tensor:
+        batch_size, height, width = self._validate_inputs(
+            hidden_states, text_hidden_states, text_attention_mask, ref_image_hidden_states, return_dict
+        )
+
+        (
+            temb,
+            text_hidden_states,
+            text_attention_mask,
+            img_tokens,
+            img_mask,
+            img_len,
+            context_rotary_emb,
+            noise_rotary_emb,
+            rotary_emb,
+            encoder_seq_lengths,
+            seq_lengths,
+        ) = self._prepare_embeddings(
+            hidden_states,
+            timestep,
+            text_hidden_states,
+            text_attention_mask,
+            freqs_cis,
+            batch_size,
+            height,
+            width,
+            ar_image_hidden_states,
+            ar_image_attention_mask,
+        )
+
+        text_hidden_states, img_tokens = self._apply_refiners(
+            text_hidden_states,
+            text_attention_mask,
+            context_rotary_emb,
+            img_tokens,
+            img_mask,
+            noise_rotary_emb,
+            temb,
+        )
+
+        max_seq_len = max(seq_lengths)
+        attention_mask = hidden_states.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
+        joint_hidden_states = hidden_states.new_zeros(batch_size, max_seq_len, self.config.hidden_size)
+        for i, (encoder_seq_len, seq_len) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+            attention_mask[i, :seq_len] = True
+            joint_hidden_states[i, :encoder_seq_len] = text_hidden_states[i, :encoder_seq_len]
+            joint_hidden_states[i, encoder_seq_len : encoder_seq_len + img_len] = img_tokens[i, :img_len]
+
+        hidden_states = self._apply_transformer_layers(joint_hidden_states, attention_mask, rotary_emb, temb)
+
+        hidden_states = self.norm_out(hidden_states, temb)
+
+        p = self.config.patch_size
+        img_hidden_states = torch.stack(
+            [
+                hidden_states[i, encoder_seq_len : encoder_seq_len + img_len]
+                for i, encoder_seq_len in enumerate(encoder_seq_lengths)
+            ],
+            dim=0,
+        )
+        output = rearrange(
+            img_hidden_states,
+            "b (h w) (p1 p2 c) -> b c (h p1) (w p2)",
+            h=height // p,
+            w=width // p,
+            p1=p,
+            p2=p,
+        )
+        return output

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -0,0 +1,405 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any, ClassVar
+
+import torch
+from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL
+from diffusers.utils.torch_utils import randn_tensor
+from torch import nn
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+
+from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
+
+from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel
+from .rope_real import RotaryPosEmbedReal
+from .schedulers import FlowMatchEulerDiscreteScheduler
+
+logger = init_logger(__name__)
+
+
+class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
+    """
+    MammothModa2 DiT + VAE generation stage (non-autoregressive).
+
+    This stage expects "image condition token hidden states" from the upstream AR stage,
+    and outputs image tensors via diffusion transformer + VAE decode.
+
+    """
+
+    _dit_modules: ClassVar[list[str]] = ["gen_transformer"]
+    _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
+    _vae_modules: ClassVar[list[str]] = ["gen_vae"]
+
+    have_multimodal_outputs = True
+
+    # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "llm_model.": None,
+            "gen_tokenizer.": None,
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        del prefix
+
+        hf_config = vllm_config.model_config.hf_config
+        if not isinstance(hf_config, Mammothmoda2Config):
+            raise TypeError(f"Expected Mammothmoda2Config, got {type(hf_config)}")
+
+        self.config = hf_config
+
+        # --- Build DiT / VAE modules (names must match checkpoint keys) ---
+        if self.config.gen_vae_config is None or self.config.gen_dit_config is None:
+            raise ValueError("Mammothmoda2Config.gen_vae_config / gen_dit_config must not be None")
+
+        self.gen_vae = AutoencoderKL.from_config(self.config.gen_vae_config)
+        self.gen_transformer = Transformer2DModel.from_config(self.config.gen_dit_config)
+
+        # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
+        llm_hidden_size = 0
+        text_config = self.config.get_text_config()
+        if text_config is None:
+            logger.warning("No text config; failed to infer llm_hidden_size.")
+        elif not hasattr(text_config, "hidden_size"):
+            logger.warning("Text config exists, but has no hidden_size attribute; failed to infer llm_hidden_size.")
+        else:
+            llm_hidden_size = int(text_config.hidden_size or 0)
+        if llm_hidden_size <= 0:
+            raise ValueError(
+                "Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.text_config.hidden_size"
+            )
+        self._reinit_caption_embedder(llm_hidden_size)
+
+        # Optional image condition Q-Former. Preview stores it as a standalone
+        # module; Dev stores it under the DiT timestep/caption embedder.
+        llm_model_type = getattr(self.config.llm_config, "model_type", "")
+        refiner_config = self.config.gen_image_condition_refiner_config
+        if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl":
+            dit_hidden_size = int(self.gen_transformer.hidden_size)
+            self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner(
+                hidden_size=llm_hidden_size,
+                output_hidden_size=dit_hidden_size,
+                num_heads=max(1, dit_hidden_size // 128),
+                **refiner_config,
+            )
+            self.gen_image_condition_refiner = None
+        elif refiner_config is not None:
+            self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
+                hidden_size=llm_hidden_size,
+                **refiner_config,
+            )
+        else:
+            self.gen_image_condition_refiner = None
+
+        # Precompute rotary freqs for diffusion transformer
+        # IMPORTANT: follow upstream mammothmoda: use top-level `config.gen_axes_*`
+        # (the checkpoint's `gen_dit_config.axes_lens` can be as small as 1024,
+        # which is insufficient for vLLM dummy-run/cudagraph warmup).
+        self.gen_freqs_cis = RotaryPosEmbedReal.get_freqs_real(
+            tuple(self.config.gen_axes_dim_rope),
+            tuple(self.config.gen_axes_lens),
+            theta=10000,
+        )
+
+        # vLLM PP interface compatibility
+        self.make_empty_intermediate_tensors = lambda: None
+
+        self._llm_hidden_size = llm_hidden_size
+
+    def _reinit_caption_embedder(self, in_features: int) -> None:
+        # Align with upstream Mammothmoda2Model's `reinit_caption_embedder`:
+        # Use Qwen2RMSNorm(in_features) + Linear(in_features -> out_features).
+        out_features = int(getattr(self.gen_transformer, "hidden_size", 0) or self.gen_transformer.config.hidden_size)
+        self.gen_transformer.time_caption_embed.caption_embedder = nn.Sequential(
+            Qwen2RMSNorm(in_features, eps=1e-5),
+            nn.Linear(in_features, out_features, bias=True),
+        )
+
+    def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
+        if num_reqs <= 0:
+            raise ValueError(f"num_reqs must be positive, got {num_reqs}")
+        if num_reqs > 1:
+            raise NotImplementedError(
+                f"get_dummy_runtime_additional_information does not support num_reqs > 1, got {num_reqs}"
+            )
+        text_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32)
+        image_prompt_embeds = torch.zeros((1, self._llm_hidden_size), dtype=torch.float32)
+        negative_prompt_embeds = torch.zeros((0, self._llm_hidden_size), dtype=torch.float32)
+        info = {
+            "text_prompt_embeds": text_prompt_embeds,
+            "image_prompt_embeds": image_prompt_embeds,
+            "negative_prompt_embeds": negative_prompt_embeds,
+            "negative_prompt_attention_mask": [],
+            "image_height": [512],
+            "image_width": [512],
+            "text_guidance_scale": [1.0],
+            "cfg_range": [0.0, 1.0],
+            "num_inference_steps": [1],
+        }
+        return [info for _ in range(num_reqs)]
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        # DiT stage does not consume token embeddings; return a dummy tensor.
+        try:
+            dtype = next(self.parameters()).dtype
+        except StopIteration:
+            dtype = torch.float32
+        return torch.zeros(
+            (input_ids.numel(), self._llm_hidden_size),
+            device=input_ids.device,
+            dtype=dtype,
+        )
+
+    def _split_ar_conditions(
+        self,
+        *,
+        full_hidden_states: torch.Tensor,
+        full_token_ids: list[int],
+        answer_start_index: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split AR-stage hidden states into text / image condition embeds.
+
+        The token ids that distinguish question (text) tokens, generated visual
+        tokens, and multi-modal placeholder tokens are read from the model config
+        (``gen_vocab_start_index`` and the vision placeholder token ids), so the
+        caller no longer needs to pass them. Mirrors the masking the bespoke
+        MammothModa2 example performed via ar2dit.
+        """
+        gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index)
+        visual_ids = [
+            int(self.config.image_token_id),
+            int(self.config.video_token_id),
+            int(self.config.vision_start_token_id),
+            int(self.config.vision_end_token_id),
+        ]
+
+        device = full_hidden_states.device
+        token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
+        positions = torch.arange(token_ids.shape[0], device=device)
+        questions_mask = positions < answer_start_index
+        answers_mask = ~questions_mask
+        gen_token_mask = token_ids >= gen_vocab_start_index
+        visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device))
+        text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
+        image_mask = answers_mask & gen_token_mask
+
+        text_cond = full_hidden_states[text_mask].to(dtype=torch.float32).contiguous()
+        image_cond = full_hidden_states[image_mask].to(dtype=torch.float32).contiguous()
+        return text_cond, image_cond
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        *,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,  # noqa: ARG002
+    ) -> OmniOutput:
+        runtime_addi = kwargs.get("runtime_additional_information", None)
+        info = runtime_addi[0]
+
+        # Sampling knobs are declared in vllm_omni/model_extras/mammothmodal2_preview.py
+        # and routed via extra_body -> sampling_params.extra_args (surfaced here as
+        # ``sampling_extra_args``). Fall back to runtime_additional_information for the
+        # legacy bespoke-example path during the transition.
+        extra_args_list = kwargs.get("sampling_extra_args") or []
+        extra_args = extra_args_list[0] if extra_args_list else {}
+        text_guidance_scale = float(extra_args.get("text_guidance_scale", info["text_guidance_scale"][0]))
+        cfg_range_val = extra_args.get("cfg_range", info["cfg_range"])
+        cfg_range = float(cfg_range_val[0]), float(cfg_range_val[1])
+        num_inference_steps = int(extra_args.get("num_inference_steps", info["num_inference_steps"][0]))
+
+        negative_cond = info.get("negative_prompt_embeds")
+        negative_attention_mask = info.get("negative_prompt_attention_mask")
+        image_hw = info["image_height"][0], info["image_width"][0]
+
+        # Split the AR hidden states into text / image conditions. The token ids that
+        # drive the split are sourced from the model config (see _split_ar_conditions),
+        # formerly supplied by the bespoke example via additional_information. Legacy
+        # fallback: ar2dit may have already produced the split conditions.
+        if "text_prompt_embeds" in info:
+            text_cond = info["text_prompt_embeds"]
+            image_cond = info["image_prompt_embeds"]
+        else:
+            text_cond, image_cond = self._split_ar_conditions(
+                full_hidden_states=info["full_hidden_states"],
+                full_token_ids=info["full_token_ids"],
+                answer_start_index=int(info["answer_start_index"][0]),
+            )
+
+        # Move to model device/dtype.
+        model_device = next(self.parameters()).device
+        if self.gen_image_condition_refiner is not None:
+            target_dtype = next(self.gen_image_condition_refiner.parameters()).dtype
+        else:
+            target_dtype = next(self.gen_transformer.parameters()).dtype
+
+        def _ensure_2d(x: torch.Tensor, name: str) -> torch.Tensor:
+            if x.ndim == 3 and x.shape[0] == 1:
+                x = x[0]
+            if x.ndim != 2:
+                raise ValueError(f"Expected {name} to be 2D [T,H], got shape={tuple(x.shape)}")
+            return x
+
+        text_cond = _ensure_2d(text_cond, "text_prompt_embeds")
+        image_cond = _ensure_2d(image_cond, "image_prompt_embeds")
+        if image_cond.shape[0] == 0:
+            answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :]
+            raise ValueError(
+                "MammothModa2 AR stage produced no visual-token hidden states; "
+                "the DiT stage requires at least one generated visual token. "
+                f"Generated token ids: {answer_token_ids[:32]}"
+            )
+        text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
+        image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
+
+        text_embeds = text_cond.unsqueeze(0)  # [1, T_text, H]
+        text_attention_mask = torch.ones(
+            (1, text_embeds.shape[1]),
+            dtype=torch.bool,
+            device=text_embeds.device,
+        )
+
+        image_embeds = image_cond.unsqueeze(0)  # [1, T_img, H]
+        image_attention_mask = torch.ones(
+            (1, image_embeds.shape[1]),
+            dtype=torch.bool,
+            device=image_embeds.device,
+        )
+
+        # Apply optional refiner ONLY on image condition tokens.
+        if self.gen_image_condition_refiner is not None and image_embeds.shape[1] > 0:
+            image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_attention_mask.bool())
+            image_attention_mask = torch.ones(
+                image_embeds.shape[:2],
+                dtype=torch.bool,
+                device=image_embeds.device,
+            )
+
+        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
+        if nested_image_embedder is None:
+            prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
+            prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+            ar_image_embeds = None
+            ar_image_attention_mask = None
+        else:
+            prompt_embeds = text_embeds
+            prompt_attention_mask = text_attention_mask
+            ar_image_embeds = image_embeds
+            ar_image_attention_mask = image_attention_mask
+
+        # Prepare negative prompt (for CFG). If none provided, fall back to unconditional.
+        negative_prompt_embeds = None
+        negative_prompt_attention_mask = None
+        if text_guidance_scale > 1.0:
+            if negative_cond is not None:
+                negative_cond = _ensure_2d(negative_cond, "negative_prompt_embeds")
+                negative_prompt_embeds = (
+                    negative_cond.to(device=model_device, dtype=target_dtype, non_blocking=True)
+                    .contiguous()
+                    .unsqueeze(0)
+                )
+                if isinstance(negative_attention_mask, torch.Tensor):
+                    neg_mask = negative_attention_mask
+                elif isinstance(negative_attention_mask, list):
+                    neg_mask = torch.tensor(negative_attention_mask, dtype=torch.bool)
+                else:
+                    neg_mask = None
+                if neg_mask is None:
+                    negative_prompt_attention_mask = torch.ones(
+                        (1, negative_prompt_embeds.shape[1]),
+                        dtype=torch.bool,
+                        device=negative_prompt_embeds.device,
+                    )
+                else:
+                    neg_mask = neg_mask.to(device=negative_prompt_embeds.device, dtype=torch.bool)
+                    if neg_mask.ndim == 1:
+                        neg_mask = neg_mask.unsqueeze(0)
+                    negative_prompt_attention_mask = neg_mask
+            else:
+                hidden_size = int(prompt_embeds.shape[-1])
+                negative_prompt_embeds = torch.zeros(
+                    (1, 0, hidden_size),
+                    dtype=target_dtype,
+                    device=prompt_embeds.device,
+                )
+                negative_prompt_attention_mask = torch.zeros(
+                    (1, 0),
+                    dtype=torch.bool,
+                    device=prompt_embeds.device,
+                )
+
+        # Output image size (px), passed from stage input processor.
+        height, width = image_hw
+        if height <= 0 or width <= 0:
+            raise ValueError(f"Invalid image size: {height}x{width}")
+        if height % 16 != 0 or width % 16 != 0:
+            raise ValueError(f"Image size must be multiples of 16, got {height}x{width}")
+        vae_scale_factor = 16
+
+        latent_channels = int(self.gen_transformer.config.in_channels)
+        shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        latents = randn_tensor(shape, device=prompt_embeds.device, dtype=prompt_embeds.dtype)
+
+        scheduler = FlowMatchEulerDiscreteScheduler()
+
+        scheduler.set_timesteps(
+            num_inference_steps=num_inference_steps,
+            device=prompt_embeds.device,
+            num_tokens=latents.shape[-2] * latents.shape[-1],
+        )
+
+        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
+        total_steps = max(1, len(scheduler.timesteps))
+        for i, t in enumerate(scheduler.timesteps):
+            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            model_pred = self.gen_transformer(
+                hidden_states=latents,
+                timestep=timestep,
+                text_hidden_states=prompt_embeds,
+                text_attention_mask=prompt_attention_mask,
+                ref_image_hidden_states=None,
+                ar_image_hidden_states=ar_image_embeds,
+                ar_image_attention_mask=ar_image_attention_mask,
+                freqs_cis=self.gen_freqs_cis,
+            )
+            guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0
+            if guidance_scale > 1.0 and negative_prompt_embeds is not None:
+                model_pred_uncond = self.gen_transformer(
+                    hidden_states=latents,
+                    timestep=timestep,
+                    text_hidden_states=negative_prompt_embeds,
+                    text_attention_mask=negative_prompt_attention_mask,
+                    ref_image_hidden_states=None,
+                    freqs_cis=self.gen_freqs_cis,
+                )
+                model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond)
+            latents = scheduler.step(model_pred, t, latents, return_dict=False)[0]
+            latents = latents.to(dtype=prompt_embeds.dtype)
+
+        # VAE decode
+        if self.gen_vae.config.scaling_factor is not None:
+            latents = latents / self.gen_vae.config.scaling_factor
+        if self.gen_vae.config.shift_factor is not None:
+            latents = latents + self.gen_vae.config.shift_factor
+        image = self.gen_vae.decode(latents, return_dict=False)[0]
+
+        return OmniOutput(
+            text_hidden_states=inputs_embeds,  # placeholder, not used by runner
+            multimodal_outputs=image,
+            intermediate_tensors=None,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:  # noqa: ARG002
+        return None
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

vllm_omni/diffusion/models/mammoth_moda2/rope_real.py

@@ -0,0 +1,252 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+import torch
+from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotary_pos_embed
+from einops import repeat
+from torch import nn
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor:
+    """
+    Apply real-valued rotary embeddings to input tensor.
+
+    Args:
+        x: Input tensor of shape [..., seq_len, num_heads, dim] or [..., seq_len, dim]
+        freqs_cos: Cosine frequencies of shape [batch, seq_len, dim] or [seq_len, dim]
+        freqs_sin: Sine frequencies of shape [batch, seq_len, dim] or [seq_len, dim]
+
+    Returns:
+        Tensor with rotary embeddings applied
+    """
+    # x: [batch, seq_len, num_heads, dim] or [batch, seq_len, dim]
+    # freqs_cos: [batch, seq_len, dim] or [seq_len, dim]
+    # freqs_sin: [batch, seq_len, dim] or [seq_len, dim]
+
+    x_shape = x.shape
+    if len(x_shape) == 4:
+        batch, seq_len, num_heads, dim = x_shape
+        x_reshaped = x.view(batch, seq_len, num_heads, dim // 2, 2)
+    elif len(x_shape) == 3:
+        batch, seq_len, dim = x_shape
+        num_heads = None
+        x_reshaped = x.view(batch, seq_len, dim // 2, 2)
+    else:
+        raise ValueError(f"Unsupported x shape: {x.shape}")
+
+    # freqs_cos/sin: [batch, seq_len, dim] or [seq_len, dim]
+    # Expand freqs_cos/sin to [batch, seq_len, dim] if needed
+    if freqs_cos.dim() == 2:
+        # [seq_len, dim] -> [1, seq_len, dim]
+        freqs_cos = freqs_cos.unsqueeze(0)
+        freqs_sin = freqs_sin.unsqueeze(0)
+    if freqs_cos.shape[0] == 1 and batch > 1:
+        freqs_cos = freqs_cos.expand(batch, -1, -1)
+        freqs_sin = freqs_sin.expand(batch, -1, -1)
+
+    # Reshape freqs to [batch, seq_len, dim//2, 2]
+    freqs_cos_reshaped = freqs_cos.view(batch, seq_len, dim // 2, 2)
+    freqs_sin_reshaped = freqs_sin.view(batch, seq_len, dim // 2, 2)
+
+    cos_1 = freqs_cos_reshaped[..., 0]  # [batch, seq_len, dim//2]
+    cos_2 = freqs_cos_reshaped[..., 1]  # [batch, seq_len, dim//2]
+    sin_1 = freqs_sin_reshaped[..., 0]  # [batch, seq_len, dim//2]
+    sin_2 = freqs_sin_reshaped[..., 1]  # [batch, seq_len, dim//2]
+
+    # Broadcast cos/sin to match x_reshaped
+    if len(x_shape) == 4:
+        # [batch, seq_len, 1, dim//2]
+        cos_1 = cos_1.unsqueeze(2)
+        cos_2 = cos_2.unsqueeze(2)
+        sin_1 = sin_1.unsqueeze(2)
+        sin_2 = sin_2.unsqueeze(2)
+
+    x1 = x_reshaped[..., 0]  # [..., seq_len, num_heads, dim//2] or [..., seq_len, dim//2]
+    x2 = x_reshaped[..., 1]  # same
+
+    out1 = x1 * cos_1 - x2 * sin_1
+    out2 = x1 * sin_2 + x2 * cos_2
+
+    out = torch.stack([out1, out2], dim=-1)
+    return out.view(*x_shape)
+
+
+def get_1d_rotary_pos_embed_real(
+    dim: int,
+    pos: np.ndarray | int,
+    theta: float = 10000.0,
+    linear_factor: float = 1.0,
+    ntk_factor: float = 1.0,
+    freqs_dtype: torch.dtype = torch.float32,
+):
+    freqs_cos, freqs_sin = _get_1d_rotary_pos_embed(
+        dim,
+        pos,
+        theta=theta,
+        use_real=True,
+        linear_factor=linear_factor,
+        ntk_factor=ntk_factor,
+        repeat_interleave_real=True,
+        freqs_dtype=freqs_dtype,
+    )
+    return freqs_cos, freqs_sin
+
+
+class RotaryPosEmbedReal(nn.Module):
+    def __init__(
+        self, theta: int, axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], patch_size: int = 2
+    ):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        self.axes_lens = axes_lens
+        self.patch_size = patch_size
+
+    @staticmethod
+    def get_freqs_real(
+        axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        freqs_real = []
+        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
+        for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
+            cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype)
+            freqs_real.append((cos_emb, sin_emb))
+        return freqs_real
+
+    def _get_freqs_real(
+        self, freqs_real: list[tuple[torch.Tensor, torch.Tensor]], ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = ids.device
+        if ids.device.type == "mps":
+            ids = ids.to("cpu")
+
+        cos_result = []
+        sin_result = []
+        for i in range(len(self.axes_dim)):
+            freqs_cos, freqs_sin = freqs_real[i]
+            freqs_cos = freqs_cos.to(ids.device)
+            freqs_sin = freqs_sin.to(ids.device)
+            index = ids[:, :, i : i + 1].repeat(1, 1, freqs_cos.shape[-1]).to(torch.int64)
+            cos_result.append(torch.gather(freqs_cos.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+            sin_result.append(torch.gather(freqs_sin.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
+
+        combined_cos = torch.cat(cos_result, dim=-1).to(device)
+        combined_sin = torch.cat(sin_result, dim=-1).to(device)
+        return combined_cos, combined_sin
+
+    def forward(
+        self, freqs_real, attention_mask, l_effective_ref_img_len, l_effective_img_len, ref_img_sizes, img_sizes, device
+    ):
+        batch_size = len(attention_mask)
+        p = self.patch_size
+
+        encoder_seq_len = attention_mask.shape[1]
+        l_effective_cap_len = attention_mask.sum(dim=1).tolist()
+        l_effective_cap_len = [int(len) for len in l_effective_cap_len]
+        seq_lengths = [
+            int(cap_len + sum(ref_img_len) + img_len)
+            for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len)
+        ]
+
+        max_seq_len = max(seq_lengths)
+        max_ref_img_len = max([int(sum(ref_img_len)) for ref_img_len in l_effective_ref_img_len])
+        max_img_len = int(max(l_effective_img_len))
+
+        # Create position IDs
+        position_ids = torch.zeros(batch_size, int(max_seq_len), 3, dtype=torch.int32, device=device)
+
+        for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)):
+            # add text position ids
+            position_ids[i, :cap_seq_len] = repeat(
+                torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3"
+            )
+            pe_shift = cap_seq_len
+            pe_shift_len = cap_seq_len
+
+            if ref_img_sizes[i] is not None:
+                for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]):
+                    H, W = ref_img_size
+                    ref_H_tokens, ref_W_tokens = H // p, W // p
+                    assert ref_H_tokens * ref_W_tokens == ref_img_len
+                    # add image position ids
+
+                    row_ids = repeat(
+                        torch.arange(ref_H_tokens, dtype=torch.int32, device=device), "h -> h w", w=ref_W_tokens
+                    ).flatten()
+                    col_ids = repeat(
+                        torch.arange(ref_W_tokens, dtype=torch.int32, device=device), "w -> h w", h=ref_H_tokens
+                    ).flatten()
+                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift
+                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids
+                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids
+
+                    pe_shift += max(ref_H_tokens, ref_W_tokens)
+                    pe_shift_len += ref_img_len
+
+            H, W = img_sizes[i]
+            H_tokens, W_tokens = H // p, W // p
+            assert H_tokens * W_tokens == l_effective_img_len[i]
+
+            row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten()
+            col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten()
+
+            assert pe_shift_len + l_effective_img_len[i] == seq_len
+            position_ids[i, pe_shift_len:seq_len, 0] = pe_shift
+            position_ids[i, pe_shift_len:seq_len, 1] = row_ids
+            position_ids[i, pe_shift_len:seq_len, 2] = col_ids
+
+        # Get combined rotary embeddings (real version)
+        freqs_cos, freqs_sin = self._get_freqs_real(freqs_real, position_ids)
+
+        # create separate rotary embeddings for captions and images
+        cap_freqs_cos = torch.zeros(
+            batch_size, encoder_seq_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype
+        )
+        cap_freqs_sin = torch.zeros(
+            batch_size, encoder_seq_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype
+        )
+        ref_img_freqs_cos = torch.zeros(
+            batch_size, max_ref_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype
+        )
+        ref_img_freqs_sin = torch.zeros(
+            batch_size, max_ref_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype
+        )
+        img_freqs_cos = torch.zeros(batch_size, max_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype)
+        img_freqs_sin = torch.zeros(batch_size, max_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype)
+
+        for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(
+            zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)
+        ):
+            cap_freqs_cos[i, :cap_seq_len] = freqs_cos[i, :cap_seq_len]
+            cap_freqs_sin[i, :cap_seq_len] = freqs_sin[i, :cap_seq_len]
+            ref_img_freqs_cos[i, : sum(ref_img_len)] = freqs_cos[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
+            ref_img_freqs_sin[i, : sum(ref_img_len)] = freqs_sin[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
+            img_freqs_cos[i, :img_len] = freqs_cos[
+                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
+            ]
+            img_freqs_sin[i, :img_len] = freqs_sin[
+                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
+            ]
+
+        return (
+            (cap_freqs_cos, cap_freqs_sin),
+            (ref_img_freqs_cos, ref_img_freqs_sin),
+            (img_freqs_cos, img_freqs_sin),
+            (freqs_cos, freqs_sin),
+            l_effective_cap_len,
+            seq_lengths,
+        )

vllm_omni/diffusion/models/mammoth_moda2/schedulers.py

@@ -0,0 +1,137 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by ByteDance Ltd. on 2025-09-30.
+#
+# Original file was released under Apache-2.0, with the full license text
+# available at https://www.apache.org/licenses/LICENSE-2.0
+#
+# This modified file is released under the same license.
+#
+# --- Upstream header preserved below ---
+#
+# Copyright 2024 Stability AI, Katherine Crowson and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+
+@dataclass
+class FlowMatchEulerDiscreteSchedulerOutput:
+    prev_sample: torch.FloatTensor
+
+
+class FlowMatchEulerDiscreteScheduler:
+    order = 1
+
+    def __init__(self, num_train_timesteps: int = 1000, dynamic_time_shift: bool = True):
+        self.num_train_timesteps = int(num_train_timesteps)
+        self.dynamic_time_shift = bool(dynamic_time_shift)
+
+        timesteps = torch.linspace(0, 1, self.num_train_timesteps + 1, dtype=torch.float32)[:-1]
+        self.timesteps = timesteps
+        self._timesteps = torch.cat([timesteps, torch.ones(1, dtype=timesteps.dtype)])
+
+        self._step_index: int | None = None
+        self._begin_index: int | None = None
+
+    @property
+    def step_index(self) -> int | None:
+        return self._step_index
+
+    @property
+    def begin_index(self) -> int | None:
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0) -> None:
+        self._begin_index = int(begin_index)
+
+    def index_for_timestep(self, timestep: torch.Tensor, schedule_timesteps: torch.Tensor | None = None) -> int:
+        schedule_timesteps = self._timesteps if schedule_timesteps is None else schedule_timesteps
+        indices = (schedule_timesteps == timestep).nonzero()
+        pos = 1 if len(indices) > 1 else 0
+        return int(indices[pos].item())
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int | None = None,
+        device: str | torch.device | None = None,
+        timesteps: list[float] | None = None,
+        num_tokens: int | None = None,
+    ) -> None:
+        if timesteps is None:
+            if num_inference_steps is None:
+                raise ValueError("`num_inference_steps` must be provided when `timesteps` is None.")
+            timesteps_np = np.linspace(0, 1, num_inference_steps + 1, dtype=np.float32)[:-1]
+            if self.dynamic_time_shift and num_tokens is not None:
+                m = np.sqrt(float(num_tokens)) / 40.0
+                timesteps_np = timesteps_np / (m - m * timesteps_np + timesteps_np)
+        else:
+            timesteps_np = np.asarray(timesteps, dtype=np.float32)
+
+        timesteps_t = torch.from_numpy(timesteps_np).to(dtype=torch.float32, device=device)
+        self.timesteps = timesteps_t
+        self._timesteps = torch.cat([timesteps_t, torch.ones(1, device=timesteps_t.device, dtype=timesteps_t.dtype)])
+
+        self._step_index = None
+        self._begin_index = None
+
+    def _init_step_index(self, timestep: float | torch.Tensor) -> None:
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            else:
+                timestep = torch.tensor(timestep, device=self.timesteps.device, dtype=self.timesteps.dtype)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.FloatTensor,
+        timestep: float | torch.FloatTensor,
+        sample: torch.FloatTensor,
+        generator: torch.Generator | None = None,  # noqa: ARG002 - kept for API compatibility
+        return_dict: bool = True,
+    ) -> FlowMatchEulerDiscreteSchedulerOutput | tuple[torch.FloatTensor]:
+        if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)):
+            raise ValueError(
+                "Integer timesteps (e.g. a loop index from enumerate) are not supported; "
+                "pass a value from `scheduler.timesteps` instead.",
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+        assert self._step_index is not None
+
+        sample_fp32 = sample.to(torch.float32)
+        t = self._timesteps[self._step_index]
+        t_next = self._timesteps[self._step_index + 1]
+
+        prev_sample = sample_fp32 + (t_next - t) * model_output
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+        return FlowMatchEulerDiscreteSchedulerOutput(prev_sample=prev_sample)
+
+    def __len__(self) -> int:
+        return self.num_train_timesteps

vllm_omni/entrypoints/cli/__init__.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """CLI helpers for vLLM-Omni entrypoints."""
 
 # To ensure patch imports work properly, disable unused import checks
@@ -6,8 +9,9 @@
 from vllm_omni.benchmarks.patch import patch
 # isort: on
 
+from vllm_omni.entrypoints.cli.benchmark.omni_duplex_eval import OmniDuplexEvalSubcommand
 from vllm_omni.entrypoints.cli.benchmark.serve import OmniBenchmarkServingSubcommand
 
 from .serve import OmniServeCommand
 
-__all__ = ["OmniServeCommand", "OmniBenchmarkServingSubcommand"]
+__all__ = ["OmniServeCommand", "OmniBenchmarkServingSubcommand", "OmniDuplexEvalSubcommand"]

vllm_omni/entrypoints/cli/benchmark/omni_duplex_eval.py

@@ -0,0 +1,134 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import argparse
+import asyncio
+import concurrent.futures
+import json
+from pathlib import Path
+
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_dataset import DEFAULT_DATASET, load_samples
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_eval import evaluate_sample, summarize_scores
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_judge import DuplexJudge
+from vllm_omni.benchmarks.duplex.omni_duplex_eval_runner import generate_sample
+from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase
+
+
+def _common(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--dataset", default=DEFAULT_DATASET)
+    parser.add_argument("--split", default="all")
+    parser.add_argument("--family", choices=("all", "rtd", "pr"), default="all")
+    parser.add_argument("--media-root")
+    parser.add_argument("--limit", type=int)
+    parser.add_argument("--ids", nargs="*")
+
+
+def add_cli_args(parser: argparse.ArgumentParser) -> None:
+    actions = parser.add_subparsers(dest="action", required=True)
+    generate = actions.add_parser("generate")
+    _common(generate)
+    generate.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
+    generate.add_argument("--model", required=True)
+    generate.add_argument("--ref-audio", required=True)
+    generate.add_argument("--response-root", required=True)
+    generate.add_argument("--fps", type=float, default=1.0)
+    generate.add_argument("--mix", choices=("question",), default="question")
+    generate.add_argument("--pace", choices=("realtime", "as-fast-as-possible"), default="realtime")
+    generate.add_argument("--clock", choices=("media",), default="media")
+    generate.add_argument("--concurrency", type=int, default=1)
+    generate.add_argument("--overwrite", action="store_true")
+    evaluate = actions.add_parser("evaluate")
+    _common(evaluate)
+    evaluate.add_argument("--response-root", required=True)
+    evaluate.add_argument("--score-root", required=True)
+    evaluate.add_argument("--judge-base-url", default="http://127.0.0.1:8000")
+    evaluate.add_argument("--judge-model", required=True)
+    evaluate.add_argument("--judge-api-key", default="EMPTY")
+    evaluate.add_argument("--judge-video-mode", choices=("video_url", "frame-sample"), default="video_url")
+    evaluate.add_argument("--judge-fps", type=int, default=2)
+    evaluate.add_argument("--window-size", type=float, default=10.0)
+    evaluate.add_argument("--allow-invalid-clock", action="store_true")
+    evaluate.add_argument("--eval-workers", type=int, default=1)
+    evaluate.add_argument("--overwrite", action="store_true")
+    summarize = actions.add_parser("summarize")
+    summarize.add_argument("--score-root", required=True)
+
+
+def run(args: argparse.Namespace) -> int:
+    if args.action == "summarize":
+        print(json.dumps(summarize_scores(args.score_root), ensure_ascii=False, indent=2))
+        return 0
+
+    samples = load_samples(
+        args.dataset,
+        split=args.split,
+        family=args.family,
+        media_root=args.media_root,
+        limit=args.limit,
+        ids=args.ids,
+    )
+    if args.action == "generate":
+        if args.concurrency < 1:
+            raise ValueError("--concurrency must be at least 1")
+
+        async def generate() -> None:
+            semaphore = asyncio.Semaphore(args.concurrency)
+
+            async def generate_one(sample) -> None:
+                async with semaphore:
+                    await generate_sample(
+                        sample,
+                        url=args.url,
+                        model=args.model,
+                        ref_audio=args.ref_audio,
+                        output_root=args.response_root,
+                        fps=args.fps,
+                        mix=args.mix,
+                        pace=args.pace,
+                        clock=args.clock,
+                        overwrite=args.overwrite,
+                    )
+
+            await asyncio.gather(*(generate_one(sample) for sample in samples))
+
+        asyncio.run(generate())
+        return 0
+
+    if args.eval_workers < 1:
+        raise ValueError("--eval-workers must be at least 1")
+    judge = DuplexJudge(args.judge_base_url, args.judge_model, api_key=args.judge_api_key)
+
+    def evaluate_one(sample) -> None:
+        response_path = Path(args.response_root) / sample.split / f"{sample.id}.json"
+        score_path = Path(args.score_root) / sample.split / f"{sample.id}.json"
+        if score_path.exists() and not args.overwrite:
+            return
+        evaluate_sample(
+            sample,
+            response_path,
+            score_path,
+            judge,
+            judge_fps=args.judge_fps,
+            judge_video_mode=args.judge_video_mode,
+            window_size=args.window_size,
+            allow_invalid_clock=args.allow_invalid_clock,
+        )
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=args.eval_workers) as executor:
+        list(executor.map(evaluate_one, samples))
+    return 0
+
+
+class OmniDuplexEvalSubcommand(OmniBenchmarkSubcommandBase):
+    """Run the Omni-DuplexEval generation or scoring workflow."""
+
+    name = "omni-duplex-eval"
+    help = "Generate, evaluate, or summarize Omni-DuplexEval artifacts."
+
+    @classmethod
+    def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
+        add_cli_args(parser)
+
+    @staticmethod
+    def cmd(args: argparse.Namespace) -> None:
+        run(args)

vllm_omni/experimental/fullduplex/client.py

@@ -435,6 +435,7 @@ class RealtimeDuplexClient:
         self.events = RealtimeEventCollector()
         self._ws: Any = None
         self._reader_task: asyncio.Task[None] | None = None
+        self._media_clock_ms = 0
 
     async def __aenter__(self) -> RealtimeDuplexClient:
         self._ws = await websockets.connect(
@@ -463,6 +464,7 @@ class RealtimeDuplexClient:
                     continue
                 event = json.loads(raw)
                 if isinstance(event, dict):
+                    event.setdefault("_media_clock_ms", self._media_clock_ms)
                     self.events.add(event)
         except ConnectionClosed:
             return
@@ -559,21 +561,36 @@ class RealtimeDuplexClient:
             PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_ms // 1000,
             PCM16_BYTES_PER_SAMPLE,
         )
+        await self.stream_av_units(
+            ((pcm16[offset : offset + chunk_bytes], None) for offset in range(0, len(pcm16), chunk_bytes)),
+            realtime=realtime,
+        )
+
+    async def stream_av_units(
+        self,
+        units: Any,
+        *,
+        realtime: bool = True,
+    ) -> None:
+        """Stream PCM16 units, optionally attaching a JPEG to each unit."""
         audio_end_ms = 0
-        for offset in range(0, len(pcm16), chunk_bytes):
-            chunk = pcm16[offset : offset + chunk_bytes]
+        for chunk, frame in units:
+            if not chunk:
+                continue
             duration_ms = len(chunk) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
             audio_end_ms += duration_ms
-            await self.send(
-                {
-                    "type": "input_audio_buffer.append",
-                    "audio": base64.b64encode(chunk).decode("ascii"),
-                    "input_audio_format": "pcm16",
-                    "sample_rate_hz": PCM16_SAMPLE_RATE,
-                    "duration_ms": duration_ms,
-                    "audio_end_ms": audio_end_ms,
-                }
-            )
+            self._media_clock_ms = audio_end_ms
+            event: dict[str, object] = {
+                "type": "input_audio_buffer.append",
+                "audio": base64.b64encode(chunk).decode("ascii"),
+                "input_audio_format": "pcm16",
+                "sample_rate_hz": PCM16_SAMPLE_RATE,
+                "duration_ms": duration_ms,
+                "audio_end_ms": audio_end_ms,
+            }
+            if frame is not None:
+                event["video_frames"] = [base64.b64encode(frame).decode("ascii") if isinstance(frame, bytes) else frame]
+            await self.send(event)
             if realtime:
                 await asyncio.sleep(duration_ms / 1000)
 

vllm_omni/model_executor/models/mammoth_moda2/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# NOTE: Do not import model classes in this file. Importing any
+# submodule in this package triggers __init__.py execution, and
+# both the model registry and pipeline registry import submodules
+# directly — heavy imports here would be loaded as a side effect
+# even though nothing depends on these re-exports.

vllm_omni/model_executor/models/mammoth_moda2/pipeline.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""MammothModa2 pipeline topology (frozen).
+
+Stage 0: AR  — multimodal understanding + latent generation
+Stage 1: DiT — latent → image
+
+For text/image understanding tasks (text output only), use MAMMOTH_MODA2_AR_PIPELINE.
+"""
+
+from vllm_omni.config.stage_config import (
+    PipelineConfig,
+    StageExecutionType,
+    StagePipelineConfig,
+)
+
+_PROC = "vllm_omni.model_executor.stage_input_processors.mammoth_moda2"
+
+MAMMOTH_MODA2_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2",
+    default_deploy_config_name="mammoth_moda2.yaml",
+    model_arch="MammothModa2ForConditionalGeneration",
+    hf_architectures=("Mammothmoda2Model", "MammothModa2ForConditionalGeneration"),
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=False,
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="latent",
+        ),
+        StagePipelineConfig(
+            stage_id=1,
+            model_stage="dit",
+            execution_type=StageExecutionType.LLM_GENERATION,
+            input_sources=(0,),
+            final_output=True,
+            final_output_type="image",
+            owns_tokenizer=False,
+            requires_multimodal_data=False,
+            engine_output_type="image",
+            custom_process_input_func=f"{_PROC}.ar2dit",
+        ),
+    ),
+)
+
+# Single-stage AR variant for understanding / summarization tasks
+MAMMOTH_MODA2_AR_PIPELINE = PipelineConfig(
+    model_type="mammoth_moda2_ar",
+    default_deploy_config_name="mammoth_moda2_ar.yaml",
+    model_arch="MammothModa2ForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="ar",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            final_output=True,
+            final_output_type="text",
+            owns_tokenizer=True,
+            requires_multimodal_data=True,
+            engine_output_type="text",
+        ),
+    ),
+)

vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -0,0 +1,15 @@
+"""
+Compatibility shim.
+
+The MammothModa2 DiT implementation lives under `vllm_omni.diffusion` to align
+with other ARDiT structured models. We keep this module path so existing
+OmniModelRegistry entries (and downstream code) keep working.
+"""
+
+from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (  # noqa: F401
+    MammothModa2DiTPipeline,
+)
+
+__all__ = [
+    "MammothModa2DiTPipeline",
+]

vllm_omni/model_executor/models/registry.py

@@ -114,6 +114,41 @@ _OMNI_MODELS = {
         "omnivoice",
         "OmniVoiceModel",
     ),
+    "MammothModa2Qwen2ForCausalLM": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen2ForCausalLM",
+    ),
+    "MammothModa2ARForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ARForConditionalGeneration",
+    ),
+    "MammothModa2Qwen3ARForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen3ARForConditionalGeneration",
+    ),
+    "MammothModa2Qwen3ForCausalLM": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen3ForCausalLM",
+    ),
+    "MammothModa2DiTPipeline": (
+        "mammoth_moda2",
+        "pipeline_mammothmoda2_dit",
+        "MammothModa2DiTPipeline",
+    ),
+    "MammothModa2ForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ForConditionalGeneration",
+    ),
+    "Mammothmoda2Model": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ForConditionalGeneration",
+    ),
     "Qwen3TTSForConditionalGeneration": (
         "qwen3_tts",
         "qwen3_tts_talker",

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -0,0 +1,115 @@
+"""Stage input processor for MammothModa2 (AR -> DiT)."""
+
+from collections.abc import Mapping
+from typing import Any
+
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+
+def _as_dict(prompt: Any) -> dict[str, Any]:
+    """Coerce an original-stage prompt to a dict.
+
+    It may arrive as a dict, a NamedTuple/object, or a bare string depending on
+    the calling flow (the shared text_to_image example vs the bespoke script).
+    """
+    if isinstance(prompt, dict):
+        return prompt
+    if hasattr(prompt, "_asdict"):
+        return prompt._asdict()
+    if hasattr(prompt, "__dict__"):
+        return vars(prompt)
+    return {}
+
+
+def _coerce_dim(value: Any, default: int) -> int:
+    try:
+        iv = int(value)
+    except (TypeError, ValueError):
+        return default
+    return iv if iv > 0 else default
+
+
+def ar2dit(
+    source_outputs: list[Any],
+    prompts: OmniTokensPrompt | TextPrompt | list[OmniTokensPrompt | TextPrompt] | None = None,
+    _requires_multimodal_data: bool = False,
+) -> list[OmniTokensPrompt]:
+    """Convert AR stage outputs to DiT stage inputs."""
+    ar_outputs = source_outputs
+
+    # The shared text_to_image example forwards a single prompt (not a list); normalize
+    # so a lone dict isn't iterated as its keys. Mirrors glm_image.ar2diffusion.
+    if not isinstance(prompts, list):
+        prompts = [prompts] if prompts is not None else [{}]
+
+    dit_inputs: list[OmniTokensPrompt] = []
+    for i, ar_output in enumerate(ar_outputs):
+        prompt_dict = _as_dict(prompts[i] if i < len(prompts) else {})
+        addi_info = prompt_dict.get("additional_information") or {}
+        mm_kwargs = prompt_dict.get("mm_processor_kwargs") or {}
+
+        # Image size: prefer mm_processor_kwargs target_h/target_w (set by the serving
+        # layer), fall back to additional_information, then a 1024 default.
+        image_height = _coerce_dim(
+            mm_kwargs.get("target_h"),
+            _coerce_dim((addi_info.get("image_height") or [None])[0], 1024),
+        )
+        image_width = _coerce_dim(
+            mm_kwargs.get("target_w"),
+            _coerce_dim((addi_info.get("image_width") or [None])[0], 1024),
+        )
+
+        # Sampling knobs arrive on the DiT stage via extra_body -> extra_args; these are
+        # defensive fallbacks (defaults mirror the former bespoke script's argparse).
+        text_guidance_scale = (addi_info.get("text_guidance_scale") or [9.0])[0]
+        cfg_range = addi_info.get("cfg_range") or [0.0, 1.0]
+        num_inference_steps = (addi_info.get("num_inference_steps") or [50])[0]
+
+        prompt_token_ids = ar_output.prompt_token_ids
+        # exclude the last token because it has no corresponding hidden state
+        completion_output = ar_output.outputs[0]
+        gen_token_ids = completion_output.cumulative_token_ids[:-1]
+        full_token_ids = prompt_token_ids + gen_token_ids
+
+        mm_output = getattr(completion_output, "multimodal_output", None)
+        if not isinstance(mm_output, Mapping) or "latent" not in mm_output:
+            raise ValueError(
+                "AR stage output missing latent multimodal output. "
+                f"request_id={getattr(ar_output, 'request_id', None)}, "
+                f"completion_has_mm={hasattr(completion_output, 'multimodal_output')}"
+            )
+        full_hidden_states = mm_output["latent"]
+        hidden_total = int(full_hidden_states.shape[0])
+        assert hidden_total == len(prompt_token_ids) + len(gen_token_ids), (
+            f"Hidden states length mismatch: expected {len(prompt_token_ids) + len(gen_token_ids)}, got {hidden_total}"
+        )
+
+        # The text/image condition split is performed in the DiT pipeline, which sources
+        # the distinguishing token ids (gen_vocab_start_index, vision placeholder ids)
+        # from the model config. Pass through the raw AR hidden states + token ids and
+        # the question/answer boundary so the pipeline can reconstruct the masks.
+        additional_information = {
+            # float32 so the tensor crosses the stage boundary (the serializer uses
+            # numpy, which has no bf16); the DiT re-casts to the model dtype.
+            "full_hidden_states": full_hidden_states.float().contiguous(),
+            "full_token_ids": full_token_ids,
+            "answer_start_index": [len(prompt_token_ids)],
+            "image_height": [int(image_height)],
+            "image_width": [int(image_width)],
+            "text_guidance_scale": [float(text_guidance_scale)],
+            "cfg_range": [float(cfg_range[0]), float(cfg_range[1])],
+            "num_inference_steps": [int(num_inference_steps)],
+        }
+
+        dit_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0],
+                additional_information=additional_information,
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+
+    return dit_inputs

vllm_omni/model_extras/mammothmodal2_preview.py

@@ -0,0 +1,101 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from typing import Any
+
+# AR image-grid patch size (pixels per visual token). Mirrors the constant the
+# bespoke MammothModa2 example used to convert image dimensions -> AR grid dims.
+_PATCH_SIZE = 16
+
+# Shared by Preview and Dev (see ``t2i_generation_config.json``).
+_EOL_TOKEN_ID = 152064
+_VISUAL_TOKEN_START_ID = 152072
+_VISUAL_TOKEN_END_ID = 168456
+
+_AR_SYSTEM_PROMPT = "You are a helpful image generator."
+
+
+def build_text_to_image_prompt(
+    prompt: str,
+    negative_prompt: str | None,
+    height: int | None = None,
+    width: int | None = None,
+) -> dict[str, Any]:
+    """Build the MammothModa2 AR-stage prompt for text-to-image generation.
+
+    Reproduces the prompt string and the structural ``additional_information``
+    that the former bespoke example (``run_mammothmoda2_t2i.py``) constructed,
+    deriving the AR image grid from ``height`` / ``width``.
+
+    Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``,
+    ``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``.
+    The structural token ids are included in ``additional_information`` so the
+    AR sampler can apply the same constraints for both Preview and Dev.
+
+    MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale``
+    and has no explicit negative-prompt path, so ``negative_prompt`` is accepted
+    for signature compatibility but not injected.
+    """
+    h = height or 1024
+    w = width or 1024
+    ar_height = h // _PATCH_SIZE
+    ar_width = w // _PATCH_SIZE
+
+    return {
+        "prompt": (
+            f"<|im_start|>system\n{_AR_SYSTEM_PROMPT}<|im_end|>\n"
+            f"<|im_start|>user\n{prompt}<|im_end|>\n"
+            "<|im_start|>assistant\n"
+            f"<|image start|>{ar_width}*{ar_height}<|image token|>"
+        ),
+        "additional_information": {
+            "omni_task": ["t2i"],
+            "ar_width": [ar_width],
+            "ar_height": [ar_height],
+            "image_height": [h],
+            "image_width": [w],
+            "eol_token_id": [_EOL_TOKEN_ID],
+            "visual_token_start_id": [_VISUAL_TOKEN_START_ID],
+            "visual_token_end_id": [_VISUAL_TOKEN_END_ID],
+        },
+    }
+
+
+MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS = frozenset(
+    {
+        "text_guidance_scale",
+        "cfg_range",
+        # MammothModa2's DiT stage consumes inputs via the kwargs interface rather
+        # than OmniDiffusionRequest, so the standard --num-inference-steps flag does
+        # not reach it; it is routed through extra_body like the CFG knobs.
+        "num_inference_steps",
+    }
+)
+
+
+def build_x_to_text_prompt(
+    model: str,
+    prompt: str,
+    has_image: bool,
+) -> tuple[dict[str, Any], list[int] | None]:
+    """Build MammothModa2's chat prompt for text-output tasks."""
+    del model
+    vision = "<|vision_start|><|image_pad|><|vision_end|>" if has_image else ""
+    return (
+        {
+            "prompt": (
+                "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+                f"<|im_start|>user\n{vision}{prompt}<|im_end|>\n"
+                "<|im_start|>assistant\n"
+            ),
+            "modalities": ["text"],
+            "additional_information": {"omni_task": ["chat"]},
+        },
+        None,
+    )
+
+
+MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS = frozenset()
+MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES = True

vllm_omni/model_extras/registry.py

@@ -36,6 +36,17 @@ from vllm_omni.model_extras.ltx2 import (
     ltx_preserves_reference_image_size,
     ltx_transformer_config_subfolder,
 )
+from vllm_omni.model_extras.mammothmodal2_preview import (
+    MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS,
+    MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS,
+    MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
+)
+from vllm_omni.model_extras.mammothmodal2_preview import (
+    build_text_to_image_prompt as build_mammothmoda2_text_to_image_prompt,
+)
+from vllm_omni.model_extras.mammothmodal2_preview import (
+    build_x_to_text_prompt as build_mammothmoda2_x_to_text_prompt,
+)
 from vllm_omni.model_extras.ming_flash_omni import (
     MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
     MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
@@ -113,6 +124,7 @@ def default_x_to_text_prompt(
 _X_TO_TEXT_SPECS: dict[str, XToTextPromptBuilder] = {
     "bagel": build_bagel_x_to_text_prompt,
     "hunyuan_image3": build_hunyuan_x_to_text_prompt,
+    "mammoth_moda2": build_mammothmoda2_x_to_text_prompt,
 }
 
 
@@ -127,6 +139,8 @@ def get_x_to_text_model_family(model: str) -> str:
         return "bagel"
     if model_type == "hunyuan_image_3_moe" or any("hunyuanimage3" in value for value in architectures):
         return "hunyuan_image3"
+    if "mammoth" in model_type or any("mammothmoda2" in value for value in architectures):
+        return "mammoth_moda2"
     return "generic"
 
 
@@ -217,6 +231,12 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
         "extra_output_params": VACE_EXTRA_OUTPUT_PARAMS,
         "image_to_video_prompt_builder": build_vace_image_to_video_prompt,
     },
+    "MammothModa2DiTPipeline": {
+        "extra_body_params": MAMMOTHMODA2_PREVIEW_EXTRA_BODY_PARAMS,
+        "extra_output_params": MAMMOTHMODA2_PREVIEW_EXTRA_OUTPUT_PARAMS,
+        "init_extra_args_for_non_diffusion_stages": MAMMOTHMODA2_PREVIEW_INIT_EXTRA_ARGS_FOR_NON_DIFFUSION_STAGES,
+        "text_to_image_prompt_builder": build_mammothmoda2_text_to_image_prompt,
+    },
     "MingImagePipeline": {
         "extra_body_params": MING_FLASH_OMNI_EXTRA_BODY_PARAMS,
         "extra_output_params": MING_FLASH_OMNI_EXTRA_OUTPUT_PARAMS,
@@ -230,6 +250,12 @@ for model_class_name in ("LTX2Pipeline", "LTX2TwoStagePipeline"):
     _EXTRA_SPECS[model_class_name]["transformer_config_subfolder_resolver"] = ltx_transformer_config_subfolder
 
 
+# Multi-stage discovery reports the top-level wrapper rather than its DiT
+# submodule, so both names must resolve to the same request builders.
+_EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
+_EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
+
+
 def _get_spec(model_class_name: str | None) -> dict[str, Any] | None:
     if not model_class_name:
         return None

vllm_omni/platforms/npu/platform.py

@@ -179,10 +179,15 @@ class NPUOmniPlatform(OmniPlatform, NPUPlatform):
                 )
                 backend_upper = "FLASH_ATTN"
 
-            if backend_upper == "FLASH_ATTN" and find_spec("mindiesd"):
-                # The NPU FLASH_ATTN backend imports mindiesd lazily at first
-                # forward, but CANN snapshots the custom-op registry at the
-                # first custom-op regInfo lookup in the process (e.g. a
+            if backend_upper in ("FLASH_ATTN", "RAINFUSION_ATTN") and find_spec("mindiesd"):
+                # Eager-import mindiesd only for backends that actually reach
+                # mindiesd kernels: FLASH_ATTN directly, and RAINFUSION_ATTN
+                # via its dense FlashAttention fallback (used before
+                # start_step and on any layer without a sparsifiable video
+                # segment). Other backends (e.g. TORCH_SDPA) never touch
+                # mindiesd, so a broken optional install must not block them.
+                # CANN snapshots the custom-op registry at the first
+                # custom-op regInfo lookup in the process (e.g. a
                 # vllm-ascend custom op during model load/warmup). Import
                 # mindiesd here so its env.py prepends the mindiesd vendor
                 # dirs (aie_ascendc etc.) to ASCEND_CUSTOM_OPP_PATH before

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -831,9 +831,9 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin
                 inputs_embeds = None
 
             # -------------------------------------- Omni-new -------------------------------------------------
-            # Some generation-stage models require model-specific runtime
-            # information (such as image size and conditioning embeddings)
-            # even during the dummy profiling run that vLLM uses to
+            # Some generation-stage models (e.g. MammothModa2DiTPipeline) require
+            # model-specific runtime information (such as image size and conditioning
+            # embeddings) even during the dummy profiling run that vLLM uses to
             # estimate KV-cache capacity.  get_dummy_runtime_additional_information
             # provides placeholder values of the correct shape so that the profiling
             # run does not raise an error due to missing inputs.

vllm_omni/tokenizers/mammoth_moda2_tokenizer.py

@@ -0,0 +1,401 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by ByteDance Ltd. on 2025-09-30.
+#
+# Original file was released under Apache-2.0, with the full license text
+# available at https://www.apache.org/licenses/LICENSE-2.0
+#
+# This modified file is released under the same license.
+#
+# --- Upstream header preserved below ---
+#
+# Copyright 2025 The Qwen Team and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import base64
+import os
+import unicodedata
+from collections.abc import Collection
+
+import tiktoken
+from loguru import logger
+from transformers import AddedToken, PreTrainedTokenizer
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "mammothu.tiktoken",
+    "special_tokens_file": "mammothu_vision_tokens.txt",
+}
+
+PAT_STR = (
+    r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?"""
+    r"""[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+)
+ENDOFTEXT = "<|endoftext|>"
+IMSTART = "<|im_start|>"
+IMEND = "<|im_end|>"
+# as the default behavior is changed to allow special tokens in
+# regular texts, the surface forms of special tokens need to be
+# as different as possible to minimize the impact
+QWEN_SPECIAL_TOKENS = (
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<tool_call>",
+    "</tool_call>",
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    "<|fim_pad|>",
+    "<|repo_name|>",
+    "<|file_sep|>",
+)
+
+# align to qwen2.5 tokenizer length (151846)
+EXTRAS = [f"<|extra_{i}|>" for i in range(181)]  # 205 - 19[len(QWEN_SPECIAL_TOKENS)] - 5
+# align to qwen2.5 embedding size (152064)
+EXTRAS += [f"<|extra_margin_{i}|>" for i in range(152064 - 151846)]
+# append new token in gen embedding range
+EXTRAS += ["<|endofline|>", "<|endoffile|>", "<|gen_placeholder|>", "<|useless token|>", "<|beginoftext|>"]
+EXTRAS = tuple(EXTRAS)
+# changed to use actual index to avoid misconfiguration with vocabulary expansion
+SPECIAL_START_ID = 151643
+
+
+def _load_tiktoken_bpe(tiktoken_bpe_file: str) -> dict[bytes, int]:
+    with open(tiktoken_bpe_file, "rb") as f:
+        contents = f.read()
+    return {
+        base64.b64decode(token): int(rank) for token, rank in (line.split() for line in contents.splitlines() if line)
+    }
+
+
+class MammothUTokenizer(PreTrainedTokenizer):
+    """MammothU tokenizer."""
+
+    vocab_files_names = VOCAB_FILES_NAMES
+
+    def __init__(
+        self,
+        vocab_file: str,
+        special_tokens_file: str,
+        errors: str = "replace",
+        bos_token: str = "<|beginoftext|>",
+        eos_token: str = "<|endoftext|>",
+        pad_token: str = "<|endoftext|>",
+        img_token: str = "<|image token|>",
+        boi_token: str = "<|image start|>",
+        eoi_token: str = "<|image end|>",
+        eol_token: str = "<|endofline|>",
+        eof_token: str = "<|endoffile|>",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        # how to handle errors in decoding UTF-8 byte sequences
+        # use ignore if you are in streaming inference
+        self.errors = errors
+        self.mergeable_ranks = _load_tiktoken_bpe(vocab_file)
+        with open(special_tokens_file) as f:
+            vision_tokens = [t.strip() for t in f.readlines() if len(t.strip()) > 0]
+        SPECIAL_TOKENS = tuple(
+            enumerate(
+                (
+                    (
+                        ENDOFTEXT,
+                        IMSTART,
+                        IMEND,
+                    )
+                    + QWEN_SPECIAL_TOKENS
+                    + EXTRAS
+                    + tuple(vision_tokens)
+                ),
+                start=SPECIAL_START_ID,
+            )
+        )
+
+        self.special_tokens = {token: index for index, token in SPECIAL_TOKENS}
+        self.special_tokens_set = set(t for _, t in SPECIAL_TOKENS)
+
+        enc = tiktoken.Encoding(
+            "mammothu",
+            pat_str=PAT_STR,
+            mergeable_ranks=self.mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+
+        assert len(self.mergeable_ranks) + len(self.special_tokens) == enc.n_vocab, (
+            f"{len(self.mergeable_ranks) + len(self.special_tokens)} != {enc.n_vocab} in encoding"
+        )
+
+        self.decoder = {v: k for k, v in self.mergeable_ranks.items()}
+        self.decoder.update({v: k for k, v in self.special_tokens.items()})
+        self.tokenizer = enc
+        self.eod_id = self.tokenizer.eot_token
+        self.bos_token = bos_token
+        self.eos_token = eos_token
+        self.pad_token = pad_token
+        self.img_token = img_token
+        self.boi_token = boi_token
+        self.eoi_token = eoi_token
+        self.eol_token = eol_token
+        self.eof_token = eof_token
+        self.image_content_token = "<|image_pad|>"  # come from Qwen2.5-VL
+        self.gen_image_token = "<|gen_image_pad|>"
+        self.gen_image_placeholder_token = "<|gen_placeholder|>"
+        self.visual_tokens = ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"]
+        self.visual_tokens_ids = [self.get_vocab()[token] for token in self.visual_tokens]
+
+        self.vision_range = (self.get_vocab()[self.boi_token], self.tokenizer.n_vocab - 1)
+        logger.info(f"MammothUTokeniser Vision range: {self.vision_range}")
+
+    def __getstate__(self):
+        # for pickle lovers
+        state = self.__dict__.copy()
+        del state["tokenizer"]
+        return state
+
+    def __setstate__(self, state):
+        # tokenizer is not python native; don't pass it; rebuild it
+        self.__dict__.update(state)
+        enc = tiktoken.Encoding(
+            "mammothu",
+            pat_str=PAT_STR,
+            mergeable_ranks=self.mergeable_ranks,
+            special_tokens=self.special_tokens,
+        )
+        self.tokenizer = enc
+
+    def __len__(self) -> int:
+        return self.tokenizer.n_vocab
+
+    def get_vocab(self) -> dict[bytes | str, int]:
+        vocab = self.mergeable_ranks.copy()
+        vocab.update(self.special_tokens)
+        return vocab
+
+    @property
+    def gen_placeholder_id(self):
+        return self.get_vocab()[self.gen_image_placeholder_token]
+
+    def convert_tokens_to_ids(self, tokens: bytes | str | list[bytes | str]) -> list[int]:
+        if isinstance(tokens, (str, bytes)):
+            if tokens in self.special_tokens:
+                return self.special_tokens[tokens]
+            else:
+                return self.mergeable_ranks.get(tokens)
+
+        ids = []
+        for token in tokens:
+            if token in self.special_tokens:
+                ids.append(self.special_tokens[token])
+            else:
+                ids.append(self.mergeable_ranks.get(token))
+        return ids
+
+    def _add_tokens(
+        self,
+        new_tokens: list[str] | list[AddedToken],
+        special_tokens: bool = False,
+    ) -> int:
+        if not special_tokens and new_tokens:
+            raise ValueError("Adding regular tokens is not supported")
+
+        added_tokens = 0
+        for token in new_tokens:
+            surface_form = token.content if isinstance(token, AddedToken) else token
+            if surface_form in self.special_tokens_set:
+                # Token already exists in our special tokens set
+                added_tokens += 1
+            else:
+                logger.warning(f"Token {surface_form} is not in the predefined special tokens set and cannot be added")
+
+        return added_tokens
+
+    def add_special_tokens(self, special_tokens_dict: dict[str, str | AddedToken]) -> int:
+        """
+        Add special tokens to the tokenizer and update the special tokens mapping.
+        Only adds tokens that are already in the special_tokens_set.
+
+        Args:
+            special_tokens_dict: dictionary of special tokens to add.
+                The key is the token type and the value is the token to add.
+
+        Returns:
+            Number of tokens added to the vocabulary.
+        """
+        added_tokens = 0
+
+        for token_type, token in special_tokens_dict.items():
+            if token_type == "additional_special_tokens" and isinstance(token, list):
+                added_tokens += self._add_tokens(token, special_tokens=True)
+            else:
+                token_value = token.content if isinstance(token, AddedToken) else token
+                if token_value in self.special_tokens_set:
+                    setattr(self, token_type, token_value)
+                    added_tokens += 1
+                else:
+                    logger.warning(
+                        f"Token {token_value} is not in the predefined special tokens set and cannot be added"
+                    )
+
+        return added_tokens
+
+    def save_vocabulary(self, save_directory: str, **kwargs) -> tuple[str]:
+        """
+        Save only the vocabulary of the tokenizer (vocabulary).
+
+        Returns:
+            `tuple(str)`: Paths to the files saved.
+        """
+        regular_file_path = os.path.join(save_directory, self.vocab_files_names["vocab_file"])
+        with open(regular_file_path, "w", encoding="utf8") as w:
+            for k, v in self.mergeable_ranks.items():
+                line = base64.b64encode(k).decode("utf8") + " " + str(v) + "\n"
+                w.write(line)
+
+        excluded_special_tokens = set(
+            (
+                ENDOFTEXT,
+                IMSTART,
+                IMEND,
+            )
+            + EXTRAS
+        )
+        special_file_path = os.path.join(save_directory, self.vocab_files_names["special_tokens_file"])
+        with open(special_file_path, "w", encoding="utf8") as w:
+            for k in self.special_tokens:
+                if k not in excluded_special_tokens:
+                    print(k, file=w)
+
+        return (regular_file_path, special_file_path)
+
+    def tokenize(
+        self,
+        text: str,
+        allowed_special: set | str = "all",
+        disallowed_special: Collection | str = (),
+        **kwargs,
+    ) -> list[bytes | str]:
+        """
+        Converts a string in a sequence of tokens.
+
+        Args:
+            text (`str`):
+                The sequence to be encoded.
+            allowed_special (`Literal["all"]` or `set`):
+                The surface forms of the tokens to be encoded as special tokens in regular texts.
+                Default to "all".
+            disallowed_special (`Literal["all"]` or `Collection`):
+                The surface forms of the tokens that should not be in regular texts and trigger errors.
+                Default to an empty tuple.
+
+            kwargs (additional keyword arguments, *optional*):
+                Will be passed to the underlying model specific encode method.
+
+        Returns:
+            `list[bytes|str]`: The list of tokens.
+        """
+        tokens = []
+        text = unicodedata.normalize("NFC", text)
+
+        # this implementation takes a detour: text -> token id -> token surface forms
+        for t in self.tokenizer.encode(text, allowed_special=allowed_special, disallowed_special=disallowed_special):
+            tokens.append(self.decoder[t])
+
+        return tokens
+
+    def convert_tokens_to_string(self, tokens: list[bytes | str]) -> str:
+        """
+        Converts a sequence of tokens in a single string.
+        """
+        text = ""
+        temp = b""
+        for t in tokens:
+            if isinstance(t, str):
+                if temp:
+                    text += temp.decode("utf-8", errors=self.errors)
+                    temp = b""
+                text += t
+            elif isinstance(t, bytes):
+                temp += t
+            else:
+                raise TypeError("token should only be of type types or str")
+        if temp:
+            text += temp.decode("utf-8", errors=self.errors)
+        return text
+
+    @property
+    def vocab_size(self) -> int:
+        return self.tokenizer.n_vocab
+
+    def _convert_id_to_token(self, index: int) -> bytes | str:
+        """Converts an id to a token, special tokens included"""
+        if index in self.decoder:
+            return self.decoder[index]
+        raise ValueError("unknown ids")
+
+    def _convert_token_to_id(self, token: bytes | str) -> int:
+        """Converts a token to an id using the vocab, special tokens included"""
+        if token in self.special_tokens:
+            return self.special_tokens[token]
+        if token in self.mergeable_ranks:
+            return self.mergeable_ranks[token]
+        raise ValueError("unknown token")
+
+    def _decode(
+        self,
+        token_ids: int | list[int],
+        skip_special_tokens: bool = False,
+        errors: str | None = None,
+        **kwargs,
+    ) -> str:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+
+        if skip_special_tokens:
+            token_ids = [i for i in token_ids if i < self.eod_id]
+
+        return self.tokenizer.decode(token_ids, errors=errors or self.errors)
+
+    def bytes_to_str(self, byte_tokens: dict) -> str:
+        """Convert byte tokens to string representation.
+
+        Args:
+            byte_tokens: A dictionary where keys are byte objects and values are integers,
+                        or a single byte object.
+
+        Returns:
+            If input is a dictionary, returns a new dictionary with byte keys converted to strings.
+            If input is a single byte object, returns the string representation.
+        """
+        if isinstance(byte_tokens, dict):
+            return {k.decode("utf-8", errors=self.errors): v for k, v in byte_tokens.items() if isinstance(k, bytes)}
+        if isinstance(byte_tokens, bytes):
+            return byte_tokens.decode("utf-8", errors=self.errors)
+        return byte_tokens

vllm_omni/transformers_utils/configs/__init__.py

@@ -11,6 +11,10 @@ import importlib
 
 _CLASS_TO_MODULE: dict[str, str] = {
     "HiggsAudioV3Config": "vllm_omni.transformers_utils.configs.higgs_audio_v3",
+    "Mammothmoda2Config": "vllm_omni.transformers_utils.configs.mammoth_moda2",
+    "Mammothmoda2Qwen2_5_VLConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2",
+    "Mammothmoda2Qwen2_5_VLTextConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2",
+    "Mammothmoda2Qwen2_5_VLVisionConfig": "vllm_omni.transformers_utils.configs.mammoth_moda2",
     "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech",
     "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
     "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech",
@@ -32,6 +36,10 @@ _CLASS_TO_MODULE: dict[str, str] = {
 
 __all__ = [
     "HiggsAudioV3Config",
+    "Mammothmoda2Config",
+    "Mammothmoda2Qwen2_5_VLConfig",
+    "Mammothmoda2Qwen2_5_VLTextConfig",
+    "Mammothmoda2Qwen2_5_VLVisionConfig",
     "FishSpeechConfig",
     "FishSpeechSlowARConfig",
     "FishSpeechFastARConfig",
@@ -72,6 +80,7 @@ from vllm_omni.transformers_utils.configs import dots_tts as _dots_tts  # noqa:
 from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import glm_tts as _glm_tts  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import higgs_audio_v3 as _higgs_audio_v3  # noqa: F401, E402
+from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import ming_flash_omni as _ming_flash_omni  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import minimax_music3 as _minimax_music3  # noqa: F401, E402
 from vllm_omni.transformers_utils.configs import omnivoice as _omnivoice  # noqa: F401, E402

vllm_omni/transformers_utils/configs/mammoth_moda2.py

@@ -0,0 +1,377 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import ClassVar, Literal
+
+from transformers import AutoConfig, AutoTokenizer, PretrainedConfig
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
+    Qwen2_5_VLConfig,
+    Qwen2_5_VLTextConfig,
+    Qwen2_5_VLVisionConfig,
+)
+from transformers.models.qwen3_vl.configuration_qwen3_vl import (
+    Qwen3VLConfig,
+    Qwen3VLTextConfig,
+    Qwen3VLVisionConfig,
+)
+
+from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer
+
+__all__ = [
+    "Mammothmoda2Config",
+    "Mammothmoda2Qwen2_5_VLConfig",
+    "Mammothmoda2Qwen2_5_VLTextConfig",
+    "Mammothmoda2Qwen2_5_VLVisionConfig",
+    "Mammothmoda2Qwen3VLConfig",
+    "Mammothmoda2Qwen3VLTextConfig",
+    "Mammothmoda2Qwen3VLVisionConfig",
+]
+
+
+class Mammothmoda2Qwen2_5_VLVisionConfig(Qwen2_5_VLVisionConfig):
+    model_type = "mammothmoda2_qwen2_5_vl_vision"
+    base_config_key = "vision_config"
+
+    def __init__(
+        self,
+        depth: int = 32,
+        hidden_size: int = 3584,
+        hidden_act: str = "silu",
+        intermediate_size: int = 3420,
+        num_heads: int = 16,
+        in_channels: int = 3,
+        patch_size: int = 14,
+        spatial_merge_size: int = 2,
+        temporal_patch_size: int = 2,
+        tokens_per_second: int = 4,
+        window_size: int = 112,
+        out_hidden_size: int = 3584,
+        fullatt_block_indexes: list[int] | None = None,
+        initializer_range: float = 0.02,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            depth=depth,
+            hidden_size=hidden_size,
+            hidden_act=hidden_act,
+            intermediate_size=intermediate_size,
+            num_heads=num_heads,
+            in_channels=in_channels,
+            patch_size=patch_size,
+            spatial_merge_size=spatial_merge_size,
+            temporal_patch_size=temporal_patch_size,
+            tokens_per_second=tokens_per_second,
+            window_size=window_size,
+            out_hidden_size=out_hidden_size,
+            fullatt_block_indexes=fullatt_block_indexes or [7, 15, 23, 31],
+            initializer_range=initializer_range,
+            **kwargs,
+        )
+
+
+class Mammothmoda2Qwen2_5_VLTextConfig(Qwen2_5_VLTextConfig):
+    model_type = "mammothmoda2_qwen2_5_vl_text"
+    base_config_key = "text_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    base_model_tp_plan = Qwen2_5_VLTextConfig.base_model_tp_plan
+    base_model_pp_plan = Qwen2_5_VLTextConfig.base_model_pp_plan
+
+    def __init__(
+        self,
+        vocab_size: int = 152064,
+        hidden_size: int = 8192,
+        intermediate_size: int = 29568,
+        num_hidden_layers: int = 80,
+        num_attention_heads: int = 64,
+        num_key_value_heads: int | None = 8,
+        hidden_act: str = "silu",
+        max_position_embeddings: int = 32768,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-05,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = False,
+        rope_theta: float = 1000000.0,
+        use_sliding_window: bool = False,
+        sliding_window: int = 4096,
+        max_window_layers: int = 80,
+        layer_types: list[str] | None = None,
+        attention_dropout: float = 0.0,
+        rope_scaling: dict | None = None,
+        image_token_id: int | None = None,
+        video_token_id: int | None = None,
+        extra_gen_vocab: bool = True,
+        gen_vocab_size: int = 32800,
+        gen_vocab_start_index: int | None = None,
+        moe_type: str = "ffn",
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            intermediate_size=intermediate_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_key_value_heads,
+            hidden_act=hidden_act,
+            max_position_embeddings=max_position_embeddings,
+            initializer_range=initializer_range,
+            rms_norm_eps=rms_norm_eps,
+            use_cache=use_cache,
+            tie_word_embeddings=tie_word_embeddings,
+            rope_theta=rope_theta,
+            use_sliding_window=use_sliding_window,
+            sliding_window=sliding_window,
+            max_window_layers=max_window_layers,
+            layer_types=layer_types,
+            attention_dropout=attention_dropout,
+            rope_scaling=rope_scaling,
+            **kwargs,
+        )
+
+        self.extra_gen_vocab = extra_gen_vocab
+        self.gen_vocab_size = gen_vocab_size
+        self.moe_type = moe_type
+        if gen_vocab_start_index is None:
+            self.gen_vocab_start_index = (
+                self.vocab_size if self.extra_gen_vocab else self.vocab_size - self.gen_vocab_size
+            )
+        else:
+            self.gen_vocab_start_index = gen_vocab_start_index
+
+        # NOTE: vLLM uses `hf_text_config.vocab_size` for sampling parameter validation
+        # (e.g., allowed_token_ids). Although MammothModa2's gen vocab is implemented via
+        # independent gen_embed/gen_head, the overall vocab size should still cover the
+        # gen vocab token ID range from the perspective of "output logits dimension".
+        if self.extra_gen_vocab:
+            self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size)
+
+        # Extra token IDs for multi-modal placeholders.
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+
+
+class Mammothmoda2Qwen2_5_VLConfig(Qwen2_5_VLConfig):
+    """Combined configuration: text_config + vision_config."""
+
+    model_type = "mammothmoda2_qwen2_5_vl"
+    sub_configs = {
+        "vision_config": Mammothmoda2Qwen2_5_VLVisionConfig,
+        "text_config": Mammothmoda2Qwen2_5_VLTextConfig,
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config: dict | PretrainedConfig | None = None,
+        vision_config: dict | PretrainedConfig | None = None,
+        image_token_id: int = 151655,
+        video_token_id: int = 151656,
+        vision_start_token_id: int = 151652,
+        vision_end_token_id: int = 151653,
+        extra_gen_vocab: bool = True,
+        gen_vocab_size: int = 32800,
+        gen_vocab_start_index: int | None = None,
+        moe_type: str = "ffn",
+        **kwargs,
+    ) -> None:
+        if isinstance(vision_config, dict):
+            vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            vision_config = self.sub_configs["vision_config"]()
+
+        text_extra_kwargs = {
+            "extra_gen_vocab": extra_gen_vocab,
+            "gen_vocab_size": gen_vocab_size,
+            "moe_type": moe_type,
+            "gen_vocab_start_index": gen_vocab_start_index,
+        }
+        if isinstance(text_config, dict):
+            for key, val in text_extra_kwargs.items():
+                text_config.setdefault(key, val)
+            text_config = self.sub_configs["text_config"](**text_config)
+        elif text_config is None:
+            text_config = self.sub_configs["text_config"](**{**text_extra_kwargs, **kwargs})
+        elif isinstance(text_config, PretrainedConfig):
+            text_config = text_config
+
+        super().__init__(
+            text_config=text_config,
+            vision_config=vision_config,
+            image_token_id=image_token_id,
+            video_token_id=video_token_id,
+            vision_start_token_id=vision_start_token_id,
+            vision_end_token_id=vision_end_token_id,
+            **kwargs,
+        )
+
+        if not hasattr(self, "text_config"):
+            self.text_config = text_config
+        if not hasattr(self, "vision_config"):
+            self.vision_config = vision_config
+
+        self.image_token_id = image_token_id
+        self.video_token_id = video_token_id
+        self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab)
+        self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size)
+        self.moe_type = getattr(self.text_config, "moe_type", moe_type)
+        self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index)
+        self.tokenizer_class = "MammothUTokenizer"
+
+
+class Mammothmoda2Qwen3VLVisionConfig(Qwen3VLVisionConfig):
+    model_type = "mammothmoda2_qwen3_vl_vision"
+    base_config_key = "vision_config"
+
+
+class Mammothmoda2Qwen3VLTextConfig(Qwen3VLTextConfig):
+    model_type = "mammothmoda2_qwen3_vl_text"
+    base_config_key = "text_config"
+
+    def __init__(
+        self,
+        extra_gen_vocab: bool = True,
+        gen_vocab_size: int = 32800,
+        gen_vocab_start_index: int | None = None,
+        moe_type: str = "ffn",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.base_vocab_size = int(self.vocab_size)
+        self.extra_gen_vocab = extra_gen_vocab
+        self.gen_vocab_size = gen_vocab_size
+        self.gen_vocab_start_index = self.vocab_size if gen_vocab_start_index is None else gen_vocab_start_index
+        self.moe_type = moe_type
+        if self.extra_gen_vocab:
+            self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size)
+
+
+class Mammothmoda2Qwen3VLConfig(Qwen3VLConfig):
+    """MammothModa2-Dev AR config backed by Qwen3-VL."""
+
+    model_type = "mammothmoda2_qwen3_vl"
+    sub_configs = {
+        "vision_config": Mammothmoda2Qwen3VLVisionConfig,
+        "text_config": Mammothmoda2Qwen3VLTextConfig,
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config: dict | PretrainedConfig | None = None,
+        vision_config: dict | PretrainedConfig | None = None,
+        extra_gen_vocab: bool = True,
+        gen_vocab_size: int = 32800,
+        gen_vocab_start_index: int | None = None,
+        moe_type: str = "ffn",
+        **kwargs,
+    ) -> None:
+        if isinstance(vision_config, dict):
+            vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            vision_config = self.sub_configs["vision_config"]()
+        extras = {
+            "extra_gen_vocab": extra_gen_vocab,
+            "gen_vocab_size": gen_vocab_size,
+            "gen_vocab_start_index": gen_vocab_start_index,
+            "moe_type": moe_type,
+        }
+        if isinstance(text_config, dict):
+            for key, value in extras.items():
+                text_config.setdefault(key, value)
+            text_config = self.sub_configs["text_config"](**text_config)
+        elif text_config is None:
+            text_config = self.sub_configs["text_config"](**extras)
+        super().__init__(text_config=text_config, vision_config=vision_config, **kwargs)
+        self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab)
+        self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size)
+        self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index)
+        self.moe_type = getattr(self.text_config, "moe_type", moe_type)
+        self.tokenizer_class = "MammothUTokenizer"
+
+
+class Mammothmoda2Config(PretrainedConfig):
+    """Top-level MammothModa2 composition configuration"""
+
+    model_type = "mammothmoda2"
+    is_composition = True
+    sub_configs: ClassVar = {"llm_config": AutoConfig}
+
+    def __init__(
+        self,
+        *,
+        llm_config: dict | None = None,
+        gen_vae_config: dict | None = None,
+        gen_dit_config: dict | None = None,
+        gen_condition_mode: Literal["text", "image", "text_image"] = "image",
+        gen_image_condition_refiner_config: dict | None = None,
+        gen_axes_dim_rope: list[int] | None = None,
+        gen_axes_lens: list[int] | None = None,
+        gen_transport_config: dict | None = None,
+        initializer_range: float = 0.02,
+        architectures: list[str] | None = None,
+        **kwargs,
+    ) -> None:
+        # Must set llm_config before super().__init__() because the parent's
+        # validation calls get_text_config() which accesses self.llm_config
+        self.llm_config = AutoConfig.for_model(**llm_config) if llm_config is not None else None
+        self.gen_vae_config = gen_vae_config
+        self.gen_dit_config = gen_dit_config
+
+        self.gen_condition_mode = gen_condition_mode
+        self.gen_image_condition_refiner_config = gen_image_condition_refiner_config
+        self.gen_axes_dim_rope = gen_axes_dim_rope or [40, 40, 40]
+        self.gen_axes_lens = gen_axes_lens or [10000, 10000, 10000]
+        self.gen_transport_config = gen_transport_config or {}
+        self.initializer_range = initializer_range
+        self.tokenizer_class = "MammothUTokenizer"
+        self.architectures = ["Mammothmoda2Model"]
+
+        super().__init__(**kwargs)
+
+    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:  # noqa: ARG002
+        if self.llm_config is None:
+            return None
+        # llm_config is a Mammothmoda2Qwen2_5_VLConfig which has nested text_config
+        return self.llm_config.text_config
+
+    def _require_llm_config(self) -> PretrainedConfig:
+        if self.llm_config is None:
+            raise AttributeError("Mammothmoda2Config.llm_config is None")
+        return self.llm_config
+
+    # ---- Proxy attrs for vLLM multimodal/mrope helpers ----
+    # vllm_omni/model_executor/layers/mrope.py expects these fields on `hf_config`.
+    # In MammothModa2, they live in the nested `llm_config` (VL config), so we
+    # expose them here to make the top-level composition config compatible.
+    @property
+    def vision_config(self):
+        return self._require_llm_config().vision_config
+
+    @property
+    def image_token_id(self) -> int:
+        return int(self._require_llm_config().image_token_id)
+
+    @property
+    def video_token_id(self) -> int:
+        return int(self._require_llm_config().video_token_id)
+
+    @property
+    def vision_start_token_id(self) -> int:
+        return int(self._require_llm_config().vision_start_token_id)
+
+    @property
+    def vision_end_token_id(self) -> int:
+        return int(self._require_llm_config().vision_end_token_id)
+
+
+# Register model_type -> config class for AutoConfig
+AutoConfig.register(Mammothmoda2Config.model_type, Mammothmoda2Config)
+AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig)
+AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig)
+AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig)
+AutoConfig.register(Mammothmoda2Qwen3VLConfig.model_type, Mammothmoda2Qwen3VLConfig)
+AutoConfig.register(Mammothmoda2Qwen3VLTextConfig.model_type, Mammothmoda2Qwen3VLTextConfig)
+AutoConfig.register(Mammothmoda2Qwen3VLVisionConfig.model_type, Mammothmoda2Qwen3VLVisionConfig)
+# Register tokenizer_type -> the configs & AutoTokenizer
+AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer)
+AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer)
+AutoTokenizer.register(config_class=Mammothmoda2Qwen3VLConfig, slow_tokenizer_class=MammothUTokenizer)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -792,9 +792,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
 
-            # Some generation-stage models require model-specific runtime
-            # information (such as image size and conditioning embeddings)
-            # even during the dummy profiling run that vLLM uses to
+            # Some generation-stage models (e.g. MammothModa2DiTPipeline) require
+            # model-specific runtime information (such as image size and conditioning
+            # embeddings) even during the dummy profiling run that vLLM uses to
             # estimate KV-cache capacity.  get_dummy_runtime_additional_information
             # provides placeholder values of the correct shape so that the profiling
             # run does not raise an error due to missing inputs.

vllm_omni/worker/gpu_model_runner.py

@@ -1296,9 +1296,13 @@ class OmniGPUModelRunner(GPUModelRunner):
         per_req_runtime_info = []
         for req_id in self.input_batch.req_ids:
             req_state = self.requests.get(req_id)
-            # Per-request runtime metadata is stamped with the number of tokens
-            # already generated so models can apply step-dependent constraints
-            # (for example end-of-row tokens during image-grid decoding).
+            # MammothModa2 AR grid constraint: the model must emit a special
+            # end-of-line (EOL) token at the end of each image row.  To determine
+            # whether the current decoding step falls on a row boundary, the
+            # constraint logic (see MammothModa2ARForConditionalGeneration.
+            # _apply_t2i_token_constraints) computes:
+            #   column_id = generated_len % (ar_width + 1)
+            # and forces the EOL token when column_id == ar_width.
             generated_len = len(req_state.output_token_ids) if req_state is not None else 0
             info = self.model_intermediate_buffer.get(req_id, {})
             if info: