[feat] Extend disaggregated MiniMax-H3 stage with VAE encoders

docs/design/architecture_overview.md

@@ -131,13 +131,13 @@ for the validated batching, attention, and quantization combinations.
 
 ### MiniMax-H3: conditioning, denoising, and decode
 
-MiniMax-H3 illustrates logical component and phase boundaries between text
-encoding, task-specific denoising, and media decoding:
-`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
-these components in one diffusion stage. These boundaries are useful for
-profiling and offload; independent placement would require a future
-multi-stage `PipelineConfig`. The stage mapping follows the implementation and
-serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
+MiniMax-H3 combines text and reference-media encoding, task-specific denoising,
+and media decoding. Standard serving co-locates these components in one diffusion
+stage. The opt-in
+[disaggregated topology](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md)
+runs Qwen and the video/audio VAE encoders in Stage 0, then sends their conditioning
+to a Stage 1 containing the DiTs and VAE decoders. The stage mapping follows the
+implementation and serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
 notes that a detailed technical report is forthcoming.
 
 #### Model architecture by component
@@ -149,16 +149,15 @@ notes that a detailed technical report is forthcoming.
 | **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
 
 Both DiT partitions remain behind one task-selection point; `task` chooses
-which partition runs. In the current implementation, the three logical
-components execute inside one diffusion stage. This is different from
-Qwen3-Omni's three independently scheduled stages, even though both models
-produce synchronized audio and video-like outputs.
+which partition runs.
 
 #### Representative component execution
 
+The following table describes the standard single-stage deployment.
+
 | Logical component | Batching | Attention and execution | Parallelism | Quantization |
 | --- | --- | --- | --- | --- |
-| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
+| **Text Encoder** | Follows the diffusion request scheduler | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
 | **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
 | **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
 

recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md

@@ -1,9 +1,10 @@
-# MiniMax H3 disaggregated text encoder
+# MiniMax H3 disaggregated encoder stage
 
-This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
-its hidden states and token-role metadata to an encoder-free diffusion stage.
-The standard MiniMax H3 recipes remain single-stage and continue to load the
-text encoder inside the diffusion pipeline.
+This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
+audio WVAE encoder together in a vLLM stage. It sends their conditioning
+payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
+remain single-stage and continue to run all components inside the diffusion
+pipeline.
 
 ## Start the server
 
@@ -23,6 +24,14 @@ deployment override for the available hardware. Diffusion quantization,
 layerwise offload, distributed layerwise offload, VAE parallelism, and USP
 settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
 
+Stage 0's `tensor_parallel_size` defines one shared rank group for all three
+encoder roles; it does not choose each role's execution strategy. The required
+`hf_overrides.minimax_h3_encoder_components` mapping declares those strategies
+explicitly: `text_encoder` uses tensor parallelism, `video_vae` uses patch
+parallelism, and `audio_vae` runs on the group leader. All three role entries
+are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
+the configuration does not create independent role schedulers or world sizes.
+
 For example, this five-GPU topology assigns one GPU to the encoder and four
 to the diffusion stage. `--stage-overrides` keeps placement and parallelism
 scoped to the owning stage rather than broadcasting an override to both:
@@ -57,8 +66,12 @@ The Stage 1 VAE patch-parallel options remain independent of offload and
 quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
 hardware-qualified profiles before combining these options.
 
-Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
-text-encoder weights. This H3 topology explicitly keeps its single-replica
+Stage 1 sets `model_loaded.text_encoder: false` and
+`model_loaded.vae_encoder: false`. It loads only the VAE decoders and skips
+tokenizer, processor, and text-encoder downloads. It requires complete text and
+media conditioning from Stage 0, with no fallback to local encoding.
+
+This H3 topology explicitly keeps its single-replica
 diffusion stage inline, avoiding serialization of decoded video through a
 subprocess. It expects both stages to run in one deployment; cross-node payload
 transport is outside this configuration.

recipes/MiniMaxAI/MiniMax-H3.md

@@ -23,6 +23,12 @@ One vLLM-Omni diffusion stage can load both DiTs while instantiating the
 tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only
 once. Requests select the DiT with `extra_params.task`.
 
+This single-stage recipe uses the default `model_loaded.text_encoder: true` and
+`model_loaded.vae_encoder: true`. Precomputed text conditioning can replace the
+local Qwen call. Setting `model_loaded.text_encoder: false` while leaving
+`model_loaded.vae_encoder: true` requires precomputed text conditioning;
+reference media is still encoded locally.
+
 The generated MP4 contains H.264 video and synchronized stereo audio.
 
 ## Prerequisites
@@ -35,7 +41,7 @@ hf auth login
 export MODEL=MiniMaxAI/MiniMax-H3
 ```
 
