[Perf] Support continuous batching for MammothModa2 (#7954)

Signed-off-by: nagisa-kun <1434936049@qq.com>
Signed-off-by: nagisa <1434936049@qq.com>

recipes/MammothModa2/MammothModa2.md

@@ -18,12 +18,24 @@ The generic example formats the AR prompt, drives the AR → DiT stage pipeline,
 and forwards MammothModa2-specific generation parameters through the
 pipeline-declared `extra_body` contract.
 
-MammothModa2's DiT stage runs in the shared diffusion runtime in request mode.
+MammothModa2's DiT stage runs in the shared diffusion runtime. The default
-The first integration intentionally supports one request and one image per
+`mammoth_moda2.yaml` uses request-level batching with up to eight compatible
-forward only (`max_num_seqs: 1`, `num_outputs_per_prompt: 1`). Request-level
+requests (`max_num_seqs: 8`) and one image per request
-batching, step execution, continuous batching, compilation, quantization,
+(`num_outputs_per_prompt: 1`). TeaCache and Cache-DiT are optional alternative
-parallelism, and offload are not enabled by this recipe. TeaCache acceleration
+cache backends in request mode. Enable TeaCache with `--cache-backend tea_cache`
-is supported for the DiT stage.
+as shown below, or enable Cache-DiT using the commented settings in that YAML.
+Compilation, DiT quantization, parallelism, and offload are not enabled by
+these presets.
+
+To use step execution and continuous batching, select
+[`mammoth_moda2_step.yaml`](../../vllm_omni/deploy/mammoth_moda2_step.yaml)
+with `--deploy-config vllm_omni/deploy/mammoth_moda2_step.yaml` in the
+text-to-image commands below. This preset inherits the same stage placement
+and capacity, enables step execution for Stage 1, and disables request-batch
+admission waiting and diffusion cache acceleration. Compatible requests can
+join between denoising steps and finish independently, including requests
+with different inference-step counts. Step mode cannot be combined with
+TeaCache, Cache-DiT, or other diffusion cache backends.
 
 Image size, seed, guidance, and denoising steps use the standard diffusion
 request fields. `cfg_range` remains a MammothModa2-specific `extra_body`

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -33,7 +33,10 @@ from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     get_mammoth_moda2_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched.step_scheduler import StepScheduler
+from vllm_omni.diffusion.worker.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
@@ -90,12 +93,12 @@ def test_root_weight_source_forwards_revision() -> None:
     assert _root_weight_source(config).revision == "rev-7"
 
 
-def test_pipeline_declares_native_components_and_batch_request_mode() -> None:
+def test_pipeline_declares_native_components_and_batch_modes() -> 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 True
-    assert MammothModa2DiTPipeline.supports_step_execution is False
+    assert MammothModa2DiTPipeline.supports_step_execution is True
 
 
 def test_mammoth_postprocess_denormalizes_nonnegative_raw_vae_output() -> None:
@@ -342,6 +345,21 @@ class _FakeTransformer(nn.Module):
         return torch.zeros_like(hidden_states)
 
 
+class _FakeImageRefiner(nn.Module):
+    def __init__(self, num_queries: int) -> None:
+        super().__init__()
+        self.anchor = nn.Parameter(torch.zeros(1))
+        self.num_queries = num_queries
+
+    def forward(self, image_embeds, attention_mask):
+        assert attention_mask.shape == image_embeds.shape[:2]
+        return image_embeds.new_zeros(
+            image_embeds.shape[0],
+            self.num_queries,
+            image_embeds.shape[-1],
+        )
+
+
 @dataclass
 class _FakeVaeConfig:
     scaling_factor: float | None = None
@@ -590,6 +608,25 @@ def test_pre_process_registers_batch_compatibility_key() -> None:
     assert req.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48, 7)
 
 
