From b308e196c97f8f8406a73e7afe9a11cbb1a2a7d4 Mon Sep 17 00:00:00 2001 From: BeatSeat <137616995+BeatSeat@users.noreply.github.com> Date: Fri, 25 Sep 2026 10:34:59 +0800 Subject: [PATCH] [Model] Support request-level batching for MammothModa2 DiT on shared runtime (#7476) Signed-off-by: BeatSeat Co-authored-by: amy-why-3459 --- tests/config/test_omni_config.py | 2 +- .../test_pipeline_mammothmoda2_dit.py | 480 ++++++++++++- vllm_omni/deploy/mammoth_moda2.yaml | 3 +- .../mammoth_moda2/mammothmoda2_dit_model.py | 29 +- .../pipeline_mammothmoda2_dit.py | 652 +++++++++++++----- vllm_omni/diffusion/registry.py | 1 + 6 files changed, 992 insertions(+), 175 deletions(-) diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index 77366d66b..26e6f9e9c 100644 --- a/tests/config/test_omni_config.py +++ b/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} 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 b49ada83e..6ef4266c1 100644 --- a/tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py +++ b/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) diff --git a/vllm_omni/deploy/mammoth_moda2.yaml b/vllm_omni/deploy/mammoth_moda2.yaml index e0b652915..0c30c2bbf 100644 --- a/vllm_omni/deploy/mammoth_moda2.yaml +++ b/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 diff --git a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py b/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py index 51ea8f431..415a17ed8 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/mammothmoda2_dit_model.py +++ b/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) 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 2cc2cdcab..c7d0e7310 100644 --- a/vllm_omni/diffusion/models/mammoth_moda2/pipeline_mammothmoda2_dit.py +++ b/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) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index b52ccd075..ff6c82899 100644 --- a/vllm_omni/diffusion/registry.py +++ b/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", }