-The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
+By default, the pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
 `Ref2VA/transformer/**`. It does not download or load the diffusers-format
 `transformer`, `transformer_ref`, or `vae` weights at the repository root, nor
 duplicate Ref2VA copies of shared components.

tests/config/test_config_factory.py

@@ -1558,7 +1558,10 @@ def test_minimax_h3_resolvers_reach_production_engine_args(self, tmp_path: Path,
         monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
 
         deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
-        assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
+        assert deploy.stages[1].engine_extras["model_loaded"] == {
+            "text_encoder": False,
+            "vae_encoder": False,
+        }
         stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
         resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
 
@@ -1589,6 +1592,11 @@ def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
         assert stages[1].yaml_runtime["num_replicas"] == 1
         assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False}
         assert stages[0].yaml_engine_args["max_num_seqs"] == 1
+        assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
+            "text_encoder": {"parallel_mode": "tp"},
+            "video_vae": {"parallel_mode": "patch"},
+            "audio_vae": {"parallel_mode": "leader"},
+        }
         assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
         assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
         parallel = stages[1].yaml_engine_args["parallel_config"]
@@ -1600,6 +1608,7 @@ def test_minimax_h3_disaggregated_defaults_match_validated_topology(self):
 
         turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
         turbo_stages = merge_pipeline_deploy(pipeline, turbo)
+        assert turbo_stages[0].yaml_engine_args["hf_overrides"] == stages[0].yaml_engine_args["hf_overrides"]
         turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
         assert turbo_sampling["num_inference_steps"] == 5
         assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}

