[Perf][Diffusion] Add configurable H3 VAE stacked tiling

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

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)

tests/diffusion/kernels/mot/test_mot_gemm.py

@@ -1,95 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-
-from __future__ import annotations
-
-import pytest
-import torch
-
-from vllm_omni.diffusion.layers.mot.ops import mot_gemm
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-_CONFIG = {
-    "BLOCK_SIZE_M": 16,
-    "BLOCK_SIZE_N": 64,
-    "BLOCK_SIZE_K": 32,
-    "GROUP_SIZE_M": 8,
-    "num_warps": 4,
-    "num_stages": 2,
-}
-
-
-class _RecordingKernel:
-    def __init__(self) -> None:
-        self.calls: list[tuple[tuple[object, ...], dict[str, object]]] = []
-
-    def __getitem__(self, grid):
-        del grid
-        return self._launch
-
-    def _launch(self, *args, **kwargs) -> None:
-        self.calls.append((args, kwargs))
-
-
-@pytest.mark.parametrize(
-    ("bias_text", "bias_vae"),
-    [(None, None), ("text", None), (None, "vae"), ("text", "vae")],
-    ids=["no-bias", "text-only", "vae-only", "both"],
-)
-def test_invoke_mot_gemm_accepts_independent_biases(
-    bias_text: str | None,
-    bias_vae: str | None,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    kernel = _RecordingKernel()
-    monkeypatch.setattr(mot_gemm, "mot_unified_gemm_kernel", kernel)
-
-    A = torch.ones((3, 4), dtype=torch.float16)
-    B_text = torch.ones((4, 5), dtype=torch.float16)
-    B_vae = torch.ones((4, 5), dtype=torch.float16)
-    C = torch.empty((3, 5), dtype=torch.float16)
-    text_indices = torch.tensor([0, 2], dtype=torch.long)
-    vae_indices = torch.tensor([1], dtype=torch.long)
-    text_bias = torch.ones(5, dtype=torch.float16) if bias_text else None
-    vae_bias = torch.ones(5, dtype=torch.float16) if bias_vae else None
-
-    mot_gemm.invoke_mot_gemm(
-        A=A,
-        B_text=B_text,
-        B_vae=B_vae,
-        C=C,
-        bias_text=text_bias,
-        bias_vae=vae_bias,
-        text_indices=text_indices,
-        vae_indices=vae_indices,
-        A_scale=None,
-        B_text_scale=None,
-        B_vae_scale=None,
-        use_fp8_w8a8=False,
-        use_int8_w8a8=False,
-        use_int8_w8a16=False,
-        use_int4_w4a16=False,
-        A_per_channel_quant=False,
-        B_per_channel_quant=False,
-        config=_CONFIG,
-    )
-
-    assert len(kernel.calls) == 1
-    args, kwargs = kernel.calls[0]
-    assert kwargs["HAS_BIAS_TEXT"] is (text_bias is not None)
-    assert kwargs["HAS_BIAS_VAE"] is (vae_bias is not None)
-    assert "HAS_BIAS" not in kwargs
-
-    launched_text_bias, launched_vae_bias = args[4:6]
-    if text_bias is None and vae_bias is None:
-        assert launched_text_bias == launched_vae_bias == 0
-    elif text_bias is None:
-        assert launched_text_bias is vae_bias
-        assert launched_vae_bias is vae_bias
-    elif vae_bias is None:
-        assert launched_text_bias is text_bias
-        assert launched_vae_bias is text_bias
-    else:
-        assert launched_text_bias is text_bias
-        assert launched_vae_bias is vae_bias

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -513,57 +513,6 @@ def test_mot_o_proj(
             )
 
 
-@pytest.mark.parametrize(
-    ("bias_text", "bias_vae"),
-    [(True, False), (False, True)],
-    ids=["text-only", "vae-only"],
-)
-def test_mot_row_one_sided_bias(bias_text: bool, bias_vae: bool):
-    """The fused row projection must support a bias on only one expert."""
-    dcfg = _parse_dtype("w16a16_bf16")
-    K = 32
-    N = 64
-    M = 5
-
-    with set_current_vllm_config(VllmConfig()):
-        text_linear = RowParallelLinear(
-            K,
-            N,
-            bias=bias_text,
-            input_is_parallel=True,
-            params_dtype=dcfg.torch_dtype,
-            disable_tp=True,
-        ).cuda()
-        vae_linear = RowParallelLinear(
-            K,
-            N,
-            bias=bias_vae,
-            input_is_parallel=True,
-            params_dtype=dcfg.torch_dtype,
-            disable_tp=True,
-        ).cuda()
-        mot_linear = MoTRowParallelLinear(
-            K,
-            N,
-            bias=bias_text,
-            vae_bias=bias_vae,
-            input_is_parallel=True,
-            params_dtype=dcfg.torch_dtype,
-            disable_tp=True,
-        ).cuda()
-
-        _sync_weights(text_linear, vae_linear, mot_linear)
-        text_idx = torch.tensor([0, 2], dtype=torch.long, device="cuda")
-        vae_idx = torch.tensor([1, 3, 4], dtype=torch.long, device="cuda")
-        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
-
-        with torch.no_grad():
-            ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear)
-            mot_out, _ = mot_linear(x, text_idx, vae_idx)
-
-        _check_and_report(ref, mot_out, f"Row one-sided bias text={bias_text} vae={bias_vae}")
-
-
 # =========================================================================
 #  Test: und-mode (text_indices=None) — falls back to standard forward
 # =========================================================================

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
 

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

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()

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

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

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

