[Model] Support request-level batching for MammothModa2 DiT on shared runtime (#7476)

Signed-off-by: BeatSeat <wendavid552@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/config/test_omni_config.py

@@ -127,7 +127,7 @@ def test_mammothmoda2_diffusion_stage_projects_native_backend_config() -> None:
     assert stage.diffusion_config.model_class_name == "MammothModa2DiTPipeline"
     assert stage.diffusion_config.model == "/models/MammothModa2-Preview"
     assert stage.diffusion_config.step_execution is False
-    assert stage.scheduler_config.max_num_seqs == 1
+    assert stage.scheduler_config.max_num_seqs == 8
     assert stage.connector_config.omni_kv_config == {"need_recv_cache": False}
 
 

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -4,6 +4,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import cache_dit
 import pytest
@@ -26,7 +27,10 @@ from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
 from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     MammothModa2DiTPipeline,
     _build_mammoth_config,
+    _MammothRequest,
+    _pad_cond_sequence,
     _root_weight_source,
+    get_mammoth_moda2_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -62,9 +66,7 @@ def _od_config() -> OmniDiffusionConfig:
 
 def test_build_mammoth_config_uses_shared_transformer_projection() -> None:
     config = _build_mammoth_config(_od_config())
-    assert config.model_type == "mammothmoda2"
-    assert config.get_text_config().hidden_size == 8
-    assert config.gen_dit_config == {"hidden_size": 8, "in_channels": 4}
+    assert config.gen_dit_config["hidden_size"] == 8
 
 
 def test_build_mammoth_config_rejects_empty_shared_projection() -> None:
@@ -88,11 +90,11 @@ def test_root_weight_source_forwards_revision() -> None:
     assert _root_weight_source(config).revision == "rev-7"
 
 
-def test_pipeline_declares_native_components_and_single_request_mode_only() -> None:
+def test_pipeline_declares_native_components_and_batch_request_mode() -> None:
     assert MammothModa2DiTPipeline._dit_modules == ["gen_transformer"]
     assert MammothModa2DiTPipeline._encoder_modules == ["gen_image_condition_refiner"]
     assert MammothModa2DiTPipeline._vae_modules == ["gen_vae"]
-    assert MammothModa2DiTPipeline.supports_request_batch is False
+    assert MammothModa2DiTPipeline.supports_request_batch is True
     assert MammothModa2DiTPipeline.supports_step_execution is False
 
 
@@ -216,13 +218,17 @@ def test_parse_request_standard_fields_precede_request_level_fallbacks() -> None
     assert parsed.num_inference_steps == 7
 
 
-def test_parse_request_rejects_multiple_requests_and_outputs() -> None:
+def test_parse_request_parses_each_request_and_rejects_multi_output() -> None:
     pipeline = _pipeline_shell()
     batch = _batch()
     batch.requests.append(_batch(request_id="req-8").requests[0])
-    with pytest.raises(ValueError, match="exactly one request"):
-        pipeline._parse_request(batch)
-    with pytest.raises(ValueError, match="num_outputs_per_prompt=1"):
+    req0 = pipeline._parse_request(batch, index=0)
+    req1 = pipeline._parse_request(batch, index=1)
+    assert req0.request_id == "req-7"
+    assert req1.request_id == "req-8"
+    assert req0.index == 0
+    assert req1.index == 1
+    with pytest.raises(ValueError, match="num_outputs_per_prompt == 1"):
         pipeline._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2)))
 
 
@@ -347,10 +353,14 @@ class _FakeVae(nn.Module):
         super().__init__()
         self.anchor = nn.Parameter(torch.zeros(1))
         self.config = _FakeVaeConfig()
+        self.last_decode_dtype: torch.dtype | None = None
 
     def decode(self, latents, return_dict=False):
         assert return_dict is False
-        return (torch.zeros(1, 3, 32, 48, dtype=latents.dtype),)
+        self.last_decode_dtype = latents.dtype
+        b = latents.shape[0]
+        base = latents[:, :1, :1, :1]
+        return (torch.zeros(b, 3, 32, 48, dtype=latents.dtype) + base,)
 
 
 class _FakeScheduler:
@@ -392,8 +402,9 @@ def test_forward_returns_diffusion_output_with_request_sampling(mocker) -> None:
         )
     )
 
-    assert isinstance(result, DiffusionOutput)
-    assert result.output.shape == (1, 3, 32, 48)
+    assert len(result) == 1
+    assert isinstance(result[0], DiffusionOutput)
+    assert result[0].output.shape == (1, 3, 32, 48)
     assert captured["shape"] == (1, 4, 4, 6)
     assert captured["seed"] == 42
     assert scheduler.requested_steps == 2
@@ -527,7 +538,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Reference run for phase 2: identical request while hooks were never
     # installed on this pipeline.
     assert not pipeline.is_cache_dit_enabled()
-    reference = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    reference = pipeline.forward(_batch(sampling=unguided_sampling))[0].output
     assert seq_lens == [cond_tokens] * 8
     assert torch.isfinite(reference.float()).all()
 
@@ -551,7 +562,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Phase 2: no-CFG request disables the hooks and must reproduce the
     # never-enabled reference bit-for-bit.
     seq_lens.clear()
-    unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    unguided = pipeline.forward(_batch(sampling=unguided_sampling))[0].output
     assert seq_lens == [cond_tokens] * 8
     assert not pipeline.is_cache_dit_enabled()
     assert not BlockAdapter.is_cached(pipeline.gen_transformer)
@@ -560,7 +571,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Phase 3: CFG with a new step count and image size re-installs the hooks
     # on a fresh context; output shape follows the new request geometry.
     seq_lens.clear()
-    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output
+    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64))[0].output
     assert seq_lens == [cond_tokens, uncond_tokens] * 10
     assert pipeline.is_cache_dit_enabled()
     assert BlockAdapter.is_cached(pipeline.gen_transformer)
@@ -570,3 +581,442 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     assert torch.isfinite(regenerated.float()).all()
 
     hook.remove()