+def test_step_pre_process_allows_different_step_counts() -> None:
+    config = _od_config()
+    config.step_execution = True
+    pre_process = get_mammoth_moda2_pre_process_func(config)
+    first = _batch(
+        request_id="short",
+        sampling=OmniDiffusionSamplingParams(height=32, width=48, num_inference_steps=2),
+    ).requests[0]
+    second = _batch(
+        request_id="long",
+        sampling=OmniDiffusionSamplingParams(height=32, width=48, num_inference_steps=7),
+    ).requests[0]
+
+    pre_process(first)
+    pre_process(second)
+
+    assert first.batch_compatibility_key == second.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48)
+
+
 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]
@@ -1020,3 +1057,367 @@ def test_admission_and_inference_threshold_resolution_matches(llm_cfg_patch: dic
     text_cond, image_cond = pipeline._split_request_conditions(parsed_at)
     assert text_cond.shape == (1, 8)
     assert image_cond.shape == (1, 8)
+
+
+def _step_state(batch: DiffusionRequestBatch) -> StepRequestState:
+    request = batch.requests[0]
+    return StepRequestState(
+        request_id=request.request_id,
+        sampling=request.sampling_params,
+        prompt=request.prompt,
+    )
+
+
+@pytest.mark.parametrize("step_execution", [False, True])
+@pytest.mark.parametrize(
+    ("standard_steps", "extra_args", "extra_info", "expected_steps"),
+    [
+        (None, {}, {}, 50),
+        (None, {"num_inference_steps": 30}, {}, 30),
+        (50, {"num_inference_steps": 30}, {}, 30),
+        (None, {}, {"num_inference_steps": [20]}, 20),
+        (10, {}, {"num_inference_steps": [20]}, 10),
+    ],
+)
+def test_pre_process_normalizes_steps_before_scheduler_admission(
+    step_execution: bool,
+    standard_steps: int | None,
+    extra_args: dict,
+    extra_info: dict,
+    expected_steps: int,
+) -> None:
+    config = _od_config()
+    config.step_execution = step_execution
+    request = _batch(
+        sampling=OmniDiffusionSamplingParams(num_inference_steps=standard_steps, extra_args=extra_args),
+    ).requests[0]
+    request.prompt["additional_information"].update(extra_info)
+    request = get_mammoth_moda2_pre_process_func(config)(request)
+
+    assert request.sampling_params.num_inference_steps == expected_steps
+    parsed = _pipeline_shell()._parse_request(DiffusionRequestBatch([request]))
+    assert parsed.num_inference_steps == expected_steps
+    if step_execution:
+        scheduler = StepScheduler()
+        assert scheduler.add_request(request) == request.request_id
+        assert scheduler._request_progress[request.request_id].total_steps == expected_steps
+    else:
+        assert request.batch_compatibility_key[-1] == expected_steps
+
+
+def test_step_protocol_matches_request_mode() -> 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)
+    batch = _batch(
+        sampling=OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=42,
+            guidance_scale=1.0,
+            num_inference_steps=2,
+        )
+    )
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        request_output = pipeline.forward(batch)[0]
+        state = pipeline.prepare_encode(_step_state(batch))
+        while not state.denoise_completed:
+            input_batch = InputBatch.make_batch([state])
+            noise_pred = pipeline.denoise_step(input_batch, states=[state])
+            pipeline.step_scheduler(state, noise_pred)
+        step_output = pipeline.post_decode(state)
+
+    torch.testing.assert_close(step_output.output, request_output.output)
+    assert state.step_index == 2
+
+
+@pytest.mark.parametrize(
+    ("guidance_scales", "expected_predictions", "expected_calls"),
+    [
+        ((1.0, 4.0), (2.0, 5.0), 2),
+        ((4.0, 1.0), (5.0, 2.0), 2),
+        ((1.0, 1.0), (2.0, 2.0), 1),
+    ],
+)
+def test_step_batch_handles_cfg_and_non_cfg_requests(
+    guidance_scales: tuple[float, float],
+    expected_predictions: tuple[float, float],
+    expected_calls: int,
+) -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    config = _od_config()
+    config.step_execution = True
+    pre_process = get_mammoth_moda2_pre_process_func(config)
+    scheduler = StepScheduler()
+    requests = [
+        pre_process(
+            _batch(
+                request_id=f"cfg-{index}",
+                sampling=OmniDiffusionSamplingParams(
+                    height=32,
+                    width=48,
+                    seed=index,
+                    guidance_scale=4.0,
+                    num_inference_steps=2,
+                    extra_args={"text_guidance_scale": scale},
+                ),
+            ).requests[0]
+        )
+        for index, scale in enumerate(guidance_scales)
+    ]
+    assert scheduler._build_sampling_params_key(requests[0]) == scheduler._build_sampling_params_key(requests[1])
+
+    def predict(*, hidden_states, text_hidden_states, **kwargs):
+        return torch.full_like(hidden_states, 2.0 if text_hidden_states.shape[1] else 1.0)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch.object(pipeline.gen_transformer, "forward", side_effect=predict) as transformer_forward,
+    ):
+        states = [pipeline.prepare_encode(_step_state(DiffusionRequestBatch([request]))) for request in requests]
+        predictions = pipeline.denoise_step(InputBatch.make_batch(states), states=states)
+
+    for prediction, expected in zip(predictions, expected_predictions):
+        torch.testing.assert_close(prediction, torch.full_like(prediction, expected))
+    assert transformer_forward.call_count == expected_calls
+
+
+@pytest.mark.parametrize(
+    ("sampling", "extra_info", "expected_guidance"),
+    [
+        (OmniDiffusionSamplingParams(num_inference_steps=2), {}, 9.0),
+        (OmniDiffusionSamplingParams(num_inference_steps=2), {"text_guidance_scale": [3.5]}, 3.5),
+    ],
+)
+def test_step_protocol_resolves_guidance_like_request_mode(
+    sampling: OmniDiffusionSamplingParams,
+    extra_info: dict,
+    expected_guidance: float,
+) -> None:
+    """Guidance from defaults or additional_information must survive step mode.
+
+    prepare_encode() must parse the already-initialized step state directly
+    instead of re-wrapping its sampling params in a fresh OmniDiffusionRequest,
+    whose __post_init__ promotes the auto-filled 1.0 guidance_scale to an
+    explicitly-provided value and silently disables model-default CFG.
+    """
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    batch = _batch(sampling=sampling)
+    batch.prompts[0]["additional_information"].update(extra_info)
+    assert sampling.guidance_scale_provided is False
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        request_spec = pipeline._parse_request(batch)
+        request_output = pipeline.forward(batch)[0]
+        state = pipeline.prepare_encode(_step_state(batch))
+        while not state.denoise_completed:
+            noise_pred = pipeline.denoise_step(InputBatch.make_batch([state]), states=[state])
+            pipeline.step_scheduler(state, noise_pred)
+        step_output = pipeline.post_decode(state)
+
+    assert request_spec.text_guidance_scale == expected_guidance
+    assert state.extra["mammoth_text_guidance_scale"] == expected_guidance
+    assert state.negative_prompt_embeds is not None
+    assert sampling.guidance_scale_provided is False
+    torch.testing.assert_close(step_output.output, request_output.output)
+
+
+def test_refiner_output_uses_query_length_mask_in_request_and_step_modes() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = _FakeImageRefiner(num_queries=3)
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    batch = _batch(
+        sampling=OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=42,
+            guidance_scale=1.0,
+            num_inference_steps=1,
+        )
+    )
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        assert len(pipeline.forward(batch)) == 1
+        state = pipeline.prepare_encode(_step_state(batch))
+
+    assert state.prompt_embeds is not None
+    assert state.prompt_embeds_mask is not None
+    assert state.prompt_embeds.shape[1] == 5
+    assert state.prompt_embeds_mask.shape == state.prompt_embeds.shape[:2]
+    assert state.prompt_embeds_mask.all()
+
+
+def test_step_protocol_keeps_request_schedulers_and_progress_independent() -> 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)
+    short = _step_state(
+        _batch(
+            request_id="short",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=1,
+                guidance_scale=1.0,
+                num_inference_steps=1,
+            ),
+        )
+    )
+    long = _step_state(
+        _batch(
+            request_id="long",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=2,
+                guidance_scale=1.0,
+                num_inference_steps=2,
+            ),
+        )
+    )
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", side_effect=lambda: _FakeScheduler()),
+        patch(f"{module}.randn_tensor", side_effect=lambda shape, **kwargs: torch.zeros(shape, dtype=kwargs["dtype"])),
+    ):
+        pipeline.prepare_encode(short)
+        pipeline.prepare_encode(long)
+        predictions = pipeline.denoise_step(InputBatch.make_batch([short, long]), states=[short, long])
+        pipeline.step_scheduler(short, predictions[:1])
+        pipeline.step_scheduler(long, predictions[1:])
+
+    assert short.scheduler is not long.scheduler
+    assert short.denoise_completed
+    assert not long.denoise_completed
+    assert short.step_index == long.step_index == 1
+
+
+class _InputDependentTransformer(_FakeTransformer):
+    """Expose dropped masks, timesteps, conditioning and latent rows in outputs."""
+
+    def forward(self, *, hidden_states, timestep, text_hidden_states, text_attention_mask, **kwargs):
+        self.calls += 1
+        conditioning = (text_hidden_states * text_attention_mask.unsqueeze(-1)).sum(dim=(1, 2)) / 100
+        return hidden_states * 0.125 + timestep[:, None, None, None] + conditioning[:, None, None, None]
+
+
+def _step_test_pipeline() -> MammothModa2DiTPipeline:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _InputDependentTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    return pipeline
+
+
+def _step_test_request(request_id: str, *, seed: int = 11, steps: int = 4, text_tokens: int = 2):
+    return _batch(
+        request_id=request_id,
+        prompt={
+            "prompt": "",
+            "additional_information": {
+                "full_hidden_states": torch.arange((text_tokens + 2) * 8, dtype=torch.float32).reshape(-1, 8) / 10,
+                "full_token_ids": [10] * text_tokens + [100, 101],
+                "answer_start_index": text_tokens,
+            },
+        },
+        sampling=OmniDiffusionSamplingParams(
+            height=32,
+            width=48,
+            seed=seed,
+            guidance_scale=4.0,
+            num_inference_steps=steps,
+            extra_args={"cfg_range": [0.25, 0.75]},
+        ),
+    )
+
+
+def test_step_protocol_seed_conditioning_and_request_parity() -> None:
+    pipeline = _step_test_pipeline()
+    outputs = []
+    for seed, text_tokens in ((11, 2), (11, 2), (12, 2), (11, 5)):
+        batch = _step_test_request("solo", seed=seed, text_tokens=text_tokens)
+        state = pipeline.prepare_encode(_step_state(batch))
+        expected_noise = torch.randn(state.latents.shape, generator=torch.Generator().manual_seed(seed))
+        torch.testing.assert_close(state.latents, expected_noise, rtol=0, atol=0)
+        while not state.denoise_completed:
+            prediction = pipeline.denoise_step(InputBatch.make_batch([state]), states=[state])
+            pipeline.step_scheduler(state, prediction)
+        output = pipeline.post_decode(state).output
+        reference = pipeline.forward(batch)[0].output
+        torch.testing.assert_close(output, reference)
+        outputs.append(output)
+    torch.testing.assert_close(outputs[0], outputs[1], rtol=0, atol=0)
+    assert not torch.allclose(outputs[0], outputs[2])
+    assert not torch.allclose(outputs[0], outputs[3])
+
+
+def test_step_protocol_handles_arrival_reordering_and_retirement() -> None:
+    pipeline = _step_test_pipeline()
+    long = pipeline.prepare_encode(_step_state(_step_test_request("long")))
+    solo_long = pipeline.prepare_encode(_step_state(_step_test_request("solo-long")))
+    solo_short = pipeline.prepare_encode(_step_state(_step_test_request("solo-short", seed=22, steps=2, text_tokens=5)))
+    short = None
+    cached_batch = None
+    completed = {}
+    retired_latents = None
+
+    # Match runner ordering: newly admitted requests precede existing requests.
+    for tick in range(4):
+        if tick == 1:
+            short = pipeline.prepare_encode(_step_state(_step_test_request("short", seed=22, steps=2, text_tokens=5)))
+            assert short.scheduler is not long.scheduler
+        active = [short, long] if tick in (1, 2) else [long]
+        references = [solo_short, solo_long] if tick in (1, 2) else [solo_long]
+        cached_batch = InputBatch.make_batch(active, cached_batch=cached_batch)
+        if tick == 1:
+            assert long.current_timestep != short.current_timestep
+            assert cached_batch.prompt_embeds_mask.sum(dim=1).tolist() == [7, 4]
+        predictions = pipeline.denoise_step(cached_batch, states=active)
+        for index, (state, reference) in enumerate(zip(active, references)):
+            solo_prediction = pipeline.denoise_step(InputBatch.make_batch([reference]), states=[reference])
+            torch.testing.assert_close(predictions[index : index + 1], solo_prediction)
+            pipeline.step_scheduler(state, predictions[index : index + 1])
+            pipeline.step_scheduler(reference, solo_prediction)
+            torch.testing.assert_close(state.latents, reference.latents)
+            assert state.step_index == reference.step_index == state.scheduler.step_index
+            if state.denoise_completed:
+                completed[state.request_id] = pipeline.post_decode(state).output
+                torch.testing.assert_close(completed[state.request_id], pipeline.post_decode(reference).output)
+        if tick == 2:
+            assert short.denoise_completed and not long.denoise_completed
+            assert list(completed) == ["short"]
+            retired_latents = short.latents.clone()
+
+    assert list(completed) == ["short", "long"]
+    assert (short.step_index, long.step_index) == (2, 4)
+    torch.testing.assert_close(short.latents, retired_latents, rtol=0, atol=0)
+    assert not torch.allclose(completed["short"], completed["long"])

vllm_omni/deploy/mammoth_moda2.yaml

@@ -19,10 +19,12 @@ stages:
     request_batch_max_wait_ms: 20
     gpu_memory_utilization: 0.3
     enforce_eager: true
+    step_execution: false
     trust_remote_code: true
     enable_prefix_caching: false
     # Optional Cache-DiT acceleration for the DiT stage (~2x faster denoise,
-    # training-free). Uncomment to enable; guidance=1.0 requests bypass it.
+    # training-free). Request mode only; uncomment to enable.
+    # guidance=1.0 requests bypass it. Use mammoth_moda2_step.yaml for step mode.
     # cache_backend: cache_dit
     # cache_config:
     #   Fn_compute_blocks: 1

vllm_omni/deploy/mammoth_moda2_step.yaml

@@ -0,0 +1,10 @@
+# MammothModa2 AR→DiT with step execution and continuous batching.
+
+base_config: mammoth_moda2.yaml
+
+stages:
+  - stage_id: 1
+    step_execution: true
+    request_batch_max_wait_ms: 0
+    # Step mode does not support diffusion cache backends.
+    cache_backend: none

vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py

@@ -2,7 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
 from dataclasses import dataclass
 from typing import ClassVar
 
@@ -25,7 +25,9 @@ from vllm_omni.diffusion.layers.norm import RMSNorm
 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.input_batch import InputBatch
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.diffusion.worker.utils import StepRequestState
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.transformers_utils.configs.mammoth_moda2 import Mammothmoda2Config
 
@@ -283,21 +285,28 @@ def get_mammoth_moda2_pre_process_func(od_config: OmniDiffusionConfig | None = N
 
     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
+    entering the scheduler queue. Request mode groups by output geometry and
-    batch_compatibility_key based on output geometry and inference steps.
+    inference steps; step mode omits the step count so independently progressing
+    requests can share compatible denoising ticks.
     """
     gen_vocab_start_index = _resolve_gen_vocab_start_index(od_config)
+    step_execution = bool(getattr(od_config, "step_execution", False))
 
     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 = (
+        # Scheduler admission and output metrics must use the model's resolved step count.
-            "mammoth_moda2_dit",
+        request.sampling_params.num_inference_steps = num_inference_steps
-            height,
+        if step_execution:
-            width,
+            request.batch_compatibility_key = ("mammoth_moda2_dit", height, width)
-            num_inference_steps,
+        else:
-        )
+            request.batch_compatibility_key = (
+                "mammoth_moda2_dit",
+                height,
+                width,
+                num_inference_steps,
+            )
         return request
 
     return pre_process_func
@@ -342,7 +351,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
     _vae_modules: ClassVar[list[str]] = ["gen_vae"]
 
     supports_request_batch = True
-    supports_step_execution = False
+    supports_step_execution = True
 
     # Load only gen_* weights; ignore llm_model.* to prevent loading the entire LLM backbone in the DiT stage.
     hf_to_vllm_mapper = WeightsMapper(
@@ -487,16 +496,23 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
 
     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
+        return self._parse_request_inputs(prompt, request.sampling_params, request.request_id, index)
+
+    def _parse_request_inputs(
+        self,
+        prompt: dict,
+        sampling: OmniDiffusionSamplingParams | None,
+        request_id: str,
+        index: int = 0,
+    ) -> _MammothRequest:
         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 request.is_dummy_run():
+        if OmniDiffusionRequest.is_dummy_run_request_id(request_id):
             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)
@@ -662,20 +678,22 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             generators.append(gen)
         return generators
 
-    def _denoise_group(
+    def _model_device_and_dtype(self) -> tuple[torch.device, torch.dtype]:
-        self,
+        model_device = next(self.parameters()).device
-        specs: list[_MammothRequest],
+        if self.gen_image_condition_refiner is not None:
-        conds: list[_RequestConditioning],
+            return model_device, next(self.gen_image_condition_refiner.parameters()).dtype
-        model_device: torch.device,
+        return model_device, next(self.gen_transformer.parameters()).dtype
-    ) -> 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,
+    def _collate_conditioning(
-        # so concatenate first and right-pad the combined sequence.
+        self,
+        conds: list[_RequestConditioning],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """Collate per-request conditions into transformer inputs.
+
+        Returns ``(prompt_embeds, prompt_attention_mask, ar_image_embeds,
+        ar_image_attention_mask)``. 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:
             image_embeds, image_mask = _pad_cond_sequence(
@@ -698,9 +716,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 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
+            return prompt_embeds, prompt_attention_mask, None, None
-            ar_image_attention_mask = None
+        if nested_image_embedder is not 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],
@@ -709,13 +726,60 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 [c.image_embeds for c in conds],
                 [c.image_mask for c in conds],
             )
