[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/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**           |     ❌     |     ✅      |           ❌           |       ✅        |         ❌         |          ❌          |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |

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:

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/config/test_config_factory.py

@@ -1314,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/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/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/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,
@@ -168,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/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/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/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: