[Fix] Restore MammothModa2 support

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

tests/config/test_config_factory.py

@@ -1314,6 +1314,8 @@ stages:
     @pytest.mark.parametrize(
         ("deploy_name", "pipeline_name", "stage_count", "final_output_type", "declared_pipeline"),
         [
+            ("mammoth_moda2.yaml", "mammoth_moda2", 2, "image", "mammoth_moda2"),
+            ("mammoth_moda2_ar.yaml", "mammoth_moda2_ar", 1, "text", "mammoth_moda2_ar"),
             ("omnivoice.yaml", "omnivoice", 1, "audio", "omnivoice"),
             ("mimo_audio.yaml", "mimo_audio", 2, "audio", None),
             ("step_audio_2.yaml", "step_audio_2", 2, "audio", None),

tests/e2e/offline_inference/test_mammoth_moda2_expansion.py

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

tests/model_executor/models/mammoth_moda2/test_mammoth_moda2_config.py

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

tests/model_executor/models/test_kv_cache_scale_mapper.py

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

tests/model_extras/test_model_extras.py

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

tests/worker/test_omni_connector_mixin.py

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