-        else:
+            return prompt_embeds, prompt_attention_mask, ar_image_embeds, ar_image_attention_mask
-            prompt_embeds, prompt_attention_mask = _pad_cond_sequence(
+        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_embeds, c.image_embeds], dim=1) for c in conds],
-                [torch.cat([c.text_mask, c.image_mask], dim=1) for c in conds],
+            [torch.cat([c.text_mask, c.image_mask], dim=1) for c in conds],
-            )
+        )
-            ar_image_embeds = None
+        return prompt_embeds, prompt_attention_mask, None, None
-            ar_image_attention_mask = None
+
+    def _init_latents(
+        self,
+        specs: list[_MammothRequest],
+        height: int,
+        width: int,
+        model_device: torch.device,
+        target_dtype: torch.dtype,
+    ) -> torch.Tensor:
+        vae_scale_factor = 16
+        latent_channels = int(self.gen_transformer.config.in_channels)
+        shape = (len(specs), latent_channels, 2 * height // vae_scale_factor, 2 * width // vae_scale_factor)
+        single_shape = (1, *shape[1:])
+        generators = self._make_latent_generators(specs, model_device)
+        if generators is None:
+            return randn_tensor(shape, device=model_device, dtype=target_dtype)
+        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))
+        return torch.cat(noise_list, dim=0)
+
+    @staticmethod
+    def _make_scheduler(num_inference_steps: int, latents: torch.Tensor) -> FlowMatchEulerDiscreteScheduler:
+        scheduler = FlowMatchEulerDiscreteScheduler()
+        scheduler.set_timesteps(
+            num_inference_steps=num_inference_steps,
+            device=latents.device,
+            num_tokens=latents.shape[-2] * latents.shape[-1],
+        )
+        return scheduler
+
+    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
+
+        prompt_embeds, prompt_attention_mask, ar_image_embeds, ar_image_attention_mask = self._collate_conditioning(
+            conds
+        )
 
         # Empty unconditional prompt for classifier-free guidance, shared by
         # all rows; rows that never take the CFG branch select their cond
