diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md
index 804655d00..acd3e87c1 100644
--- a/docs/models/supported_models.md
+++ b/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` |
diff --git a/examples/offline_inference/mammothmodal2_preview/README.md b/examples/offline_inference/mammothmodal2_preview/README.md
new file mode 100644
index 000000000..44b20c76a
--- /dev/null
+++ b/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
+```
diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py
new file mode 100644
index 000000000..e29f907f2
--- /dev/null
+++ b/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()
diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py
new file mode 100644
index 000000000..1f10fb288
--- /dev/null
+++ b/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 ``_0``, ``_1``, …
+ """
+ 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()
diff --git a/tests/e2e/offline_inference/test_mammoth_moda2.py b/tests/e2e/offline_inference/test_mammoth_moda2.py
new file mode 100644
index 000000000..5293b5ed1
--- /dev/null
+++ b/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: /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()
diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py
index 25ea24b19..c9d7fcf2c 100644
--- a/vllm_omni/__init__.py
+++ b/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
diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py
new file mode 100644
index 000000000..ed42e24ca
--- /dev/null
+++ b/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
diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
new file mode 100644
index 000000000..e55338963
--- /dev/null
+++ b/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)
diff --git a/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py b/vllm_omni/diffusion/models/mammoth_moda2/rope_real.py
new file mode 100644
index 000000000..d16181a69
--- /dev/null
+++ b/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,
+ )
diff --git a/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py b/vllm_omni/diffusion/models/mammoth_moda2/schedulers.py
new file mode 100644
index 000000000..ac73301ce
--- /dev/null
+++ b/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
diff --git a/vllm_omni/model_executor/models/mammoth_moda2/__init__.py b/vllm_omni/model_executor/models/mammoth_moda2/__init__.py
new file mode 100644
index 000000000..80e7c2263
--- /dev/null
+++ b/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",
+]
diff --git a/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py b/vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py
new file mode 100644
index 000000000..be56a14a6
--- /dev/null
+++ b/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}")
diff --git a/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/model_executor/models/mammoth_moda2/pipeline_mammothmoda2_dit.py
new file mode 100644
index 000000000..afe38f015
--- /dev/null
+++ b/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",
+]
diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py
index add27a7b1..ec879448f 100644
--- a/vllm_omni/model_executor/models/registry.py
+++ b/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(
{
**{
diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml
new file mode 100644
index 000000000..505548ecc
--- /dev/null
+++ b/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
diff --git a/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml b/vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml
new file mode 100644
index 000000000..f95ed2b3d
--- /dev/null
+++ b/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
diff --git a/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py b/vllm_omni/model_executor/stage_input_processors/mammoth_moda2.py
new file mode 100644
index 000000000..1239aafaf
--- /dev/null
+++ b/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
diff --git a/vllm_omni/tokenizers/__init__.py b/vllm_omni/tokenizers/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py b/vllm_omni/tokenizers/mammoth_moda2_tokenizer.py
new file mode 100644
index 000000000..a419b2a92
--- /dev/null
+++ b/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|>",
+ "",
+ "",
+ "<|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
diff --git a/vllm_omni/transformers_utils/__init__.py b/vllm_omni/transformers_utils/__init__.py
new file mode 100644
index 000000000..208f01a7c
--- /dev/null
+++ b/vllm_omni/transformers_utils/__init__.py
@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py
new file mode 100644
index 000000000..136a680a3
--- /dev/null
+++ b/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
diff --git a/vllm_omni/transformers_utils/configs/mammoth_moda2.py b/vllm_omni/transformers_utils/configs/mammoth_moda2.py
new file mode 100644
index 000000000..e8a4bd6dd
--- /dev/null
+++ b/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)
diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py
index 3629cb2a9..67e530190 100644
--- a/vllm_omni/worker/gpu_generation_model_runner.py
+++ b/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:
diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py
index 058cba444..966d6e5b8 100644
--- a/vllm_omni/worker/gpu_model_runner.py
+++ b/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"]