+
+
+def test_pre_process_registers_batch_compatibility_key() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch().requests[0]
+    pre_process(req)
+    assert req.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48, 7)
+
+
+def test_pre_process_rejects_missing_ar_conditions() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(request_id="req-bad-ar", prompt={"prompt": "test"}).requests[0]
+    with pytest.raises(ValueError, match="Missing additional_information AR conditions.*req-bad-ar"):
+        pre_process(req)
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "message"),
+    [
+        (0, 32, "Invalid image size.*req-dim"),
+        (32, -1, "Invalid image size.*req-dim"),
+        (30, 32, "multiples of 16.*req-dim"),
+        (32, 31, "multiples of 16.*req-dim"),
+    ],
+)
+def test_pre_process_rejects_invalid_dimensions(height, width, message) -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    batch = _batch(request_id="req-dim")
+    batch.prompts[0].update(height=height, width=width)
+    with pytest.raises(ValueError, match=message):
+        pre_process(batch.requests[0])
+
+
+def test_pre_process_rejects_non_positive_steps() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(
+        request_id="req-zero-steps",
+        sampling=OmniDiffusionSamplingParams(num_inference_steps=0),
+    ).requests[0]
+    with pytest.raises(ValueError, match="num_inference_steps must be positive.*req-zero-steps"):
+        pre_process(req)
+
+
+def test_pre_process_allows_dummy_run() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, num_inference_steps=20),
+    ).requests[0]
+    pre_process(req)
+    assert req.batch_compatibility_key == ("mammoth_moda2_dit", 512, 512, 20)
+
+
+def test_forward_casts_latents_to_vae_dtype() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    vae = _FakeVae()
+    vae.anchor = nn.Parameter(torch.zeros(1, dtype=torch.float16))
+    pipeline.gen_vae = vae
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+    ):
+        result = pipeline.forward(_batch())
+
+    assert len(result) == 1
+    assert vae.last_decode_dtype == torch.float16
+
+
+def test_group_requests_by_geometry_and_steps() -> None:
+    def make_req(idx: int, h: int, w: int, s: int) -> _MammothRequest:
+        return _MammothRequest(
+            index=idx,
+            request_id=f"r{idx}",
+            full_hidden_states=torch.zeros(1, 8),
+            full_token_ids=[1],
+            answer_start_index=0,
+            height=h,
+            width=w,
+            text_guidance_scale=1.0,
+            cfg_range=(0.0, 1.0),
+            num_inference_steps=s,
+            seed=None,
+            generator=None,
+            generator_device=None,
+        )
+
+    reqs = [make_req(0, 32, 32, 5), make_req(1, 32, 32, 5), make_req(2, 64, 64, 5)]
+    groups = MammothModa2DiTPipeline._group_requests(reqs)
+    assert sorted(groups) == [[0, 1], [2]]
+
+
+def test_pad_cond_sequence_right_pads() -> None:
+    embeds = [torch.ones(1, 2, 8), torch.ones(1, 4, 8) * 2]
+    masks = [torch.ones(1, 2, dtype=torch.bool), torch.ones(1, 4, dtype=torch.bool)]
+    padded, mask = _pad_cond_sequence(embeds, masks)
+    assert padded.shape == (2, 4, 8)
+    assert mask.tolist() == [[True, True, False, False], [True, True, True, True]]
+
+
+def test_forward_batched_matches_solo_output() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    def fake_randn_tensor(shape, *, generator, device, dtype):
+        if isinstance(generator, list):
+            tensors = [
+                torch.full((1, *shape[1:]), float(g.initial_seed() if g else 0), device=device, dtype=dtype)
+                for g in generator
+            ]
+            return torch.cat(tensors, dim=0)
+        val = float(generator.initial_seed() if generator else 0)
+        return torch.full(shape, val, device=device, dtype=dtype)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(f"{module}.randn_tensor", side_effect=fake_randn_tensor),
+    ):
+        req1 = _batch(
+            request_id="r1",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=10, guidance_scale=1.0, num_inference_steps=2
+            ),
+        ).requests[0]
+        req2 = _batch(
+            request_id="r2",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=20, guidance_scale=1.0, num_inference_steps=2
+            ),
+        ).requests[0]
+
+        batched_out = pipeline.forward(DiffusionRequestBatch([req1, req2]))
+        solo1 = pipeline.forward(DiffusionRequestBatch([req1]))
+        solo2 = pipeline.forward(DiffusionRequestBatch([req2]))
+
+    assert len(batched_out) == 2
+    torch.testing.assert_close(batched_out[0].output, solo1[0].output)
+    torch.testing.assert_close(batched_out[1].output, solo2[0].output)
+    assert not torch.equal(batched_out[0].output, batched_out[1].output)
+
+
+def test_refiner_output_length_preserved_when_input_len_differs_from_queries() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+
+    class _MockRefiner(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(1))
+
+        def forward(self, features, mask=None):
+            b, _, d = features.shape
+            return torch.zeros(b, 8, d, device=features.device, dtype=features.dtype)
+
+    pipeline.gen_image_condition_refiner = _MockRefiner()
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    captured_cond_lens = []
+
+    def fake_forward(*, hidden_states, text_hidden_states, **kwargs):
+        captured_cond_lens.append(text_hidden_states.shape[1])
+        return torch.zeros_like(hidden_states)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+        patch.object(pipeline.gen_transformer, "forward", fake_forward),
+    ):
+        prompt = {
+            "prompt": "",
+            "additional_information": {
+                "full_hidden_states": torch.zeros(6, 8),
+                "full_token_ids": [10, 11, 100, 101, 102, 103],
+                "answer_start_index": 2,
+            },
+        }
+        batch = _batch(
+            prompt=prompt,
+            sampling=OmniDiffusionSamplingParams(height=32, width=48, guidance_scale=1.0, num_inference_steps=1),
+        )
+        pipeline.forward(batch)
+
+    # 2 text tokens + 8 refined queries = 10 total conditioning tokens seen by transformer
+    assert captured_cond_lens == [10]
+
+
+def test_nested_image_embedder_compacts_unequal_text_lengths() -> None:
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+        Transformer2DModel,
+    )
+
+    model = object.__new__(Transformer2DModel)
+    nn.Module.__init__(model)
+    object.__setattr__(model, "_internal_dict", SimpleNamespace(patch_size=2))
+
+    dim = 8
+    batch_size = 2
+    max_text_len = 4
+
+    class _MockTimeCaptionEmbed(nn.Module):
+        def forward(self, timestep, text_hidden_states, dtype):
+            return torch.zeros(batch_size, dim), text_hidden_states
+
+        def image_embedder(self, feats, mask=None):
+            return torch.full((batch_size, 3, dim), 99.0)
+
+    model.time_caption_embed = _MockTimeCaptionEmbed()
+    model.x_embedder = nn.Identity()
+
+    def mock_rope(freqs, mask, *args):
+        return None, None, None, None, mask.sum(dim=-1).tolist(), []
+
+    model.rope_embedder = mock_rope
+
+    hidden_states = torch.zeros(batch_size, 4, 4, 4)
+    timestep = torch.tensor([1.0, 1.0])
+    text_hidden_states = torch.zeros(batch_size, max_text_len, dim)
+    text_hidden_states[0, 0, :] = 1.0
+    text_hidden_states[0, 1, :] = 2.0
+    text_hidden_states[1, 0, :] = 10.0
+    text_hidden_states[1, 1, :] = 20.0
+    text_hidden_states[1, 2, :] = 30.0
+    text_hidden_states[1, 3, :] = 40.0
+
+    text_attention_mask = torch.tensor(
+        [
+            [True, True, False, False],
+            [True, True, True, True],
+        ]
+    )
+
+    ar_image_hidden_states = torch.zeros(batch_size, 5, dim)
+
+    res = model._prepare_embeddings(
+        hidden_states=hidden_states,
+        timestep=timestep,
+        text_hidden_states=text_hidden_states,
+        text_attention_mask=text_attention_mask,
+        freqs_cis=torch.zeros(1),
+        batch_size=batch_size,
+        height=4,
+        width=4,
+        ar_image_hidden_states=ar_image_hidden_states,
+    )
+
+    compacted_text = res[1]
+    compacted_mask = res[2]
+    encoder_seq_lengths = res[9]
+
+    assert encoder_seq_lengths == [5, 7]
+    assert compacted_text.shape == (2, 7, dim)
+    assert compacted_text[0, 0, 0].item() == 1.0
+    assert compacted_text[0, 1, 0].item() == 2.0
+    assert compacted_text[0, 2, 0].item() == 99.0
+    assert compacted_text[0, 3, 0].item() == 99.0
+    assert compacted_text[0, 4, 0].item() == 99.0
+    assert compacted_mask[0].tolist() == [True, True, True, True, True, False, False]
+    assert compacted_mask[1].tolist() == [True, True, True, True, True, True, True]
+
+
+def test_pre_process_rejects_missing_visual_tokens_and_multi_output() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+
+    multi_out = _batch(
+        request_id="req-multi",
+        sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2),
+    ).requests[0]
+    with pytest.raises(ValueError, match="num_outputs_per_prompt == 1.*req-multi"):
+        pre_process(multi_out)
+
+    no_vis_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 11, 12],
+            "answer_start_index": 1,
+        },
+    }
+    no_vis_req = _batch(request_id="req-no-vis", prompt=no_vis_prompt).requests[0]
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-no-vis"):
+        pre_process(no_vis_req)
+
+    valid_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 100, 101],
+            "answer_start_index": 1,
+        },
+    }
+    valid_req = _batch(request_id="req-valid", prompt=valid_prompt).requests[0]
+    pre_process(valid_req)
+    assert valid_req.batch_compatibility_key is not None
+
+
+def test_cfg_precomputation_uses_cpu_booleans_without_device_sync() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    raw_any = torch.Tensor.any
+    tensor_any_called = []
+
+    def wrapped_any(self, *args, **kwargs):
+        tensor_any_called.append(self)
+        return raw_any(self, *args, **kwargs)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+        patch.object(torch.Tensor, "any", wrapped_any),
+    ):
+        req1 = _batch(
+            request_id="r1",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=1,
+                guidance_scale=4.0,
+                num_inference_steps=5,
+                extra_args={"cfg_range": [0.2, 0.8]},
+            ),
+        ).requests[0]
+        req2 = _batch(
+            request_id="r2",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=2, guidance_scale=1.0, num_inference_steps=5
+            ),
+        ).requests[0]
+
+        pipeline.forward(DiffusionRequestBatch([req1, req2]))
+
+    assert len(tensor_any_called) == 0
+
+
+def test_dummy_request_splits_valid_ar_conditions_without_error() -> None:
+    """Startup warmup dummy run must split into valid non-empty AR conditions."""
+    pipeline = _pipeline_shell()
+    batch = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, seed=1, guidance_scale=0.0, num_inference_steps=2),
+    )
+    parsed = pipeline._parse_request(batch)
+    text_cond, image_cond = pipeline._split_request_conditions(parsed)
+    assert text_cond.shape == (1, pipeline._llm_hidden_size)
+    assert image_cond.shape == (1, pipeline._llm_hidden_size)
+
+
+@pytest.mark.parametrize(
+    ("llm_cfg_patch", "expected_threshold"),
+    [
+        # Parent-level llm_config threshold
+        ({"gen_vocab_start_index": 200}, 200),
+        # Nested text_config threshold
+        ({"text_config": {"model_type": "mammothmoda2_qwen2_5_vl_text", "gen_vocab_start_index": 350}}, 350),
+        # Default-derived threshold (fallback to vocab_size / 152064)
+        ({}, 152064),
+    ],
+)
+def test_admission_and_inference_threshold_resolution_matches(llm_cfg_patch: dict, expected_threshold: int) -> None:
+    """Verify identical threshold resolution for parent-level, nested, and default configs."""
+    raw = _raw_config()
+    raw["llm_config"] = {"model_type": "mammothmoda2_qwen2_5_vl", **llm_cfg_patch}
+    od_cfg = OmniDiffusionConfig(
+        model="/models/MammothModa2-Preview",
+        model_class_name="MammothModa2DiTPipeline",
+        tf_model_config=TransformerConfig.from_dict(raw),
+    )
+    # 1. Check admission preprocessor resolution
+    pre_process = get_mammoth_moda2_pre_process_func(od_cfg)
+
+    # Token just below threshold -> rejected
+    below_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(2, 8),
+            "full_token_ids": [10, expected_threshold - 1],
+            "answer_start_index": 1,
+        },
+    }
+    below_req = _batch(request_id="req-below", prompt=below_prompt).requests[0]
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-below"):
+        pre_process(below_req)
+
+    # Token at or above threshold -> accepted
+    at_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(2, 8),
+            "full_token_ids": [10, expected_threshold],
+            "answer_start_index": 1,
+        },
+    }
+    at_req = _batch(request_id="req-at", prompt=at_prompt).requests[0]
+    pre_process(at_req)
+    assert at_req.batch_compatibility_key is not None
+
+    # 2. Check inference pipeline model config resolution
+    pipeline = _pipeline_shell()
+    pipeline.config = _build_mammoth_config(od_cfg)
+
+    # Token below threshold rejected by pipeline
+    parsed_below = pipeline._parse_request(_batch(request_id="req-inf-below", prompt=below_prompt))
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-inf-below"):
+        pipeline._split_request_conditions(parsed_below)
+
+    # Token at threshold accepted by pipeline
+    parsed_at = pipeline._parse_request(_batch(request_id="req-inf-at", prompt=at_prompt))
+    text_cond, image_cond = pipeline._split_request_conditions(parsed_at)
+    assert text_cond.shape == (1, 8)
+    assert image_cond.shape == (1, 8)