@@ -736,27 +800,8 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
                 device=model_device,
             )
 
-        vae_scale_factor = 16
+        latents = self._init_latents(specs, height, width, model_device, target_dtype)
-        latent_channels = int(self.gen_transformer.config.in_channels)
+        scheduler = self._make_scheduler(num_inference_steps, latents)
-        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=num_inference_steps,
-            device=model_device,
-            num_tokens=latents.shape[-2] * latents.shape[-1],
-        )
 
         scale_vec = torch.tensor(
             [s.text_guidance_scale for s in specs],
@@ -853,12 +898,7 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
         # 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
+        model_device, target_dtype = self._model_device_and_dtype()
-        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
@@ -882,6 +922,153 @@ class MammothModa2DiTPipeline(nn.Module, SupportsComponentDiscovery):
             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 prepare_encode(self, state: StepRequestState, **kwargs) -> StepRequestState:
+        del kwargs
+        prompt = state.prompt if isinstance(state.prompt, dict) else {}
+        spec = self._parse_request_inputs(prompt, state.sampling, state.request_id)
+        text_cond, image_cond = self._split_request_conditions(spec)
+
+        model_device, target_dtype = self._model_device_and_dtype()
+        cond = self._build_conditioning(text_cond, image_cond, model_device, target_dtype)
+        prompt_embeds, prompt_attention_mask, ar_image_embeds, ar_image_attention_mask = self._collate_conditioning(
+            [cond]
+        )
+
+        # InputBatch requires consistent negative fields for CFG and non-CFG requests.
+        # denoise_step decides whether the current batch needs an unconditional pass.
+        negative_prompt_embeds = prompt_embeds.new_zeros((1, 0, int(prompt_embeds.shape[-1])))
+        negative_prompt_attention_mask = torch.zeros(
+            (1, 0),
+            dtype=torch.bool,
+            device=model_device,
+        )
+
+        latents = self._init_latents([spec], spec.height, spec.width, model_device, target_dtype)
+        scheduler = self._make_scheduler(spec.num_inference_steps, latents)
+
+        state.prompt_embeds = prompt_embeds
+        state.prompt_embeds_mask = prompt_attention_mask
+        state.negative_prompt_embeds = negative_prompt_embeds
+        state.negative_prompt_embeds_mask = negative_prompt_attention_mask
+        state.latents = latents
+        state.timesteps = scheduler.timesteps
+        # Mid-denoise resume (sampling_params.step_index > 0) is not supported;
+        # like other step pipelines, encoding always restarts from step 0.
+        state.step_index = 0
+        state.scheduler = scheduler
+        state.do_true_cfg = False
+        state.txt_seq_lens = [int(prompt_embeds.shape[1])]
+        state.negative_txt_seq_lens = [0]
+        state.extra.update(
+            {
+                "mammoth_ar_image_embeds": ar_image_embeds,
+                "mammoth_ar_image_attention_mask": ar_image_attention_mask,
+                "mammoth_text_guidance_scale": spec.text_guidance_scale,
+                "mammoth_cfg_range": spec.cfg_range,
+            }
+        )
+        return state
+
+    def denoise_step(
+        self,
+        input_batch: InputBatch,
+        *,
+        states: Sequence[StepRequestState] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del kwargs
+        batch_states = list(states if states is not None else input_batch.states)
+        if input_batch.prompt_embeds is None or input_batch.prompt_embeds_mask is None:
+            raise ValueError("MammothModa2 step batch is missing prompt conditioning")
+        if len(batch_states) != int(input_batch.latents.shape[0]):
+            raise ValueError("MammothModa2 supports one latent row per step request")
+
+        ar_embeds = [state.extra["mammoth_ar_image_embeds"] for state in batch_states]
+        ar_masks = [state.extra["mammoth_ar_image_attention_mask"] for state in batch_states]
+        if all(value is None for value in ar_embeds):
+            ar_image_embeds = None
+            ar_image_attention_mask = None
+        elif all(value is not None for value in ar_embeds) and all(value is not None for value in ar_masks):
+            ar_image_embeds, ar_image_attention_mask = _pad_cond_sequence(
+                [value for value in ar_embeds if value is not None],
+                [value for value in ar_masks if value is not None],
+            )
+        else:
+            raise ValueError("Mixed MammothModa2 AR image-conditioning layouts in one step batch")
+
+        batch_size = int(input_batch.latents.shape[0])
+        timesteps = input_batch.timesteps.reshape(-1).expand(batch_size).to(input_batch.latents.dtype)
+        model_pred = self.gen_transformer(
+            hidden_states=input_batch.latents,
+            timestep=timesteps,
+            text_hidden_states=input_batch.prompt_embeds,
+            text_attention_mask=input_batch.prompt_embeds_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_scales = [float(state.extra["mammoth_text_guidance_scale"]) for state in batch_states]
+        active_cfg = []
+        for state, guidance_scale in zip(batch_states, guidance_scales):
+            cfg_start, cfg_end = state.extra["mammoth_cfg_range"]
+            fraction = state.step_index / max(1, state.total_steps)
+            active_cfg.append(guidance_scale > 1.0 and cfg_start <= fraction <= cfg_end)
+        if not any(active_cfg):
+            return model_pred
+        if input_batch.negative_prompt_embeds is None or input_batch.negative_prompt_embeds_mask is None:
+            raise ValueError("MammothModa2 CFG batch is missing negative conditioning")
+
+        model_pred_uncond = self.gen_transformer(
+            hidden_states=input_batch.latents,
+            timestep=timesteps,
+            text_hidden_states=input_batch.negative_prompt_embeds,
+            text_attention_mask=input_batch.negative_prompt_embeds_mask,
+            ref_image_hidden_states=None,
+            freqs_cis=self.gen_freqs_cis,
+        )
+        scale = model_pred.new_tensor(guidance_scales).view(batch_size, 1, 1, 1)
+        blended = torch.lerp(model_pred_uncond, model_pred, scale)
+        if all(active_cfg):
+            return blended
+        active_tensor = torch.tensor(
+            active_cfg,
+            device=model_pred.device,
+            dtype=torch.bool,
+        ).view(batch_size, 1, 1, 1)
+        return torch.where(active_tensor, blended, model_pred)
+
+    def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs) -> None:
+        del kwargs
+        if (
+            state.scheduler is None
+            or state.latents is None
+            or state.current_timestep is None
+            or state.prompt_embeds is None
+        ):
+            raise ValueError(f"MammothModa2 scheduler state is incomplete for request {state.request_id}")
+        state.latents = state.scheduler.step(
+            noise_pred,
+            state.current_timestep,
+            state.latents,
+            return_dict=False,
+        )[0].to(dtype=state.prompt_embeds.dtype)
+        state.step_index += 1
+
+    def post_decode(self, state: StepRequestState, **kwargs) -> DiffusionOutput:
+        del kwargs
+        if state.latents is None:
+            raise ValueError(f"MammothModa2 has no final latents for request {state.request_id}")
+        latents = state.latents
+        if self.gen_vae.config.scaling_factor is not None:
+            latents = latents / self.gen_vae.config.scaling_factor
+        if self.gen_vae.config.shift_factor is not None:
+            latents = latents + self.gen_vae.config.shift_factor
+        vae_dtype = next(self.gen_vae.parameters()).dtype
+        image = self.gen_vae.decode(latents.to(dtype=vae_dtype), return_dict=False)[0]
+        return DiffusionOutput(output=image)
+
     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)