diff --git a/docs/user_guide/diffusion/parallelism/vae_parallelism.md b/docs/user_guide/diffusion/parallelism/vae_parallelism.md index ca0c428e1..205b44f5d 100644 --- a/docs/user_guide/diffusion/parallelism/vae_parallelism.md +++ b/docs/user_guide/diffusion/parallelism/vae_parallelism.md @@ -120,10 +120,18 @@ Additional requirements: | Parameter | Type | Default | Description | |-----------|------|---------|-------------| | `vae_use_tiling` | bool | False | Must be set to `True` when using VAE patch parallelism. | +| `vae_stack_tiling` | `"auto"` \| `"true"` \| `"false"` | `"false"` | For VAEs with native support (currently MiniMax H3), batch local decoder tiles into fewer model calls. `"auto"` enables it when the request has at least two tiles per VAE-parallel rank. | !!! note "Automatic VAE Tiling" When `vae_patch_parallel_size > 1` and the model has a distributed VAE (`DistributedVaeMixin`), the system automatically sets `vae_use_tiling=True` if not already enabled. +!!! tip "MiniMax H3 stacked tiling" + MiniMax H3 supports lossless native stacked tiling. Enable it with + `--vae-stack-tiling auto` (or `true`) to batch each rank's spatial tiles + into one decoder call per temporal window. `auto` keeps the sequential + path for shapes that do not provide multiple tiles per rank, while + `false` remains the default to avoid increasing the VAE activation peak. + --- ## Spatially-Sharded Decode (Wan) diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index bd59ee94b..c621a7192 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -614,6 +614,54 @@ def test_initialize_model_sets_current_diffusion_config_during_model_constructio assert get_current_diffusion_config_or_none() is None +def test_configure_vae_stack_tiling_enables_tiling_and_native_capability(): + from vllm_omni.diffusion.registry import _configure_vae_stack_tiling + + calls = [] + vae = SimpleNamespace( + supports_stack_tiling=True, + set_stack_tiling=lambda mode: calls.append(mode), + ) + model = SimpleNamespace(vae=vae) + od_config = SimpleNamespace( + model_class_name="MiniMaxH3Pipeline", + vae_stack_tiling="auto", + vae_use_tiling=False, + ) + + _configure_vae_stack_tiling(model, od_config) + + assert od_config.vae_use_tiling is True + assert calls == ["auto"] + + +def test_configure_vae_stack_tiling_rejects_unsupported_vae(): + from vllm_omni.diffusion.registry import _configure_vae_stack_tiling + + od_config = SimpleNamespace( + model_class_name="UnsupportedPipeline", + vae_stack_tiling="true", + vae_use_tiling=False, + ) + + with pytest.raises(ValueError, match="does not declare stacked-tiling support"): + _configure_vae_stack_tiling(SimpleNamespace(vae=SimpleNamespace()), od_config) + + +def test_configure_vae_stack_tiling_auto_falls_back_for_unsupported_vae(): + from vllm_omni.diffusion.registry import _configure_vae_stack_tiling + + od_config = SimpleNamespace( + model_class_name="UnsupportedPipeline", + vae_stack_tiling="auto", + vae_use_tiling=False, + ) + + _configure_vae_stack_tiling(SimpleNamespace(vae=SimpleNamespace()), od_config) + + assert od_config.vae_use_tiling is False + + def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch): import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod diff --git a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py index ef04019f6..8a132babc 100644 --- a/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py +++ b/tests/diffusion/models/minimax_h3/test_minimax_h3_vae_tiling.py @@ -30,6 +30,7 @@ class _FakeCheckpointModel: def __init__(self): self.vae_ratio = RATIO self.parallel_tiling = True + self.stack_tiling = False self.decoded_with = [] def split_tiles(self, input_len, is_decoder=False): @@ -57,6 +58,7 @@ def _vae(parallel_size, state): vae.model = _FakeCheckpointModel() vae.remote = remote vae.parallel_size = parallel_size + vae.stack_tiling_mode = "false" return vae @@ -89,6 +91,33 @@ def test_tile_count_uses_the_decoder_grid_not_the_encoder_one(h, w, decoder_tile assert as_encoder == encoder_tiles != decoder_tiles +@pytest.mark.parametrize( + ("parallel_size", "h", "w", "expected"), + [ + (1, 48, 84, True), + (2, 48, 84, True), + (4, 48, 84, True), + (8, 48, 84, True), + (16, 48, 84, False), + (4, 16, 16, False), + ], +) +def test_auto_stack_tiling_requires_multiple_tiles_per_rank(parallel_size, h, w, expected): + vae = _vae(parallel_size, {}) + vae.set_stack_tiling("auto") + + assert MiniMaxH3VideoVAE._should_stack_decode_tiles(vae, _latent(h, w)) is expected + + +@pytest.mark.parametrize(("mode", "expected"), [(False, False), (True, True), ("false", False), ("true", True)]) +def test_stack_tiling_configuration_accepts_bool_and_string_modes(mode, expected): + vae = _vae(1, {}) + + vae.set_stack_tiling(mode) + + assert vae.model.stack_tiling is expected + + def test_rank_local_tiling_restores_the_group_state(): state = {"sp_size": 4, "sp_rank": 2, "sp_enabled": True, "sp_process_group": "pg"} vae = _vae(4, state) @@ -156,6 +185,7 @@ def _dispatch_vae(parallel_size): seen["sp_enabled"] = state["sp_enabled"] seen["sp_process_group"] = state["sp_process_group"] seen["parallel_tiling"] = vae.model.parallel_tiling + seen["stack_tiling"] = vae.model.stack_tiling # the guard never inspects the decoded shape, so keep it small return torch.zeros(1, 3, 2, 4, 4) @@ -212,3 +242,15 @@ def test_decode_latent_leaves_the_group_state_alone_when_the_grid_is_large_enoug assert seen["sp_enabled"] is True assert seen["sp_process_group"] is state["sp_process_group"] assert seen["parallel_tiling"] is True + + +def test_decode_latent_restores_native_stack_tiling_state(): + vae, _, seen = _dispatch_vae(4) + vae.set_stack_tiling("auto") + # The adapter must not leak its per-request decision into the remote VAE. + vae.model.stack_tiling = True + + MiniMaxH3VideoVAE.decode_latent(vae, _latent(16, 16)) + + assert seen["stack_tiling"] is False + assert vae.model.stack_tiling is True diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 71a6e00a5..aba576777 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -395,6 +395,28 @@ def test_serve_cli_accepts_diffusion_pipeline_profiler_flag(): assert stage_cfg["engine_args"]["enable_diffusion_pipeline_profiler"] is True +def test_serve_cli_forwards_vae_stack_tiling(): + parser = TrackingArgumentParser() + subparsers = parser.add_subparsers(dest="command") + OmniServeCommand().subparser_init(subparsers) + + args = parser.parse_args( + [ + "serve", + "MiniMaxAI/MiniMax-H3", + "--omni", + "--vae-stack-tiling", + "auto", + ] + ) + + explicit_kwargs = args.get_explicit_kwargs_dict() + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(explicit_kwargs)[0] + + assert args.vae_stack_tiling == "auto" + assert stage_cfg["engine_args"]["vae_stack_tiling"] == "auto" + + def test_serve_cli_forwards_distilled_lora_to_diffusion_stage(): """Ensure startup distilled LoRA options reach the online diffusion stage.""" parser = TrackingArgumentParser() diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 2ef6bf276..1dc32fbcc 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -711,6 +711,7 @@ class _DiffusionConfigProjection: fa_deterministic: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_stack_tiling: Literal["auto", "true", "false"] = "false" mask_strategy_file_path: str | None = None skip_time_steps: int = 15 VSA_sparsity: float = 0.0 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 6bfcd7e25..e37983b90 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping from dataclasses import asdict, dataclass, field, fields from enum import Enum from pathlib import Path -from typing import Any, NamedTuple +from typing import Any, Literal, NamedTuple from transformers import PretrainedConfig from vllm.logger import init_logger @@ -462,6 +462,7 @@ class StageDeployConfig: step_execution: bool | None = None vae_use_slicing: bool | None = None vae_use_tiling: bool | None = None + vae_stack_tiling: Literal["auto", "true", "false"] | None = None boundary_ratio: float | None = None flow_shift: float | None = None diffusion_kv_cache_dtype: str | None = None diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index d069d95d1..fc97d7d06 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -8,7 +8,7 @@ import random from collections.abc import Callable, Mapping from dataclasses import dataclass, field, fields from enum import Enum -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import diffusers import huggingface_hub @@ -831,6 +831,7 @@ class OmniDiffusionConfig: # VAE memory optimization parameters vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_stack_tiling: Literal["auto", "true", "false"] = "false" # STA (Sliding Tile Attention) parameters mask_strategy_file_path: str | None = None diff --git a/vllm_omni/diffusion/models/minimax_h3/vae.py b/vllm_omni/diffusion/models/minimax_h3/vae.py index 2e830fa04..5644bada3 100644 --- a/vllm_omni/diffusion/models/minimax_h3/vae.py +++ b/vllm_omni/diffusion/models/minimax_h3/vae.py @@ -9,7 +9,7 @@ import json from collections.abc import Iterator from contextlib import AbstractContextManager, contextmanager, nullcontext from pathlib import Path -from typing import Any +from typing import Any, ClassVar import torch import torch.distributed as dist @@ -115,6 +115,8 @@ class _AudioVAEDeterminismContext(AbstractContextManager): class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): """Adapter around the checkpoint's native parallel-tiled video VAE.""" + supports_stack_tiling: ClassVar[bool] = True + def __init__( self, component_path: str, @@ -153,6 +155,30 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): self.use_slicing = False self.parallel_size = 1 self.device_module = torch.get_device_module() + self.stack_tiling_mode = "false" + + def set_stack_tiling(self, mode: str | bool) -> None: + """Configure native decoder tile batching. + + ``auto`` is resolved per decode because the useful batch size depends + on both the request's tile grid and the VAE parallel group size. + """ + if not hasattr(self.model, "stack_tiling"): + raise RuntimeError("MiniMax H3 video VAE remote model does not expose stack_tiling") + if isinstance(mode, bool): + mode = "true" if mode else "false" + if mode not in ("auto", "true", "false"): + raise ValueError(f"invalid VAE stacked-tiling mode: {mode!r}") + self.stack_tiling_mode = mode + self.model.stack_tiling = mode == "true" + + def _should_stack_decode_tiles(self, latent: torch.Tensor) -> bool: + if self.stack_tiling_mode != "auto": + return self.stack_tiling_mode == "true" + # Stacking is worthwhile when every rank has more than one tile in + # the common round-robin distribution. The tile grid is global and + # therefore identical on all ranks. + return self._decoder_tile_count(latent) >= 2 * self.parallel_size def load_to_device(self) -> None: if self._stager is not None: @@ -388,7 +414,18 @@ class MiniMaxH3VideoVAE(nn.Module, DistributedVaeMixin): tiling_context = nullcontext() with tiling_context: - decoded = self.model.decode_base(latent * std + mean) + stack_tiling_mode = getattr(self, "stack_tiling_mode", "false") + if stack_tiling_mode == "false" and not hasattr(self.model, "stack_tiling"): + # Keep the default path compatible with older H3 snapshots + # that predate the native stacked-tiling attribute. + decoded = self.model.decode_base(latent * std + mean) + else: + previous_stack_tiling = self.model.stack_tiling + self.model.stack_tiling = self._should_stack_decode_tiles(latent) + try: + decoded = self.model.decode_base(latent * std + mean) + finally: + self.model.stack_tiling = previous_stack_tiling frames = self.model.processor.revert_tensor(decoded) if frames.ndim == 4: frames = frames.unsqueeze(0).transpose(1, 2) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index ed6a816a9..f7df51855 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -387,6 +387,38 @@ def _prepare_diffusion_quant_config( configure_quant_config(quant_config, model_class) +def _configure_vae_stack_tiling(model: nn.Module, od_config: OmniDiffusionConfig) -> None: + """Configure optional native batching of VAE decoder tiles.""" + mode = getattr(od_config, "vae_stack_tiling", "false") + if isinstance(mode, bool): + mode = "true" if mode else "false" + if mode not in ("auto", "true", "false"): + raise ValueError(f"invalid vae_stack_tiling mode: {mode!r}") + if mode == "false": + return + + vae = getattr(model, "vae", None) + if vae is None or not getattr(vae, "supports_stack_tiling", False): + if mode == "auto": + logger.info( + "VAE stacked tiling auto mode is unavailable for %s; using sequential tiles.", + od_config.model_class_name, + ) + return + raise ValueError( + f"vae_stack_tiling was requested, but {od_config.model_class_name} does not declare stacked-tiling support" + ) + setter = getattr(vae, "set_stack_tiling", None) + if not callable(setter): + raise ValueError("vae_stack_tiling was requested, but the VAE does not provide set_stack_tiling()") + + if not od_config.vae_use_tiling: + logger.info("vae_stack_tiling requires vae_use_tiling; automatically enabling it.") + od_config.vae_use_tiling = True + setter(mode) + logger.info("VAE stacked tiling mode set to %s for %s.", mode, od_config.model_class_name) + + def initialize_model( od_config: OmniDiffusionConfig, ) -> nn.Module: @@ -428,6 +460,8 @@ def initialize_model( ) od_config.vae_use_tiling = True + _configure_vae_stack_tiling(model, od_config) + # Configure VAE memory optimization settings from config if hasattr(model, "vae") and hasattr(model.vae, "use_slicing"): model.vae.use_slicing = od_config.vae_use_slicing diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 0eb18c686..e00765463 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -7,7 +7,7 @@ import os import tempfile from collections.abc import Callable from dataclasses import dataclass, field, fields -from typing import Any, cast +from typing import Any, Literal, cast from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs from vllm.logger import init_logger @@ -566,6 +566,7 @@ class OrchestratorArgs: step_execution: bool = False vae_use_slicing: bool = False vae_use_tiling: bool = False + vae_stack_tiling: Literal["auto", "true", "false"] = "false" enable_multithread_weight_load: bool = True num_weight_load_threads: int = 4 enable_cpu_offload: bool = False diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 9e894039f..1d0add076 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1132,6 +1132,7 @@ class AsyncOmniEngine: "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), "vae_use_slicing": kwargs.get("vae_use_slicing", False), "vae_use_tiling": kwargs.get("vae_use_tiling", False), + "vae_stack_tiling": kwargs.get("vae_stack_tiling", "false"), "cache_backend": cache_backend, "cache_config": cache_config, "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index f4c890816..702287570 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -686,6 +686,15 @@ class OmniServeCommand(CLISubcommand): action="store_true", help="Enable VAE tiling for memory optimization (useful for mitigating OOM issues).", ) + omni_config_group.add_argument( + "--vae-stack-tiling", + choices=("auto", "true", "false"), + default="false", + help=( + "Batch native VAE tiles into fewer model calls. 'auto' enables it " + "when each VAE-parallel rank has multiple tiles (default: false)." + ), + ) # Parallel weight loading (faster diffusion startup) omni_config_group.add_argument(