vllm_omni/deploy/mammoth_moda2.yaml

@@ -15,7 +15,8 @@ stages:
 
   - stage_id: 1
     devices: "0"
-    max_num_seqs: 1
+    max_num_seqs: 8
+    request_batch_max_wait_ms: 20
     gpu_memory_utilization: 0.3
     enforce_eager: true
     trust_remote_code: true

vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py

@@ -631,8 +631,33 @@ class Transformer2DModel(ModelMixin, ConfigMixin):
                 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)
+            # When batch requests have unequal text lengths, right-padded text_hidden_states
+            # followed by concatenated image_hidden_states places padding in the middle of
+            # the row: [valid text, pad, valid image]. Since downstream rope_embedder and
+            # joint_hidden_states prefix-pack using encoder_seq_lengths (sum of valid mask),
+            # middle padding causes image queries to be truncated.
+            # Compact each row so valid tokens are placed contiguously at the start.
+            if batch_size > 1 and not text_attention_mask.all():
+                text_lens = text_attention_mask.sum(dim=-1).tolist()
+                img_lens = image_attention_mask.sum(dim=-1).tolist()
+                max_valid_len = max(t + img for t, img in zip(text_lens, img_lens))
+                compact_hidden_states = text_hidden_states.new_zeros(
+                    batch_size, max_valid_len, text_hidden_states.shape[-1]
+                )
+                compact_attention_mask = torch.zeros((batch_size, max_valid_len), dtype=torch.bool, device=device)
+                for i in range(batch_size):
+                    t_len = text_lens[i]
+                    i_len = img_lens[i]
+                    if t_len > 0:
+                        compact_hidden_states[i, :t_len] = text_hidden_states[i, :t_len]
+                    if i_len > 0:
+                        compact_hidden_states[i, t_len : t_len + i_len] = image_hidden_states[i, :i_len]
+                    compact_attention_mask[i, : t_len + i_len] = True
+                text_hidden_states = compact_hidden_states
+                text_attention_mask = compact_attention_mask
+            else:
+                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)

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -24,7 +24,9 @@ from vllm_omni.diffusion.data import DiffusionCacheConfig, DiffusionOutput, Omni
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
+from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
 from .mammothmoda2_dit_model import SimpleQFormerImageRefiner, Transformer2DModel
