[Model] Support MammothModa2-Dev (#5411)

Signed-off-by: princepride <wangzhipeng628@gmail.com>

docs/models/supported_models.md

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

examples/offline_inference/text_to_image/text_to_image.py

@@ -556,9 +556,21 @@ def main():
         elif init_non_diffusion and hasattr(params, "extra_args"):
             if params.extra_args is None:
                 params.extra_args = {}
+            params.extra_args.update(diffusion_params.extra_args or {})
             if args.seed is not None and hasattr(params, "seed"):
                 params.seed = args.seed
 
+            # MammothModa2's AR stage emits one visual token per grid cell,
+            # one EOL token per row, and one final look-ahead token whose hidden
+            # state is unavailable. Size the first stage from the prompt metadata
+            # instead of SamplingParams' default of 16 tokens.
+            prompt_info = prompt_dict.get("additional_information", {})
+            if idx == 0 and prompt_info.get("omni_task") == ["t2i"]:
+                ar_width = int(prompt_info.get("ar_width", [0])[0])
+                ar_height = int(prompt_info.get("ar_height", [0])[0])
+                if ar_width > 0 and ar_height > 0:
+                    params.max_tokens = ar_height * (ar_width + 1) + 1
+
     if not diffusion_replaced and len(sampling_params_list) == 1:
         sampling_params_list = [diffusion_params]
 

recipes/MammothModa2/MammothModa2.md

@@ -1,13 +1,12 @@
-# MammothModa2-Preview
+# MammothModa2
 
-> MammothModa2-Preview generation and understanding through shared offline examples
+> MammothModa2-Preview and MammothModa2-Dev unified understanding and generation
 
 ## Summary
 
 - Vendor: ByteDance Research
-- Model: `bytedance-research/MammothModa2-Preview`
-- Task: Text-to-image generation (AR → DiT two-stage pipeline), plus
-  image-to-text understanding (AR stage)
+- Models: `bytedance-research/MammothModa2-Preview`, `bytedance-research/MammothModa2-Dev`
+- Tasks: Preview and Dev text-to-image (AR → DiT); Dev text/image understanding
 - Mode: Offline inference
 - Maintainer: Community
 
@@ -29,7 +28,9 @@ flags. Image size uses the standard `--height` / `--width` flags.
 
 - Upstream model:
   [`bytedance-research/MammothModa2-Preview`](https://huggingface.co/bytedance-research/MammothModa2-Preview)
-- Related text-to-image example:
+- Dev model:
+  [`bytedance-research/MammothModa2-Dev`](https://huggingface.co/bytedance-research/MammothModa2-Dev)
+- Related offline example:
   [`examples/offline_inference/text_to_image/text_to_image.py`](../../examples/offline_inference/text_to_image/text_to_image.py)
 - Related T2T/I2T example:
   [`examples/offline_inference/x_to_text/x_to_text.py`](../../examples/offline_inference/x_to_text/x_to_text.py)
@@ -131,10 +132,24 @@ ls -lh mammoth_t2i.png
 python -c "from PIL import Image; print(Image.open('mammoth_t2i.png').size)"
 ```
 
-### Image-to-Text (I2T, simplest)
+## MammothModa2-Dev unified inference
 
-For the simplest offline image-understanding run, use a short inline Python
-snippet from the repository root:
+MammothModa2-Dev uses a Qwen3-VL AR backbone, while MammothModa2-Preview uses
+Qwen2.5-VL. vLLM-Omni selects the matching implementation from the nested
+`llm_config.model_type`; no checkpoint edits or `trust_remote_code` flag are
+required.
+
+Text-to-text and image-to-text use the AR-only deploy. Text-to-image loads the
+Qwen3 generation experts (`gen_mlp`), extra visual vocabulary and image head,
+then sends the generated visual tokens and hidden states to the DiT stage.
+
+Download the checkpoint:
+
+```bash
+hf download bytedance-research/MammothModa2-Dev --local-dir ./MammothModa2-Dev
+```
+
+Run image-to-text inference from the repository root:
 
 ```python
 from PIL import Image
@@ -142,35 +157,44 @@ from vllm import SamplingParams
 from vllm.multimodal.image import convert_image_mode
 from vllm_omni import Omni
 
-def main():
-    prompt = (
-        "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
-        "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
-        "Summarize this image.<|im_end|>\n"
-        "<|im_start|>assistant\n"
-    )
+prompt = (
+    "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
+    "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
+    "Summarize this image.<|im_end|>\n"
+    "<|im_start|>assistant\n"
+)
 
-    omni = Omni(
-        model="./MammothModa2-Preview",
-        deploy_config="vllm_omni/deploy/mammoth_moda2_ar.yaml",
-    )
-    try:
-        outputs = list(
-            omni.generate(
-                [{
-                    "prompt": prompt,
-                    "multi_modal_data": {"image": convert_image_mode(Image.open("./image.png"), "RGB")},
-                    "additional_information": {"omni_task": ["chat"]},
-                }],
-                [SamplingParams(temperature=0.2, top_p=0.9, top_k=-1, max_tokens=512, seed=42)],
-            )
+omni = Omni(
+    model="./MammothModa2-Dev",
+    deploy_config="vllm_omni/deploy/mammoth_moda2_ar.yaml",
+)
+try:
+    outputs = list(
+        omni.generate(
+            [{
+                "prompt": prompt,
+                "multi_modal_data": {
+                    "image": convert_image_mode(Image.open("./image.png"), "RGB"),
+                },
+                "additional_information": {"omni_task": ["chat"]},
+            }],
+            [SamplingParams(
+                temperature=0.2,
+                top_p=0.9,
+                top_k=-1,
+                max_tokens=512,
+                seed=42,
+            )],
         )
-    finally:
-        omni.close()
+    )
+finally:
+    omni.close()
 
-    ro = getattr(outputs[-1], "request_output", outputs[-1])
-    print(ro.outputs[0].text.strip())
-
-if __name__ == "__main__":
-    main()
+request_output = getattr(outputs[-1], "request_output", outputs[-1])
+print(request_output.outputs[0].text.strip())
 ```
+
+The Dev checkpoint is approximately 47.55 GiB on disk. In the verified AR-only
+run, loaded model weights used approximately 16.97 GiB of GPU memory before KV
+and encoder caches. Allow additional GPU memory for those caches and the input
+image.

recipes/README.md

@@ -48,7 +48,7 @@ recipes/
 | [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
-| [`MammothModa2/MammothModa2-Preview.md`](./MammothModa2/MammothModa2-Preview.md) | Text-to-image with the shared offline image example (AR → DiT) | 1x L40S 48GB / 1x ≥40GB GPU |
+| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
 | [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
 | [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |

tests/model_extras/test_model_extras.py

@@ -533,15 +533,29 @@ def test_declared_extra_args_apply_to_existing_sampling_params() -> None:
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_mammothmoda2_extra_registry_declares_request_and_response_params() -> None:
-    assert get_extra_body_params("MammothModa2DiTPipeline") == frozenset(
-        {
-            "text_guidance_scale",
-            "cfg_range",
-            "num_inference_steps",
-        }
+    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",
+        negative_prompt=None,
+        height=256,
+        width=256,
     )
-    assert get_extra_output_params("MammothModa2DiTPipeline") == frozenset()
-    assert should_init_extra_args_for_non_diffusion_stages("MammothModa2DiTPipeline") is True
+    assert wrapper_prompt["additional_information"]["omni_task"] == ["t2i"]
 
 
 @pytest.mark.core_model
@@ -568,5 +582,8 @@ def test_mammothmoda2_text_to_image_prompt_builder() -> None:
             "ar_height": [32],
             "image_height": [512],
             "image_width": [768],
+            "eol_token_id": [152064],
+            "visual_token_start_id": [152072],
+            "visual_token_end_id": [168456],
         },
     }

tests/unit/mammoth_moda2/test_mammoth_moda2_config.py

@@ -16,8 +16,12 @@ 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]
+
 
 @pytest.mark.cpu
 class TestMammothmoda2Config:
@@ -89,3 +93,30 @@ class TestMammothmoda2Qwen2_5_VLConfig:
         assert hasattr(text_config, "vocab_size")
         assert hasattr(text_config, "image_token_id")
         assert hasattr(text_config, "video_token_id")
+
+
+@pytest.mark.cpu
+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]

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -189,6 +189,7 @@ class SimpleQFormerImageRefiner(nn.Module):
     def __init__(
         self,
         hidden_size: int,
+        output_hidden_size: int | None = None,
         num_queries: int = 128,
         num_layers: int = 2,
         num_heads: int | None = None,
@@ -196,6 +197,8 @@ class SimpleQFormerImageRefiner(nn.Module):
         norm_eps: float = 1e-5,
     ) -> None:
         super().__init__()
+        input_hidden_size = hidden_size
+        hidden_size = output_hidden_size or hidden_size
         self.hidden_size = hidden_size
         self.num_queries = num_queries
         # ensure num_heads divides hidden_size
@@ -203,8 +206,8 @@ class SimpleQFormerImageRefiner(nn.Module):
             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),
+            Qwen2RMSNorm(input_hidden_size, eps=norm_eps),
+            nn.Linear(input_hidden_size, hidden_size, bias=True),
         )
 
         # Learnable query embeddings
@@ -662,11 +665,32 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         batch_size: int,
         height: int,
         width: int,
+        ar_image_hidden_states: torch.Tensor | None = None,
+        ar_image_attention_mask: torch.Tensor | None = None,
     ):
         device = hidden_states.device
         p = self.config.patch_size
 
         temb, text_hidden_states = self.time_caption_embed(timestep, text_hidden_states, hidden_states.dtype)
+        image_embedder = getattr(self.time_caption_embed, "image_embedder", None)
+        if image_embedder is not None and ar_image_hidden_states is not None:
+            if ar_image_attention_mask is None:
+                ar_image_attention_mask = torch.ones(
+                    ar_image_hidden_states.shape[:2],
+                    dtype=torch.bool,
+                    device=ar_image_hidden_states.device,
+                )
+            image_hidden_states = image_embedder(
+                ar_image_hidden_states,
+                ~ar_image_attention_mask.bool(),
+            )
+            image_attention_mask = torch.ones(
+                image_hidden_states.shape[:2],
+                dtype=torch.bool,
+                device=image_hidden_states.device,
+            )
+            text_hidden_states = torch.cat([text_hidden_states, image_hidden_states], dim=1)
+            text_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
 
         img_tokens = rearrange(hidden_states, "b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1=p, p2=p)
         img_tokens = self.x_embedder(img_tokens)
@@ -699,6 +723,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         return (
             temb,
             text_hidden_states,
+            text_attention_mask,
             img_tokens,
             img_mask,
             img_len,
@@ -740,6 +765,8 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         freqs_cis: torch.Tensor,
         text_attention_mask: torch.Tensor,
         ref_image_hidden_states: list[list[torch.Tensor]] | None = None,
+        ar_image_hidden_states: torch.Tensor | None = None,
+        ar_image_attention_mask: torch.Tensor | None = None,
         return_dict: bool = False,
     ) -> torch.Tensor:
         batch_size, height, width = self._validate_inputs(
@@ -749,6 +776,7 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
         (
             temb,
             text_hidden_states,
+            text_attention_mask,
             img_tokens,
             img_mask,
             img_len,
@@ -766,6 +794,8 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
             batch_size,
             height,
             width,
+            ar_image_hidden_states,
+            ar_image_attention_mask,
         )
 
         text_hidden_states, img_tokens = self._apply_refiners(

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -42,6 +42,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
             "llm_model.": None,
+            "gen_tokenizer.": None,
         }
     )
 
@@ -77,11 +78,23 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             )
         self._reinit_caption_embedder(llm_hidden_size)
 
-        # Optional: image condition refiner (Q-Former)
-        if self.config.gen_image_condition_refiner_config is not None:
+        # Optional image condition Q-Former. Preview stores it as a standalone
+        # module; Dev stores it under the DiT timestep/caption embedder.
+        llm_model_type = getattr(self.config.llm_config, "model_type", "")
+        refiner_config = self.config.gen_image_condition_refiner_config
+        if refiner_config is not None and llm_model_type == "mammothmoda2_qwen3_vl":
+            dit_hidden_size = int(self.gen_transformer.hidden_size)
+            self.gen_transformer.time_caption_embed.image_embedder = SimpleQFormerImageRefiner(
+                hidden_size=llm_hidden_size,
+                output_hidden_size=dit_hidden_size,
+                num_heads=max(1, dit_hidden_size // 128),
+                **refiner_config,
+            )
+            self.gen_image_condition_refiner = None
+        elif refiner_config is not None:
             self.gen_image_condition_refiner = SimpleQFormerImageRefiner(
                 hidden_size=llm_hidden_size,
-                **self.config.gen_image_condition_refiner_config,
+                **refiner_config,
             )
         else:
             self.gen_image_condition_refiner = None
@@ -237,6 +250,13 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
 
         text_cond = _ensure_2d(text_cond, "text_prompt_embeds")
         image_cond = _ensure_2d(image_cond, "image_prompt_embeds")
+        if image_cond.shape[0] == 0:
+            answer_token_ids = info.get("full_token_ids", [])[int(info.get("answer_start_index", [0])[0]) :]
+            raise ValueError(
+                "MammothModa2 AR stage produced no visual-token hidden states; "
+                "the DiT stage requires at least one generated visual token. "
+                f"Generated token ids: {answer_token_ids[:32]}"
+            )
         text_cond = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
         image_cond = image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous()
 
@@ -263,8 +283,17 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 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)
+        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
+        if nested_image_embedder is None:
+            prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
+            prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+            ar_image_embeds = None
+            ar_image_attention_mask = None
+        else:
+            prompt_embeds = text_embeds
+            prompt_attention_mask = text_attention_mask
+            ar_image_embeds = image_embeds
+            ar_image_attention_mask = image_attention_mask
 
         # Prepare negative prompt (for CFG). If none provided, fall back to unconditional.
         negative_prompt_embeds = None
@@ -337,6 +366,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 text_hidden_states=prompt_embeds,
                 text_attention_mask=prompt_attention_mask,
                 ref_image_hidden_states=None,
+                ar_image_hidden_states=ar_image_embeds,
+                ar_image_attention_mask=ar_image_attention_mask,
                 freqs_cis=self.gen_freqs_cis,
             )
             guidance_scale = text_guidance_scale if cfg_range[0] <= i / total_steps <= cfg_range[1] else 1.0

vllm_omni/model_executor/models/mammoth_moda2/mammoth_moda2.py

@@ -6,7 +6,7 @@ from typing import Any
 
 import torch
 from torch import nn
-from transformers import Qwen2Config
+from transformers import Qwen2Config, Qwen3VLProcessor
 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
@@ -30,7 +30,15 @@ from vllm.model_executor.models.qwen2_5_vl import (
     Qwen2_5_VLProcessingInfo,
 )
 from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
+from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer, Qwen3MLP
+from vllm.model_executor.models.qwen3_vl import (
+    Qwen3VLDummyInputsBuilder,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMultiModalProcessor,
+    Qwen3VLProcessingInfo,
+)
 from vllm.model_executor.models.utils import (
+    AutoWeightsLoader,
     PPMissingLayer,
     WeightsMapper,
     init_vllm_registered_model,
@@ -51,6 +59,18 @@ from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
 
+def _runtime_meta(runtime_info: dict[str, Any]) -> dict[str, Any]:
+    """Accept both legacy flat metadata and the canonical nested payload."""
+    meta = runtime_info.get("meta")
+    return meta if isinstance(meta, dict) else runtime_info
+
+
+def _first_int(value: Any) -> int:
+    if isinstance(value, (list, tuple)):
+        value = value[0]
+    return int(value)
+
+
 def moe_enable(moe_type, layer_type, layer_idx) -> bool:
     """Determine if MoE should be enabled for a specific layer type and index.
 
@@ -291,6 +311,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
         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)
+        self.org_base_vocab_size = int(getattr(config, "base_vocab_size", self.base_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))
@@ -299,6 +320,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
             self.embed_tokens = VocabParallelEmbedding(
                 self.base_vocab_size,
                 config.hidden_size,
+                org_num_embeddings=self.org_base_vocab_size,
                 quant_config=quant_config,
                 prefix=f"{prefix}.embed_tokens",
             )
@@ -322,6 +344,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
             self.lm_head = ParallelLMHead(
                 self.base_vocab_size,
                 config.hidden_size,
+                org_num_embeddings=self.org_base_vocab_size,
                 quant_config=quant_config,
                 prefix=f"{prefix}.lm_head",
             )
@@ -425,6 +448,7 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
+        deepstack_input_embeds: IntermediateTensors | None = None,
     ) -> torch.Tensor | IntermediateTensors:
         if get_pp_group().is_first_rank:
             if inputs_embeds is not None:
@@ -446,8 +470,17 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
             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)):
+        for idx, layer in enumerate(
+            islice(self.layers, self.start_layer, self.end_layer),
+            start=self.start_layer,
+        ):
             hidden_states, residual = layer(positions, hidden_states, residual, gen_token_mask)
+            # The vision encoder's deepstack_visual_indexes select which vision
+            # layers produce these features. The resulting tensors are keyed
+            # 0..N-1 and, as in Qwen3-Omni, are injected into the first N
+            # language-model layers.
+            if deepstack_input_embeds is not None and idx < len(deepstack_input_embeds):
+                hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{idx}"]
 
         if not get_pp_group().is_last_rank:
             tensors = {"hidden_states": hidden_states, "residual": residual}
@@ -532,6 +565,65 @@ class MammothModa2Qwen2ForCausalLM(nn.Module, SupportsPP):
         return loaded_params
 
 
+class Mammoth3DecoderLayer(Qwen3DecoderLayer):
+    """Qwen3 decoder layer with MammothModa2's generation MLP expert."""
+
+    def __init__(
+        self,
+        config,
+        layer_idx: int,
+        cache_config: CacheConfig | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ) -> None:
+        super().__init__(config, cache_config, quant_config, prefix)
+        self.moe_enable = moe_enable(config.moe_type, "ffn", layer_idx)
+        self.gen_mlp = (
+            Qwen3MLP(
+                hidden_size=self.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.gen_mlp",
+            )
+            if self.moe_enable
+            else 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]:
+        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)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = moe_forward(
+            hidden_states,
+            self.mlp,
+            self.gen_mlp,
+            gen_token_mask if self.moe_enable else None,
+        )
+        return hidden_states, residual
+
+
+class MammothModa2Qwen3ForCausalLM(MammothModa2Qwen2ForCausalLM):
+    """MammothModa2 generation-aware language model using Qwen3 attention."""
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__(
+            vllm_config=vllm_config,
+            prefix=prefix,
+            decoder_layer_type=Mammoth3DecoderLayer,
+        )
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     MammothModa2ARMultiModalProcessor,
     info=MammothModa2ARProcessingInfo,
@@ -604,17 +696,17 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration)
         num_reqs = int(logits.shape[0])
         for i in range(num_reqs):
             runtime_info = runtime_infos[i] if isinstance(runtime_infos[i], dict) else {}
-            meta = runtime_info.get("meta", {})
+            meta = _runtime_meta(runtime_info)
             omni_task = meta.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 = meta["ar_width"][0]
-            eol_token_id = meta["eol_token_id"][0]
-            visual_start = meta["visual_token_start_id"][0]
-            visual_end = meta["visual_token_end_id"][0]
+            ar_width = _first_int(meta["ar_width"])
+            eol_token_id = _first_int(meta["eol_token_id"])
+            visual_start = _first_int(meta["visual_token_start_id"])
+            visual_end = _first_int(meta["visual_token_end_id"])
             generated_len = runtime_info["generated_len"]
 
             row = logits[i]
@@ -682,6 +774,100 @@ class MammothModa2ARForConditionalGeneration(Qwen2_5_VLForConditionalGeneration)
         return logits
 
 
+class MammothModa2Qwen3ProcessingInfo(Qwen3VLProcessingInfo):
+    """Qwen3-VL processing using MammothModa2's nested AR configuration."""
+
+    def get_hf_config(self):
+        mammoth_cfg: Mammothmoda2Config = self.ctx.get_hf_config(Mammothmoda2Config)
+        return mammoth_cfg.llm_config
+
+    def get_hf_processor(self, **kwargs: object) -> Qwen3VLProcessor:
+        # The MammothModa2 repository does not ship tokenizer_config.json, so
+        # AutoProcessor otherwise resolves only the image processor and leaves
+        # Qwen3-VL's image placeholder undefined.
+        return self.ctx.get_hf_processor(
+            Qwen3VLProcessor,
+            use_fast=kwargs.pop("use_fast", True),
+            **kwargs,
+        )
+
+    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+        return {"image": None}
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+    Qwen3VLMultiModalProcessor,
+    info=MammothModa2Qwen3ProcessingInfo,
+    dummy_inputs=Qwen3VLDummyInputsBuilder,
+)
+class MammothModa2Qwen3ARForConditionalGeneration(Qwen3VLForConditionalGeneration):
+    """MammothModa2-Dev Qwen3-VL stage for understanding and generation."""
+
+    hf_to_vllm_mapper = WeightsMapper(
+        orig_to_new_prefix={
+            "llm_model.model.visual.": "visual.",
+            "llm_model.lm_head.": "language_model.lm_head.",
+            "llm_model.model.language_model.gen_embed_tokens.": "language_model.gen_embed_tokens.",
+            "llm_model.model.language_model.": "language_model.",
+            "llm_model.gen_head.": "language_model.gen_head.",
+            "gen_image_condition_refiner.": None,
+            "gen_transformer.": None,
+            "gen_vae.": None,
+            "gen_tokenizer.": None,
+        }
+    )
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        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,
+        )
+        super().__init__(vllm_config=ar_vllm_config, prefix=prefix)
+        text_config = ar_hf_config.text_config
+        self.language_model = MammothModa2Qwen3ForCausalLM(
+            vllm_config=ar_vllm_config.with_hf_config(text_config),
+            prefix=maybe_prefix(prefix, "language_model"),
+        )
+        self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors
+        self._last_runtime_additional_information: list[dict[str, Any]] | None = None
+
+    def forward(self, *args, **kwargs):
+        runtime_infos = kwargs.get("runtime_additional_information")
+        self._last_runtime_additional_information = runtime_infos if isinstance(runtime_infos, list) else None
+        return super().forward(*args, **kwargs)
+
+    def compute_logits(self, hidden_states: torch.Tensor):
+        logits = super().compute_logits(hidden_states)
+        if not isinstance(logits, torch.Tensor) or self._last_runtime_additional_information is None:
+            return logits
+        neg_inf = -float("inf")
+        for i, runtime_info in enumerate(self._last_runtime_additional_information):
+            meta = _runtime_meta(runtime_info) if isinstance(runtime_info, dict) else {}
+            if (meta.get("omni_task") or [None])[0] != "t2i":
+                logits[i, self.language_model.base_vocab_size :] = neg_inf
+                continue
+            ar_width = _first_int(meta["ar_width"])
+            eol_token_id = _first_int(meta["eol_token_id"])
+            visual_start = _first_int(meta["visual_token_start_id"])
+            visual_end = _first_int(meta["visual_token_end_id"])
+            row = logits[i]
+            if int(runtime_info["generated_len"]) % (ar_width + 1) == ar_width:
+                eol_logit = row[eol_token_id].clone()
+                row.fill_(neg_inf)
+                row[eol_token_id] = eol_logit
+            else:
+                row[:visual_start] = neg_inf
+                row[visual_end + 1 :] = neg_inf
+                row[eol_token_id] = neg_inf
+        return logits
+
+    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)
+
+
 @MULTIMODAL_REGISTRY.register_processor(
     MammothModa2ARMultiModalProcessor,
     info=MammothModa2ARProcessingInfo,
@@ -716,7 +902,11 @@ class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
                 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"],
+                architectures=[
+                    "MammothModa2Qwen3ARForConditionalGeneration"
+                    if getattr(cfg.llm_config, "model_type", "") == "mammothmoda2_qwen3_vl"
+                    else "MammothModa2ARForConditionalGeneration"
+                ],
             )
             self.model = self.ar
         elif self.model_stage == "dit":
@@ -750,6 +940,26 @@ class MammothModa2ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor
             return self.model.get_language_model()
         return self.model
 
+    def embed_input_ids(
+        self,
+        input_ids: torch.Tensor,
+        multimodal_embeddings=None,
+        *,
+        is_multimodal: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if isinstance(self.model, MammothModa2Qwen3ARForConditionalGeneration):
+            return self.model.embed_input_ids(
+                input_ids,
+                multimodal_embeddings=multimodal_embeddings,
+                is_multimodal=is_multimodal,
+            )
+        return SupportsMultiModal.embed_input_ids(
+            self,
+            input_ids,
+            multimodal_embeddings=multimodal_embeddings,
+            is_multimodal=is_multimodal,
+        )
+
     def get_multimodal_embeddings(self, **kwargs: object):
         # Backward compatibility: route through embed_multimodal.
         return self.embed_multimodal(**kwargs)

vllm_omni/model_executor/models/registry.py

@@ -99,6 +99,16 @@ _OMNI_MODELS = {
         "mammoth_moda2",
         "MammothModa2ARForConditionalGeneration",
     ),
+    "MammothModa2Qwen3ARForConditionalGeneration": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen3ARForConditionalGeneration",
+    ),
+    "MammothModa2Qwen3ForCausalLM": (
+        "mammoth_moda2",
+        "mammoth_moda2",
+        "MammothModa2Qwen3ForCausalLM",
+    ),
     "MammothModa2DiTPipeline": (
         "mammoth_moda2",
         "pipeline_mammothmoda2_dit",

vllm_omni/model_extras/mammothmodal2_preview.py

@@ -9,6 +9,11 @@ from typing import Any
 # bespoke MammothModa2 example used to convert image dimensions -> AR grid dims.
 _PATCH_SIZE = 16
 
+# Shared by Preview and Dev (see ``t2i_generation_config.json``).
+_EOL_TOKEN_ID = 152064
+_VISUAL_TOKEN_START_ID = 152072
+_VISUAL_TOKEN_END_ID = 168456
+
 _AR_SYSTEM_PROMPT = "You are a helpful image generator."
 
 
@@ -26,9 +31,8 @@ def build_text_to_image_prompt(
 
     Model-specific sampling knobs (``text_guidance_scale``, ``cfg_range``,
     ``num_inference_steps``) flow separately via ``extra_body`` -> ``extra_args``.
-    Config-derived token ids (``eol_token_id``, ``visual_token_start_id``,
-    ``visual_token_end_id``, ``visual_ids``) are sourced inside the AR stage
-    rather than passed by the caller -- see the ar2dit stage input processor.
+    The structural token ids are included in ``additional_information`` so the
+    AR sampler can apply the same constraints for both Preview and Dev.
 
     MammothModa2 t2i uses classifier-free guidance via ``text_guidance_scale``
     and has no explicit negative-prompt path, so ``negative_prompt`` is accepted
@@ -52,6 +56,9 @@ def build_text_to_image_prompt(
             "ar_height": [ar_height],
             "image_height": [h],
             "image_width": [w],
+            "eol_token_id": [_EOL_TOKEN_ID],
+            "visual_token_start_id": [_VISUAL_TOKEN_START_ID],
+            "visual_token_end_id": [_VISUAL_TOKEN_END_ID],
         },
     }
 

vllm_omni/model_extras/registry.py

@@ -243,6 +243,11 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
     },
 }
 
+# Multi-stage discovery reports the top-level wrapper rather than its DiT
+# submodule, so both names must resolve to the same request adapters.
+_EXTRA_SPECS["MammothModa2ForConditionalGeneration"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
+_EXTRA_SPECS["Mammothmoda2Model"] = _EXTRA_SPECS["MammothModa2DiTPipeline"]
+
 
 def _get_spec(model_class_name: str | None) -> dict[str, Any] | None:
     if not model_class_name:

vllm_omni/transformers_utils/configs/mammoth_moda2.py

@@ -8,6 +8,11 @@ from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import (
     Qwen2_5_VLTextConfig,
     Qwen2_5_VLVisionConfig,
 )
+from transformers.models.qwen3_vl.configuration_qwen3_vl import (
+    Qwen3VLConfig,
+    Qwen3VLTextConfig,
+    Qwen3VLVisionConfig,
+)
 
 from vllm_omni.tokenizers.mammoth_moda2_tokenizer import MammothUTokenizer
 
@@ -16,6 +21,9 @@ __all__ = [
     "Mammothmoda2Qwen2_5_VLConfig",
     "Mammothmoda2Qwen2_5_VLTextConfig",
     "Mammothmoda2Qwen2_5_VLVisionConfig",
+    "Mammothmoda2Qwen3VLConfig",
+    "Mammothmoda2Qwen3VLTextConfig",
+    "Mammothmoda2Qwen3VLVisionConfig",
 ]
 
 
@@ -209,6 +217,77 @@ class Mammothmoda2Qwen2_5_VLConfig(Qwen2_5_VLConfig):
         self.tokenizer_class = "MammothUTokenizer"
 
 
+class Mammothmoda2Qwen3VLVisionConfig(Qwen3VLVisionConfig):
+    model_type = "mammothmoda2_qwen3_vl_vision"
+    base_config_key = "vision_config"
+
+
+class Mammothmoda2Qwen3VLTextConfig(Qwen3VLTextConfig):
+    model_type = "mammothmoda2_qwen3_vl_text"
+    base_config_key = "text_config"
+
+    def __init__(
+        self,
+        extra_gen_vocab: bool = True,
+        gen_vocab_size: int = 32800,
+        gen_vocab_start_index: int | None = None,
+        moe_type: str = "ffn",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.base_vocab_size = int(self.vocab_size)
+        self.extra_gen_vocab = extra_gen_vocab
+        self.gen_vocab_size = gen_vocab_size
+        self.gen_vocab_start_index = self.vocab_size if gen_vocab_start_index is None else gen_vocab_start_index
+        self.moe_type = moe_type
+        if self.extra_gen_vocab:
+            self.vocab_size = int(self.gen_vocab_start_index) + int(self.gen_vocab_size)
+
+
+class Mammothmoda2Qwen3VLConfig(Qwen3VLConfig):
+    """MammothModa2-Dev AR config backed by Qwen3-VL."""
+
+    model_type = "mammothmoda2_qwen3_vl"
+    sub_configs = {
+        "vision_config": Mammothmoda2Qwen3VLVisionConfig,
+        "text_config": Mammothmoda2Qwen3VLTextConfig,
+    }
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        text_config: dict | PretrainedConfig | None = None,
+        vision_config: dict | PretrainedConfig | None = None,
+        extra_gen_vocab: bool = True,
+        gen_vocab_size: int = 32800,
+        gen_vocab_start_index: int | None = None,
+        moe_type: str = "ffn",
+        **kwargs,
+    ) -> None:
+        if isinstance(vision_config, dict):
+            vision_config = self.sub_configs["vision_config"](**vision_config)
+        elif vision_config is None:
+            vision_config = self.sub_configs["vision_config"]()
+        extras = {
+            "extra_gen_vocab": extra_gen_vocab,
+            "gen_vocab_size": gen_vocab_size,
+            "gen_vocab_start_index": gen_vocab_start_index,
+            "moe_type": moe_type,
+        }
+        if isinstance(text_config, dict):
+            for key, value in extras.items():
+                text_config.setdefault(key, value)
+            text_config = self.sub_configs["text_config"](**text_config)
+        elif text_config is None:
+            text_config = self.sub_configs["text_config"](**extras)
+        super().__init__(text_config=text_config, vision_config=vision_config, **kwargs)
+        self.extra_gen_vocab = getattr(self.text_config, "extra_gen_vocab", extra_gen_vocab)
+        self.gen_vocab_size = getattr(self.text_config, "gen_vocab_size", gen_vocab_size)
+        self.gen_vocab_start_index = getattr(self.text_config, "gen_vocab_start_index", gen_vocab_start_index)
+        self.moe_type = getattr(self.text_config, "moe_type", moe_type)
+        self.tokenizer_class = "MammothUTokenizer"
+
+
 class Mammothmoda2Config(PretrainedConfig):
     """Top-level MammothModa2 composition configuration"""
 
@@ -289,6 +368,10 @@ AutoConfig.register(Mammothmoda2Config.model_type, Mammothmoda2Config)
 AutoConfig.register(Mammothmoda2Qwen2_5_VLConfig.model_type, Mammothmoda2Qwen2_5_VLConfig)
 AutoConfig.register(Mammothmoda2Qwen2_5_VLTextConfig.model_type, Mammothmoda2Qwen2_5_VLTextConfig)
 AutoConfig.register(Mammothmoda2Qwen2_5_VLVisionConfig.model_type, Mammothmoda2Qwen2_5_VLVisionConfig)
+AutoConfig.register(Mammothmoda2Qwen3VLConfig.model_type, Mammothmoda2Qwen3VLConfig)
+AutoConfig.register(Mammothmoda2Qwen3VLTextConfig.model_type, Mammothmoda2Qwen3VLTextConfig)
+AutoConfig.register(Mammothmoda2Qwen3VLVisionConfig.model_type, Mammothmoda2Qwen3VLVisionConfig)
 # Register tokenizer_type -> the configs & AutoTokenizer
 AutoTokenizer.register(config_class=Mammothmoda2Config, slow_tokenizer_class=MammothUTokenizer)
 AutoTokenizer.register(config_class=Mammothmoda2Qwen2_5_VLConfig, slow_tokenizer_class=MammothUTokenizer)
+AutoTokenizer.register(config_class=Mammothmoda2Qwen3VLConfig, slow_tokenizer_class=MammothUTokenizer)