tests/diffusion/models/minimax_h3/test_minimax_h3_collectives.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+import time
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
+
+
+def _collective_worker(rank, rendezvous, backend):
+    from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
+    from vllm_omni.errors import OmniClientError
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
+    from vllm_omni.platforms import current_omni_platform
+
+    torch.set_num_threads(1)
+    device = torch.device("cuda", rank) if backend == "nccl" else torch.device("cpu")
+    if backend == "nccl":
+        current_omni_platform.set_device(device)
+    dist.init_process_group(
+        backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=45)
+    )
+    try:
+        module._dit_rank_world = lambda: (dist.group.WORLD, rank, 2)
+
+        class VisualCodec:
+            fail = False
+
+            def is_distributed_enabled(self):
+                return True
+
+            def collective(self):
+                value = torch.ones(1, device=device)
+                dist.all_reduce(value)
+                assert value.item() == 2
+
+            def encode_image(self, image):
+                self.collective()
+                assert image.size == (32, 32)
+                return torch.ones(1, 96)
+
+            def encode_video(self, frames):
+                self.collective()
+                assert frames.shape == (4, 32, 32, 3)
+                if self.fail and rank == 1:
+                    raise ValueError("codec failed on rank 1")
+                return torch.full((1, 96), 2.0), (1, 2, 2)
+
+        class AudioCodec:
+            def encode_waveform(self, waveform, sample_rate):
+                assert rank == 0
+                assert sample_rate == 32000
+                assert waveform.shape == (64000,)
+                return torch.full((160, 32), 3.0), 80
+
+        model = object.__new__(module.MiniMaxH3Pipeline)
+        torch.nn.Module.__init__(model)
+        model.device = device
+        model.od_config = SimpleNamespace(enable_layerwise_offload=False)
+        model.video_vae = VisualCodec()
+        model.audio_vae = AudioCodec()
+        model.partition = "combined"
+        model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
+        model.text_encoder_tp_size = 2
+        model.text_encoder_group = SimpleNamespace(
+            world_size=2, ranks=[0, 1], rank_in_group=rank, device_group=dist.group.WORLD
+        )
+        model.processor = SimpleNamespace(image_processor=SimpleNamespace(merge_size=2))
+        model.tokenizer = lambda text, **kwargs: {"input_ids": [1, 2, 3]}
+
+        def encode_text(ids, kwargs):
+            assert ids.tolist() == [1, 2, 3]
+            assert not kwargs
+            hidden = torch.ones(3, 5120, device=device, dtype=torch.bfloat16)
+            dist.all_reduce(hidden)
+            return hidden
+
+        model._encode_text_hidden = encode_text
+        sampling = OmniDiffusionSamplingParams(
+            height=32, width=32, num_frames=96, extra_args={"task": "t2va", "aspect_ratio": "1:1"}
+        )
+        text_conditioning = model._prepare_local_conditioning({"prompt": "test"}, sampling)
+        torch.testing.assert_close(text_conditioning.hidden_states, torch.full_like(text_conditioning.hidden_states, 2))
+        assert text_conditioning.visual_condition is None
+
+        media = (
+            MiniMaxH3EncoderMediaInput(
+                task="ref2va",
+                height=32,
+                width=32,
+                num_frames=107,
+                latent_t=32,
+                audio_t=178,
+                images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
+                videos=(torch.zeros(4, 32, 32, 3, dtype=torch.uint8),),
+                video_audios=(None,),
+                audios=((torch.zeros(64000), 32000),),
+            )
+            if rank == 0
+            else None
+        )
+        conditioning = model._encode_local_media(media)
+        assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video", "audio"]
+        torch.testing.assert_close(
+            conditioning.visual_condition.cpu(), torch.cat([torch.ones(1, 96), torch.full((1, 96), 2.0)])
+        )
+        torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
+        assert conditioning.audio_condition_lengths == (80,)
+
+        model.video_vae.fail = True
+        with pytest.raises(OmniClientError, match="codec failed on rank 1"):
+            model._encode_local_media(media)
+        with pytest.raises(OmniClientError, match="non-empty prompt"):
+            model._prepare_local_conditioning({"prompt": ""}, sampling)
+        model.video_vae.fail = False
+        assert model._encode_local_media(media).visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
+    finally:
+        dist.destroy_process_group()
+
+
+def _run_collectives(tmp_path, backend):
+    context = mp.spawn(_collective_worker, args=(str(tmp_path / "rendezvous"), backend), nprocs=2, join=False)
+    deadline = time.monotonic() + 150
+    try:
+        while not context.join(timeout=1):
+            if time.monotonic() > deadline:
+                pytest.fail("MiniMax H3 collectives did not finish within 150 seconds")
+    finally:
+        for process in context.processes:
+            if process.is_alive():
+                process.terminate()
+            process.join(timeout=5)
+
+
+@pytest.mark.cpu
+def test_two_rank_conditioning_and_errors_gloo(tmp_path):
+    _run_collectives(tmp_path, "gloo")
+
+
+@pytest.mark.cuda
+@pytest.mark.parallel
+@pytest.mark.skipif(
+    os.environ.get("VLLM_TEST_MINIMAX_H3_COLLECTIVES") != "1",
+    reason="set VLLM_TEST_MINIMAX_H3_COLLECTIVES=1 with two reserved GPUs",
+)
+def test_two_rank_conditioning_and_errors_nccl(tmp_path):
+    if torch.accelerator.device_count() < 2:
+        pytest.skip("requires two GPUs")
+    _run_collectives(tmp_path, "nccl")