@@ -75,8 +77,235 @@ def _root_weight_source(
     )
 
 
+def _pad_cond_sequence(
+    embeds: list[torch.Tensor],
+    masks: list[torch.Tensor],
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Right-pad per-request ``[1, T_i, H]`` embeds into ``[B, T_max, H]``.
+
+    The transformer derives per-row ``encoder_seq_lengths`` from the boolean
+    mask, so padding must sit on the right and masked positions are excluded
+    from the packed joint sequence.
+    """
+    if not embeds:
+        raise ValueError("Cannot pad an empty conditioning sequence list")
+    if len(embeds) != len(masks):
+        raise ValueError(f"Conditioning embeds/mask count mismatch: {len(embeds)} vs {len(masks)}")
+
+    # Fast-path for single-item batch: bypass padding tensor allocation
+    if len(embeds) == 1:
+        emb0, m0 = embeds[0], masks[0]
+        if emb0.ndim != 3 or emb0.shape[0] != 1:
+            raise ValueError(f"Conditioning embeds[0] must be [1, T, H], got {tuple(emb0.shape)}")
+        if m0.ndim != 2 or m0.shape != emb0.shape[:2]:
+            raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb0.shape)} vs {tuple(m0.shape)}")
+        return emb0, m0
+
+    hidden = embeds[0].shape[-1]
+    lengths = [int(e.shape[1]) for e in embeds]
+    max_len = max(lengths)
+
+    # Fast-path for homogeneous lengths: concatenate directly in C++
+    if all(length == max_len for length in lengths):
+        for i, (emb, m) in enumerate(zip(embeds, masks)):
+            if emb.ndim != 3 or emb.shape[0] != 1 or emb.shape[-1] != hidden:
+                raise ValueError(f"Conditioning embeds[{i}] must be [1, T, {hidden}], got {tuple(emb.shape)}")
+            if m.ndim != 2 or m.shape[0] != 1 or m.shape[1] != max_len:
+                raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb.shape)} vs {tuple(m.shape)}")
+        return torch.cat(embeds, dim=0), torch.cat(masks, dim=0)
+
+    out = embeds[0].new_zeros((len(embeds), max_len, hidden))
+    mask = torch.zeros((len(embeds), max_len), dtype=torch.bool, device=out.device)
+    for i, (emb, m) in enumerate(zip(embeds, masks)):
+        if emb.ndim != 3 or emb.shape[0] != 1 or emb.shape[-1] != hidden:
+            raise ValueError(f"Conditioning embeds[{i}] must be [1, T, {hidden}], got {tuple(emb.shape)}")
+        if m.ndim != 2 or m.shape[0] != 1 or m.shape[1] != emb.shape[1]:
+            raise ValueError(f"Conditioning embeds/mask length mismatch: {tuple(emb.shape)} vs {tuple(m.shape)}")
+        out[i, : emb.shape[1]] = emb[0]
+        mask[i, : m.shape[1]] = m[0].bool()
+    return out, mask
+
+
+def _resolve_request_sampling_and_dims(
+    prompt: dict,
+    sampling: OmniDiffusionSamplingParams | None,
+    request_id: str,
+) -> tuple[int, int, float, tuple[float, float], int]:
+    """Shared extraction and validation of geometry and sampling knobs."""
+    dimensions = []
+    for name in ("height", "width"):
+        value = DiffusionRequestBatch.get_prompt_field(prompt, name)
+        if value is None and sampling is not None:
+            value = getattr(sampling, name, None)
+        if value is None:
+            value = 1024
+        try:
+            dimensions.append(int(value))
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc
+    height, width = dimensions
+    if height <= 0 or width <= 0:
+        raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}")
+    if height % 16 != 0 or width % 16 != 0:
+        raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}")
+
+    info = prompt.get("additional_information")
+    request_info = info if isinstance(info, dict) else {}
+    extra_args = sampling.extra_args or {} if sampling else {}
+
+    guidance = extra_args.get("text_guidance_scale")
+    if guidance is None and sampling is not None:
+        guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None
+    if guidance is None:
+        guidance = _first_request_value(request_info.get("text_guidance_scale"))
+    if guidance is None:
+        guidance = 9.0
+    try:
+        text_guidance_scale = float(guidance)
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc
+
+    raw_num_inference_steps = extra_args.get("num_inference_steps")
+    if raw_num_inference_steps is None and sampling is not None:
+        raw_num_inference_steps = sampling.num_inference_steps
+    if raw_num_inference_steps is None:
+        raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps"))
+    if raw_num_inference_steps is None:
+        raw_num_inference_steps = 50
+    try:
+        num_inference_steps = int(raw_num_inference_steps)
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc
+    if num_inference_steps <= 0:
+        raise ValueError(f"num_inference_steps must be positive for request {request_id}")
+
+    cfg_range = extra_args.get("cfg_range")
+    if cfg_range is None:
+        cfg_range = request_info.get("cfg_range")
+    if cfg_range is None:
+        cfg_range = [0.0, 1.0]
+    if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2:
+        raise ValueError(f"cfg_range requires two values for request {request_id}")
+    try:
+        cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1])
+    except (TypeError, ValueError, OverflowError) as exc:
+        raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc
+    if not 0 <= cfg_start <= cfg_end <= 1:
+        raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}")
+
+    return height, width, text_guidance_scale, (cfg_start, cfg_end), num_inference_steps
+
+
+def _resolve_gen_vocab_start_index(od_config: OmniDiffusionConfig | None) -> int:
+    """Resolve the visual token threshold through the normalized model configuration."""
+    if od_config is not None and getattr(od_config, "tf_model_config", None) is not None:
+        try:
+            cfg = _build_mammoth_config(od_config)
+            llm_cfg = getattr(cfg, "llm_config", None)
+            if llm_cfg is not None:
+                val = getattr(llm_cfg, "gen_vocab_start_index", None)
+                if val is not None:
+                    return int(val)
+                text_cfg = getattr(llm_cfg, "text_config", None)
+                if text_cfg is not None and getattr(text_cfg, "gen_vocab_start_index", None) is not None:
+                    return int(text_cfg.gen_vocab_start_index)
+        except Exception:
+            pass
+
+        # Fallback inspection of raw dictionary if _build_mammoth_config fails (e.g. mock test objects)
+        raw_cfg = od_config.tf_model_config.to_dict()
+        if isinstance(raw_cfg, dict):
+            llm_cfg = raw_cfg.get("llm_config")
+            if isinstance(llm_cfg, dict):
+                text_cfg = llm_cfg.get("text_config")
+                if isinstance(text_cfg, dict) and text_cfg.get("gen_vocab_start_index") is not None:
+                    return int(text_cfg["gen_vocab_start_index"])
+                if llm_cfg.get("gen_vocab_start_index") is not None:
+                    return int(llm_cfg["gen_vocab_start_index"])
+    return 152064
+
+
+def _validate_request_for_admission(
+    request: OmniDiffusionRequest,
+    od_config: OmniDiffusionConfig | None = None,
+    gen_vocab_start_index: int | None = None,
+) -> tuple[int, int, int]:
+    """Validate request at admission time to fail-fast before entering scheduler queue.
+
+    Returns (height, width, num_inference_steps).
+    """
+    request_id = request.request_id
+    prompt = request.prompt if isinstance(request.prompt, dict) else {}
+    sampling = request.sampling_params
+    if sampling is not None and getattr(sampling, "num_outputs_per_prompt", 1) != 1:
+        raise ValueError(
+            f"MammothModa2 requires num_outputs_per_prompt == 1, got {sampling.num_outputs_per_prompt} "
+            f"for request {request_id}"
+        )
+    if not request.is_dummy_run():
+        info = prompt.get("additional_information")
+        if not isinstance(info, dict):
+            raise ValueError(f"Missing additional_information AR conditions for request {request_id}")
+        full_hidden_states = info.get("full_hidden_states")
+        full_token_ids = info.get("full_token_ids")
+        if not isinstance(full_hidden_states, torch.Tensor) or not isinstance(full_token_ids, list):
+            raise ValueError(f"Expected full_hidden_states tensor and full_token_ids list for request {request_id}")
+        try:
+            answer_start_index = int(info.get("answer_start_index"))
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError(f"Invalid answer_start_index for request {request_id}") from exc
+        if full_hidden_states.ndim != 2:
+            raise ValueError(f"Expected 2D full_hidden_states for request {request_id}")
+        if full_hidden_states.shape[0] != len(full_token_ids):
+            raise ValueError(f"AR hidden-state/token-count mismatch for request {request_id}")
+        if not 0 <= answer_start_index <= len(full_token_ids):
+            raise ValueError(f"answer_start_index outside token range for request {request_id}")
+        try:
+            int_tokens = [int(token_id) for token_id in full_token_ids]
+        except (TypeError, ValueError, OverflowError) as exc:
+            raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc
+
+        if gen_vocab_start_index is None:
+            gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
+        answer_tokens = int_tokens[answer_start_index:]
+        if not any(token_id >= gen_vocab_start_index for token_id in answer_tokens):
+            raise ValueError(
+                f"MammothModa2 AR stage produced no visual-token hidden states for request {request_id}; "
+                f"the DiT stage requires at least one generated visual token. Generated token ids: {answer_tokens[:32]}"
+            )
+
+    height, width, _, _, num_inference_steps = _resolve_request_sampling_and_dims(prompt, sampling, request_id)
+    return height, width, num_inference_steps
+
+
+def get_mammoth_moda2_pre_process_func(od_config: OmniDiffusionConfig | None = None):
+    """Admission preprocessor: fail-fast per-request validation and grouping.
+
+    Validates AR conditions, dimensions, and sampling knobs at admission so
+    malformed requests are rejected individually with their request_id before
+    entering the scheduler queue. Admitted requests are assigned a
+    batch_compatibility_key based on output geometry and inference steps.
+    """
+    gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
+
+    def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        height, width, num_inference_steps = _validate_request_for_admission(
+            request, od_config=od_config, gen_vocab_start_index=gen_vocab_start_index
+        )
+        request.batch_compatibility_key = (
+            "mammoth_moda2_dit",
+            height,
+            width,
+            num_inference_steps,
+        )
+        return request
+
+    return pre_process_func
+
+
 @dataclass(frozen=True)
 class _MammothRequest:
+    index: int
     request_id: str
     full_hidden_states: torch.Tensor
     full_token_ids: list[int]
@@ -88,6 +317,15 @@ class _MammothRequest:
     num_inference_steps: int
     seed: int | None
     generator: torch.Generator | list[torch.Generator] | None
+    generator_device: torch.device | str | None = None
+
+
+@dataclass(frozen=True)
+class _RequestConditioning:
+    text_embeds: torch.Tensor
+    text_mask: torch.Tensor
+    image_embeds: torch.Tensor
+    image_mask: torch.Tensor
 
 
 class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
@@ -103,7 +341,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     _encoder_modules: ClassVar[list[str]] = ["gen_image_condition_refiner"]
     _vae_modules: ClassVar[list[str]] = ["gen_vae"]
 
-    supports_request_batch = False
+    supports_request_batch = True
     supports_step_execution = False
 
     # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
@@ -234,20 +472,42 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             nn.Linear(in_features, out_features, bias=True),
         )
 
-    def _parse_request(self, req: DiffusionRequestBatch) -> _MammothRequest:
-        if req.num_reqs != 1:
-            raise ValueError("MammothModa2 diffusion requires exactly one request")
-        request_id = req.request_id
-        sampling = req.sampling_params
-        if sampling.num_outputs_per_prompt != 1:
-            raise ValueError(f"MammothModa2 requires num_outputs_per_prompt=1 for request {request_id}")
+    @staticmethod
+    def _group_requests(specs: list[_MammothRequest]) -> list[list[int]]:
+        """Group co-scheduled requests by (height, width, num_inference_steps).
 
-        prompt = req.prompts[0]
-        prompt = prompt if isinstance(prompt, dict) else {}
+        Preserves arrival order: the first-seen request of each bucket
+        determines the group's execution order.
+        """
+        buckets: dict[tuple[int, int, int], list[int]] = {}
+        for idx, spec in enumerate(specs):
+            key = (spec.height, spec.width, spec.num_inference_steps)
+            buckets.setdefault(key, []).append(idx)
+        return list(buckets.values())
+
+    def _parse_request(self, req: DiffusionRequestBatch, index: int = 0) -> _MammothRequest:
+        request = req.requests[index]
+        request_id = request.request_id
+        prompt = request.prompt if isinstance(request.prompt, dict) else {}
+        sampling = request.sampling_params
+        if sampling is not None and getattr(sampling, "num_outputs_per_prompt", 1) != 1:
+            raise ValueError(
+                f"MammothModa2 requires num_outputs_per_prompt == 1, got {sampling.num_outputs_per_prompt} "
+                f"for request {request_id}"
+            )
         info = prompt.get("additional_information")
-        if req.is_dummy_run():
-            full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32, device="cpu")
-            full_token_ids = [0, int(self.config.llm_config.gen_vocab_start_index)]
+        if request.is_dummy_run():
+            gen_start = 152064
+            if hasattr(self, "config") and hasattr(self.config, "llm_config") and self.config.llm_config is not None:
+                val = getattr(self.config.llm_config, "gen_vocab_start_index", None)
+                if val is not None:
+                    gen_start = val
+                else:
+                    text_cfg = getattr(self.config.llm_config, "text_config", None)
+                    if text_cfg is not None and getattr(text_cfg, "gen_vocab_start_index", None) is not None:
+                        gen_start = text_cfg.gen_vocab_start_index
+            full_hidden_states = torch.zeros((2, self._llm_hidden_size), dtype=torch.float32)
+            full_token_ids = [0, int(gen_start)]
             answer_start_index = 1
         else:
             if not isinstance(info, dict):
@@ -271,70 +531,18 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             except (TypeError, ValueError, OverflowError) as exc:
                 raise ValueError(f"Invalid full_token_ids for request {request_id}") from exc
 
-        dimensions = []
-        for name in ("height", "width"):
-            value = DiffusionRequestBatch.get_prompt_field(prompt, name)
-            if value is None:
-                value = getattr(sampling, name)
-            if value is None:
-                value = 1024
-            try:
-                dimensions.append(int(value))
-            except (TypeError, ValueError, OverflowError) as exc:
-                raise ValueError(f"Invalid image size {name}={value!r} for request {request_id}") from exc
-        height, width = dimensions
-        if height <= 0 or width <= 0:
-            raise ValueError(f"Invalid image size: {height}x{width} for request {request_id}")
-        if height % 16 != 0 or width % 16 != 0:
-            raise ValueError(f"Image size must be multiples of 16, got {height}x{width} for request {request_id}")
+        height, width, text_guidance_scale, cfg_range, num_inference_steps = _resolve_request_sampling_and_dims(
+            prompt, sampling, request_id
+        )
 
-        request_info = info if isinstance(info, dict) else {}
-        extra_args = sampling.extra_args or {}
-        guidance = extra_args.get("text_guidance_scale")
-        if guidance is None:
-            guidance = sampling.guidance_scale if sampling.guidance_scale_provided else None
-        if guidance is None:
-            guidance = _first_request_value(request_info.get("text_guidance_scale"))
-        if guidance is None:
-            guidance = 9.0
-        try:
-            text_guidance_scale = float(guidance)
-        except (TypeError, ValueError, OverflowError) as exc:
-            raise ValueError(f"Invalid text_guidance_scale for request {request_id}") from exc
-        raw_num_inference_steps = extra_args.get("num_inference_steps")
-        if raw_num_inference_steps is None:
-            raw_num_inference_steps = sampling.num_inference_steps
-        if raw_num_inference_steps is None:
-            raw_num_inference_steps = _first_request_value(request_info.get("num_inference_steps"))
-        if raw_num_inference_steps is None:
-            raw_num_inference_steps = 50
-        try:
-            num_inference_steps = int(raw_num_inference_steps)
-        except (TypeError, ValueError, OverflowError) as exc:
-            raise ValueError(f"Invalid num_inference_steps for request {request_id}") from exc
-        if num_inference_steps <= 0:
-            raise ValueError(f"num_inference_steps must be positive for request {request_id}")
-        cfg_range = extra_args.get("cfg_range")
-        if cfg_range is None:
-            cfg_range = request_info.get("cfg_range")
-        if cfg_range is None:
-            cfg_range = [0.0, 1.0]
-        if not isinstance(cfg_range, (list, tuple)) or len(cfg_range) != 2:
-            raise ValueError(f"cfg_range requires two values for request {request_id}")
-        try:
-            cfg_start, cfg_end = float(cfg_range[0]), float(cfg_range[1])
-        except (TypeError, ValueError, OverflowError) as exc:
-            raise ValueError(f"cfg_range requires two values convertible to floats for request {request_id}") from exc
-        if not 0 <= cfg_start <= cfg_end <= 1:
-            raise ValueError(f"cfg_range must satisfy 0 <= start <= end <= 1 for request {request_id}")
-
-        generator = sampling.generator
+        generator = sampling.generator if sampling else None
         if isinstance(generator, list) and len(generator) != 1:
             raise ValueError(
                 f"MammothModa2 single-output request mode requires exactly one generator for request {request_id}"
             )
 
         return _MammothRequest(
+            index=index,
             request_id=request_id,
             full_hidden_states=full_hidden_states,
             full_token_ids=full_token_ids,
@@ -342,10 +550,11 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             height=height,
             width=width,
             text_guidance_scale=text_guidance_scale,
-            cfg_range=(cfg_start, cfg_end),
+            cfg_range=cfg_range,
             num_inference_steps=num_inference_steps,
-            seed=sampling.seed,
+            seed=sampling.seed if sampling else None,
             generator=generator,
+            generator_device=getattr(sampling, "generator_device", None) if sampling else None,
         )
 
     def _split_ar_conditions(
@@ -363,13 +572,23 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         caller no longer needs to pass them. Mirrors the masking the bespoke
         MammothModa2 example performed via ar2dit.
         """
-        gen_vocab_start_index = int(self.config.llm_config.gen_vocab_start_index)
-        visual_ids = [
-            int(self.config.image_token_id),
-            int(self.config.video_token_id),
-            int(self.config.vision_start_token_id),
-            int(self.config.vision_end_token_id),
-        ]
+        gen_vocab_start_index = getattr(self.config.llm_config, "gen_vocab_start_index", None)
+        if gen_vocab_start_index is None:
+            text_cfg = getattr(self.config.llm_config, "text_config", None)
+            gen_vocab_start_index = getattr(text_cfg, "gen_vocab_start_index", 152064)
+        gen_vocab_start_index = int(gen_vocab_start_index)
+        cached_visual_ids = getattr(self, "_cached_visual_ids_tensor", None)
+        if cached_visual_ids is None:
+            cached_visual_ids = torch.tensor(
+                [
+                    int(self.config.image_token_id),
+                    int(self.config.video_token_id),
+                    int(self.config.vision_start_token_id),
+                    int(self.config.vision_end_token_id),
+                ],
+                dtype=torch.long,
+            )
+            self._cached_visual_ids_tensor = cached_visual_ids
 
         device = full_hidden_states.device
         token_ids = torch.tensor(full_token_ids, dtype=torch.long, device=device)
@@ -377,7 +596,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         questions_mask = positions < answer_start_index
         answers_mask = ~questions_mask
         gen_token_mask = token_ids >= gen_vocab_start_index
-        visual_token_mask = torch.isin(token_ids, torch.tensor(visual_ids, dtype=torch.long, device=device))
+        visual_ids_device = cached_visual_ids.to(device=device, non_blocking=True)
+        visual_token_mask = torch.isin(token_ids, visual_ids_device)
         text_mask = questions_mask & ~(visual_token_mask | gen_token_mask)
         image_mask = answers_mask & gen_token_mask
 
@@ -389,9 +609,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         image_cond = full_hidden_states[image_mask].contiguous()
         return text_cond, image_cond
 
-    @torch.inference_mode()
-    def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput:
-        request = self._parse_request(req)
+    def _split_request_conditions(self, request: _MammothRequest) -> tuple[torch.Tensor, torch.Tensor]:
+        """Split AR hidden states and fail fast before touching model params."""
         text_cond, image_cond = self._split_ar_conditions(
             full_hidden_states=request.full_hidden_states,
             full_token_ids=request.full_token_ids,
@@ -404,104 +623,181 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 f"request {request.request_id}; the DiT stage requires at least one generated visual token. "
                 f"Generated token ids: {answer_token_ids[:32]}"
             )
+        return text_cond, image_cond
 
-        # Move to model device/dtype.
-        model_device = next(self.parameters()).device
+    def _build_conditioning(
+        self,
+        text_cond: torch.Tensor,
+        image_cond: torch.Tensor,
+        model_device: torch.device,
+        target_dtype: torch.dtype,
+    ) -> _RequestConditioning:
+        text_embeds = text_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous().unsqueeze(0)
+        text_mask = torch.ones((1, text_embeds.shape[1]), dtype=torch.bool, device=model_device)
+        image_embeds = (
+            image_cond.to(device=model_device, dtype=target_dtype, non_blocking=True).contiguous().unsqueeze(0)
+        )
+        image_mask = torch.ones((1, image_embeds.shape[1]), dtype=torch.bool, device=model_device)
+        return _RequestConditioning(
+            text_embeds=text_embeds,
+            text_mask=text_mask,
+            image_embeds=image_embeds,
+            image_mask=image_mask,
+        )
+
+    @staticmethod
+    def _make_latent_generators(
+        specs: list[_MammothRequest],
+        fallback_device: torch.device,
+    ) -> list[torch.Generator] | None:
+        generators: list[torch.Generator] = []
+        for s in specs:
+            if s.generator is not None:
+                gen = s.generator[0] if isinstance(s.generator, list) else s.generator
+            elif s.seed is not None:
+                device = s.generator_device or fallback_device
+                gen = torch.Generator(device=device).manual_seed(s.seed)
+            else:
+                return None
+            generators.append(gen)
+        return generators
+
+    def _denoise_group(
+        self,
+        specs: list[_MammothRequest],
+        conds: list[_RequestConditioning],
+        model_device: torch.device,
+    ) -> list[torch.Tensor]:
+        batch = len(specs)
+        height = specs[0].height
+        width = specs[0].width
+        num_inference_steps = specs[0].num_inference_steps
+        target_dtype = conds[0].text_embeds.dtype
+
+        # Conditioning collation. Token order is text-then-image per request,
+        # so concatenate first and right-pad the combined sequence.
+        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
         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
-
-        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, image_mask = _pad_cond_sequence(
+                [c.image_embeds for c in conds],
+                [c.image_mask for c in conds],
+            )
+            # Apply optional refiner ONLY on image condition tokens.
+            if image_embeds.shape[1] > 0:
+                image_embeds = self.gen_image_condition_refiner(image_embeds, ~image_mask.bool())
+            refined_image_mask = torch.ones(
                 image_embeds.shape[:2],
                 dtype=torch.bool,
                 device=image_embeds.device,
             )
-
-        nested_image_embedder = getattr(self.gen_transformer.time_caption_embed, "image_embedder", None)
-        if nested_image_embedder is None:
-            prompt_embeds = torch.cat([text_embeds, image_embeds], dim=1)
-            prompt_attention_mask = torch.cat([text_attention_mask, image_attention_mask], dim=1)
+            seq_embeds = []
+            seq_masks = []
+            for i, c in enumerate(conds):
+                cur_img_embed = image_embeds[i : i + 1]
+                cur_img_mask = refined_image_mask[i : i + 1]
+                seq_embeds.append(torch.cat([c.text_embeds, cur_img_embed], dim=1))
+                seq_masks.append(torch.cat([c.text_mask, cur_img_mask], dim=1))
+            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(seq_embeds, seq_masks)
             ar_image_embeds = None
             ar_image_attention_mask = None
+        elif nested_image_embedder is not None:
+            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
+                [c.text_embeds for c in conds],
+                [c.text_mask for c in conds],
+            )
+            ar_image_embeds, ar_image_attention_mask = _pad_cond_sequence(
+                [c.image_embeds for c in conds],
+                [c.image_mask for c in conds],
+            )
         else:
