add support for MammothModa2 model (#336)

Signed-off-by: HonestDeng <2958906959@qq.com>
Signed-off-by: iwzbi <iwzbi@zju.edu.cn>
Signed-off-by: iwzbi <wzbi@zju.edu.cn>
Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: catcat <108673086+iwzbi@users.noreply.github.com>
Co-authored-by: iwzbi <iwzbi@zju.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: dsinghvi <divyanshsinghvi@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>

docs/models/supported_models.md

@@ -35,6 +35,7 @@ th {
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
+|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |

examples/offline_inference/mammothmodal2_preview/README.md

@@ -0,0 +1,32 @@
+# MammothModa2-Preview
+
+## Run examples (MammothModa2-Preview)
+
+Download model
+```bash
+hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
+```
+
+### Text-to-Image (T2I)
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+  --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+  --height 1024 \
+  --width 1024 \
+  --num-inference-steps 50 \
+  --text-guidance-scale 4.0 \
+  --out output.png
+```
+
+### Image Summary
+
+```bash
+python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
+  --model ./MammothModa2-Preview \
+  --stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
+  --question "Summarize this image." \
+  --image ./image.png
+```

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -0,0 +1,115 @@
+"""
+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 \
+    --stage-config vllm_omni/model_executor/stage_configs/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(
+        "--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
+    )
+    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")
+    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,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=args.trust_remote_code,
+    )
+    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:
+        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
+        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
+        for ro in req_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/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -0,0 +1,251 @@
+"""
+Offline inference example for MammothModa2 Text-to-Image (T2I) generation.
+This script uses the vllm_omni.Omni pipeline with a multi-stage configuration.
+
+Workflow:
+1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states.
+2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion
+   and VAE decoding to produce the final image.
+
+Example Usage:
+    uv run python examples/offline_inference/run_mammothmoda2_t2i.py \
+        --model path/to/MammothModa2-Preview \
+        --stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
+        --prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
+        --out output.png
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import logging
+import os
+from pathlib import Path
+from typing import NamedTuple
+
+import torch
+from PIL import Image
+from vllm.sampling_params import SamplingParams
+
+from vllm_omni import Omni
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+_PATCH_SIZE = 16  # AR image grid patch size (pixels per token)
+
+
+class T2IGenConfig(NamedTuple):
+    eol_token_id: int
+    visual_token_start_id: int
+    visual_token_end_id: int
+    top_k: int  # AR sampling top-k (covers the full visual generation vocabulary)
+    # Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|>
+    visual_ids: list[int]
+
+
+def load_t2i_generation_config(model_dir: str) -> T2IGenConfig:
+    """Load T2I token IDs from t2i_generation_config.json and config.json.
+
+    Supports both local directory paths and HuggingFace Hub model IDs.
+    """
+    model_path = Path(model_dir)
+
+    def _read_json(filename: str) -> dict:
+        local = model_path / filename
+        if local.exists():
+            with local.open(encoding="utf-8") as f:
+                return json.load(f)
+        # Fall back to HuggingFace Hub when model_dir is a repo ID.
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError(
+                "huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub"
+            ) from exc
+        cached = hf_hub_download(repo_id=model_dir, filename=filename)
+        with open(cached, encoding="utf-8") as f:
+            return json.load(f)
+
+    gen_cfg = _read_json("t2i_generation_config.json")
+    llm_cfg = _read_json("config.json").get("llm_config", {})
+
+    return T2IGenConfig(
+        eol_token_id=int(gen_cfg["eol_token_id"]),
+        visual_token_start_id=int(gen_cfg["visual_token_start_id"]),
+        visual_token_end_id=int(gen_cfg["visual_token_end_id"]),
+        top_k=int(gen_cfg["top_k"]),
+        visual_ids=[
+            int(llm_cfg["image_token_id"]),
+            int(llm_cfg["video_token_id"]),
+            int(llm_cfg["vision_start_token_id"]),
+            int(llm_cfg["vision_end_token_id"]),
+        ],
+    )
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.")
+    p.add_argument("--model", type=str, required=True, help="Path to the model directory.")
+    p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.")
+    p.add_argument(
+        "--prompt",
+        type=str,
+        action="append",
+        default=None,
+        help=(
+            "Text prompt for image generation. Can be provided multiple times "
+            "to generate multiple images with shared height/width/CFG settings."
+        ),
+    )
+    p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).")
+    p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).")
+    p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.")
+    p.add_argument(
+        "--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT."
+    )
+    p.add_argument(
+        "--cfg-range",
+        type=float,
+        nargs=2,
+        default=(0.0, 1.0),
+        help="Relative step range [start, end] where CFG is active.",
+    )
+    p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
+    p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
+    args = p.parse_args()
+    if not args.prompt:
+        args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
+    return args
+
+
+def tensor_to_pil(image: torch.Tensor) -> Image.Image:
+    """Convert a normalized torch tensor [-1, 1] to a PIL Image."""
+    if image.ndim == 4:
+        image = image[0]
+    image = image.detach().to("cpu")
+    image = (image / 2 + 0.5).clamp(0, 1)
+    image = (image * 255).to(torch.uint8)
+    image = image.permute(1, 2, 0).contiguous().numpy()
+    return Image.fromarray(image)
+
+
+def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
+    """Build the AR-stage prompt string including the image grid header."""
+    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 _collect_images(outputs: list) -> list[torch.Tensor]:
+    """Extract all image tensors produced by the final (DiT) stage."""
+    images: list[torch.Tensor] = []
+    for out in outputs:
+        ro_list = getattr(out, "request_output", out)
+        if not isinstance(ro_list, list):
+            ro_list = [ro_list]
+        for ro_item in ro_list:
+            for completion in getattr(ro_item, "outputs", None) or []:
+                mm = getattr(completion, "multimodal_output", None)
+                if not isinstance(mm, dict) or "image" not in mm:
+                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
+                payload = mm["image"]
+                for tensor in payload if isinstance(payload, list) else [payload]:
+                    if not isinstance(tensor, torch.Tensor):
+                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                    images.append(tensor)
+    return images
+
+
+def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]:
+    """Save image tensors to disk.
+
+    Single image: written to *out_path* exactly.
+    Multiple images: suffixed as ``<base>_0<ext>``, ``<base>_1<ext>``, …
+    """
+    if not images:
+        raise RuntimeError("No images to save.")
+    base, ext = os.path.splitext(out_path)
+    ext = ext or ".png"
+    paths = []
+    for i, tensor in enumerate(images):
+        path = out_path if len(images) == 1 else f"{base}_{i}{ext}"
+        tensor_to_pil(tensor).save(path)
+        paths.append(path)
+    return paths
+
+
+def main() -> None:
+    args = parse_args()
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+
+    if args.height <= 0 or args.width <= 0:
+        raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}")
+    if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0:
+        raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}")
+
+    ar_height = args.height // _PATCH_SIZE
+    ar_width = args.width // _PATCH_SIZE
+    gen_cfg = load_t2i_generation_config(args.model)
+    expected_grid_tokens = ar_height * (ar_width + 1)
+
+    logger.info("Initializing Omni pipeline...")
+    omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code)
+    try:
+        ar_sampling = SamplingParams(
+            temperature=1.0,
+            top_p=1.0,
+            top_k=gen_cfg.top_k,
+            max_tokens=max(1, expected_grid_tokens + 1),  # +1 for hidden state of eoi
+            detokenize=False,
+        )
+        dit_sampling = SamplingParams(
+            temperature=0.0,
+            top_p=1.0,
+            top_k=-1,
+            max_tokens=1,
+            detokenize=False,
+        )
+
+        additional_information = {
+            "omni_task": ["t2i"],
+            "ar_width": [ar_width],
+            "ar_height": [ar_height],
+            "eol_token_id": [gen_cfg.eol_token_id],
+            "visual_token_start_id": [gen_cfg.visual_token_start_id],
+            "visual_token_end_id": [gen_cfg.visual_token_end_id],
+            "image_height": [args.height],
+            "image_width": [args.width],
+            "num_inference_steps": [args.num_inference_steps],
+            "text_guidance_scale": [args.text_guidance_scale],
+            "cfg_range": [args.cfg_range[0], args.cfg_range[1]],
+            "visual_ids": gen_cfg.visual_ids,
+        }
+        inputs = [
+            {
+                "prompt": _format_prompt(p, ar_width, ar_height),
+                "additional_information": dict(additional_information),
+            }
+            for p in args.prompt
+        ]
+
+        logger.info("Starting generation...")
+        # omni.generate() returns a Generator; consume it to run the full pipeline.
+        outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling]))
+
+        logger.info("Post-processing and saving image(s)...")
+        for path in _save_images(_collect_images(outputs), args.out):
+            logger.info(f"Saved: {path}")
+    finally:
+        omni.close()
+
+
+if __name__ == "__main__":
+    main()

tests/e2e/offline_inference/test_mammoth_moda2.py

@@ -0,0 +1,215 @@
+"""
+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 weights: env var ``MAMMOTHMODA2_MODEL_PATH``
+  (default: <repo_root>/MammothModa2-Preview)
+Stage config:  env var ``MAMMOTHMODA2_T2I_STAGE_CONFIG``
+  (default: vllm_omni/model_executor/stage_configs/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.py``
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+import pytest
+import torch
+from vllm.sampling_params import SamplingParams
+
+from tests.utils import hardware_test
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+# ---------------------------------------------------------------------------
+# 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
+
+_STAGE_CONFIGS_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "model_executor" / "stage_configs"
+MODEL_PATH = os.environ.get(
+    "MAMMOTHMODA2_MODEL_PATH",
+    str(Path(__file__).resolve().parents[3] / "MammothModa2-Preview"),
+)
+T2I_STAGE_CONFIG = os.environ.get(
+    "MAMMOTHMODA2_T2I_STAGE_CONFIG",
+    str(_STAGE_CONFIGS_DIR / "mammoth_moda2.yaml"),
+)
+
+# 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(model_dir: str) -> dict:
+    cfg_path = Path(model_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
+# ---------------------------------------------------------------------------
+@pytest.mark.core_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "L4"})
+def test_mammothmoda2_t2i_e2e():
+    """
+    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``).
+    """
+    from vllm_omni import Omni
+
+    if not Path(MODEL_PATH).exists():
+        pytest.skip(f"Model weights not found at {MODEL_PATH}")
+    if not Path(T2I_STAGE_CONFIG).exists():
+        pytest.skip(f"Stage config not found at {T2I_STAGE_CONFIG}")
+
+    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(model=MODEL_PATH, stage_configs_path=T2I_STAGE_CONFIG, trust_remote_code=True)
+    try:
+        # Greedy / deterministic sampling so pixel values are reproducible.
+        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"
+
+        # Find the image tensor produced by the DiT stage and compare pixels.
+        found_image = False
+        for out in outputs:
+            ro_list = getattr(out, "request_output", 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"
+    finally:
+        omni.close()

vllm_omni/__init__.py

@@ -20,6 +20,8 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
     # Allow importing vllm_omni without vllm (e.g., documentation builds)
     patch = None  # type: ignore
 
+# Register custom configs (AutoConfig, AutoTokenizer) as early as possible.
+from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
 
 from .config import OmniModelConfig
 from .entrypoints.async_omni import AsyncOmni

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -0,0 +1,809 @@
+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,
+        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__()
+        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(hidden_size, eps=norm_eps),
+            nn.Linear(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,
+    ):
+        device = hidden_states.device
+        p = self.config.patch_size
+
+        temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype)
+
+        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,
+            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,
+        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,
+            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,
+        )
+
+        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,298 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+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.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+
+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
+
+
+class MammothModa2DiTPipeline(nn.Module):
+    """
+    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.
+
+    """
+
+    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,
+        }
+    )
+
+    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_hidden_size = int(getattr(self.config.llm_config, "hidden_size", 0) or 0)
+        if llm_hidden_size <= 0:
+            raise ValueError("Failed to infer llm hidden_size from Mammothmoda2Config.llm_config.hidden_size")
+        self._reinit_caption_embedder(llm_hidden_size)
+
+        # Optional: image condition refiner (Q-Former)
+        if self.config.gen_image_condition_refiner_config is not None:
+            self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
+                hidden_size=llm_hidden_size,
+                **self.config.gen_image_condition_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,
+        )
+
+    @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]
+        text_cond = info["text_prompt_embeds"]
+        image_cond = info["image_prompt_embeds"]
+        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]
+        text_guidance_scale = info["text_guidance_scale"][0]
+        cfg_range = info["cfg_range"][0], info["cfg_range"][1]
+        num_inference_steps = info["num_inference_steps"][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")
+        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,
+            )
+
+        prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
+        prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+
+        # 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,
+                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,250 @@
+# 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
+
+
+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.float32 if torch.backends.mps.is_available() else torch.float64
+        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,17 @@
+from transformers import AutoTokenizer
+
+from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import (  # noqa: F401 registers AutoConfig
+    Mammothmoda2Config,
+    Mammothmoda2Qwen2_5_VLConfig,
+)
+
+from .mammoth_moda2 import MammothModa2ARForConditionalGeneration
+
+# AutoConfig.register is done inside transformers_utils/configs/mammoth_moda2.py
+AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer)
+AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer)
+
+__all__ = [
+    "MammothModa2ARForConditionalGeneration",
+]

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -0,0 +1,826 @@
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable, Mapping
+from itertools import islice
+from typing import Any
+
+import torch
+from torch import nn
+from transformers import Qwen2Config
+from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
+from vllm.config import CacheConfig, VllmConfig
+from vllm.distributed import get_pp_group
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    default_weight_loader,
+    maybe_remap_kv_scale_name,
+)
+from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
+from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP
+from vllm.model_executor.models.qwen2_5_vl import (
+    Qwen2_5_VLDummyInputsBuilder,
+    Qwen2_5_VLForConditionalGeneration,
+    Qwen2_5_VLMultiModalProcessor,
+    Qwen2_5_VLProcessingInfo,
+)
+from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
+from vllm.model_executor.models.utils import (
+    PPMissingLayer,
+    WeightsMapper,
+    init_vllm_registered_model,
+    is_pp_missing_parameter,
+    make_layers,
+    maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.sequence import IntermediateTensors
+from vllm.transformers_utils.config import (
+    is_interleaved,
+    patch_rope_parameters,
+    set_default_rope_theta,
+)
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
+from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
+
+
+def moe_enable(moe_type, layer_type, layer_idx) -> bool:
+    """Determine if MoE should be enabled for a specific layer type and index.
+
+    Args:
+        moe_type (str): The MoE configuration string (e.g., "ffn", "ffn_attention-14:28").
+        layer_type (str): The type of layer being checked (e.g., "ffn", "attention").
+        layer_idx (int): The index of the current layer.
+
+    Returns:
+        bool: True if MoE should be enabled for this layer, False otherwise.
+    """
+    if ":" in moe_type:
+        # moe_type like ffn_attention-14:28
+        moe_type, layers = moe_type.split("-")
+        start, end = [int(n) for n in layers.split(":")]
+    else:
+        start, end = 0, float("inf")
+    assert moe_type in ["none", "attention", "ffn", "ffn_attention"]
+    return layer_type in moe_type and start <= layer_idx < end
+
+
+def moe_forward(
+    hidden_states: torch.Tensor,
+    und_expert: Callable[[torch.Tensor], torch.Tensor],
+    gen_expert: Callable[[torch.Tensor], torch.Tensor] | None,
+    gen_token_mask: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Perform Mixture-of-Experts (MoE) routing and forward pass.
+
+    This function routes tokens to either the understanding expert (`und_expert`) or the
+    generation expert (`gen_expert`) based on the `gen_token_mask`.
+
+    Routing Logic:
+    - If `gen_expert` is None: All tokens go to `und_expert`.
+    - If `gen_token_mask` is None or all False: All tokens go to `und_expert`.
+    - If `gen_token_mask` is all True: All tokens go to `gen_expert`.
+    - Otherwise (mixed batch):
+        - Tokens where `gen_token_mask` is True go to `gen_expert`.
+        - Tokens where `gen_token_mask` is False go to `und_expert`.
+        - Results are concatenated and reordered to match the original input order.
+
+    Args:
+        hidden_states (torch.Tensor): Input hidden states. Shape: `(batch_size, seq_len, hidden_size)`
+            or `(num_tokens, hidden_size)`.
+        und_expert (Callable): The expert module for understanding/text tokens.
+            Takes `(N, D)` tensor, returns `(N, D_out)`.
+        gen_expert (Callable | None): The expert module for generation/image tokens, or None.
+            If provided, takes `(N, D)` tensor, returns `(N, D_out)`.
+        gen_token_mask (torch.Tensor | None): Boolean mask indicating generation tokens.
+            Shape matches `hidden_states` (excluding feature dim). True for generation tokens.
+
+    Returns:
+        torch.Tensor: The processed hidden states with the same shape as input (except potentially
+        different feature dimension if `D_out != D`).
+    """
+    if gen_expert is None:
+        return und_expert(hidden_states)
+
+    if gen_token_mask is None or not gen_token_mask.any():
+        return und_expert(hidden_states)
+    if gen_token_mask.all():
+        return gen_expert(hidden_states)
+
+    if hidden_states.ndim == 2:
+        flat_hid = hidden_states
+        d_model = hidden_states.shape[-1]
+        total_tokens = hidden_states.shape[0]
+    elif hidden_states.ndim == 3:
+        d_model = hidden_states.shape[-1]
+        flat_hid = hidden_states.reshape(-1, d_model)  # (B*L, D)
+        total_tokens = flat_hid.shape[0]
+    else:
+        raise ValueError(f"Unexpected hidden_states shape: {tuple(hidden_states.shape)}")
+
+    # Validate before reshape to catch shape mismatches where numel() would
+    # coincidentally match after flattening dimensions of different sizes.
+    if gen_token_mask.numel() != total_tokens:  # type: ignore[union-attr]
+        raise ValueError(
+            "gen_token_mask shape mismatch: "
+            f"mask={tuple(gen_token_mask.shape)}, hidden_states={tuple(hidden_states.shape)}"
+        )
+    # mask: [num_tokens] or [B, L] -> flatten to [total_tokens]
+    flat_mask = gen_token_mask.reshape(-1)  # type: ignore[union-attr]
+    gen_pos = torch.where(flat_mask)[0]
+    und_pos = torch.where(~flat_mask)[0]
+    permute_order = torch.cat([gen_pos, und_pos], dim=0)
+    inverse_order = torch.argsort(permute_order)
+    gen_token_num = int(flat_mask.sum().item())
+    gen_hid, und_hid = flat_hid[permute_order].split([gen_token_num, total_tokens - gen_token_num], dim=0)
+
+    # 1.1 Generation tokens (True)
+    gen_out = gen_expert(gen_hid)  # (N_gen, D)
+
+    # 1.2 Understanding tokens (False)
+    und_out = und_expert(und_hid)  # (N_und, D)
+    out_dim = und_out.shape[-1]
+
+    merged = torch.cat([gen_out, und_out], dim=0)
+    merged = merged[inverse_order]
+
+    if hidden_states.ndim == 2:
+        return merged.view(total_tokens, out_dim).contiguous()
+    return merged.view(*hidden_states.shape[:-1], out_dim).contiguous()
+
+
+class Mammothmoda2Processor(Qwen2_5_VLProcessor):
+    """Qwen2.5-VL Processor with MammothU tokenizer."""
+
+    tokenizer_class = ("MammothUTokenizer", None)
+
+
+class MammothModa2ARProcessingInfo(Qwen2_5_VLProcessingInfo):
+    """Processes multi-modal information for MammothModa2 AR, returning the VL sub-configuration."""
+
+    def get_hf_config(self):
+        mammoth_cfg: Mammothmoda2Config = self.ctx.get_hf_config(Mammothmoda2Config)
+        llm_cfg = getattr(mammoth_cfg, "llm_config", None)
+        return llm_cfg
+
+    def get_hf_processor(self, **kwargs: object) -> Mammothmoda2Processor:
+        return self.ctx.get_hf_processor(
+            Mammothmoda2Processor,
+            use_fast=kwargs.pop("use_fast", True),
+            **kwargs,
+        )
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        # MammothModa2 currently supports only image input, not video.
+        return {"image": None}
+
+    def get_data_parser(self) -> Qwen2VLMultiModalDataParser:
+        # vLLM >=0.16 expects the parser to be provided by ProcessingInfo,
+        # not by BaseMultiModalProcessor._get_data_parser.
+        return Qwen2VLMultiModalDataParser(
+            spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
+            expected_hidden_size=self._get_expected_hidden_size(),
+        )
+
+
+class MammothModa2ARDummyInputsBuilder(Qwen2_5_VLDummyInputsBuilder):
+    """Reuse Qwen2.5-VL's dummy input generation logic."""
+
+
+class MammothModa2ARMultiModalProcessor(Qwen2_5_VLMultiModalProcessor):
+    """Reuse Qwen2.5-VL's multi-modal processing,"""
+
+
+class Mammoth2DecoderLayer(Qwen2DecoderLayer):
+    def __init__(
+        self,
+        config: Qwen2Config,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        # Must patch rope parameters on config BEFORE calling super().__init__,
+        # because Qwen2DecoderLayer.__init__ creates Qwen2Attention using the
+        # config's rope settings.  Patching afterwards would leave the attention
+        # module initialised with wrong parameters, causing shape mismatches in
+        # rotary_embedding at runtime.
+        patch_rope_parameters(config)
+        set_default_rope_theta(config, default_theta=1000000)
+        super().__init__(config, cache_config, quant_config, prefix)
+
+        self.moe_enable = moe_enable(config.moe_type, "ffn", layer_idx)
+        if self.moe_enable:
+            self.gen_mlp = Qwen2MLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.gen_mlp",
+            )
+        else:
+            self.gen_mlp = None
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+        gen_token_mask: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Self Attention
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+        hidden_states = self.self_attn(
+            positions=positions,
+            hidden_states=hidden_states,
+        )
+
+        # Fully Connected
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = moe_forward(hidden_states, self.mlp, self.gen_mlp, gen_token_mask)
+        return hidden_states, residual
+
+
+class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
+    def __init__(
+        self, *, vllm_config: VllmConfig, prefix: str = "", decoder_layer_type: type[nn.Module] = Mammoth2DecoderLayer
+    ):
+        super().__init__()
+
+        hf_config = vllm_config.model_config.hf_config
+        if hasattr(hf_config, "get_text_config"):
+            config = hf_config.get_text_config()
+        elif hasattr(hf_config, "text_config"):
+            config = hf_config.text_config
+        else:
+            config = hf_config
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        self.prefix = prefix
+
+        if is_interleaved(vllm_config.model_config.hf_text_config):
+            assert config.max_window_layers == config.num_hidden_layers, (
+                "Sliding window for some but all layers is not supported. "
+                f"This model uses sliding window but `max_window_layers` = {config.max_window_layers} "
+                f"is less than `num_hidden_layers` = {config.num_hidden_layers}. Please open an issue "
+                "to discuss this feature."
+            )
+
+        self.config = config
+        self.quant_config = quant_config
+        # NOTE: MammothModa2 supports extra generation vocabulary (for image tokens).
+        # Token ID range: [gen_vocab_start_index, gen_vocab_start_index + gen_vocab_size).
+        # vLLM sampler/processor expects "last dimension of logits == model_config.get_vocab_size()",
+        # so we output base+gen logits in compute_logits, and embeddings must accept these IDs.
+        self.extra_gen_vocab = bool(getattr(config, "extra_gen_vocab", False))
+        # Starting index for generation tokens (used for gen_token_mask).
+        self.gen_vocab_start_index = getattr(hf_config, "gen_vocab_start_index", None) or getattr(
+            config, "gen_vocab_start_index", None
+        )
+        self.gen_vocab_size = int(getattr(config, "gen_vocab_size", 0) or 0)
+
+        self.base_vocab_size = int(self.gen_vocab_start_index) if self.extra_gen_vocab else int(config.vocab_size)
+        # The configuration level (hf_text_config.vocab_size) has been extended to base+gen
+        # by the upstream config class. Use config.vocab_size as the total vocab size.
+        self.total_vocab_size = int(getattr(config, "vocab_size", self.base_vocab_size))
+
+        if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
+            self.embed_tokens = VocabParallelEmbedding(
+                self.base_vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=f"{prefix}.embed_tokens",
+            )
+        else:
+            self.embed_tokens = PPMissingLayer()
+
+        if self.extra_gen_vocab:
+            if get_pp_group().is_first_rank or (config.tie_word_embeddings and get_pp_group().is_last_rank):
+                self.gen_embed_tokens = VocabParallelEmbedding(
+                    self.gen_vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.gen_embed_tokens",
+                )
+            else:
+                self.gen_embed_tokens = PPMissingLayer()
+        else:
+            self.gen_embed_tokens = None
+
+        if get_pp_group().is_last_rank:
+            self.lm_head = ParallelLMHead(
+                self.base_vocab_size,
+                config.hidden_size,
+                quant_config=quant_config,
+                prefix=f"{prefix}.lm_head",
+            )
+        else:
+            self.lm_head = PPMissingLayer()
+
+        if self.extra_gen_vocab:
+            if get_pp_group().is_last_rank:
+                self.gen_head = ParallelLMHead(
+                    self.gen_vocab_size,
+                    config.hidden_size,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.gen_head",
+                )
+            else:
+                self.gen_head = PPMissingLayer()
+        else:
+            self.gen_head = None
+
+        self.logits_processor = LogitsProcessor(self.base_vocab_size)
+        self.gen_logits_processor = LogitsProcessor(self.gen_vocab_size) if self.extra_gen_vocab else None
+
+        decoder_layer_type = decoder_layer_type or Mammoth2DecoderLayer
+
+        def _make_decoder_layer(*, prefix: str) -> nn.Module:
+            try:
+                layer_idx = int(prefix.rsplit(".", 1)[-1])
+            except Exception:
+                layer_idx = 0
+            return decoder_layer_type(
+                config=config,
+                layer_idx=layer_idx,
+                cache_config=cache_config,
+                quant_config=quant_config,
+                prefix=prefix,
+            )
+
+        self.start_layer, self.end_layer, self.layers = make_layers(
+            config.num_hidden_layers,
+            _make_decoder_layer,
+            prefix=f"{prefix}.layers",
+        )
+
+        def _make_empty_intermediate_tensors(
+            batch_size: int,
+            dtype: torch.dtype,
+            device: torch.device,
+        ) -> IntermediateTensors:
+            return IntermediateTensors(
+                {
+                    "hidden_states": torch.zeros((batch_size, config.hidden_size), dtype=dtype, device=device),
+                    "residual": torch.zeros((batch_size, config.hidden_size), dtype=dtype, device=device),
+                    "gen_token_mask": torch.zeros((batch_size,), dtype=torch.bool, device=device),
+                }
+            )
+
+        self.make_empty_intermediate_tensors = _make_empty_intermediate_tensors
+        if get_pp_group().is_last_rank:
+            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        else:
+            self.norm = PPMissingLayer()
+
+    @property
+    def model(self) -> MammothModa2Qwen2ForCausalLM:
+        return self
+
+    def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
+        if not self.extra_gen_vocab or self.gen_embed_tokens is None:
+            return self.embed_tokens(input_ids)
+
+        gen_mask = input_ids >= int(self.gen_vocab_start_index)
+        if not gen_mask.any():
+            return self.embed_tokens(input_ids)
+        if gen_mask.all():
+            gen_ids = input_ids - int(self.gen_vocab_start_index)
+            return self.gen_embed_tokens(gen_ids)
+
+        flat_ids = input_ids.reshape(-1)
+        flat_mask = gen_mask.reshape(-1)
+        out = torch.empty(
+            (flat_ids.shape[0], self.config.hidden_size),
+            dtype=self.embed_tokens.weight.dtype,  # type: ignore[attr-defined]
+            device=flat_ids.device,
+        )
+
+        base_pos = torch.where(~flat_mask)[0]
+        gen_pos = torch.where(flat_mask)[0]
+        if base_pos.numel() > 0:
+            out[base_pos] = self.embed_tokens(flat_ids[base_pos])
+        if gen_pos.numel() > 0:
+            gen_ids = flat_ids[gen_pos] - int(self.gen_vocab_start_index)
+            out[gen_pos] = self.gen_embed_tokens(gen_ids)
+        return out.view(*input_ids.shape, -1).contiguous()
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.get_input_embeddings(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor | IntermediateTensors:
+        if get_pp_group().is_first_rank:
+            if inputs_embeds is not None:
+                hidden_states = inputs_embeds
+            else:
+                assert input_ids is not None
+                hidden_states = self.get_input_embeddings(input_ids)
+            # gen_token_mask: True indicates image generation tokens, which use gen_mlp.
+            # In vLLM v1 path, only inputs_embeds might be provided, with input_ids set to None.
+            # In this case, gen tokens cannot be distinguished by ID, falling back to und_expert.
+            if self.gen_vocab_start_index is None or input_ids is None:
+                gen_token_mask = None
+            else:
+                gen_token_mask = input_ids >= self.gen_vocab_start_index
+            residual = None
+        else:
+            assert intermediate_tensors is not None
+            hidden_states = intermediate_tensors["hidden_states"]
+            residual = intermediate_tensors["residual"]
+            gen_token_mask = intermediate_tensors.tensors.get("gen_token_mask")
+
+        for idx, layer in enumerate(islice(self.layers, self.start_layer, self.end_layer)):
+            hidden_states, residual = layer(positions, hidden_states, residual, gen_token_mask)
+
+        if not get_pp_group().is_last_rank:
+            tensors = {"hidden_states": hidden_states, "residual": residual}
+            if gen_token_mask is not None:
+                tensors["gen_token_mask"] = gen_token_mask
+            return IntermediateTensors(tensors)
+
+        hidden_states, _ = self.norm(hidden_states, residual)
+
+        return hidden_states
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        if isinstance(self.lm_head, PPMissingLayer):
+            return None
+        base_logits = self.logits_processor(self.lm_head, hidden_states)
+        if not self.extra_gen_vocab:
+            return base_logits
+        if self.gen_head is None or isinstance(self.gen_head, PPMissingLayer):
+            return base_logits
+        assert self.gen_logits_processor is not None
+        gen_logits = self.gen_logits_processor(self.gen_head, hidden_states)
+        if base_logits is None or gen_logits is None:
+            return None
+        return torch.cat([base_logits, gen_logits], dim=-1)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+            ("gate_up_proj", "gate_proj", 0),
+            ("gate_up_proj", "up_proj", 1),
+        ]
+        params_dict = dict(self.named_parameters(remove_duplicate=False))
+        loaded_params: set[str] = set()
+
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+
+            if self.quant_config is not None and (scale_name := self.quant_config.get_cache_scale(name)):
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+
+                name = name.replace(weight_name, param_name)
+
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                if name.endswith("scale"):
+                    name = maybe_remap_kv_scale_name(name, params_dict)
+                    if name is None:
+                        continue
+
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                if weight_loader == default_weight_loader:
+                    weight_loader(param, loaded_weight)
+                else:
+                    weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+
+            loaded_params.add(name)
+
+        return loaded_params
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MammothModa2ARMultiModalProcessor,
+    info=MammothModa2ARProcessingInfo,
+    dummy_inputs=MammothModa2ARDummyInputsBuilder,
+)
+class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration):
+    """Replaces the language backbone with MoE within the Qwen2_5_VLForConditionalGeneration multi-modal framework."""
+
+    have_multimodal_outputs = True
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            # Skip generation-side (DiT/VAE) weights as they do not belong to the AR stage.
+            "gen_image_condition_refiner.": None,
+            "gen_transformer.": None,
+            "gen_vae.": None,
+            # LLM backbone: checkpoint uses the llm_model.* prefix.
+            # Extra generation vocab (image tokens) weights: mapped separately to the vLLM language_model submodule.
+            "llm_model.model.language_model.gen_embed_tokens.": "language_model.gen_embed_tokens.",
+            "llm_model.gen_head.": "language_model.gen_head.",
+            "llm_model.model.language_model.": "language_model.",
+            "llm_model.model.visual.": "visual.",
+            "llm_model.lm_head.": "language_model.lm_head.",
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        # Switch hf_config to the AR sub-config to ensure the Qwen2.5-VL path receives the correct type.
+        mammoth_cfg = vllm_config.model_config.hf_config
+        ar_hf_config = getattr(mammoth_cfg, "llm_config", mammoth_cfg)
+        ar_vllm_config = vllm_config.with_hf_config(ar_hf_config, architectures=vllm_config.model_config.architectures)
+        # Initialize multi-modal components like the vision tower first.
+        super().__init__(vllm_config=ar_vllm_config, prefix=prefix)
+        # Replace with the custom MoE language model.
+        lm_hf_config = getattr(
+            ar_vllm_config.model_config.hf_config, "text_config", ar_vllm_config.model_config.hf_config
+        )
+        self.language_model = init_vllm_registered_model(
+            vllm_config=ar_vllm_config,
+            prefix=maybe_prefix(prefix, "language_model"),
+            hf_config=lm_hf_config,
+            architectures=["MammothModa2Qwen2ForCausalLM"],
+        )
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+
+        # -------- t2i (AR grid) token constraints --------
+        # Constraint logic depends on per-step sampling_metadata + runtime_additional_information.
+        # These are passed by the vllm-omni runner via kwargs, so caching them in the model is sufficient.
+        self._last_runtime_additional_information: list[dict[str, Any]] | None = None
+
+    def _apply_t2i_token_constraints(self, logits: torch.Tensor) -> torch.Tensor:
+        """Applies per-request token constraints.
+
+        - For T2I requests: constrain AR grid tokens (force EOL at row end and
+          restrict intra-row sampling to visual token range).
+        - For non-T2I (text/understanding/chat) requests: disallow sampling
+          from the extra generation vocabulary (image tokens) to prevent
+          accidentally emitting visual-token sequences.
+        """
+        if logits is None or not isinstance(logits, torch.Tensor):
+            return logits
+
+        runtime_infos = self._last_runtime_additional_information
+
+        if runtime_infos is None:
+            # There is no runtime info in dummy/profile run
+            return logits
+
+        neg_inf = -float("inf")
+        num_reqs = int(logits.shape[0])
+        for i in range(num_reqs):
+            runtime_info = runtime_infos[i] if isinstance(runtime_infos[i], dict) else {}
+            omni_task = runtime_info.get("omni_task")
+            if not isinstance(omni_task, list) or not omni_task or omni_task[0] != "t2i":
+                # Text/understanding/chat: forbid sampling from the extra gen vocab.
+                logits[i, self.language_model.base_vocab_size :] = neg_inf
+                continue
+
+            ar_width = runtime_info["ar_width"][0]
+            eol_token_id = runtime_info["eol_token_id"][0]
+            visual_start = runtime_info["visual_token_start_id"][0]
+            visual_end = runtime_info["visual_token_end_id"][0]
+            generated_len = runtime_info["generated_len"]
+
+            row = logits[i]
+            column_id = generated_len % (ar_width + 1)
+            if column_id == ar_width:
+                # End-of-row token: only allow eol.
+                eol_logit = row[eol_token_id].clone()
+                row.fill_(neg_inf)
+                row[eol_token_id] = eol_logit
+            else:
+                # Intra-row tokens: only allow visual tokens (explicitly forbid eol).
+                row[:visual_start] = neg_inf
+                row[visual_end + 1 :] = neg_inf
+                row[eol_token_id] = neg_inf
+
+        return logits
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Any,
+    ):
+        # vllm-omni runner passes sampling_metadata and runtime_additional_information
+        # in each forward step. compute_logits is called immediately after
+        # forward, so caching here enables step-by-step dynamic token constraints.
+        runtime_infos = kwargs.get("runtime_additional_information")
+        self._last_runtime_additional_information = runtime_infos if isinstance(runtime_infos, list) else None
+        hidden_states = super().forward(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+        if isinstance(hidden_states, IntermediateTensors) and not get_pp_group().is_last_rank:
+            return hidden_states
+        # NOTE: gpu_model_runner._dummy_run performs hidden_states[logit_indices] after forward.
+        # We must ensure text_hidden_states is a torch.Tensor to avoid errors when
+        # indexing (which happens if it's a list/tuple).
+        if isinstance(hidden_states, IntermediateTensors):
+            text_hidden_states = hidden_states["hidden_states"]
+            out_intermediate_tensors = hidden_states
+        elif isinstance(hidden_states, list):
+            text_hidden_states = hidden_states[0]
+            out_intermediate_tensors = None
+        else:
+            text_hidden_states = hidden_states
+            out_intermediate_tensors = None
+
+        return OmniOutput(
+            text_hidden_states=text_hidden_states,
+            multimodal_outputs={},
+            intermediate_tensors=out_intermediate_tensors,
+        )
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput):
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        logits = super().compute_logits(hidden_states)
+        if isinstance(logits, torch.Tensor):
+            logits = self._apply_t2i_token_constraints(logits)
+        return logits
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    MammothModa2ARMultiModalProcessor,
+    info=MammothModa2ARProcessingInfo,
+    dummy_inputs=MammothModa2ARDummyInputsBuilder,
+)
+class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
+    # Ensure vllm_omni/worker/gpu_model_runner.py's `extract_multimodal_outputs` follows
+    # the OmniOutput branch to retrieve text_hidden_states as a pure torch.Tensor,
+    # preventing errors in `hidden_states[logit_indices]` due to type mismatch (list/tuple).
+    have_multimodal_outputs = True
+
+    multimodal_cpu_fields = {"image_grid_thw", "video_grid_thw"}
+    merge_by_field_config = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        # Consistent with Qwen2_5OmniForConditionalGeneration: instance-level flag.
+        self.have_multimodal_outputs = True
+        self.vllm_config = vllm_config
+        cfg = vllm_config.model_config.hf_config
+        self.model_stage = vllm_config.model_config.model_stage
+        self.multimodal_config = vllm_config.model_config.multimodal_config
+
+        # For debugging/alignment with qwen2.5-omni: explicitly nullify unused stages.
+        self.ar = None
+        self.dit = None
+        self.vae = None
+
+        if self.model_stage == "ar":
+            # AR stage: multi-modal + MoE text.
+            self.ar = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "ar"),
+                hf_config=cfg.llm_config if hasattr(cfg, "llm_config") else cfg.text_config,
+                architectures=["MammothModa2ARForConditionalGeneration"],
+            )
+            self.model = self.ar
+        elif self.model_stage == "dit":
+            self.dit = init_vllm_registered_model(
+                vllm_config=vllm_config,
+                prefix=maybe_prefix(prefix, "dit"),
+                # NOTE: init_vllm_registered_model -> VllmConfig.with_hf_config requires a
+                # transformers.PretrainedConfig; however, Mammothmoda2Config.gen_dit_config
+                # is a dict (diffusers config). The DiT stage hf_config still uses the
+                # top-level Mammothmoda2Config, and the DiT module reads its own
+                # gen_dit_config / gen_vae_config dicts.
+                hf_config=cfg,
+                architectures=["MammothModa2DiTPipeline"],
+            )
+            self.model = self.dit
+        elif self.model_stage == "vae":
+            # Reserved: VAEs not implemented yet; raise explicit error.
+            raise NotImplementedError("MammothModa2 VAE stage not implemented yet.")
+        else:
+            raise ValueError(f"Unsupported model_stage: {self.model_stage}")
+
+        # Expose intermediate tensor factory for PP if provided by the submodule.
+        self.make_empty_intermediate_tensors = getattr(self.model, "make_empty_intermediate_tensors", lambda: None)
+
+    @classmethod
+    def get_placeholder_str(cls, modality: str, i: int):  # noqa: ARG003
+        return Qwen2_5_VLForConditionalGeneration.get_placeholder_str(modality, i)
+
+    def get_language_model(self) -> nn.Module:
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
+    def get_multimodal_embeddings(self, **kwargs: object):
+        # Backward compatibility: route through embed_multimodal.
+        return self.embed_multimodal(**kwargs)
+
+    def embed_multimodal(self, **kwargs: object):
+        if hasattr(self.model, "embed_multimodal"):
+            return self.model.embed_multimodal(**kwargs)
+        if hasattr(self.model, "get_multimodal_embeddings"):
+            return self.model.get_multimodal_embeddings(**kwargs)
+        return []
+
+    def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings=None) -> torch.Tensor:
+        if hasattr(self.model, "get_input_embeddings"):
+            return self.model.get_input_embeddings(input_ids, multimodal_embeddings=multimodal_embeddings)
+        # DiT stage does not consume token embeddings from `input_ids`; it uses
+        # condition embeddings passed via additional_information.
+        # However, vLLM's generation runner may still request token embeddings
+        # to populate `inputs_embeds` buffers, so we provide a dummy tensor.
+        if self.model_stage == "dit":
+            hidden_size = int(self.vllm_config.model_config.get_hidden_size())
+            try:
+                target_dtype = next(self.model.parameters()).dtype
+            except StopIteration:
+                target_dtype = self.vllm_config.model_config.dtype
+            return torch.zeros(
+                (input_ids.numel(), hidden_size),
+                device=input_ids.device,
+                dtype=target_dtype,
+            )
+        raise NotImplementedError("Underlying model does not implement get_input_embeddings")
+
+    def forward(self, *args, **kwargs) -> OmniOutput | torch.Tensor:
+        out = self.model(*args, **kwargs)
+        if isinstance(out, OmniOutput):
+            return out
+        if isinstance(out, list):
+            out = out[0]
+        return OmniOutput(text_hidden_states=out, multimodal_outputs={}, intermediate_tensors=None)
+
+    def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, *args, **kwargs):
+        if isinstance(hidden_states, OmniOutput):
+            hidden_states = hidden_states.text_hidden_states
+        if hasattr(self.model, "compute_logits"):
+            return self.model.compute_logits(hidden_states)
+        return None
+
+    def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]:
+        if self.model_stage != "dit":
+            raise RuntimeError(
+                f"get_dummy_runtime_additional_information only valid for dit stage, got {self.model_stage}"
+            )
+        if self.dit is None:
+            raise RuntimeError("dit stage model is not initialized")
+        if not hasattr(self.dit, "get_dummy_runtime_additional_information"):
+            raise AttributeError("dit model missing get_dummy_runtime_additional_information")
+        return self.dit.get_dummy_runtime_additional_information(num_reqs)
+
+    def load_weights(self, weights):
+        if self.model_stage == "ar":
+            if self.ar is None or not hasattr(self.ar, "load_weights"):
+                return set()
+            loaded = self.ar.load_weights(weights)
+            return add_prefix_to_loaded_weights(loaded, "ar")
+        if self.model_stage == "dit":
+            if self.dit is None or not hasattr(self.dit, "load_weights"):
+                return set()
+            loaded = self.dit.load_weights(weights)
+            return add_prefix_to_loaded_weights(loaded, "dit")
+        if self.model_stage == "vae":
+            return set()
+        raise ValueError(f"Unsupported model_stage: {self.model_stage}")

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

@@ -1,4 +1,8 @@
-from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry
+from vllm.model_executor.models.registry import (
+    _VLLM_MODELS,
+    _LazyRegisteredModel,
+    _ModelRegistry,
+)
 
 _OMNI_MODELS = {
     "Qwen2_5OmniForConditionalGeneration": (
@@ -48,6 +52,31 @@ _OMNI_MODELS = {
         "qwen3_omni_code2wav",
         "Qwen3OmniMoeCode2Wav",
     ),
+    "MammothModa2Qwen2ForCausalLM": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen2ForCausalLM",
+    ),
+    "MammothModa2ARForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2ARForConditionalGeneration",
+    ),
+    "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",
@@ -98,7 +127,6 @@ _VLLM_OMNI_MODELS = {
     **_OMNI_MODELS,
 }
 
-
 OmniModelRegistry = _ModelRegistry(
     {
         **{

vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml

@@ -0,0 +1,42 @@
+# Stage config for running MammothModa2 with full AR→DiT pipeline (text-to-image).
+# Stage 0: AR  (multimodal understanding + latent generation)
+# Stage 1: DiT (latent → image)
+#
+# For text/image understanding tasks (text output only), use mammoth_moda2_ar.yaml.
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0"
+      max_batch_size: 100
+    engine_args:
+      model_stage: ar
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 8192
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      enable_prefix_caching: false
+    final_output: false
+
+  - stage_id: 1
+    runtime:
+      devices: "0"
+      max_batch_size: 1
+    engine_args:
+      model_stage: dit
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      gpu_memory_utilization: 0.3
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: image
+      enable_prefix_caching: false
+    engine_input_source: [0]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mammoth_moda2.ar2dit
+    final_output: true
+    final_output_type: image

vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml

@@ -0,0 +1,24 @@
+# Stage config for running MammothModa2 with single AR stage (understanding / summarization).
+# Stage 0: AR (multimodal understanding + text generation)
+#
+# Use this config for image/video understanding tasks where only text output is needed.
+# For text-to-image generation (AR → DiT), use mammoth_moda2.yaml instead.
+
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0"
+      max_batch_size: 16
+    engine_args:
+      model_stage: ar
+      model_arch: MammothModa2ForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      max_model_len: 8192
+      gpu_memory_utilization: 0.5
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: text
+      enable_prefix_caching: false
+    final_output: true
+    final_output_type: text

vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py

@@ -0,0 +1,99 @@
+"""Stage input processor for MammothModa2 (AR -> DiT)."""
+
+from typing import Any
+
+import torch
+from vllm.inputs import TextPrompt
+
+from vllm_omni.inputs.data import OmniTokensPrompt
+
+
+def ar2dit(
+    stage_list: list[Any],
+    engine_input_source: list[int],
+    prompts: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,  # noqa: ARG001 — interface param, unused for ar2dit
+) -> list[OmniTokensPrompt]:
+    """Convert AR stage outputs to DiT stage inputs."""
+
+    source_stage_id = engine_input_source[0]
+    ar_outputs = stage_list[source_stage_id].engine_outputs
+
+    dit_inputs: list[OmniTokensPrompt] = []
+    for ar_output, prompt in zip(ar_outputs, prompts):
+        addi_info = prompt["additional_information"]
+        image_height = addi_info["image_height"][0]
+        image_width = addi_info["image_width"][0]
+        text_guidance_scale = addi_info["text_guidance_scale"][0]
+        cfg_range = addi_info["cfg_range"]
+        num_inference_steps = addi_info["num_inference_steps"][0]
+        gen_vocab_start_index = addi_info["visual_token_start_id"][0]
+        # ["<|image_pad|>", "<|video_pad|>", "<|vision_start|>", "<|vision_end|>"]
+        visual_ids = addi_info["visual_ids"]
+
+        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.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, dict) 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}"
+        )
+
+        mask_device = full_hidden_states.device
+        full_token_ids_t = torch.tensor(full_token_ids, dtype=torch.long, device=mask_device)
+        attention_mask = torch.ones_like(full_token_ids_t, dtype=torch.bool)
+
+        pos = torch.arange(full_token_ids_t.shape[0], device=mask_device)
+        answer_start_index = len(prompt_token_ids)
+        questions_mask = pos < answer_start_index
+        answers_mask = ~questions_mask
+
+        gen_token_mask = full_token_ids_t >= gen_vocab_start_index
+
+        visual_token_mask = torch.isin(
+            full_token_ids_t,
+            torch.tensor(visual_ids, dtype=torch.long, device=mask_device),
+        )
+
+        text_condition_token_mask = questions_mask & ~(visual_token_mask | gen_token_mask) & attention_mask
+        image_condition_token_mask = answers_mask & gen_token_mask & attention_mask
+
+        text_condition = full_hidden_states[text_condition_token_mask]
+        image_condition = full_hidden_states[image_condition_token_mask]
+
+        text_prompt_embeds = text_condition.to(dtype=torch.float32).contiguous()
+        image_prompt_embeds = image_condition.to(dtype=torch.float32).contiguous()
+
+        additional_information = {
+            "text_prompt_embeds": text_prompt_embeds,
+            "text_prompt_embeds_shape": list(text_prompt_embeds.shape),
+            "image_prompt_embeds": image_prompt_embeds,
+            "image_prompt_embeds_shape": list(image_prompt_embeds.shape),
+            "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/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/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm_omni/transformers_utils/configs/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Custom model configs that cannot be expressed via HuggingFace Transformers alone,
+following the same pattern as vllm.transformers_utils.configs.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+_CLASS_TO_MODULE: dict[str, str] = {
+    "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",
+}
+
+__all__ = [
+    "Mammothmoda2Config",
+    "Mammothmoda2Qwen2_5_VLConfig",
+    "Mammothmoda2Qwen2_5_VLTextConfig",
+    "Mammothmoda2Qwen2_5_VLVisionConfig",
+]
+
+
+def __getattr__(name: str):
+    if name in _CLASS_TO_MODULE:
+        module_name = _CLASS_TO_MODULE[name]
+        module = importlib.import_module(module_name)
+        return getattr(module, name)
+
+    raise AttributeError(f"module 'vllm_omni.transformers_utils.configs' has no attribute '{name}'")
+
+
+def __dir__():
+    return sorted(list(__all__))
+
+
+# Eagerly import all config modules so their AutoConfig.register() side-effects
+# run as soon as `vllm_omni.transformers_utils.configs` is imported.
+from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2  # noqa: F401, E402

vllm_omni/transformers_utils/configs/mammoth_moda2.py

@@ -0,0 +1,283 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import ClassVar, Literal
+
+from transformers import AutoConfig, PretrainedConfig
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
+    Qwen2_5_VLConfig,
+    Qwen2_5_VLTextConfig,
+    Qwen2_5_VLVisionConfig,
+)
+
+__all__ = [
+    "Mammothmoda2Config",
+    "Mammothmoda2Qwen2_5_VLConfig",
+    "Mammothmoda2Qwen2_5_VLTextConfig",
+    "Mammothmoda2Qwen2_5_VLVisionConfig",
+]
+
+
+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 V1 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 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:
+        super().__init__(**kwargs)
+        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"]
+
+    def get_text_config(self, decoder: bool = False) -> PretrainedConfig:  # noqa: ARG002
+        return self.llm_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)

vllm_omni/worker/gpu_generation_model_runner.py

@@ -657,6 +657,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 input_ids = self.input_ids.gpu[:num_tokens_padded]
                 inputs_embeds = None
 
+            # 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.
+            if hasattr(self.model, "get_dummy_runtime_additional_information"):
+                runtime_addi = self.model.get_dummy_runtime_additional_information(num_reqs)
+                model_kwargs["runtime_additional_information"] = runtime_addi
+
             if self.uses_mrope:
                 positions = self.mrope_positions.gpu[:, :num_tokens_padded]
             elif self.uses_xdrope_dim > 0:

vllm_omni/worker/gpu_model_runner.py

@@ -910,8 +910,18 @@ class OmniGPUModelRunner(GPUModelRunner):
         """Gather per-request model_intermediate_buffer in batch order."""
         per_req_runtime_info = []
         for req_id in self.input_batch.req_ids:
+            req_state = self.requests.get(req_id)
+            # 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:
+                info["generated_len"] = generated_len
                 per_req_runtime_info.append(info)
                 if "thinker_reply_part_per_request" in info:
                     q = info["thinker_reply_part_per_request"]