From b9cf74352a782513913fb301cf389eeb3c913050 Mon Sep 17 00:00:00 2001 From: hurukawa <61525444+nagisa-kunhah@users.noreply.github.com> Date: Mon, 5 Oct 2026 02:22:37 +0800 Subject: [PATCH] [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 | 24 +- .../test_pipeline_mammothmoda2_dit.py | 405 +++++++++++++++++- vllm_omni/deploy/mammoth_moda2.yaml | 4 +- vllm_omni/deploy/mammoth_moda2_step.yaml | 10 + .../pipeline_mammothmoda2_dit.py | 313 +++++++++++--- 5 files changed, 684 insertions(+), 72 deletions(-) create mode 100644 vllm_omni/deploy/mammoth_moda2_step.yaml diff --git a/recipes/MammothModa2/MammothModa2.md b/recipes/MammothModa2/MammothModa2.md index 7c1e8bfd8..536ee6e99 100644 --- a/recipes/MammothModa2/MammothModa2.md +++ b/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. -The first integration intentionally supports one request and one image per -forward only (`max_num_seqs: 1`, `num_outputs_per_prompt: 1`). Request-level -batching, step execution, continuous batching, compilation, quantization, -parallelism, and offload are not enabled by this recipe. TeaCache acceleration -is supported for the DiT stage. +MammothModa2's DiT stage runs in the shared diffusion runtime. The default +`mammoth_moda2.yaml` uses request-level batching with up to eight compatible +requests (`max_num_seqs: 8`) and one image per request +(`num_outputs_per_prompt: 1`). TeaCache and Cache-DiT are optional alternative +cache backends in request mode. Enable TeaCache with `--cache-backend tea_cache` +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` diff --git a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py b/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py index 6ef4266c1..b33d9c247 100644 --- a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py +++ b/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"]) diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml index 0c30c2bbf..e6892a5b0 100644 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ b/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 diff --git a/vllm_omni/deploy/mammoth_moda2_step.yaml b/vllm_omni/deploy/mammoth_moda2_step.yaml new file mode 100644 index 000000000..d695b0440 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py b/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py index 6a35955ae..8d14db7d2 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/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 - batch_compatibility_key based on output geometry and inference steps. + entering the scheduler queue. Request mode groups by output geometry and + 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 = ( - "mammoth_moda2_dit", - height, - width, - num_inference_steps, - ) + # Scheduler admission and output metrics must use the model's resolved step count. + request.sampling_params.num_inference_steps = num_inference_steps + if step_execution: + request.batch_compatibility_key = ("mammoth_moda2_dit", height, width) + 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( - 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 + def _model_device_and_dtype(self) -> tuple[torch.device, torch.dtype]: + model_device = next(self.parameters()).device + if self.gen_image_condition_refiner is not None: + return model_device, next(self.gen_image_condition_refiner.parameters()).dtype + return model_device, next(self.gen_transformer.parameters()).dtype - # Conditioning collation. Token order is text-then-image per request, - # so concatenate first and right-pad the combined sequence. + def _collate_conditioning( + 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 - ar_image_attention_mask = None - elif nested_image_embedder is not None: + return prompt_embeds, prompt_attention_mask, None, None + if 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: - 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 + return prompt_embeds, prompt_attention_mask, ar_image_embeds, ar_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], + ) + return prompt_embeds, prompt_attention_mask, None, 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 - latent_channels = int(self.gen_transformer.config.in_channels) - 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], - ) + latents = self._init_latents(specs, height, width, model_device, target_dtype) + scheduler = self._make_scheduler(num_inference_steps, latents) 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 - 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 - + model_device, target_dtype = self._model_device_and_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)