-            prompt_embeds = text_embeds
-            prompt_attention_mask = text_attention_mask
-            ar_image_embeds = image_embeds
-            ar_image_attention_mask = image_attention_mask
+            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
+                [torch.cat([c.text_embeds, c.image_embeds], dim=1) for c in conds],
+                [torch.cat([c.text_mask, c.image_mask], dim=1) for c in conds],
+            )
+            ar_image_embeds = None
+            ar_image_attention_mask = None
 
-        # Empty unconditional prompt for classifier-free guidance.
+        # Empty unconditional prompt for classifier-free guidance, shared by
+        # all rows; rows that never take the CFG branch select their cond
+        # prediction directly via torch.where below.
+        needs_uncond = any(s.text_guidance_scale > 1.0 for s in specs)
         negative_prompt_embeds = None
         negative_prompt_attention_mask = None
-        if request.text_guidance_scale > 1.0:
+        if needs_uncond:
             hidden_size = int(prompt_embeds.shape[-1])
             negative_prompt_embeds = torch.zeros(
-                (1, 0, hidden_size),
+                (batch, 0, hidden_size),
                 dtype=target_dtype,
-                device=prompt_embeds.device,
+                device=model_device,
             )
             negative_prompt_attention_mask = torch.zeros(
-                (1, 0),
+                (batch, 0),
                 dtype=torch.bool,
-                device=prompt_embeds.device,
+                device=model_device,
             )
 