vllm_omni/diffusion/layers/mot/ops/mot_gemm.py

@@ -238,7 +238,6 @@ def _get_mot_pointers(
     cur_stride_bk = stride_bk_vae
     cur_stride_bn = stride_bn_vae
     M_limit = M_vae
-    is_vae = 1
 
     # Text Path
     if pid < num_pid_m_text * num_pid_n:
@@ -252,7 +251,6 @@ def _get_mot_pointers(
         cur_stride_bk = stride_bk_text
         cur_stride_bn = stride_bn_text
         M_limit = M_text
-        is_vae = 0
 
     # 3. Calculate Grid coordinates(grouping)
     cur_num_pid_m = tl.cdiv(M_limit, BLOCK_SIZE_M)
@@ -285,7 +283,6 @@ def _get_mot_pointers(
         offs_n,
         n_mask,  # N-dim info
         M_limit,  # Boundary
-        is_vae,  # Selected expert kind: 0=text, 1=VAE
         cur_b_ptr,
         cur_bias_ptr,  # Selected Pointers
         cur_scale_b_ptr,  # Selected Scale Pointer
@@ -536,8 +533,7 @@ def mot_unified_gemm_kernel(
     # Quant Control
     # 0=None, 1=W8A8, 2=W8A16, 3=W4A16
     QUANT_TYPE: tl.constexpr,
-    HAS_BIAS_TEXT: tl.constexpr = False,
-    HAS_BIAS_VAE: tl.constexpr = False,
+    HAS_BIAS: tl.constexpr = False,
 ):
     pid = tl.program_id(axis=0)
 
@@ -552,7 +548,6 @@ def mot_unified_gemm_kernel(
         offs_n,
         n_mask,
         M_limit,
-        is_vae,
         cur_b_ptr,
         cur_bias_ptr,
         cur_scale_b_ptr,
@@ -680,12 +675,8 @@ def mot_unified_gemm_kernel(
     # -----------------------------------------------------------
 
     # Bias Add
-    if HAS_BIAS_TEXT or HAS_BIAS_VAE:
+    if HAS_BIAS:
         bias = tl.load(cur_bias_ptr + offs_n, mask=n_mask, other=0.0)
-        if HAS_BIAS_TEXT and not HAS_BIAS_VAE:
-            bias = tl.where(is_vae == 0, bias, 0.0)
-        elif HAS_BIAS_VAE and not HAS_BIAS_TEXT:
-            bias = tl.where(is_vae == 1, bias, 0.0)
         c = c + bias[None, :]  # reshape to (1, BLOCK_SIZE_N)
 
     # Cast C into the output dtype
@@ -841,8 +832,11 @@ def invoke_mot_gemm(
     STRIDE_BK_IS_1 = (B_text.stride(0) == 1) and (B_vae.stride(0) == 1)
     STRIDE_BN_IS_1 = (B_text.stride(1) == 1) and (B_vae.stride(1) == 1)
 
-    has_bias_text = bias_text is not None
-    has_bias_vae = bias_vae is not None
+    # bias check
+    assert (bias_text is None) == (bias_vae is None), (
+        "Bias must be provided for both Text and VAE simultaneously, or neither."
+    )
+    has_bias = bias_text is not None
 
     # --- 3. Grid Calculation ---
     def grid(META):
@@ -859,8 +853,7 @@ def invoke_mot_gemm(
             "ACCUMULATOR_DTYPE": ACCUMULATOR_DTYPE,
             "COMPUTE_DTYPE": COMPUTE_DTYPE,
             "OUTPUT_DTYPE": OUTPUT_DTYPE,
-            "HAS_BIAS_TEXT": has_bias_text,
-            "HAS_BIAS_VAE": has_bias_vae,
+            "HAS_BIAS": has_bias,
             "EVEN_K": EVEN_K,
             "EVEN_N": EVEN_N,
             "STRIDE_AK_IS_1": STRIDE_AK_IS_1,
@@ -873,17 +866,8 @@ def invoke_mot_gemm(
     p_a_scale = A_scale if A_scale is not None else 0
     p_b_text_scale = B_text_scale if B_text_scale is not None else 0
     p_b_vae_scale = B_vae_scale if B_vae_scale is not None else 0
-    # Triton requires both route-selected bias arguments to have a pointer
-    # type. Alias the missing route to the available tensor; its compile-time
-    # bias flag and runtime route mask keep it from being read for that route.
-    if bias_text is None:
-        p_bias_text = bias_vae if bias_vae is not None else 0
-    else:
-        p_bias_text = bias_text
-    if bias_vae is None:
-        p_bias_vae = bias_text if bias_text is not None else 0
-    else:
-        p_bias_vae = bias_vae
+    p_bias_text = bias_text if bias_text is not None else 0
+    p_bias_vae = bias_vae if bias_vae is not None else 0
 
     # Quantization granularity
     stride_scale_a = 1 if A_per_channel_quant else 0

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)

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

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

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),

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(