-        generator = request.generator
-        if generator is None and request.seed is not None:
-            generator = torch.Generator(device=model_device).manual_seed(request.seed)
-
-        # Output image size (px), passed from stage input processor.
-        height, width = request.height, request.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, generator=generator, device=prompt_embeds.device, dtype=prompt_embeds.dtype)
+        shape = (batch, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        single_shape = (1, latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        generators = self._make_latent_generators(specs, model_device)
+        if generators is not None:
+            noise_list = []
+            for gen in generators:
+                gen_device = gen.device if hasattr(gen, "device") else model_device
+                n = randn_tensor(single_shape, generator=gen, device=gen_device, dtype=target_dtype)
+                noise_list.append(n.to(device=model_device))
+            latents = torch.cat(noise_list, dim=0)
+        else:
+            latents = randn_tensor(shape, device=model_device, dtype=target_dtype)
 
         scheduler = FlowMatchEulerDiscreteScheduler()
-
         scheduler.set_timesteps(
-            num_inference_steps=request.num_inference_steps,
-            device=prompt_embeds.device,
+            num_inference_steps=num_inference_steps,
+            device=model_device,
             num_tokens=latents.shape[-2] * latents.shape[-1],
         )
 
-        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
+        scale_vec = torch.tensor(
+            [s.text_guidance_scale for s in specs],
+            device=model_device,
+            dtype=target_dtype,
+        ).view(batch, 1, 1, 1)
+        cfg_specs = [(s.text_guidance_scale > 1.0, float(s.cfg_range[0]), float(s.cfg_range[1])) for s in specs]
+
         total_steps = max(1, len(scheduler.timesteps))
+        any_cfg_active = any(s.text_guidance_scale > 1.0 for s in specs)
         self._prepare_cache_dit_for_request(
             num_inference_steps=total_steps,
-            cfg_active=request.text_guidance_scale > 1.0,
+            cfg_active=any_cfg_active,
         )
-        # cache-dit distinguishes the conditional vs unconditional passes
-        # purely by transformer-forward parity, so when it is active the
-        # uncond pass must run on every step. Outside cfg_range CFG is
-        # neutralized via scale=1.0 (uncond + 1.0 * (cond - uncond) == cond),
-        # matching the skip optimization's cond-only update.
         requires_paired_cfg = self._cache_dit_runtime.is_enabled and getattr(
             self, "_cache_dit_requires_paired_cfg", False
         )
+
+        # Precompute mixed-CFG active masks across timesteps to avoid CPU list
+        # comprehensions and H2D copies in the step loop.
+        precomputed_active_masks = None
+        all_active_per_step = None
+        any_active_per_step = None
+        if needs_uncond:
+            precomputed_active_masks = torch.empty((total_steps, batch, 1, 1, 1), device=model_device, dtype=torch.bool)
+            all_active_per_step = []
+            any_active_per_step = []
+            for step_idx in range(total_steps):
+                frac = step_idx / total_steps
+                step_mask = [is_active and (lo <= frac <= hi) for is_active, lo, hi in cfg_specs]
+                all_active_per_step.append(all(step_mask))
+                any_active_per_step.append(any(step_mask))
+                mask_t = torch.tensor(step_mask, device=model_device, dtype=torch.bool)
+                precomputed_active_masks[step_idx] = mask_t.view(batch, 1, 1, 1)
+
+        # Precompute timesteps in latent dtype to eliminate per-step cast overhead
+        timesteps_dtype = scheduler.timesteps.to(latents.dtype)
+
+        # Run diffusion loop (CFG supported when text_guidance_scale > 1.0)
         for i, t in enumerate(scheduler.timesteps):
-            timestep = t.expand(latents.shape[0]).to(latents.dtype)
+            timestep = timesteps_dtype[i].expand(batch)
             model_pred = self.gen_transformer(
                 hidden_states=latents,
                 timestep=timestep,
@@ -512,10 +808,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 ar_image_attention_mask=ar_image_attention_mask,
                 freqs_cis=self.gen_freqs_cis,
             )
-            guidance_scale = (
-                request.text_guidance_scale if request.cfg_range[0] <= i / total_steps <= request.cfg_range[1] else 1.0
-            )
-            run_uncond = guidance_scale > 1.0 and negative_prompt_embeds is not None
+            run_uncond = (any_active_per_step is not None and any_active_per_step[i]) if needs_uncond else False
             if requires_paired_cfg and negative_prompt_embeds is not None:
                 run_uncond = True
             if run_uncond:
@@ -527,18 +820,65 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                     ref_image_hidden_states=None,
                     freqs_cis=self.gen_freqs_cis,
                 )
-                model_pred = model_pred_uncond + guidance_scale * (model_pred - model_pred_uncond)
+                # Fused CFG blend: torch.lerp fuses (uncond + scale * (cond - uncond)) into 1 kernel
+                blended = torch.lerp(model_pred_uncond, model_pred, scale_vec)
+                if all_active_per_step is not None and all_active_per_step[i]:
+                    # Fast path: all requests in the batch are active.
+                    model_pred = blended
+                else:
+                    # Inactive rows keep their conditional prediction exactly:
+                    # gating the blend with torch.where avoids propagating
+                    # uncond-branch NaN/rounding into CFG-free rows.
+                    active_tensor = precomputed_active_masks[i]
+                    model_pred = torch.where(active_tensor, blended, model_pred)
             latents = scheduler.step(model_pred, t, latents, return_dict=False)[0]
-            latents = latents.to(dtype=prompt_embeds.dtype)
+            latents = latents.to(dtype=target_dtype)
 
-        # VAE decode
+        # VAE decode (in-place scaling to reduce peak VRAM before decode)
         if self.gen_vae.config.scaling_factor is not None:
-            latents = latents / self.gen_vae.config.scaling_factor
+            latents.div_(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]
+            latents.add_(self.gen_vae.config.shift_factor)
+        vae_dtype = next(self.gen_vae.parameters()).dtype
+        image = self.gen_vae.decode(latents.to(dtype=vae_dtype), return_dict=False)[0]  # [B, C, H, W]
+        return [image[i : i + 1] for i in range(batch)]
 
-        return DiffusionOutput(output=image)
+    @torch.inference_mode()
+    def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
+        specs = [self._parse_request(req, i) for i in range(req.num_reqs)]
+
+        # Validate AR conditions before touching model parameters so malformed
+        # requests fail fast instead of surfacing device/dtype errors.
+        raw_conds = [self._split_request_conditions(spec) for spec in specs]
+
+        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
+
+        conds = [
+            self._build_conditioning(text_cond, image_cond, model_device, target_dtype)
+            for text_cond, image_cond in raw_conds
+        ]
+
+        groups = self._group_requests(specs)
+        if len(groups) == 1 and len(groups[0]) == len(specs):
+            images = self._denoise_group(specs, conds, model_device)
+            return [DiffusionOutput(output=img) for img in images]
+
+        outputs: list[DiffusionOutput | None] = [None] * len(specs)
+        for group in groups:
+            images = self._denoise_group(
+                [specs[i] for i in group],
+                [conds[i] for i in group],
+                model_device,
+            )
+            for index, image in zip(group, images):
+                outputs[index] = DiffusionOutput(output=image)
+        if any(output is None for output in outputs):
+            raise RuntimeError("DiT batching produced no image for at least one scheduled request")
+        return [output for output in outputs if output is not None]
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         loader = AutoWeightsLoader(self)

vllm_omni/diffusion/registry.py

@@ -746,6 +746,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = {
     "SanaWmPipeline": "get_sana_wm_pre_process_func",
     "Cosmos3OmniDiffusersPipeline": "get_cosmos3_pre_process_func",
     "Cosmos3OmniPipeline": "get_cosmos3_pre_process_func",
+    "MammothModa2DiTPipeline": "get_mammoth_moda2_pre_process_func",
 }