[diffusion] model: support MOVA (#17704)

Co-authored-by: gaoyang07 <Gary1546308416AL@gmail.com>
Co-authored-by: cms42 <c@cms42.top>
Co-authored-by: cms42 <44895820+cms42@users.noreply.github.com>
Co-authored-by: Ruixiao Li <cgruixiao@outlook.com>
Co-authored-by: Li Ruixiao(SII) <80368770+Li-dongyang@users.noreply.github.com>

python/sglang/multimodal_gen/configs/models/bridges/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
+    MOVADualTowerConfig,
+)
+
+__all__ = ["MOVADualTowerConfig"]

python/sglang/multimodal_gen/configs/models/bridges/mova_dual_tower.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Configuration for MOVA dual tower bridge model."""
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def _is_conditioner_block(name: str, module) -> bool:
+    """Check if module is a ConditionalCrossAttentionBlock."""
+    return "ConditionalCrossAttentionBlock" in type(module).__name__
+
+
+@dataclass
+class MOVADualTowerArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(
+        default_factory=lambda: [_is_conditioner_block]
+    )
+
+    # Model architecture parameters
+    visual_layers: int = 40
+    audio_layers: int = 30
+    visual_hidden_dim: int = 5120
+    audio_hidden_dim: int = 1536
+    audio_fps: float = 50.0
+    head_dim: int = 128
+    interaction_strategy: str = "full"
+    apply_cross_rope: bool = True
+    apply_first_frame_bias_in_rope: bool = False
+    trainable_condition_scale: bool = False
+    pooled_adaln: bool = False
+    eps: float = 1e-6
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.visual_hidden_dim
+        self.num_attention_heads = self.visual_hidden_dim // self.head_dim
+
+
+@dataclass
+class MOVADualTowerConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MOVADualTowerArchConfig)

python/sglang/multimodal_gen/configs/models/dits/__init__.py

@@ -1,6 +1,13 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
 from sglang.multimodal_gen.configs.models.dits.hunyuanvideo import HunyuanVideoConfig
+from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
+from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
 from sglang.multimodal_gen.configs.models.dits.wanvideo import WanVideoConfig
 
-__all__ = ["HunyuanVideoConfig", "WanVideoConfig"]
+__all__ = [
+    "HunyuanVideoConfig",
+    "MOVAAudioConfig",
+    "MOVAVideoConfig",
+    "WanVideoConfig",
+]

python/sglang/multimodal_gen/configs/models/dits/mova_audio.py

@@ -0,0 +1,66 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def _is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class MOVAAudioArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_blocks])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1",
+            r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1",
+            r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1",
+        }
+    )
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    dim: int = 1536
+    in_dim: int = 128
+    ffn_dim: int = 6144
+    out_dim: int = 128
+    text_dim: int = 4096
+    freq_dim: int = 256
+    eps: float = 1e-6
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    num_heads: int = 12
+    num_layers: int = 30
+    has_image_input: bool = False
+    has_image_pos_emb: bool = False
+    has_ref_conv: bool = False
+    add_control_adapter: bool = False
+    in_dim_control_adapter: int = 24
+    seperated_timestep: bool = False
+    require_vae_embedding: bool = False
+    require_clip_embedding: bool = False
+    fuse_vae_embedding_in_latents: bool = False
+    vae_type: str = "dac"
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.num_heads
+        self.num_channels_latents = self.out_dim
+        assert (
+            not self.has_image_input
+        ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)."
+
+
+@dataclass
+class MOVAAudioConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MOVAAudioArchConfig)
+    prefix: str = "mova_audio"

python/sglang/multimodal_gen/configs/models/dits/mova_video.py

@@ -0,0 +1,65 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.models.dits.base import DiTArchConfig, DiTConfig
+
+
+def _is_blocks(n: str, m) -> bool:
+    return "blocks" in n and str.isdigit(n.split(".")[-1])
+
+
+@dataclass
+class MOVAVideoArchConfig(DiTArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_is_blocks])
+
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^text_embedding\.0\.(.*)$": r"text_embedding.fc_in.\1",
+            r"^text_embedding\.2\.(.*)$": r"text_embedding.fc_out.\1",
+            r"^time_embedding\.0\.(.*)$": r"time_embedding.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"time_embedding.fc_out.\1",
+            r"^img_emb\.proj\.1\.(.*)$": r"img_emb.fc_in.\1",
+            r"^img_emb\.proj\.3\.(.*)$": r"img_emb.fc_out.\1",
+        }
+    )
+    reverse_param_names_mapping: dict = field(default_factory=dict)
+    lora_param_names_mapping: dict = field(default_factory=dict)
+
+    dim: int = 5120
+    in_dim: int = 16
+    ffn_dim: int = 13824
+    out_dim: int = 16
+    text_dim: int = 4096
+    freq_dim: int = 256
+    eps: float = 1e-6
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    num_heads: int = 40
+    num_layers: int = 40
+    has_image_input: bool = False
+    has_image_pos_emb: bool = False
+    has_ref_conv: bool = False
+    add_control_adapter: bool = False
+    in_dim_control_adapter: int = 24
+    seperated_timestep: bool = False
+    require_vae_embedding: bool = True
+    require_clip_embedding: bool = True
+    fuse_vae_embedding_in_latents: bool = False
+
+    def __post_init__(self):
+        super().__post_init__()
+        self.hidden_size = self.dim
+        self.num_attention_heads = self.num_heads
+        self.num_channels_latents = self.out_dim
+        assert (
+            not self.has_image_input
+        ), "has_image_input must be False; it's a config from Diffsynth Studio, which means the model uses CLIP for image encoding (we don't)."
+
+
+@dataclass
+class MOVAVideoConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=MOVAVideoArchConfig)
+    prefix: str = "mova_video"

python/sglang/multimodal_gen/configs/models/vaes/__init__.py

@@ -1,9 +1,11 @@
 # Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
 
+from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig
 from sglang.multimodal_gen.configs.models.vaes.hunyuanvae import HunyuanVAEConfig
 from sglang.multimodal_gen.configs.models.vaes.wanvae import WanVAEConfig
 
 __all__ = [
+    "DacVAEConfig",
     "HunyuanVAEConfig",
     "WanVAEConfig",
 ]

python/sglang/multimodal_gen/configs/models/vaes/dac.py

@@ -0,0 +1,30 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/dac_vae.py
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+from typing import List
+
+from sglang.multimodal_gen.configs.models.base import ArchConfig, ModelConfig
+
+
+@dataclass
+class DacVAEArchConfig(ArchConfig):
+    codebook_dim: int = 8
+    codebook_size: int = 1024
+    continuous: bool = True
+    decoder_dim: int = 2048
+    decoder_rates: List[int] = field(default_factory=lambda: [8, 5, 4, 3, 2])
+    encoder_dim: int = 128
+    encoder_rates: List[int] = field(default_factory=lambda: [2, 3, 4, 5, 8])
+    hop_length: int = 3840
+    latent_dim: int = 128
+    n_codebooks: int = 9
+    quantizer_dropout: bool = False
+    sample_rate: int = 48000
+
+
+@dataclass
+class DacVAEConfig(ModelConfig):
+    arch_config: DacVAEArchConfig = field(default_factory=DacVAEArchConfig)
+    load_encoder: bool = True
+    load_decoder: bool = True

python/sglang/multimodal_gen/configs/pipeline_configs/__init__.py

@@ -20,6 +20,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan import (
     HunyuanConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
 from sglang.multimodal_gen.configs.pipeline_configs.wan import (
     SelfForcingWanT2V480PConfig,
     WanI2V480PConfig,
@@ -39,6 +40,7 @@ __all__ = [
     "Flux2FinetunedPipelineConfig",
     "PipelineConfig",
     "SlidingTileAttnConfig",
+    "MOVAPipelineConfig",
     "WanT2V480PConfig",
     "WanI2V480PConfig",
     "WanT2V720PConfig",

python/sglang/multimodal_gen/configs/pipeline_configs/mova.py

@@ -0,0 +1,185 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+MOVA pipeline configuration.
+"""
+
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from sglang.multimodal_gen.configs.models.dits import MOVAAudioConfig, MOVAVideoConfig
+from sglang.multimodal_gen.configs.models.encoders import T5Config
+from sglang.multimodal_gen.configs.models.vaes import DacVAEConfig, WanVAEConfig
+from sglang.multimodal_gen.configs.pipeline_configs.base import (
+    ModelTaskType,
+    PipelineConfig,
+)
+from sglang.multimodal_gen.configs.pipeline_configs.wan import t5_postprocess_text
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class MOVAPipelineConfig(PipelineConfig):
+    """Configuration for MOVA (text+image -> video+audio) pipelines."""
+
+    task_type: ModelTaskType = ModelTaskType.T2V
+
+    # Model configs
+    dit_config: MOVAVideoConfig = field(default_factory=MOVAVideoConfig)
+    audio_dit_config: MOVAAudioConfig = field(default_factory=MOVAAudioConfig)
+
+    # Video VAE (Wan) + Audio VAE (DAC)
+    vae_config: WanVAEConfig = field(default_factory=WanVAEConfig)
+    audio_vae_config: DacVAEConfig = field(default_factory=DacVAEConfig)
+    audio_vae_precision: str = "fp32"
+
+    # Text encoder (UMT5 compatible)
+    text_encoder_configs: tuple = field(default_factory=lambda: (T5Config(),))
+    postprocess_text_funcs: tuple = field(
+        default_factory=lambda: (t5_postprocess_text,)
+    )
+
+    # MOVA specific
+    audio_vae_type: str = "dac"
+    boundary_ratio: float | None = 0.9
+
+    # temporal alignment: MOVA expects (num_frames - 1) % 4 == 0
+    time_division_factor: int = 4
+    time_division_remainder: int = 1
+
+    def _center_crop_and_resize(
+        self, image: torch.Tensor | Image.Image, target_height: int, target_width: int
+    ) -> torch.Tensor | Image.Image:
+        if not isinstance(image, (Image.Image, torch.Tensor)):
+            raise TypeError(f"Unsupported image type: {type(image)}")
+        if isinstance(image, Image.Image):
+            image = torch.from_numpy(np.array(image))
+
+        if image.ndim == 2:
+            image = image[..., None]
+
+        if not image.dtype.is_floating_point:
+            image = image.to(torch.float32).div(255.0)
+
+        if image.ndim == 3:
+            if image.shape[0] in (1, 3, 4) and image.shape[-1] not in (1, 3, 4):
+                image = image.unsqueeze(0)
+            else:
+                image = image.permute(2, 0, 1).unsqueeze(0)
+        elif image.ndim == 4:
+            if image.shape[1] not in (1, 3, 4) and image.shape[-1] in (1, 3, 4):
+                image = image.permute(0, 3, 1, 2)
+
+        image_height, image_width = image.shape[-2], image.shape[-1]
+        if image_height == target_height and image_width == target_width:
+            return image
+
+        logger.info(
+            "Center cropping and resizing image to %dx%d", target_width, target_height
+        )
+
+        if image_height * target_width < image_width * target_height:
+            cropped_width = (image_height * target_width) // target_height
+            left = (image_width - cropped_width) // 2
+            image = image[..., :, left : left + cropped_width]
+        else:
+            cropped_height = (image_width * target_height) // target_width
+            top = (image_height - cropped_height) // 2
+            image = image[..., top : top + cropped_height, :]
+
+        image = F.interpolate(
+            image,
+            size=(target_height, target_width),
+            mode="bilinear",
+            align_corners=False,
+            antialias=True,
+        )
+        return image
+
+    def adjust_num_frames(self, num_frames: int) -> int:
+        if num_frames is None:
+            return num_frames
+        if num_frames % self.time_division_factor != self.time_division_remainder:
+            adjusted = (
+                (num_frames + self.time_division_factor - 1)
+                // self.time_division_factor
+                * self.time_division_factor
+                + self.time_division_remainder
+            )
+            logger.warning(
+                "`num_frames` (%s) is not compatible with MOVA temporal constraints. "
+                "Rounding to %s.",
+                num_frames,
+                adjusted,
+            )
+            return adjusted
+        return num_frames
+
+    def preprocess_condition_image(
+        self, image, target_width, target_height, _vae_image_processor
+    ):
+        image = self._center_crop_and_resize(image, target_height, target_width)
+        return image, (target_width, target_height)
+
+    def prepare_latent_shape(self, batch, batch_size, num_frames):
+        spatial = self.vae_config.arch_config.spatial_compression_ratio
+        length = (num_frames - 1) // self.time_division_factor + 1
+        shape = (
+            batch_size,
+            self.dit_config.arch_config.out_dim,
+            length,
+            batch.height // spatial,
+            batch.width // spatial,
+        )
+        return shape
+
+    def prepare_audio_latent_shape(self, batch_size, num_samples, audio_vae):
+        latent_T = (num_samples + audio_vae.hop_length - 1) // audio_vae.hop_length
+        return (batch_size, audio_vae.latent_dim, latent_T)
+
+    def normalize_video_latents(self, latents: torch.Tensor, video_vae) -> torch.Tensor:
+        latents_mean = getattr(video_vae.config, "latents_mean", None)
+        latents_std = getattr(video_vae.config, "latents_std", None)
+        if latents_mean is None or latents_std is None:
+            return latents
+        mean = torch.tensor(
+            latents_mean, device=latents.device, dtype=latents.dtype
+        ).view(1, video_vae.config.z_dim, 1, 1, 1)
+        inv_std = (
+            1.0 / torch.tensor(latents_std, device=latents.device, dtype=latents.dtype)
+        ).view(1, video_vae.config.z_dim, 1, 1, 1)
+        return (latents - mean) * inv_std
+
+    def denormalize_video_latents(
+        self, latents: torch.Tensor, video_vae
+    ) -> torch.Tensor:
+        latents_mean = getattr(video_vae.config, "latents_mean", None)
+        latents_std = getattr(video_vae.config, "latents_std", None)
+        if latents_mean is None or latents_std is None:
+            return latents
+        mean = torch.tensor(
+            latents_mean, device=latents.device, dtype=latents.dtype
+        ).view(1, video_vae.config.z_dim, 1, 1, 1)
+        std = torch.tensor(
+            latents_std, device=latents.device, dtype=latents.dtype
+        ).view(1, video_vae.config.z_dim, 1, 1, 1)
+        return latents * std + mean
+
+
+@dataclass
+class MOVA360PConfig(MOVAPipelineConfig):
+    """Configuration for MOVA 360P (text+image -> video+audio) pipelines."""
+
+    max_area: int = 352 * 640
+
+
+@dataclass
+class MOVA720PConfig(MOVAPipelineConfig):
+    """Configuration for MOVA 720P (text+image -> video+audio) pipelines."""
+
+    max_area: int = 720 * 1280

python/sglang/multimodal_gen/configs/sample/mova.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
+
+
+@dataclass
+class MOVASamplingParams(SamplingParams):
+    # Video parameters (MOVA defaults)
+    height: int = 352
+    width: int = 640
+    num_frames: int = 193
+    fps: int = 24
+
+    # Denoising stage
+    guidance_scale: float = 5.0
+    num_inference_steps: int = 50
+    sigma_shift: float = 5.0
+    visual_shift: float = 5.0
+    audio_shift: float = 5.0
+
+    negative_prompt: str = (
+        "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，"
+        "整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，"
+        "画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，"
+        "静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+    )
+
+
+@dataclass
+class MOVA_360P_SamplingParams(MOVASamplingParams):
+    # Video parameters (MOVA 360P)
+    height: int = 352
+    width: int = 640
+
+    # MOVA 360P supported resolutions
+    supported_resolutions: list[tuple[int, int]] = field(
+        default_factory=lambda: [
+            (352, 640),
+            (640, 352),
+        ]
+    )
+
+
+@dataclass
+class MOVA_720P_SamplingParams(MOVASamplingParams):
+    # Video parameters (MOVA 720P)
+    height: int = 720
+    width: int = 1280
+
+    # MOVA 720P supported resolutions
+    supported_resolutions: list[tuple[int, int]] = field(
+        default_factory=lambda: [
+            (720, 1280),
+            (1280, 720),
+        ]
+    )

python/sglang/multimodal_gen/registry.py

@@ -46,6 +46,10 @@ from sglang.multimodal_gen.configs.pipeline_configs.glm_image import (
     GlmImagePipelineConfig,
 )
 from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
+from sglang.multimodal_gen.configs.pipeline_configs.mova import (
+    MOVA360PConfig,
+    MOVA720PConfig,
+)
 from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
     QwenImageEditPipelineConfig,
     QwenImageEditPlus_2511_PipelineConfig,
@@ -72,6 +76,10 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
     HunyuanSamplingParams,
 )
 from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams
+from sglang.multimodal_gen.configs.sample.mova import (
+    MOVA_360P_SamplingParams,
+    MOVA_720P_SamplingParams,
+)
 from sglang.multimodal_gen.configs.sample.qwenimage import (
     QwenImage2512SamplingParams,
     QwenImageEditPlusSamplingParams,
@@ -543,6 +551,21 @@ def _register_configs():
             "FastVideo/FastWan2.1-T2V-1.3B-Diffusers",
         ],
     )
+    # MOVA
+    register_configs(
+        sampling_param_cls=MOVA_360P_SamplingParams,
+        pipeline_config_cls=MOVA360PConfig,
+        model_detectors=[
+            lambda hf_id: "mova" in hf_id.lower() and "360p" in hf_id.lower()
+        ],
+    )
+    register_configs(
+        sampling_param_cls=MOVA_720P_SamplingParams,
+        pipeline_config_cls=MOVA720PConfig,
+        model_detectors=[
+            lambda hf_id: "mova" in hf_id.lower() and "720p" in hf_id.lower()
+        ],
+    )
     # FLUX
     register_configs(
         sampling_param_cls=FluxSamplingParams,

python/sglang/multimodal_gen/runtime/loader/component_loader.py

@@ -268,18 +268,31 @@ class ComponentLoader(ABC):
             "processor": (AutoProcessorLoader, "transformers"),
             "vision_language_encoder": (VisionLanguageEncoderLoader, "transformers"),
         }
-
         # Loaders for audio/video specific components that might vary
         av_module_loaders = {
+            "audio_dit": (TransformerLoader, "diffusers"),
             "audio_vae": (VAELoader, "diffusers"),
-            "vocoder": (VocoderLoader, "diffusers"),
             "connectors": (AdapterLoader, "diffusers"),
+            "dual_tower_bridge": (BridgeLoader, "diffusers"),
+            "video_dit": (TransformerLoader, "diffusers"),
+            "video_vae": (VAELoader, "diffusers"),
+            "vocoder": (VocoderLoader, "diffusers"),
         }
 
         # NOTE(FlamingoPg): special for LTX-2 models
         if module_type == "vocoder" or module_type == "connectors":
             transformers_or_diffusers = "diffusers"
 
+        # NOTE(CloudRipple): special for MOVA models
+        # TODO(CloudRipple): remove most of these special cases after unifying the loading logic
+        if module_type in [
+            "audio_vae",
+            "audio_dit",
+            "dual_tower_bridge",
+            "video_dit",
+        ]:
+            transformers_or_diffusers = "diffusers"
+
         if module_type in module_loaders:
             loader_cls, expected_library = module_loaders[module_type]
             # Assert that the library matches what's expected for this module type
@@ -616,7 +629,7 @@ class TokenizerLoader(ComponentLoader):
 
 
 class VAELoader(ComponentLoader):
-    """Loader for VAE."""
+    """Shared loader for (video/audio) VAE modules."""
 
     def should_offload(
         self, server_args: ServerArgs, model_config: ModelConfig | None = None
@@ -636,17 +649,20 @@ class VAELoader(ComponentLoader):
         server_args.model_paths[module_name] = component_model_path
 
         logger.debug("HF model config: %s", config)
-        if module_name == "audio_vae":
-            vae_config = server_args.pipeline_config.audio_vae_config
-            vae_precision = server_args.pipeline_config.audio_vae_precision
+        if module_name in ("vae", "video_vae"):
+            pipeline_vae_config_attr = "vae_config"
+            pipeline_vae_precision = "vae_precision"
+        elif module_name in ("audio_vae",):
+            pipeline_vae_config_attr = "audio_vae_config"
+            pipeline_vae_precision = "audio_vae_precision"
         else:
-            vae_config = server_args.pipeline_config.vae_config
-            vae_precision = server_args.pipeline_config.vae_precision
-
+            raise ValueError(f"Unsupported module name for VAE loader: {module_name}")
+        vae_config = getattr(server_args.pipeline_config, pipeline_vae_config_attr)
+        vae_precision = getattr(server_args.pipeline_config, pipeline_vae_precision)
         vae_config.update_model_arch(config)
-
-        # NOTE: some post init logics are only available after updated with config
-        vae_config.post_init()
+        if hasattr(vae_config, "post_init"):
+            # NOTE: some post init logics are only available after updated with config
+            vae_config.post_init()
 
         should_offload = self.should_offload(server_args)
         target_device = self.target_device(should_offload)
@@ -685,6 +701,16 @@ class VAELoader(ComponentLoader):
         ), f"Found {len(safetensors_list)} safetensors files in {component_model_path}"
         loaded = safetensors_load_file(safetensors_list[0])
         vae.load_state_dict(loaded, strict=False)
+
+        state_keys = set(vae.state_dict().keys())
+        loaded_keys = set(loaded.keys())
+        missing_keys = sorted(state_keys - loaded_keys)
+        unexpected_keys = sorted(loaded_keys - state_keys)
+        if missing_keys:
+            logger.warning("VAE missing keys: %s", missing_keys)
+        if unexpected_keys:
+            logger.warning("VAE unexpected keys: %s", unexpected_keys)
+
         return vae.eval()
 
 
@@ -753,11 +779,96 @@ class VocoderLoader(ComponentLoader):
         return vocoder.eval()
 
 
-class TransformerLoader(ComponentLoader):
-    """Loader for transformer."""
+class BridgeLoader(ComponentLoader):
+    """Loader for MOVA dual tower bridge with FSDP support."""
+
+    pipeline_bridge_config_attr: str = "bridge_config"
 
     def load_customized(
-        self, component_model_path: str, server_args: ServerArgs, *args
+        self, component_model_path: str, server_args: ServerArgs, module_name: str
+    ):
+        config = get_diffusers_component_config(model_path=component_model_path)
+        hf_config = deepcopy(config)
+        class_name = config.pop("_class_name", None)
+        if class_name is None:
+            raise ValueError(
+                "Model config does not contain a _class_name attribute. "
+                "Only diffusers format is supported."
+            )
+        server_args.model_paths[module_name] = component_model_path
+
+        # Try to get bridge config from pipeline config, fallback to creating one
+        bridge_config = getattr(
+            server_args.pipeline_config, self.pipeline_bridge_config_attr, None
+        )
+        if bridge_config is not None:
+            bridge_config.update_model_arch(config)
+        else:
+            # Create a minimal config from hf_config
+            from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
+                MOVADualTowerConfig,
+            )
+
+            bridge_config = MOVADualTowerConfig()
+            bridge_config.update_model_arch(config)
+
+        model_cls, _ = ModelRegistry.resolve_model_cls(class_name)
+
+        # Find all safetensors files
+        safetensors_list = _list_safetensors_files(component_model_path)
+        if not safetensors_list:
+            raise ValueError(f"No safetensors files found in {component_model_path}")
+
+        default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+
+        logger.info(
+            "Loading %s from %s safetensors files, default_dtype: %s",
+            class_name,
+            len(safetensors_list),
+            default_dtype,
+        )
+
+        # Check if FSDP loading is available
+        if (
+            server_args.hsdp_shard_dim is not None
+            and hasattr(model_cls, "_fsdp_shard_conditions")
+            and model_cls._fsdp_shard_conditions
+        ):
+            # Load with FSDP support
+            model = maybe_load_fsdp_model(
+                model_cls=model_cls,
+                init_params={"config": bridge_config, "hf_config": hf_config},
+                weight_dir_list=safetensors_list,
+                device=get_local_torch_device(),
+                hsdp_replicate_dim=server_args.hsdp_replicate_dim,
+                hsdp_shard_dim=server_args.hsdp_shard_dim,
+                cpu_offload=server_args.dit_cpu_offload,
+                pin_cpu_memory=server_args.pin_cpu_memory,
+                fsdp_inference=server_args.use_fsdp_inference,
+                default_dtype=default_dtype,
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.float32,
+                output_dtype=None,
+                strict=False,
+            )
+        else:
+            # Fallback to simple loading (for non-FSDP or legacy models)
+            model = model_cls.from_pretrained(
+                component_model_path, torch_dtype=default_dtype
+            )
+            model = model.to(device=get_local_torch_device(), dtype=default_dtype)
+
+        total_params = sum(p.numel() for p in model.parameters())
+        logger.info("Loaded bridge model with %.2fM parameters", total_params / 1e6)
+
+        return model.eval()
+
+
+class TransformerLoader(ComponentLoader):
+    """Shared loader for (video/audio) DiT transformers."""
+
+    def load_customized(
+        self, component_model_path: str, server_args: ServerArgs, module_name: str
     ):
         """Load the transformer based on the model path, and inference args."""
         config = get_diffusers_component_config(model_path=component_model_path)
@@ -769,10 +880,17 @@ class TransformerLoader(ComponentLoader):
                 "Only diffusers format is supported."
             )
 
-        server_args.model_paths["transformer"] = component_model_path
+        module_name = _normalize_module_type(module_name)
+        server_args.model_paths[module_name] = component_model_path
 
+        if module_name in ("transformer", "video_dit"):
+            pipeline_dit_config_attr = "dit_config"
+        elif module_name in ("audio_dit",):
+            pipeline_dit_config_attr = "audio_dit_config"
+        else:
+            raise ValueError(f"Invalid module name: {module_name}")
         # Config from Diffusers supersedes sgl_diffusion's model config
-        dit_config = server_args.pipeline_config.dit_config
+        dit_config = getattr(server_args.pipeline_config, pipeline_dit_config_attr)
         dit_config.update_model_arch(config)
 
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -98,6 +98,9 @@ class GPUWorker:
                 [
                     self.pipeline.get_module("transformer"),
                     self.pipeline.get_module("transformer_2"),
+                    self.pipeline.get_module("video_dit"),
+                    self.pipeline.get_module("video_dit_2"),
+                    self.pipeline.get_module("audio_dit"),
                 ],
             ):
                 if isinstance(dit, OffloadableDiTMixin):
@@ -164,6 +167,8 @@ class GPUWorker:
             if isinstance(result, Req):
                 output_batch = OutputBatch(
                     output=result.output,
+                    audio=getattr(result, "audio", None),
+                    audio_sample_rate=getattr(result, "audio_sample_rate", None),
                     timings=result.timings,
                     trajectory_timesteps=getattr(result, "trajectory_timesteps", None),
                     trajectory_latents=getattr(result, "trajectory_latents", None),

python/sglang/multimodal_gen/runtime/models/bridges/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from sglang.multimodal_gen.runtime.models.bridges.mova_dual_tower import (
+    DualTowerConditionalBridge,
+)
+
+__all__ = ["DualTowerConditionalBridge"]

python/sglang/multimodal_gen/runtime/models/bridges/mova_dual_tower.py

@@ -0,0 +1,671 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: mossVG/mova/diffusion/models/interactionv2.py
+
+
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from sglang.multimodal_gen.configs.models.bridges.mova_dual_tower import (
+    MOVADualTowerConfig,
+)
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.attention import USPAttention
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    RMSNorm,
+    tensor_parallel_rms_norm,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
+    apply_flashinfer_rope_qk_inplace,
+)
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+@torch.no_grad()
+def compute_rope_cos_sin(
+    position_ids: torch.Tensor,
+    head_dim: int,
+    base: float = 10000.0,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute RoPE cos/sin embeddings for given position IDs.
+
+    This is a functional implementation that doesn't require storing buffers,
+    making it compatible with FSDP meta device initialization.
+
+    Args:
+        position_ids: Position IDs tensor [B, L] or [1, L]
+        head_dim: Dimension of each attention head
+        base: RoPE base frequency (default: 10000.0)
+        device: Target device
+        dtype: Output dtype
+
+    Returns:
+        (cos, sin): Each with shape [B, L, head_dim]
+    """
+    device = device or position_ids.device
+    dtype = dtype or torch.float32
+
+    # Compute inverse frequencies
+    inv_freq = 1.0 / (
+        base
+        ** (torch.arange(0, head_dim, 2, dtype=torch.float32, device=device) / head_dim)
+    )
+
+    # Expand for batch computation: [B, L] -> [B, 1, L] @ [1, head_dim/2, 1] -> [B, head_dim/2, L]
+    inv_freq_expanded = inv_freq[None, :, None].expand(position_ids.shape[0], -1, 1)
+    position_ids_expanded = position_ids[:, None, :].float()
+
+    # Compute frequencies: [B, head_dim/2, L] -> [B, L, head_dim/2]
+    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+
+    # Double the frequencies for full head_dim: [B, L, head_dim]
+    emb = torch.cat((freqs, freqs), dim=-1)
+
+    cos = emb.cos().to(dtype=dtype)
+    sin = emb.sin().to(dtype=dtype)
+
+    return cos, sin
+
+
+class PerFrameAttentionPooling(nn.Module):
+    """Per-frame multi-head attention pooling.
+
+    Flattens the input sequence [B, L, D] and grid size (T, H, W).
+    Performs single-query attention pooling on the H*W tokens for each time frame.
+    Output shape: [B, T, D].
+    """
+
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        assert dim % num_heads == 0, "dim must be divisible by num_heads"
+        self.dim = dim
+        self.num_heads = num_heads
+
+        self.probe = nn.Parameter(torch.randn(1, 1, dim))
+        nn.init.normal_(self.probe, std=0.02)
+
+        self.attention = nn.MultiheadAttention(
+            embed_dim=dim, num_heads=num_heads, batch_first=True
+        )
+        self.layernorm = nn.LayerNorm(dim, eps=eps)
+
+    def forward(self, x: torch.Tensor, grid_size: Tuple[int, int, int]) -> torch.Tensor:
+        """Forward pass.
+
+        Args:
+            x: Input tensor of shape [B, L, D], where L = T * H * W.
+            grid_size: Tuple of (T, H, W).
+
+        Returns:
+            Pooled tensor of shape [B, T, D].
+        """
+        B, L, D = x.shape
+        T, H, W = grid_size
+        assert (
+            D == self.dim
+        ), f"Input dimension D={D} does not match module dim={self.dim}"
+        assert L == T * H * W, f"Flattened length L={L} does not match T*H*W={T*H*W}"
+
+        S = H * W
+        x_bt_s_d = x.view(B, T, S, D).contiguous().view(B * T, S, D)  # [B*T, S, D]
+        probe = self.probe.expand(B * T, -1, -1)  # [B*T, 1, D]
+
+        pooled_bt_1_d = self.attention(probe, x_bt_s_d, x_bt_s_d, need_weights=False)[0]
+        pooled_bt_d = pooled_bt_1_d.squeeze(1)  # [B*T, D]
+
+        pooled = pooled_bt_d.view(B, T, D)
+        pooled = self.layernorm(pooled)
+        return pooled
+
+
+class CrossModalInteractionController:
+    """Strategy class to control dual-tower interaction.
+
+    Manages the interaction mapping between Visual DiT (e.g., 30 layers)
+    and Audio DiT (e.g., 30 layers).
+    """
+
+    def __init__(self, visual_layers: int = 30, audio_layers: int = 30):
+        self.visual_layers = visual_layers
+        self.audio_layers = audio_layers
+        self.min_layers = min(visual_layers, audio_layers)
+
+    def get_interaction_layers(
+        self, strategy: str = "shallow_focus"
+    ) -> Dict[str, List[Tuple[int, int]]]:
+        """Gets the mapping relationship of interaction layers."""
+        if strategy == "shallow_focus":
+            num_interact = min(10, self.min_layers // 3)
+            interact_layers = list(range(0, num_interact))
+        elif strategy == "distributed":
+            step = 3
+            interact_layers = list(range(0, self.min_layers, step))
+        elif strategy == "progressive":
+            shallow = list(range(0, min(8, self.min_layers)))
+            if self.min_layers > 8:
+                deep = list(range(8, self.min_layers, 3))
+                interact_layers = shallow + deep
+            else:
+                interact_layers = shallow
+        elif strategy == "custom":
+            interact_layers = [0, 2, 4, 6, 8, 12, 16, 20]
+            interact_layers = [i for i in interact_layers if i < self.min_layers]
+        elif strategy == "full":
+            interact_layers = list(range(0, self.min_layers))
+        else:
+            raise ValueError(f"Unknown interaction strategy: {strategy}")
+
+        mapping = {
+            "v2a": [(i, i) for i in interact_layers],
+            "a2v": [(i, i) for i in interact_layers],
+        }
+        return mapping
+
+    def should_interact(
+        self, layer_idx: int, direction: str, interaction_mapping: Dict
+    ) -> bool:
+        """Determines if the specified layer needs to interact."""
+        if direction not in interaction_mapping:
+            return False
+        return any(src == layer_idx for src, _ in interaction_mapping[direction])
+
+
+class ConditionalCrossAttention(nn.Module):
+    """
+    Cross-modal attention for dual-tower bridge with Tensor Parallel support.
+
+    This module handles attention between video and audio hidden states,
+    which have different sequence lengths.
+    """
+
+    def __init__(self, dim: int, kv_dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.q_dim = dim
+        self.kv_dim = kv_dim
+        self.num_heads = num_heads
+        self.head_dim = self.q_dim // num_heads
+
+        self.tp_size = get_tp_world_size()
+        if self.num_heads % self.tp_size != 0:
+            raise ValueError(
+                f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})."
+            )
+        self.num_heads_per_rank = self.num_heads // self.tp_size
+
+        # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output.
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.k = ColumnParallelLinear(kv_dim, dim, bias=True, gather_output=False)
+        self.v = ColumnParallelLinear(kv_dim, dim, bias=True, gather_output=False)
+        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+        self.attn = USPAttention(
+            num_heads=self.num_heads_per_rank,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=None,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ):
+        ctx = y
+        q, _ = self.q(x)
+        k, _ = self.k(ctx)
+        v, _ = self.v(ctx)
+
+        # RMSNorm over sharded hidden dimension
+        if self.tp_size > 1:
+            q = tensor_parallel_rms_norm(q, self.norm_q)
+            k = tensor_parallel_rms_norm(k, self.norm_k)
+        else:
+            q = self.norm_q(q)
+            k = self.norm_k(k)
+
+        if x_freqs is not None:
+            x_cos, x_sin = x_freqs
+            q_view = rearrange(q, "b l (h d) -> b l h d", d=self.head_dim)
+            x_cos = x_cos.to(q_view.dtype).to(q_view.device).squeeze(0)
+            x_sin = x_sin.to(q_view.dtype).to(q_view.device).squeeze(0)
+            # FlashInfer expects cos_sin_cache with shape [seqlen, head_dim],
+            # where the first half is cos and the second half is sin, each with
+            # head_dim//2 elements. Since compute_rope_cos_sin duplicates the
+            # frequencies (cat((freqs, freqs))), we only take the first half.
+            half_dim = self.head_dim // 2
+            cos_sin_cache = torch.cat(
+                [
+                    x_cos[:, :half_dim].to(dtype=torch.float32).contiguous(),
+                    x_sin[:, :half_dim].to(dtype=torch.float32).contiguous(),
+                ],
+                dim=-1,
+            )
+            q_view, _ = apply_flashinfer_rope_qk_inplace(
+                q_view, q_view.clone(), cos_sin_cache, is_neox=True
+            )
+            q = rearrange(q_view, "b l h d -> b l (h d)")
+
+        if y_freqs is not None:
+            y_cos, y_sin = y_freqs
+            k_view = rearrange(k, "b l (h d) -> b l h d", d=self.head_dim)
+            y_cos = y_cos.to(k_view.dtype).to(k_view.device).squeeze(0)
+            y_sin = y_sin.to(k_view.dtype).to(k_view.device).squeeze(0)
+            # FlashInfer expects cos_sin_cache with shape [seqlen, head_dim],
+            # where the first half is cos and the second half is sin, each with
+            # head_dim//2 elements. Since compute_rope_cos_sin duplicates the
+            # frequencies (cat((freqs, freqs))), we only take the first half.
+            half_dim = self.head_dim // 2
+            cos_sin_cache = torch.cat(
+                [
+                    y_cos[:, :half_dim].to(dtype=torch.float32).contiguous(),
+                    y_sin[:, :half_dim].to(dtype=torch.float32).contiguous(),
+                ],
+                dim=-1,
+            )
+            k_view, _ = apply_flashinfer_rope_qk_inplace(
+                k_view, k_view.clone(), cos_sin_cache, is_neox=True
+            )
+            k = rearrange(k_view, "b l h d -> b l (h d)")
+
+        q = rearrange(q, "b l (h d) -> b l h d", h=self.num_heads_per_rank)
+        k = rearrange(k, "b l (h d) -> b l h d", h=self.num_heads_per_rank)
+        v = rearrange(v, "b l (h d) -> b l h d", h=self.num_heads_per_rank)
+
+        x = self.attn(q, k, v)
+        x = rearrange(x, "b l h d -> b l (h d)")
+        x, _ = self.o(x)
+        return x
+
+
+class AdaLayerNorm(nn.Module):
+    """
+    Norm layer modified to incorporate timestep embeddings.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_embeddings: Optional[int] = None,
+        output_dim: Optional[int] = None,
+        norm_elementwise_affine: bool = False,
+        norm_eps: float = 1e-5,
+        chunk_dim: int = 0,
+    ):
+        super().__init__()
+
+        self.chunk_dim = chunk_dim
+        output_dim = output_dim or embedding_dim * 2
+
+        if num_embeddings is not None:
+            self.emb = nn.Embedding(num_embeddings, embedding_dim)
+        else:
+            self.emb = None
+
+        self.silu = nn.SiLU()
+        self.linear = ReplicatedLinear(embedding_dim, output_dim)
+        self.norm = nn.LayerNorm(output_dim // 2, norm_eps, norm_elementwise_affine)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        timestep: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.emb is not None:
+            temb = self.emb(timestep)
+
+        temb, _ = self.linear(self.silu(temb))
+
+        if self.chunk_dim == 2:
+            scale, shift = temb.chunk(2, dim=2)
+        elif self.chunk_dim == 1:
+            shift, scale = temb.chunk(2, dim=1)
+            shift = shift[:, None, :]
+            scale = scale[:, None, :]
+        else:
+            scale, shift = temb.chunk(2, dim=0)
+
+        x = self.norm(x) * (1 + scale) + shift
+        return x
+
+
+class ConditionalCrossAttentionBlock(nn.Module):
+    """A wrapper block for ConditionalCrossAttention that applies LayerNorm to the condition input y."""
+
+    def __init__(
+        self,
+        dim: int,
+        kv_dim: int,
+        num_heads: int,
+        eps: float = 1e-6,
+        pooled_adaln: bool = False,
+    ):
+        super().__init__()
+        self.y_norm = nn.LayerNorm(kv_dim, eps=eps)
+        self.inner = ConditionalCrossAttention(
+            dim=dim, kv_dim=kv_dim, num_heads=num_heads, eps=eps
+        )
+        self.pooled_adaln = pooled_adaln
+        if pooled_adaln:
+            self.per_frame_pooling = PerFrameAttentionPooling(
+                kv_dim, num_heads=num_heads, eps=eps
+            )
+            self.adaln = AdaLayerNorm(kv_dim, output_dim=dim * 2, chunk_dim=2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: torch.Tensor,
+        x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        video_grid_size: Optional[Tuple[int, int, int]] = None,
+    ) -> torch.Tensor:
+        if self.pooled_adaln:
+            assert video_grid_size is not None, "video_grid_size cannot be None"
+            pooled_y = self.per_frame_pooling(y, video_grid_size)
+            if pooled_y.shape[1] != x.shape[1]:
+                pooled_y = F.interpolate(
+                    pooled_y.permute(0, 2, 1),
+                    size=x.shape[1],
+                    mode="linear",
+                    align_corners=False,
+                ).permute(0, 2, 1)
+            x = self.adaln(x, temb=pooled_y)
+        y = self.y_norm(y)
+        return self.inner(x=x, y=y, x_freqs=x_freqs, y_freqs=y_freqs)
+
+
+class DualTowerConditionalBridge(
+    CachableDiT,
+    OffloadableDiTMixin,
+):
+    """Dual-tower conditional bridge module v2 (SGLang optimized version).
+
+    Implements the correct architecture:
+    1. Audio latents -> Audio DiT -> Audio hidden states [B, L, 1536].
+    2. Visual latents -> Visual DiT -> Visual hidden states [B, L, 5120].
+    3. Cross-attention interaction between the hidden states of the two DiTs.
+    """
+
+    _fsdp_shard_conditions = MOVADualTowerConfig()._fsdp_shard_conditions
+    _compile_conditions = MOVADualTowerConfig()._compile_conditions
+    _supported_attention_backends = MOVADualTowerConfig()._supported_attention_backends
+    param_names_mapping = MOVADualTowerConfig().param_names_mapping
+    reverse_param_names_mapping = MOVADualTowerConfig().reverse_param_names_mapping
+    lora_param_names_mapping = MOVADualTowerConfig().lora_param_names_mapping
+
+    def __init__(
+        self,
+        config: MOVADualTowerConfig | None = None,
+        hf_config: dict[str, Any] | None = None,
+        # Fallback parameters for from_pretrained compatibility
+        visual_layers: int = 40,
+        audio_layers: int = 30,
+        visual_hidden_dim: int = 5120,
+        audio_hidden_dim: int = 1536,
+        audio_fps: float = 50.0,
+        head_dim: int = 128,
+        interaction_strategy: str = "full",
+        apply_cross_rope: bool = True,
+        apply_first_frame_bias_in_rope: bool = False,
+        trainable_condition_scale: bool = False,
+        pooled_adaln: bool = False,
+    ):
+        super().__init__(config=config, hf_config=hf_config)
+
+        # Use config if provided, otherwise use individual parameters
+        if config is not None:
+            visual_layers = config.visual_layers
+            audio_layers = config.audio_layers
+            visual_hidden_dim = config.visual_hidden_dim
+            audio_hidden_dim = config.audio_hidden_dim
+            audio_fps = config.audio_fps
+            head_dim = config.head_dim
+            interaction_strategy = config.interaction_strategy
+            apply_cross_rope = config.apply_cross_rope
+            apply_first_frame_bias_in_rope = config.apply_first_frame_bias_in_rope
+            trainable_condition_scale = config.trainable_condition_scale
+            pooled_adaln = config.pooled_adaln
+
+        self.visual_hidden_dim = visual_hidden_dim
+        self.audio_hidden_dim = audio_hidden_dim
+        self.audio_fps = audio_fps
+        self.head_dim = head_dim
+        self.apply_cross_rope = apply_cross_rope
+        self.apply_first_frame_bias_in_rope = apply_first_frame_bias_in_rope
+        self.trainable_condition_scale = trainable_condition_scale
+        self.pooled_adaln = pooled_adaln
+
+        if self.trainable_condition_scale:
+            self.condition_scale = nn.Parameter(
+                torch.tensor([1.0], dtype=torch.float32)
+            )
+        else:
+            self.condition_scale = 1.0
+
+        self.controller = CrossModalInteractionController(visual_layers, audio_layers)
+        self.interaction_mapping = self.controller.get_interaction_layers(
+            interaction_strategy
+        )
+
+        # Cross-modal attention modules - interaction at DiT hidden states level
+        self.audio_to_video_conditioners = nn.ModuleDict()
+        self.video_to_audio_conditioners = nn.ModuleDict()
+
+        self.rope_base = 10000.0  # RoPE base frequency hardcode. adapted from original mova implementation.
+
+        # Audio DiT hidden states conditioning Video DiT
+        for v_layer, _ in self.interaction_mapping["a2v"]:
+            self.audio_to_video_conditioners[str(v_layer)] = (
+                ConditionalCrossAttentionBlock(
+                    dim=visual_hidden_dim,
+                    kv_dim=audio_hidden_dim,
+                    num_heads=visual_hidden_dim // head_dim,
+                    pooled_adaln=False,
+                )
+            )
+
+        # Visual DiT hidden states conditioning Audio DiT
+        for a_layer, _ in self.interaction_mapping["v2a"]:
+            self.video_to_audio_conditioners[str(a_layer)] = (
+                ConditionalCrossAttentionBlock(
+                    dim=audio_hidden_dim,
+                    kv_dim=visual_hidden_dim,
+                    num_heads=audio_hidden_dim // head_dim,
+                    pooled_adaln=self.pooled_adaln,
+                )
+            )
+
+        # Required attributes for CachableDiT/BaseDiT
+        self.hidden_size = visual_hidden_dim
+        self.num_attention_heads = visual_hidden_dim // head_dim
+        self.num_channels_latents = (
+            visual_hidden_dim  # Bridge doesn't output latents, but required by BaseDiT
+        )
+        self.layer_names = [
+            "audio_to_video_conditioners",
+            "video_to_audio_conditioners",
+        ]
+        self.__post_init__()
+
+    @torch.no_grad()
+    def build_aligned_freqs(
+        self,
+        video_fps: float,
+        grid_size: Tuple[int, int, int],
+        audio_steps: int,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
+        """Generates aligned RoPE (cos, sin) based on video FPS, grid size, and audio length.
+
+        Uses functional RoPE computation to avoid FSDP meta device issues.
+
+        Args:
+            video_fps: FPS of the video.
+            grid_size: Tuple of (f_v, h, w).
+            audio_steps: Length of the audio sequence.
+            device: Target device.
+            dtype: Output dtype.
+
+        Returns:
+            A tuple of ((cos_v, sin_v), (cos_a, sin_a)).
+        """
+        f_v, h, w = grid_size
+        L_v = f_v * h * w
+        L_a = int(audio_steps)
+
+        device = device or next(self.parameters()).device
+        dtype = dtype or torch.float32
+
+        # Audio positions: 0, 1, 2, ..., L_a-1
+        audio_pos = torch.arange(L_a, device=device, dtype=torch.float32).unsqueeze(0)
+
+        # Video positions: Align video frames to audio step units
+        if self.apply_first_frame_bias_in_rope:
+            video_effective_fps = float(video_fps) / 4.0
+            if f_v > 0:
+                t_starts = torch.zeros((f_v,), device=device, dtype=torch.float32)
+                if f_v > 1:
+                    t_starts[1:] = (1.0 / float(video_fps)) + torch.arange(
+                        f_v - 1, device=device, dtype=torch.float32
+                    ) * (1.0 / video_effective_fps)
+            else:
+                t_starts = torch.zeros((0,), device=device, dtype=torch.float32)
+            video_pos_per_frame = t_starts * float(self.audio_fps)
+        else:
+            scale = float(self.audio_fps) / float(video_fps / 4.0)
+            video_pos_per_frame = (
+                torch.arange(f_v, device=device, dtype=torch.float32) * scale
+            )
+
+        video_pos = video_pos_per_frame.repeat_interleave(h * w).unsqueeze(0)
+
+        # Use functional RoPE to compute cos/sin
+        cos_v, sin_v = compute_rope_cos_sin(
+            video_pos, self.head_dim, base=self.rope_base, device=device, dtype=dtype
+        )
+        cos_a, sin_a = compute_rope_cos_sin(
+            audio_pos, self.head_dim, base=self.rope_base, device=device, dtype=dtype
+        )
+
+        return (cos_v, sin_v), (cos_a, sin_a)
+
+    def should_interact(self, layer_idx: int, direction: str) -> bool:
+        return self.controller.should_interact(
+            layer_idx, direction, self.interaction_mapping
+        )
+
+    def apply_conditional_control(
+        self,
+        layer_idx: int,
+        direction: str,
+        primary_hidden_states: torch.Tensor,
+        condition_hidden_states: torch.Tensor,
+        x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        condition_scale: Optional[float] = None,
+        video_grid_size: Optional[Tuple[int, int, int]] = None,
+    ) -> torch.Tensor:
+        """Applies conditional control at the DiT hidden states level."""
+        if not self.controller.should_interact(
+            layer_idx, direction, self.interaction_mapping
+        ):
+            return primary_hidden_states
+
+        if direction == "a2v":
+            conditioner = self.audio_to_video_conditioners[str(layer_idx)]
+        elif direction == "v2a":
+            conditioner = self.video_to_audio_conditioners[str(layer_idx)]
+        else:
+            raise ValueError(f"Invalid direction: {direction}")
+
+        conditioned_features = conditioner(
+            x=primary_hidden_states,
+            y=condition_hidden_states,
+            x_freqs=x_freqs,
+            y_freqs=y_freqs,
+            video_grid_size=video_grid_size,
+        )
+
+        if self.trainable_condition_scale and condition_scale is not None:
+            logger.warning(
+                "The current model has a trainable condition_scale, but condition_scale "
+                "was passed externally. Ignoring the trainable condition_scale and "
+                "using the external condition_scale=%s.",
+                condition_scale,
+            )
+
+        scale = condition_scale if condition_scale is not None else self.condition_scale
+
+        primary_hidden_states = primary_hidden_states + conditioned_features * scale
+
+        return primary_hidden_states
+
+    def forward(
+        self,
+        layer_idx: int,
+        visual_hidden_states: torch.Tensor,
+        audio_hidden_states: torch.Tensor,
+        *,
+        x_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        y_freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        a2v_condition_scale: Optional[float] = None,
+        v2a_condition_scale: Optional[float] = None,
+        condition_scale: Optional[float] = None,
+        video_grid_size: Optional[Tuple[int, int, int]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Performs bidirectional conditional control for both visual and audio towers."""
+        visual_conditioned = self.apply_conditional_control(
+            layer_idx=layer_idx,
+            direction="a2v",
+            primary_hidden_states=visual_hidden_states,
+            condition_hidden_states=audio_hidden_states,
+            x_freqs=x_freqs,
+            y_freqs=y_freqs,
+            condition_scale=(
+                a2v_condition_scale
+                if a2v_condition_scale is not None
+                else condition_scale
+            ),
+            video_grid_size=video_grid_size,
+        )
+
+        audio_conditioned = self.apply_conditional_control(
+            layer_idx=layer_idx,
+            direction="v2a",
+            primary_hidden_states=audio_hidden_states,
+            condition_hidden_states=visual_hidden_states,
+            x_freqs=y_freqs,
+            y_freqs=x_freqs,
+            condition_scale=(
+                v2a_condition_scale
+                if v2a_condition_scale is not None
+                else condition_scale
+            ),
+            video_grid_size=video_grid_size,
+        )
+
+        return visual_conditioned, audio_conditioned
+
+
+EntryClass = DualTowerConditionalBridge

python/sglang/multimodal_gen/runtime/models/dits/mova_audio_dit.py

@@ -0,0 +1,306 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/wan_audio_dit.py
+# SPDX-License-Identifier: Apache-2.0
+#
+# NOTE: This module reuses common functions from mova_video_dit.py to reduce code duplication.
+# Audio-specific functions (precompute_freqs_cis_1d, legacy_precompute_freqs_cis_1d) are kept here.
+
+import math
+from typing import Any, Optional, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.distributed.tensor import DTensor
+
+from sglang.multimodal_gen.configs.models.dits.mova_audio import MOVAAudioConfig
+from sglang.multimodal_gen.runtime.layers.linear import ReplicatedLinear
+from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+
+# Reuse common functions and classes from mova_video_dit
+from .mova_video_dit import DiTBlock, precompute_freqs_cis, sinusoidal_embedding_1d
+
+
+# Audio-specific positional encoding functions
+def legacy_precompute_freqs_cis_1d(
+    dim: int,
+    end: int = 16384,
+    theta: float = 10000.0,
+    base_tps=4.0,
+    target_tps=44100 / 2048,
+):
+    s = float(base_tps) / float(target_tps)
+    # 1d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta, s)
+    # No positional encoding is applied to the remaining dimensions
+    no_freqs_cis = precompute_freqs_cis(dim // 3, end, theta, s)
+    no_freqs_cis = torch.ones_like(no_freqs_cis)
+    return f_freqs_cis, no_freqs_cis, no_freqs_cis
+
+
+def precompute_freqs_cis_1d(dim: int, end: int = 16384, theta: float = 10000.0):
+    f_freqs_cis = precompute_freqs_cis(dim, end, theta)
+    return f_freqs_cis.chunk(3, dim=-1)
+
+
+class Head(nn.Module):
+    def __init__(
+        self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float
+    ):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size))
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        if len(t_mod.shape) == 3:
+            shift, scale = (
+                self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device)
+                + t_mod.unsqueeze(2)
+            ).chunk(2, dim=2)
+            x, _ = self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2))
+        else:
+            # NOTE: t_mod was originally [B, C]. This works correctly with broadcasting when B=1, but it won't match [1, 2, C] when B > 1.
+            shift, scale = (
+                self.modulation.to(dtype=t_mod.dtype, device=t_mod.device)
+                + t_mod.unsqueeze(1)
+            ).chunk(2, dim=1)
+            x, _ = self.head(self.norm(x) * (1 + scale) + shift)
+        return x
+
+
+class Conv1dLocalIsland(nn.Conv1d):
+    """Inherits from Conv1d and overrides forward.
+
+    - Parameters remain as DTensors (optimizer consistency is maintained).
+    - In the forward pass, x, weight, and bias are aggregated as Replicate,
+      and then local convolution is performed via to_local.
+    - The output is then redistributed as a DTensor (default is Replicate,
+      placements can be customized).
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        if isinstance(input, DTensor):
+            x_local = input.to_local()  # type: ignore[attr-defined]
+            w_local = self.weight.to_local()  # type: ignore[attr-defined]
+            b_local = (
+                self.bias.to_local() if self.bias is not None else None  # type: ignore[attr-defined]
+            )
+
+            return self._conv_forward(x_local, w_local, b_local)
+        else:
+            return super().forward(input)
+
+
+class WanAudioModel(CachableDiT, OffloadableDiTMixin):
+    _fsdp_shard_conditions = MOVAAudioConfig()._fsdp_shard_conditions
+    _compile_conditions = MOVAAudioConfig()._compile_conditions
+    _supported_attention_backends = MOVAAudioConfig()._supported_attention_backends
+    param_names_mapping = MOVAAudioConfig().param_names_mapping
+    reverse_param_names_mapping = MOVAAudioConfig().reverse_param_names_mapping
+    lora_param_names_mapping = MOVAAudioConfig().lora_param_names_mapping
+
+    def __init__(self, config: MOVAAudioConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        # Extract parameters from config
+        dim = config.dim
+        in_dim = config.in_dim
+        ffn_dim = config.ffn_dim
+        out_dim = config.out_dim
+        text_dim = config.text_dim
+        freq_dim = config.freq_dim
+        eps = config.eps
+        patch_size = config.patch_size
+        num_heads = config.num_heads
+        num_layers = config.num_layers
+        has_image_pos_emb = config.has_image_pos_emb
+        has_ref_conv = config.has_ref_conv
+        add_control_adapter = config.add_control_adapter
+        in_dim_control_adapter = config.in_dim_control_adapter
+        seperated_timestep = config.seperated_timestep
+        require_vae_embedding = config.require_vae_embedding
+        require_clip_embedding = config.require_clip_embedding
+        fuse_vae_embedding_in_latents = config.fuse_vae_embedding_in_latents
+        vae_type = config.vae_type
+
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.patch_size = patch_size
+        self.seperated_timestep = seperated_timestep
+        self.require_vae_embedding = require_vae_embedding
+        self.require_clip_embedding = require_clip_embedding
+        self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents
+        self.vae_type = vae_type
+        # self.patch_embedding = nn.Conv3d(
+        #     in_dim, dim, kernel_size=patch_size, stride=patch_size)
+        self.patch_embedding = Conv1dLocalIsland(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.text_embedding = MLP(
+            text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh"
+        )
+        self.time_embedding = MLP(freq_dim, dim, output_dim=dim, act_type="silu")
+        # Preserve state_dict keys (time_projection.1.weight/bias).
+        self.time_projection = nn.Sequential(nn.SiLU(), ReplicatedLinear(dim, dim * 6))
+        self.blocks = nn.ModuleList(
+            [DiTBlock(dim, num_heads, ffn_dim, eps) for _ in range(num_layers)]
+        )
+        self.head = Head(dim, out_dim, patch_size, eps)
+        self.num_heads = num_heads
+        self.freqs = None
+        self.img_pos_emb = None
+        if has_ref_conv:
+            self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2))
+        self.has_image_pos_emb = has_image_pos_emb
+        self.has_ref_conv = has_ref_conv
+        self.hidden_size = dim
+        self.num_attention_heads = num_heads
+        self.num_channels_latents = out_dim
+        self.layer_names = ["blocks"]
+        self.cnt = 0
+        self.teacache_thresh = 0
+        self.coefficients = []
+        self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = None
+        self.previous_resiual = None
+        self.previous_e0_even = None
+        self.previous_e0_odd = None
+        self.previous_residual_even = None
+        self.previous_residual_odd = None
+        self.is_even = False
+        self.should_calc_even = True
+        self.should_calc_odd = True
+        self.accumulated_rel_l1_distance_even = 0
+        self.accumulated_rel_l1_distance_odd = 0
+        self.__post_init__()
+        if add_control_adapter:
+            from .wan_video_camera_controller import SimpleAdapter
+
+            self.control_adapter = SimpleAdapter(
+                in_dim_control_adapter,
+                dim,
+                kernel_size=patch_size[1:],
+                stride=patch_size[1:],
+            )
+        else:
+            self.control_adapter = None
+
+    def _init_freqs(self):
+        if self.freqs is not None:
+            return
+        head_dim = self.dim // self.num_heads
+        if self.vae_type == "dac":
+            self.freqs = precompute_freqs_cis_1d(head_dim)
+        else:
+            raise ValueError(f"Invalid VAE type: {self.vae_type}")
+
+    def patchify(
+        self,
+        x: torch.Tensor,
+        control_camera_latents_input: Optional[torch.Tensor] = None,
+    ):
+        x = self.patch_embedding(x)
+        if (
+            self.control_adapter is not None
+            and control_camera_latents_input is not None
+        ):
+            y_camera = self.control_adapter(control_camera_latents_input)
+            if isinstance(x, list):
+                x = [u + v for u, v in zip(x, y_camera)]
+                x = x[0].unsqueeze(0)
+            else:
+                if isinstance(y_camera, list):
+                    x = x + y_camera[0]
+                else:
+                    x = x + y_camera
+        grid_size = x.shape[2:]
+        x = rearrange(x, "b c f -> b f c").contiguous()
+        return x, grid_size  # x, grid_size: (f)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: tuple[int]):
+        return rearrange(
+            x, "b f (p c) -> b c (f p)", f=grid_size[0], p=self.patch_size[0]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
+        guidance=None,
+        use_gradient_checkpointing: bool = False,
+        use_gradient_checkpointing_offload: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        # MOVA audio uses x/context naming historically.
+        x = hidden_states
+        context = (
+            encoder_hidden_states[0]
+            if isinstance(encoder_hidden_states, list)
+            else encoder_hidden_states
+        )
+
+        t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep))
+        t_proj, _ = self.time_projection(t)
+        t_mod = t_proj.unflatten(1, (6, self.dim))
+        context = self.text_embedding(context)
+
+        x, (f,) = self.patchify(x)
+
+        freqs = (
+            torch.cat(
+                [
+                    self.freqs[0][:f].view(f, -1).expand(f, -1),
+                    self.freqs[1][:f].view(f, -1).expand(f, -1),
+                    self.freqs[2][:f].view(f, -1).expand(f, -1),
+                ],
+                dim=-1,
+            )
+            .reshape(f, 1, -1)
+            .to(x.device)
+        )
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+
+            return custom_forward
+
+        for block in self.blocks:
+            if self.training and use_gradient_checkpointing:
+                if use_gradient_checkpointing_offload:
+                    with torch.autograd.graph.save_on_cpu():
+                        x = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            x,
+                            context,
+                            t_mod,
+                            freqs,
+                            use_reentrant=False,
+                        )
+                else:
+                    x = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        x,
+                        context,
+                        t_mod,
+                        freqs,
+                        use_reentrant=False,
+                    )
+            else:
+                x = block(x, context, t_mod, freqs)
+
+        x = self.head(x, t)
+        x = self.unpatchify(x, (f,))
+        return x
+
+
+EntryClass = WanAudioModel

python/sglang/multimodal_gen/runtime/models/dits/mova_video_dit.py

@@ -0,0 +1,570 @@
+# Copied and adapted from: mossVG/mova/diffusion/models/wan_video_dit.py
+# SPDX-License-Identifier: Apache-2.0
+#
+# NOTE: This module shares common functions (sinusoidal_embedding_1d, precompute_freqs_cis, etc.)
+# with wanvideo.py. These functions are kept here for MOVA-specific model architecture,
+# but could be refactored to a common module in the future.
+
+import math
+from typing import Any, Tuple
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.distributed.tensor import DTensor
+
+from sglang.multimodal_gen.configs.models.dits.mova_video import MOVAVideoConfig
+from sglang.multimodal_gen.runtime.distributed import get_tp_world_size
+from sglang.multimodal_gen.runtime.layers.attention import LocalAttention, USPAttention
+
+# Reuse SGLang's optimized RMSNorm instead of torch.nn.RMSNorm or custom SlowRMSNorm
+from sglang.multimodal_gen.runtime.layers.layernorm import (
+    RMSNorm,
+    tensor_parallel_rms_norm,
+)
+from sglang.multimodal_gen.runtime.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from sglang.multimodal_gen.runtime.layers.mlp import MLP
+from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+# @torch.compile(fullgraph=True)
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return x * (1 + scale) + shift
+
+
+def sinusoidal_embedding_1d(dim, position):
+    sinusoid = torch.outer(
+        position.type(torch.float64),
+        torch.pow(
+            10000,
+            -torch.arange(dim // 2, dtype=torch.float64, device=position.device).div(
+                dim // 2
+            ),
+        ),
+    )
+    x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+    return x.to(position.dtype)
+
+
+def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0):
+    # 3d rope precompute
+    f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta)
+    h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta)
+    return f_freqs_cis, h_freqs_cis, w_freqs_cis
+
+
+def precompute_freqs_cis(
+    dim: int, end: int = 1024, theta: float = 10000.0, s: float = 1.0
+):
+    # 1d rope precompute
+    # Note: s parameter is used for audio-specific scaling (e.g., tps adjustment)
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim))
+    pos = torch.arange(end, dtype=torch.float64, device=freqs.device) * s
+    freqs = torch.outer(pos, freqs)
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def rope_apply(x, freqs, num_heads):
+    x = rearrange(x, "b s (n d) -> b s n d", n=num_heads)
+    x_out = torch.view_as_complex(
+        x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)
+    )
+    x_out = torch.view_as_real(x_out * freqs).flatten(2)
+    return x_out.to(x.dtype)
+
+
+def rope_apply_head_dim(x, freqs, head_dim):
+    x = rearrange(x, "b s (n d) -> b s n d", d=head_dim)
+    x_out = torch.view_as_complex(
+        x.to(torch.float64).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)
+    )
+    # print(f"{x_out.shape = }, {freqs.shape = }")
+    x_out = torch.view_as_real(x_out * freqs).flatten(2)
+    return x_out.to(x.dtype)
+
+
+class SelfAttention(nn.Module):
+    """
+    Self-Attention module for MOVA DiT with Sequence Parallelism support.
+
+    SP is handled at the pipeline level (latents are pre-sharded before DiT forward).
+    USPAttention internally handles the all-to-all communication for distributed attention.
+    Input x should already be the local shard [B, S_local, D] when SP is enabled.
+    """
+
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.tp_size = get_tp_world_size()
+        if self.num_heads % self.tp_size != 0:
+            raise ValueError(
+                f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})."
+            )
+        self.num_heads_per_rank = self.num_heads // self.tp_size
+
+        # TP strategy: shard Q/K/V over heads (column-parallel), then row-parallel output.
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+        self.attn = USPAttention(
+            # Local heads per TP rank.
+            num_heads=self.num_heads_per_rank,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=None,
+        )
+
+    def forward(self, x, freqs):
+        """
+        Forward pass for self-attention.
+
+        Args:
+            x: Input tensor [B, S_local, D] - already sharded by SP when SP > 1
+            freqs: RoPE frequencies [S_local, 1, head_dim] - should match x's sequence length
+
+        Returns:
+            Output tensor [B, S_local, D]
+        """
+        if isinstance(freqs, DTensor):
+            freqs = freqs.to_local()
+
+        # Compute Q, K, V on local sequence
+        q, _ = self.q(x)
+        k, _ = self.k(x)
+        v, _ = self.v(x)
+
+        # RMSNorm over sharded hidden dimension.
+        if self.tp_size > 1:
+            q = tensor_parallel_rms_norm(q, self.norm_q)
+            k = tensor_parallel_rms_norm(k, self.norm_k)
+        else:
+            q = self.norm_q(q)
+            k = self.norm_k(k)
+
+        # Apply RoPE
+        q = rope_apply_head_dim(q, freqs, self.head_dim)
+        k = rope_apply_head_dim(k, freqs, self.head_dim)
+
+        # USPAttention expects [B, S_local, H, D] format
+        q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
+        k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
+        v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
+
+        # USPAttention handles SP communication internally
+        out = self.attn(q, k, v)
+        out = rearrange(out, "b s n d -> b s (n d)")
+
+        out, _ = self.o(out)
+        return out
+
+
+class CrossAttention(nn.Module):
+    """
+    Cross-Attention module for MOVA DiT.
+
+    Cross-attention does NOT require SP communication because:
+    - Query comes from the main sequence (already sharded by SP)
+    - Key/Value come from context (text embeddings, which are replicated across all ranks)
+
+    Uses LocalAttention instead of USPAttention for efficiency.
+    """
+
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-6):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.tp_size = get_tp_world_size()
+        if self.num_heads % self.tp_size != 0:
+            raise ValueError(
+                f"num_heads ({self.num_heads}) must be divisible by tp_size ({self.tp_size})."
+            )
+        self.num_heads_per_rank = self.num_heads // self.tp_size
+
+        self.q = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.k = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.v = ColumnParallelLinear(dim, dim, bias=True, gather_output=False)
+        self.o = RowParallelLinear(dim, dim, bias=True, input_is_parallel=True)
+        self.norm_q = RMSNorm(dim, eps=eps)
+        self.norm_k = RMSNorm(dim, eps=eps)
+
+        # Use LocalAttention for cross-attention (no SP communication needed)
+        self.attn = LocalAttention(
+            num_heads=self.num_heads_per_rank,
+            head_size=self.head_dim,
+            causal=False,
+            softmax_scale=None,
+        )
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor):
+        """
+        Forward pass for cross-attention.
+
+        Args:
+            x: Query tensor [B, S_local, D] - the main sequence (sharded by SP)
+            y: Context tensor [B, S_ctx, D] - text/image embeddings (replicated)
+
+        Returns:
+            Output tensor [B, S_local, D]
+        """
+        ctx = y
+
+        q, _ = self.q(x)
+        k, _ = self.k(ctx)
+        v, _ = self.v(ctx)
+
+        if self.tp_size > 1:
+            q = tensor_parallel_rms_norm(q, self.norm_q)
+            k = tensor_parallel_rms_norm(k, self.norm_k)
+        else:
+            q = self.norm_q(q)
+            k = self.norm_k(k)
+
+        q = rearrange(q, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
+        k = rearrange(k, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
+        v = rearrange(v, "b s (n d) -> b s n d", n=self.num_heads_per_rank)
+        x = self.attn(q, k, v)
+        x = rearrange(x, "b s n d -> b s (n d)")
+        x, _ = self.o(x)
+        return x
+
+
+class GateModule(nn.Module):
+    def __init__(
+        self,
+    ):
+        super().__init__()
+
+    def forward(self, x, gate, residual):
+        return x + gate * residual
+
+
+class DiTBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        ffn_dim: int,
+        eps: float = 1e-6,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.ffn_dim = ffn_dim
+
+        self.self_attn = SelfAttention(dim, num_heads, eps)
+        self.cross_attn = CrossAttention(dim, num_heads, eps)
+        self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        self.norm3 = nn.LayerNorm(dim, eps=eps)
+        self.ffn = MLP(dim, ffn_dim, output_dim=dim, act_type="gelu_pytorch_tanh")
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+        self.gate = GateModule()
+
+    def forward(self, x, context, t_mod, freqs):
+        has_seq = len(t_mod.shape) == 4
+        chunk_dim = 2 if has_seq else 1
+        # msa: multi-head self-attention  mlp: multi-layer perceptron
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod
+        ).chunk(6, dim=chunk_dim)
+        if has_seq:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                shift_msa.squeeze(2),
+                scale_msa.squeeze(2),
+                gate_msa.squeeze(2),
+                shift_mlp.squeeze(2),
+                scale_mlp.squeeze(2),
+                gate_mlp.squeeze(2),
+            )
+        input_x = modulate(self.norm1(x), shift_msa, scale_msa)
+        x = self.gate(x, gate_msa, self.self_attn(input_x, freqs))
+        x = x + self.cross_attn(self.norm3(x), context)
+        input_x = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = self.gate(x, gate_mlp, self.ffn(input_x))
+        return x
+
+
+class Head(nn.Module):
+    def __init__(
+        self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float
+    ):
+        super().__init__()
+        self.dim = dim
+        self.patch_size = patch_size
+        self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False)
+        # Output dim is small for MOVA; replicate to avoid TP shape coupling.
+        self.head = ReplicatedLinear(dim, out_dim * math.prod(patch_size))
+        self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+    def forward(self, x, t_mod):
+        if len(t_mod.shape) == 3:
+            shift, scale = (
+                self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device)
+                + t_mod.unsqueeze(2)
+            ).chunk(2, dim=2)
+            x, _ = self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2))
+        else:
+            shift, scale = (
+                self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod
+            ).chunk(2, dim=1)
+            x, _ = self.head(self.norm(x) * (1 + scale) + shift)
+        return x
+
+
+class Conv3dLocalIsland(nn.Conv3d):
+    """
+    Inherits from Conv3d and overrides the forward method.
+
+    Key behaviors:
+    - Parameters are kept as DTensor to maintain optimizer consistency.
+    - The forward pass aggregates input, weight, and bias into a Replicate state,
+      then performs the convolution locally using to_local().
+    - The output is then redistributed as a DTensor (defaults to Replicate,
+      but placements can be customized).
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, input):
+        if isinstance(input, DTensor):
+            # NOTE: DTensor typing stubs are incomplete; at runtime DTensor has
+            # to_local() and parameters may also be DTensor.
+            x_local = input.to_local()  # type: ignore[attr-defined]
+            w_local = self.weight.to_local()  # type: ignore[attr-defined]
+            b_local = (
+                self.bias.to_local() if self.bias is not None else None  # type: ignore[attr-defined]
+            )
+
+            return self._conv_forward(x_local, w_local, b_local)
+        else:
+            return super().forward(input)
+
+
+class WanModel(CachableDiT, OffloadableDiTMixin):
+    _fsdp_shard_conditions = MOVAVideoConfig()._fsdp_shard_conditions
+    _compile_conditions = MOVAVideoConfig()._compile_conditions
+    _supported_attention_backends = MOVAVideoConfig()._supported_attention_backends
+    param_names_mapping = MOVAVideoConfig().param_names_mapping
+    reverse_param_names_mapping = MOVAVideoConfig().reverse_param_names_mapping
+    lora_param_names_mapping = MOVAVideoConfig().lora_param_names_mapping
+
+    def __init__(self, config: MOVAVideoConfig, hf_config: dict[str, Any]) -> None:
+        super().__init__(config=config, hf_config=hf_config)
+
+        # Extract parameters from config
+        dim = config.dim
+        in_dim = config.in_dim
+        ffn_dim = config.ffn_dim
+        out_dim = config.out_dim
+        text_dim = config.text_dim
+        freq_dim = config.freq_dim
+        eps = config.eps
+        patch_size = config.patch_size
+        num_heads = config.num_heads
+        num_layers = config.num_layers
+        has_image_pos_emb = config.has_image_pos_emb
+        has_ref_conv = config.has_ref_conv
+        add_control_adapter = config.add_control_adapter
+        in_dim_control_adapter = config.in_dim_control_adapter
+        seperated_timestep = config.seperated_timestep
+        require_vae_embedding = config.require_vae_embedding
+        require_clip_embedding = config.require_clip_embedding
+        fuse_vae_embedding_in_latents = config.fuse_vae_embedding_in_latents
+
+        self.dim = dim
+        self.freq_dim = freq_dim
+        self.patch_size = patch_size
+        self.seperated_timestep = seperated_timestep
+        self.require_vae_embedding = require_vae_embedding
+        self.require_clip_embedding = require_clip_embedding
+        self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents
+
+        self.patch_embedding = Conv3dLocalIsland(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size
+        )
+        self.text_embedding = MLP(
+            text_dim, dim, output_dim=dim, act_type="gelu_pytorch_tanh"
+        )
+        self.time_embedding = MLP(freq_dim, dim, output_dim=dim, act_type="silu")
+        # Preserve state_dict keys (time_projection.1.weight/bias).
+        self.time_projection = nn.Sequential(nn.SiLU(), ReplicatedLinear(dim, dim * 6))
+        self.blocks = nn.ModuleList(
+            [DiTBlock(dim, num_heads, ffn_dim, eps) for _ in range(num_layers)]
+        )
+        self.head = Head(dim, out_dim, patch_size, eps)
+        self.num_heads = num_heads
+        self.freqs = None
+
+        if has_ref_conv:
+            self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2))
+        self.has_image_pos_emb = has_image_pos_emb
+        self.has_ref_conv = has_ref_conv
+        self.hidden_size = dim
+        self.num_attention_heads = num_heads
+        self.num_channels_latents = out_dim
+        self.layer_names = ["blocks"]
+        self.cnt = 0
+        self.teacache_thresh = 0
+        self.coefficients = []
+        self.accumulated_rel_l1_distance = 0
+        self.previous_modulated_input = None
+        self.previous_resiual = None
+        self.previous_e0_even = None
+        self.previous_e0_odd = None
+        self.previous_residual_even = None
+        self.previous_residual_odd = None
+        self.is_even = False
+        self.should_calc_even = True
+        self.should_calc_odd = True
+        self.accumulated_rel_l1_distance_even = 0
+        self.accumulated_rel_l1_distance_odd = 0
+        self.__post_init__()
+        if add_control_adapter:
+            from .wan_video_camera_controller import SimpleAdapter
+
+            self.control_adapter = SimpleAdapter(
+                in_dim_control_adapter,
+                dim,
+                kernel_size=patch_size[1:],
+                stride=patch_size[1:],
+            )
+        else:
+            self.control_adapter = None
+
+    def _init_freqs(self):
+        if self.freqs is not None:
+            return
+        head_dim = self.dim // self.num_heads
+        self.freqs = precompute_freqs_cis_3d(head_dim)
+
+    def patchify(
+        self, x: torch.Tensor, control_camera_latents_input: torch.Tensor | None = None
+    ):
+        # NOTE(dhyu): avoid slow_conv
+        x = x.contiguous(memory_format=torch.channels_last_3d)
+        x = self.patch_embedding(x)
+        if (
+            self.control_adapter is not None
+            and control_camera_latents_input is not None
+        ):
+            y_camera = self.control_adapter(control_camera_latents_input)
+            if isinstance(x, list):
+                x = [u + v for u, v in zip(x, y_camera)]
+                x = x[0].unsqueeze(0)
+            else:
+                # Some adapters may return a list even when x is a Tensor.
+                if isinstance(y_camera, list):
+                    x = x + y_camera[0]
+                else:
+                    x = x + y_camera
+        grid_size = x.shape[2:]
+        x = rearrange(x, "b c f h w -> b (f h w) c").contiguous()
+        return x, grid_size  # x, grid_size: (f, h, w)
+
+    def unpatchify(self, x: torch.Tensor, grid_size: tuple[int, int, int]):
+        return rearrange(
+            x,
+            "b (f h w) (x y z c) -> b c (f x) (h y) (w z)",
+            f=grid_size[0],
+            h=grid_size[1],
+            w=grid_size[2],
+            x=self.patch_size[0],
+            y=self.patch_size[1],
+            z=self.patch_size[2],
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | list[torch.Tensor],
+        timestep: torch.LongTensor,
+        encoder_hidden_states_image: torch.Tensor | list[torch.Tensor],
+        y: torch.Tensor,
+        guidance=None,
+        use_gradient_checkpointing: bool = False,
+        use_gradient_checkpointing_offload: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        # MOVA code historically uses x/context/y/clip_feature naming.
+        x = hidden_states
+        context = (
+            encoder_hidden_states[0]
+            if isinstance(encoder_hidden_states, list)
+            else encoder_hidden_states
+        )
+        t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep))
+        t_proj, _ = self.time_projection(t)
+        t_mod = t_proj.unflatten(1, (6, self.dim))
+        context = self.text_embedding(context)
+
+        x, (f, h, w) = self.patchify(x)
+
+        freqs = (
+            torch.cat(
+                [
+                    self.freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+                    self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+                    self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1),
+                ],
+                dim=-1,
+            )
+            .reshape(f * h * w, 1, -1)
+            .to(x.device)
+        )
+
+        def create_custom_forward(module):
+            def custom_forward(*inputs):
+                return module(*inputs)
+
+            return custom_forward
+
+        for block in self.blocks:
+            if self.training and use_gradient_checkpointing:
+                if use_gradient_checkpointing_offload:
+                    with torch.autograd.graph.save_on_cpu():
+                        x = torch.utils.checkpoint.checkpoint(
+                            create_custom_forward(block),
+                            x,
+                            context,
+                            t_mod,
+                            freqs,
+                            use_reentrant=False,
+                        )
+                else:
+                    x = torch.utils.checkpoint.checkpoint(
+                        create_custom_forward(block),
+                        x,
+                        context,
+                        t_mod,
+                        freqs,
+                        use_reentrant=False,
+                    )
+            else:
+                x = block(x, context, t_mod, freqs)
+
+        x = self.head(x, t)
+        x = self.unpatchify(x, (f, h, w))
+        return x
+
+
+EntryClass = WanModel

python/sglang/multimodal_gen/runtime/models/model_stages/mova.py

@@ -0,0 +1,911 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+MOVA-specific pipeline stages.
+
+Sequence Parallelism (SP) Support:
+- Video latents are sharded along the sequence dimension (T*H*W) after patchify
+- Audio latents are sharded along the sequence dimension (L) after patchify
+- USPAttention handles all-to-all communication internally
+- Latents are gathered before unpatchify to restore full sequence
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+import os
+from collections.abc import Iterable
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+from tqdm.auto import tqdm
+
+from sglang.multimodal_gen.runtime.distributed import (
+    get_local_torch_device,
+    get_world_group,
+)
+from sglang.multimodal_gen.runtime.distributed.communication_op import (
+    cfg_model_parallel_all_reduce,
+    sequence_model_parallel_all_gather,
+)
+from sglang.multimodal_gen.runtime.distributed.parallel_state import (
+    get_cfg_group,
+    get_classifier_free_guidance_rank,
+    get_sp_parallel_rank,
+    get_sp_world_size,
+)
+from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context
+
+# Both audio and video DiT use the same sinusoidal_embedding_1d function
+# Import from mova_video_dit where it's defined (mova_audio_dit re-exports it)
+from sglang.multimodal_gen.runtime.models.dits.mova_video_dit import (
+    sinusoidal_embedding_1d,
+)
+
+# Create aliases for backward compatibility
+video_sinusoidal_embedding_1d = sinusoidal_embedding_1d
+audio_sinusoidal_embedding_1d = sinusoidal_embedding_1d
+from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
+from sglang.multimodal_gen.runtime.pipelines_core.stages.base import (
+    PipelineStage,
+    StageParallelismType,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import (
+    _ensure_tensor_decode_output,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    StageValidators as V,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
+    VerificationResult,
+)
+from sglang.multimodal_gen.runtime.platforms import current_platform
+from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args
+from sglang.multimodal_gen.runtime.utils.layerwise_offload import OffloadableDiTMixin
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+from sglang.multimodal_gen.runtime.utils.perf_logger import StageProfiler
+from sglang.multimodal_gen.runtime.utils.profiler import SGLDiffusionProfiler
+from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
+
+logger = init_logger(__name__)
+
+
+class MOVALatentPreparationStage(PipelineStage):
+    """Prepare video/audio noise latents for MOVA."""
+
+    def __init__(self, audio_vae, require_vae_embedding: bool = True) -> None:
+        super().__init__()
+        self.audio_vae = audio_vae
+        self.require_vae_embedding = require_vae_embedding
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        batch_size = batch.batch_size
+        num_frames = batch.num_frames
+        if num_frames is None:
+            raise ValueError("num_frames is required for MOVA")
+
+        audio_num_samples = int(self.audio_vae.sample_rate * num_frames / batch.fps)
+
+        video_shape = server_args.pipeline_config.prepare_latent_shape(
+            batch, batch_size, num_frames
+        )
+        audio_shape = server_args.pipeline_config.prepare_audio_latent_shape(
+            batch_size, audio_num_samples, self.audio_vae
+        )
+
+        device = get_local_torch_device()
+        generator = batch.generator
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+
+        dit_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
+        batch.latents = randn_tensor(
+            video_shape, generator=generator, device=device, dtype=dit_dtype
+        )
+        batch.audio_latents = randn_tensor(
+            audio_shape, generator=generator, device=device, dtype=dit_dtype
+        )
+
+        if batch.image_latent is not None:
+            batch.y = batch.image_latent.to(device=device, dtype=dit_dtype)
+        elif self.require_vae_embedding:
+            raise ValueError("MOVA requires reference image latents for denoising")
+        return batch
+
+
+class MOVATimestepPreparationStage(PipelineStage):
+    """Prepare paired timesteps for MOVA."""
+
+    def __init__(self, scheduler) -> None:
+        super().__init__()
+        self.scheduler = scheduler
+
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        self.scheduler.set_timesteps(
+            batch.num_inference_steps,
+            denoising_strength=1.0,
+            shift=getattr(batch, "sigma_shift", self.scheduler.shift),
+        )
+        self.scheduler.set_pair_postprocess_by_name(
+            "dual_sigma_shift",
+            visual_shift=getattr(batch, "visual_shift", 5.0),
+            audio_shift=getattr(batch, "audio_shift", 5.0),
+        )
+        paired = self.scheduler.get_pairs()
+        batch.paired_timesteps = paired
+        batch.timesteps = paired
+        return batch
+
+
+class MOVADenoisingStage(PipelineStage):
+    """Run MOVA dual-tower denoising loop."""
+
+    def __init__(self, video_dit, video_dit_2, audio_dit, dual_tower_bridge, scheduler):
+        super().__init__()
+        self.video_dit = video_dit
+        self.video_dit_2 = video_dit_2
+        self.audio_dit = audio_dit
+        self.dual_tower_bridge = dual_tower_bridge
+        self.scheduler = scheduler
+        self._cache_dit_enabled = False
+        self._cached_num_steps = None
+        self._torch_compiled = False
+
+    @property
+    def parallelism_type(self) -> StageParallelismType:
+        if get_global_server_args().enable_cfg_parallel:
+            return StageParallelismType.CFG_PARALLEL
+        return StageParallelismType.REPLICATED
+
+    def _predict(
+        self,
+        visual_dit,
+        visual_latents,
+        audio_latents,
+        y,
+        context,
+        timestep,
+        audio_timestep,
+        video_fps,
+        timestep_index: int,
+        attn_metadata,
+        forward_batch: Req | None = None,
+    ):
+        # Set forward context for distributed attention (USPAttention)
+        with set_forward_context(
+            current_timestep=timestep_index,
+            attn_metadata=attn_metadata,
+            forward_batch=forward_batch,
+        ):
+            return self.inference_single_step(
+                visual_dit=visual_dit,
+                visual_latents=visual_latents,
+                audio_latents=audio_latents,
+                y=y,
+                context=context,
+                timestep=timestep,
+                audio_timestep=audio_timestep,
+                video_fps=video_fps,
+            )
+
+    def _cfg_combine(self, pos, neg, guidance_scale, cfg_rank, enable_cfg_parallel):
+        if not enable_cfg_parallel:
+            return neg + guidance_scale * (pos - neg)
+        if cfg_rank == 0:
+            partial = guidance_scale * pos
+        else:
+            partial = (1 - guidance_scale) * neg
+        return cfg_model_parallel_all_reduce(partial)
+
+    def compile_module_with_torch_compile(self, module, server_args: ServerArgs):
+        if not server_args.enable_torch_compile or module is None:
+            return module
+        if not hasattr(module, "forward"):
+            return module
+        try:
+            import torch._inductor.config as _inductor_cfg
+
+            _inductor_cfg.reorder_for_compute_comm_overlap = True
+        except ImportError:
+            pass
+        mode = os.environ.get("SGLANG_TORCH_COMPILE_MODE", "max-autotune-no-cudagraphs")
+        logger.info("Compiling %s with mode: %s", module.__class__.__name__, mode)
+        compiled_forward = torch.compile(getattr(module, "forward"), mode=mode)
+        setattr(module, "forward", compiled_forward)
+        return module
+
+    def _maybe_compile_dits(self, server_args: ServerArgs):
+        if self._torch_compiled or not server_args.enable_torch_compile:
+            return
+        for module in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]):
+            self.compile_module_with_torch_compile(module, server_args)
+        self._torch_compiled = True
+
+    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        """Verify denoising stage inputs."""
+        result = VerificationResult()
+        result.add_check("y", batch.y, V.is_tensor)
+        result.add_check("paired_timesteps", batch.paired_timesteps, V.is_tensor)
+        result.add_check("latents", batch.latents, V.is_tensor)
+        result.add_check("audio_latents", batch.audio_latents, V.is_tensor)
+        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
+        result.add_check(
+            "negative_prompt_embeds",
+            batch.negative_prompt_embeds,
+            lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x),
+        )
+        result.add_check(
+            "num_inference_steps", batch.num_inference_steps, V.positive_int
+        )
+        result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float)
+        result.add_check(
+            "guidance_rescale", batch.guidance_rescale, V.non_negative_float
+        )
+        result.add_check(
+            "do_classifier_free_guidance",
+            batch.do_classifier_free_guidance,
+            V.bool_value,
+        )
+        return result
+
+    def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
+        """Verify denoising stage outputs."""
+        result = VerificationResult()
+        result.add_check("latents", batch.latents, V.is_tensor)
+        result.add_check("audio_latents", batch.audio_latents, V.is_tensor)
+        return result
+
+    def progress_bar(
+        self, iterable: Iterable | None = None, total: int | None = None
+    ) -> tqdm:
+        """
+        Create a progress bar for the denoising process.
+        """
+        local_rank = get_world_group().local_rank
+        disable = local_rank != 0
+        return tqdm(iterable=iterable, total=total, disable=disable)
+
+    def step_profile(self):
+        profiler = SGLDiffusionProfiler.get_instance()
+        if profiler:
+            profiler.step_denoising_step()
+
+    def rescale_noise_cfg(
+        self, noise_cfg, noise_pred_text, guidance_rescale=0.0
+    ) -> torch.Tensor:
+        """
+        Rescale noise prediction according to guidance_rescale.
+
+        Based on findings of "Common Diffusion Noise Schedules and Sample Steps are Flawed"
+        (https://arxiv.org/pdf/2305.08891.pdf), Section 3.4.
+        """
+        std_text = noise_pred_text.std(
+            dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+        )
+        std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True)
+        noise_pred_rescaled = noise_cfg * (std_text / std_cfg)
+        noise_cfg = (
+            guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg
+        )
+        return noise_cfg
+
+    def prepare_extra_func_kwargs(self, func, kwargs) -> dict[str, object]:
+        if not kwargs:
+            return {}
+
+        if isinstance(func, functools.partial) and func.args:
+            func = getattr(func.args[0], "_original_forward", func)
+
+        target_func = inspect.unwrap(func)
+        params = inspect.signature(target_func).parameters
+        return {k: v for k, v in kwargs.items() if k in params}
+
+    def _build_attn_metadata(
+        self, i: int, batch: Req, server_args: ServerArgs
+    ) -> object | None:
+        return None
+
+    def _manage_device_placement(
+        self,
+        model_to_use: torch.nn.Module | None,
+        model_to_offload: torch.nn.Module | None,
+        server_args: ServerArgs,
+    ):
+        if not server_args.dit_cpu_offload:
+            return
+
+        if (
+            model_to_offload is not None
+            and next(model_to_offload.parameters()).device.type == "cuda"
+        ):
+            model_to_offload.to("cpu")
+
+        if (
+            model_to_use is not None
+            and next(model_to_use.parameters()).device.type == "cpu"
+        ):
+            model_to_use.to(get_local_torch_device())
+
+    def _select_visual_dit(
+        self, timestep: float, boundary_ratio: float | None, server_args: ServerArgs
+    ):
+        if boundary_ratio is None or self.video_dit_2 is None:
+            self._manage_device_placement(self.video_dit, None, server_args)
+            return self.video_dit
+
+        boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
+        if timestep >= boundary_timestep:
+            current_model = self.video_dit
+            model_to_offload = self.video_dit_2
+        else:
+            current_model = self.video_dit_2
+            model_to_offload = self.video_dit
+
+        self._manage_device_placement(current_model, model_to_offload, server_args)
+        return current_model
+
+    def _apply_guidance_rescale(
+        self,
+        noise_pred,
+        noise_pred_text,
+        guidance_rescale,
+        cfg_rank,
+        enable_cfg_parallel,
+    ):
+        if guidance_rescale <= 0.0:
+            return noise_pred
+        if enable_cfg_parallel:
+            std_cfg = noise_pred.std(dim=list(range(1, noise_pred.ndim)), keepdim=True)
+            if cfg_rank == 0:
+                assert noise_pred_text is not None
+                std_text = noise_pred_text.std(
+                    dim=list(range(1, noise_pred_text.ndim)), keepdim=True
+                )
+            else:
+                std_text = torch.empty_like(std_cfg)
+            std_text = get_cfg_group().broadcast(std_text, src=0)
+            noise_pred_rescaled = noise_pred * (std_text / std_cfg)
+            return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * (
+                noise_pred
+            )
+        return self.rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale)
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> Req:
+        self._maybe_compile_dits(server_args)
+        self._manage_device_placement(self.audio_dit, None, server_args)
+
+        paired_timesteps = batch.paired_timesteps
+        if paired_timesteps is None:
+            raise ValueError("paired_timesteps must be set for MOVA")
+
+        y = batch.y if batch.y is not None else batch.image_latent
+        if getattr(self.video_dit, "require_vae_embedding", False) and y is None:
+            raise ValueError("MOVA requires reference image latents for denoising")
+
+        boundary_ratio = server_args.pipeline_config.boundary_ratio
+        total_steps = paired_timesteps.shape[0]
+        cfg_rank = get_classifier_free_guidance_rank()
+        enable_cfg_parallel = server_args.enable_cfg_parallel
+
+        is_warmup = batch.is_warmup
+        extra_step_kwargs = self.prepare_extra_func_kwargs(
+            self.scheduler.step_from_to,
+            getattr(batch, "extra_step_kwargs", None) or {},
+        )
+
+        timings = getattr(batch, "timings", None)
+        perf_dump_path_provided = getattr(batch, "perf_dump_path", None) is not None
+
+        with self.progress_bar(total=total_steps) as progress_bar:
+            for idx_step in range(total_steps):
+                with StageProfiler(
+                    f"denoising_step_{idx_step}",
+                    logger=logger,
+                    timings=timings,
+                    perf_dump_path_provided=perf_dump_path_provided,
+                ):
+                    pair_t = paired_timesteps[idx_step]
+                    if getattr(pair_t, "shape", None) == (2,):
+                        timestep, audio_timestep = pair_t
+                    else:
+                        timestep = pair_t
+                        audio_timestep = pair_t
+
+                    cur_visual_dit = self._select_visual_dit(
+                        timestep.item(), boundary_ratio, server_args
+                    )
+
+                    timestep = timestep.unsqueeze(0).to(device=get_local_torch_device())
+                    audio_timestep = audio_timestep.unsqueeze(0).to(
+                        device=get_local_torch_device()
+                    )
+
+                    attn_metadata = self._build_attn_metadata(
+                        idx_step, batch, server_args
+                    )
+
+                    if not batch.do_classifier_free_guidance:
+                        visual_noise_pred, audio_noise_pred = self._predict(
+                            cur_visual_dit,
+                            batch.latents,
+                            batch.audio_latents,
+                            y,
+                            batch.prompt_embeds[0],
+                            timestep,
+                            audio_timestep,
+                            batch.fps,
+                            idx_step,
+                            attn_metadata,
+                            batch,
+                        )
+                    else:
+                        if enable_cfg_parallel:
+                            if cfg_rank == 0:
+                                pos = self._predict(
+                                    cur_visual_dit,
+                                    batch.latents,
+                                    batch.audio_latents,
+                                    y,
+                                    batch.prompt_embeds[0],
+                                    timestep,
+                                    audio_timestep,
+                                    batch.fps,
+                                    idx_step,
+                                    attn_metadata,
+                                    batch,
+                                )
+                                neg = (None, None)
+                            else:
+                                pos = (None, None)
+                                neg = self._predict(
+                                    cur_visual_dit,
+                                    batch.latents,
+                                    batch.audio_latents,
+                                    y,
+                                    batch.negative_prompt_embeds[0],
+                                    timestep,
+                                    audio_timestep,
+                                    batch.fps,
+                                    idx_step,
+                                    attn_metadata,
+                                    batch,
+                                )
+                        else:
+                            pos = self._predict(
+                                cur_visual_dit,
+                                batch.latents,
+                                batch.audio_latents,
+                                y,
+                                batch.prompt_embeds[0],
+                                timestep,
+                                audio_timestep,
+                                batch.fps,
+                                idx_step,
+                                attn_metadata,
+                                batch,
+                            )
+                            neg = self._predict(
+                                cur_visual_dit,
+                                batch.latents,
+                                batch.audio_latents,
+                                y,
+                                batch.negative_prompt_embeds[0],
+                                timestep,
+                                audio_timestep,
+                                batch.fps,
+                                idx_step,
+                                attn_metadata,
+                                batch,
+                            )
+
+                            visual_noise_pred = self._cfg_combine(
+                                pos[0] if pos[0] is not None else neg[0],
+                                neg[0] if neg[0] is not None else pos[0],
+                                batch.guidance_scale,
+                                cfg_rank,
+                                enable_cfg_parallel,
+                            )
+                            audio_noise_pred = self._cfg_combine(
+                                pos[1] if pos[1] is not None else neg[1],
+                                neg[1] if neg[1] is not None else pos[1],
+                                batch.guidance_scale,
+                                cfg_rank,
+                                enable_cfg_parallel,
+                            )
+
+                            if batch.guidance_rescale > 0.0:
+                                visual_noise_pred = self._apply_guidance_rescale(
+                                    visual_noise_pred,
+                                    pos[0] if pos[0] is not None else None,
+                                    batch.guidance_rescale,
+                                    cfg_rank,
+                                    enable_cfg_parallel,
+                                )
+                                audio_noise_pred = self._apply_guidance_rescale(
+                                    audio_noise_pred,
+                                    pos[1] if pos[1] is not None else None,
+                                    batch.guidance_rescale,
+                                    cfg_rank,
+                                    enable_cfg_parallel,
+                                )
+
+                        if idx_step + 1 < total_steps:
+                            next_pair_t = paired_timesteps[idx_step + 1]
+                            if getattr(next_pair_t, "shape", None) == (2,):
+                                next_timestep, next_audio_timestep = next_pair_t
+                            else:
+                                next_timestep = next_pair_t
+                                next_audio_timestep = next_pair_t
+                        else:
+                            next_timestep = None
+                            next_audio_timestep = None
+
+                        batch.latents = self.scheduler.step_from_to(
+                            visual_noise_pred,
+                            timestep,
+                            next_timestep,
+                            batch.latents,
+                            **extra_step_kwargs,
+                        )
+                        batch.audio_latents = self.scheduler.step_from_to(
+                            audio_noise_pred,
+                            audio_timestep,
+                            next_audio_timestep,
+                            batch.audio_latents,
+                            **extra_step_kwargs,
+                        )
+
+                    if progress_bar is not None:
+                        progress_bar.update()
+                    if not is_warmup and hasattr(self, "step_profile"):
+                        self.step_profile()
+
+        for dit in filter(None, [self.video_dit, self.video_dit_2, self.audio_dit]):
+            if isinstance(dit, OffloadableDiTMixin):
+                dit.prepare_for_next_denoise()
+
+        return batch
+
+    def _shard_sequence_for_sp(
+        self, x: torch.Tensor, dim: int = 1
+    ) -> tuple[torch.Tensor, int]:
+        """
+        Shard tensor along sequence dimension for Sequence Parallelism.
+
+        Args:
+            x: Input tensor
+            dim: Dimension to shard along
+
+        Returns:
+            (sharded_tensor, pad_len)
+        """
+        sp_size = get_sp_world_size()
+        if sp_size <= 1:
+            return x, 0
+
+        sp_rank = get_sp_parallel_rank()
+        seq_len = x.shape[dim]
+
+        # Pad if needed
+        pad_len = (sp_size - (seq_len % sp_size)) % sp_size
+        if pad_len > 0:
+            pad_shape = list(x.shape)
+            pad_shape[dim] = pad_len
+            pad = torch.zeros(pad_shape, dtype=x.dtype, device=x.device)
+            x = torch.cat([x, pad], dim=dim)
+
+        # Shard
+        chunk_size = x.shape[dim] // sp_size
+        start = sp_rank * chunk_size
+        end = start + chunk_size
+        idx = [slice(None)] * x.dim()
+        idx[dim] = slice(start, end)
+        return x[tuple(idx)], pad_len
+
+    def _gather_sequence_from_sp(
+        self, x: torch.Tensor, pad_len: int, dim: int = 1
+    ) -> torch.Tensor:
+        """
+        Gather tensor along sequence dimension after Sequence Parallelism.
+
+        Args:
+            x: Sharded tensor
+            pad_len: Padding length that was added during sharding
+            dim: Dimension to gather along
+
+        Returns:
+            Gathered tensor with padding removed
+        """
+        sp_size = get_sp_world_size()
+        if sp_size <= 1:
+            return x
+
+        gathered = sequence_model_parallel_all_gather(x, dim=dim)
+        if pad_len > 0:
+            idx = [slice(None)] * gathered.dim()
+            idx[dim] = slice(0, gathered.shape[dim] - pad_len)
+            gathered = gathered[tuple(idx)]
+        return gathered
+
+    def inference_single_step(
+        self,
+        visual_dit,
+        visual_latents: torch.Tensor,
+        audio_latents: torch.Tensor,
+        y,
+        context: torch.Tensor,
+        timestep: torch.Tensor,
+        audio_timestep: torch.Tensor,
+        video_fps: float,
+    ):
+        """
+        Single inference step for MOVA dual-tower denoising.
+
+        Supports Sequence Parallelism (SP):
+        - After patchify, sequences are sharded across SP ranks
+        - USPAttention handles distributed attention communication
+        - Before unpatchify, sequences are gathered back
+        """
+        model_dtype = visual_dit.time_embedding.fc_in.weight.dtype
+        device = visual_latents.device
+
+        visual_context = context.to(device=device, dtype=model_dtype)
+        audio_context = context.to(device=device, dtype=model_dtype)
+        with torch.autocast(
+            device_type=current_platform.device_type, dtype=torch.float32
+        ):
+            visual_t = visual_dit.time_embedding(
+                video_sinusoidal_embedding_1d(visual_dit.freq_dim, timestep)
+            )
+            visual_t_mod, _ = visual_dit.time_projection(visual_t)
+            visual_t_mod = visual_t_mod.unflatten(1, (6, visual_dit.dim))
+
+            audio_t = self.audio_dit.time_embedding(
+                audio_sinusoidal_embedding_1d(self.audio_dit.freq_dim, audio_timestep)
+            )
+            audio_t_mod, _ = self.audio_dit.time_projection(audio_t)
+            audio_t_mod = audio_t_mod.unflatten(1, (6, self.audio_dit.dim))
+
+        visual_t = visual_t.to(model_dtype)
+        visual_t_mod = visual_t_mod.to(model_dtype)
+        audio_t = audio_t.to(model_dtype)
+        audio_t_mod = audio_t_mod.to(model_dtype)
+
+        visual_context_emb = visual_dit.text_embedding(visual_context)
+        audio_context_emb = self.audio_dit.text_embedding(audio_context)
+
+        visual_x = visual_latents.to(model_dtype)
+        audio_x = audio_latents.to(model_dtype)
+
+        if getattr(visual_dit, "require_vae_embedding", False):
+            visual_x = torch.cat([visual_x, y], dim=1)
+
+        # Patchify visual latents
+        visual_x, (t, h, w) = visual_dit.patchify(visual_x)
+        grid_size = (t, h, w)
+        full_visual_seq_len = t * h * w
+
+        # Build visual freqs for full sequence
+        visual_dit._init_freqs()
+        visual_freqs = tuple(freq.to(visual_x.device) for freq in visual_dit.freqs)
+        visual_freqs = (
+            torch.cat(
+                [
+                    visual_freqs[0][:t].view(t, 1, 1, -1).expand(t, h, w, -1),
+                    visual_freqs[1][:h].view(1, h, 1, -1).expand(t, h, w, -1),
+                    visual_freqs[2][:w].view(1, 1, w, -1).expand(t, h, w, -1),
+                ],
+                dim=-1,
+            )
+            .reshape(full_visual_seq_len, 1, -1)
+            .to(visual_x.device)
+        )
+
+        # Patchify audio latents
+        audio_x, (f,) = self.audio_dit.patchify(audio_x, None)
+        full_audio_seq_len = f
+
+        # Build audio freqs for full sequence
+        self.audio_dit._init_freqs()
+        audio_freqs = (
+            torch.cat(
+                [
+                    self.audio_dit.freqs[0][:f].view(f, -1).expand(f, -1),
+                    self.audio_dit.freqs[1][:f].view(f, -1).expand(f, -1),
+                    self.audio_dit.freqs[2][:f].view(f, -1).expand(f, -1),
+                ],
+                dim=-1,
+            )
+            .reshape(full_audio_seq_len, 1, -1)
+            .to(audio_x.device)
+        )
+
+        # Shard sequences for SP
+        visual_x, visual_pad_len = self._shard_sequence_for_sp(visual_x, dim=1)
+        audio_x, audio_pad_len = self._shard_sequence_for_sp(audio_x, dim=1)
+
+        # Shard freqs to match local sequence length
+        visual_freqs, _ = self._shard_sequence_for_sp(visual_freqs, dim=0)
+        audio_freqs, _ = self._shard_sequence_for_sp(audio_freqs, dim=0)
+
+        # Forward through dual-tower DiT
+        visual_x, audio_x = self.forward_dual_tower_dit(
+            visual_dit=visual_dit,
+            visual_x=visual_x,
+            audio_x=audio_x,
+            visual_context=visual_context_emb,
+            audio_context=audio_context_emb,
+            visual_t_mod=visual_t_mod,
+            audio_t_mod=audio_t_mod,
+            visual_freqs=visual_freqs,
+            audio_freqs=audio_freqs,
+            grid_size=grid_size,
+            video_fps=video_fps,
+            full_visual_seq_len=full_visual_seq_len,
+            full_audio_seq_len=full_audio_seq_len,
+        )
+
+        # Gather sequences back from SP before head/unpatchify
+        visual_x = self._gather_sequence_from_sp(visual_x, visual_pad_len, dim=1)
+        audio_x = self._gather_sequence_from_sp(audio_x, audio_pad_len, dim=1)
+
+        visual_output = visual_dit.head(visual_x, visual_t)
+        visual_output = visual_dit.unpatchify(visual_output, grid_size)
+
+        audio_output = self.audio_dit.head(audio_x, audio_t)
+        audio_output = self.audio_dit.unpatchify(audio_output, (f,))
+
+        return visual_output.float(), audio_output.float()
+
+    def forward_dual_tower_dit(
+        self,
+        visual_dit,
+        visual_x: torch.Tensor,
+        audio_x: torch.Tensor,
+        visual_context: torch.Tensor,
+        audio_context: torch.Tensor,
+        visual_t_mod: torch.Tensor,
+        audio_t_mod: torch.Tensor,
+        visual_freqs: torch.Tensor,
+        audio_freqs: torch.Tensor,
+        grid_size: tuple[int, int, int],
+        video_fps: float,
+        full_visual_seq_len: int,
+        full_audio_seq_len: int,
+        condition_scale: float | None = 1.0,
+        a2v_condition_scale: float | None = None,
+        v2a_condition_scale: float | None = None,
+    ):
+        """
+        Forward pass through dual-tower DiT with cross-modal interaction.
+
+        Sequence Parallelism (SP) Support:
+        - visual_x and audio_x are already sharded along sequence dimension
+        - visual_freqs and audio_freqs match the local sequence length
+        - USPAttention in self-attention handles distributed communication
+        - LocalAttention in cross-attention operates on local sequence vs replicated context
+        - Cross-modal attention (dual_tower_bridge) uses LocalAttention (no SP communication)
+
+        Args:
+            full_visual_seq_len: Full visual sequence length before SP sharding
+            full_audio_seq_len: Full audio sequence length before SP sharding
+        """
+        min_layers = min(len(visual_dit.blocks), len(self.audio_dit.blocks))
+        visual_layers = len(visual_dit.blocks)
+        sp_size = get_sp_world_size()
+
+        # Build RoPE frequencies for cross-attention if needed (only used when SP == 1)
+        # When SP > 1, we rebuild freqs inside the loop after gathering full sequences
+        visual_rope_cos_sin, audio_rope_cos_sin = (
+            self.dual_tower_bridge.build_aligned_freqs(
+                video_fps=video_fps,
+                grid_size=grid_size,
+                audio_steps=full_audio_seq_len,
+                device=visual_x.device,
+                dtype=visual_x.dtype,
+            )
+        )
+        if visual_rope_cos_sin is not None:
+            visual_rope_cos_sin = [
+                self._shard_sequence_for_sp(rope_cos_sin, dim=1)[0]
+                for rope_cos_sin in visual_rope_cos_sin
+            ]
+        if audio_rope_cos_sin is not None:
+            audio_rope_cos_sin = [
+                self._shard_sequence_for_sp(rope_cos_sin, dim=1)[0]
+                for rope_cos_sin in audio_rope_cos_sin
+            ]
+
+        for layer_idx in range(min_layers):
+            visual_block = visual_dit.blocks[layer_idx]
+            audio_block = self.audio_dit.blocks[layer_idx]
+
+            # Cross-modal interaction via dual tower bridge
+            # Bridge operations (PerFrameAttentionPooling, RoPE) expect full sequences
+            # When SP is enabled, we need to gather before bridge and shard after
+            if self.dual_tower_bridge.should_interact(layer_idx, "a2v"):
+                visual_x, audio_x = self.dual_tower_bridge(
+                    layer_idx,
+                    visual_x,
+                    audio_x,
+                    x_freqs=visual_rope_cos_sin,
+                    y_freqs=audio_rope_cos_sin,
+                    a2v_condition_scale=a2v_condition_scale,
+                    v2a_condition_scale=v2a_condition_scale,
+                    condition_scale=condition_scale,
+                    video_grid_size=grid_size,
+                )
+
+            # Self-attention and FFN in DiT blocks
+            visual_x = visual_block(
+                visual_x, visual_context, visual_t_mod, visual_freqs
+            )
+            audio_x = audio_block(audio_x, audio_context, audio_t_mod, audio_freqs)
+
+        # Process remaining visual layers (if visual has more layers than audio)
+        for layer_idx in range(min_layers, visual_layers):
+            visual_block = visual_dit.blocks[layer_idx]
+            visual_x = visual_block(
+                visual_x, visual_context, visual_t_mod, visual_freqs
+            )
+
+        return visual_x, audio_x
+
+
+class MOVADecodingStage(PipelineStage):
+    """Decode video and audio outputs for MOVA."""
+
+    def __init__(self, video_vae, audio_vae) -> None:
+        super().__init__()
+        self.video_vae = video_vae
+        self.audio_vae = audio_vae
+
+    @property
+    def parallelism_type(self) -> StageParallelismType:
+        if get_global_server_args().enable_cfg_parallel:
+            return StageParallelismType.MAIN_RANK_ONLY
+        return StageParallelismType.REPLICATED
+
+    @torch.no_grad()
+    def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
+        self.video_vae = self.video_vae.to(get_local_torch_device())
+        self.audio_vae = self.audio_vae.to(get_local_torch_device())
+
+        video_latents = server_args.pipeline_config.denormalize_video_latents(
+            batch.latents, self.video_vae
+        )
+
+        vae_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision]
+        vae_autocast_enabled = (
+            vae_dtype != torch.float32
+        ) and not server_args.disable_autocast
+
+        with torch.autocast(
+            device_type=current_platform.device_type,
+            dtype=vae_dtype,
+            enabled=vae_autocast_enabled,
+        ):
+            if server_args.pipeline_config.vae_tiling:
+                self.video_vae.enable_tiling()
+            if not vae_autocast_enabled:
+                video_latents = video_latents.to(vae_dtype)
+            decode_output = self.video_vae.decode(video_latents)
+            video = _ensure_tensor_decode_output(decode_output)
+
+        video = (video / 2 + 0.5).clamp(0, 1)
+
+        with torch.autocast(
+            device_type=current_platform.device_type, dtype=torch.float32
+        ):
+            audio = self.audio_vae.decode(batch.audio_latents)
+        output_batch = OutputBatch(
+            output=video,
+            audio=audio,
+            audio_sample_rate=getattr(self.audio_vae, "sample_rate", None),
+            timings=batch.timings,
+        )
+        return output_batch

python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py

@@ -0,0 +1,501 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copied and adapted from: mossVG/mova/diffusion/schedulers/flow_match.py and flow_match_pair.py
+
+from __future__ import annotations
+
+import math
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin
+
+from sglang.multimodal_gen.runtime.models.schedulers.base import BaseScheduler
+
+
+class FlowMatchScheduler(BaseScheduler):
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self.order = 1
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.exponential_shift = exponential_shift
+        self.exponential_shift_mu = exponential_shift_mu
+        self.shift_terminal = shift_terminal
+        self.train_timesteps = None
+        self.train_sigmas = None
+        self.set_timesteps(num_train_timesteps)
+        self.set_timesteps(num_inference_steps)
+        BaseScheduler.__init__(self)
+
+    def set_shift(self, shift: float) -> None:
+        self.shift = shift
+
+    def set_timesteps(
+        self,
+        num_inference_steps=100,
+        denoising_strength=1.0,
+        training=False,
+        shift=None,
+        dynamic_shift_len=None,
+    ):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = (
+            self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        )
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps + 1
+            )[:-1]
+        else:
+            self.sigmas = torch.linspace(
+                sigma_start, self.sigma_min, num_inference_steps
+            )
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        if self.exponential_shift:
+            mu = (
+                self.calculate_shift(dynamic_shift_len)
+                if dynamic_shift_len is not None
+                else self.exponential_shift_mu
+            )
+            self.sigmas = math.exp(mu) / (math.exp(mu) + (1 / self.sigmas - 1))
+        else:
+            self.sigmas = (
+                self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+            )
+        if self.shift_terminal is not None:
+            one_minus_z = 1 - self.sigmas
+            scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+            self.sigmas = 1 - (one_minus_z / scale_factor)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        # 第一次设置 train_timesteps
+        if self.train_timesteps is None:
+            self.train_timesteps = self.timesteps
+            self.train_sigmas = self.sigmas
+        if training:
+            x = self.timesteps
+            y = torch.exp(
+                -2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2
+            )
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+            self.training = True
+        else:
+            self.training = False
+
+    def scale_model_input(self, sample: torch.Tensor, timestep: int | None = None):
+        return sample
+
+    def step(self, model_output, timestep, sample, to_final=False, **kwargs):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def return_to_timestep(self, timestep, sample, sample_stablized):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        model_output = (sample - sample_stablized) / sigma
+        return model_output
+
+    def add_noise(self, original_samples, noise, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin(
+            (self.timesteps - timestep.to(self.timesteps.device)).abs()
+        )
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+
+    def calculate_shift(
+        self,
+        image_seq_len,
+        base_seq_len: int = 256,
+        max_seq_len: int = 8192,
+        base_shift: float = 0.5,
+        max_shift: float = 0.9,
+    ):
+        m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+        b = base_shift - m * base_seq_len
+        mu = image_seq_len * m + b
+        return mu
+
+
+class FlowMatchPairScheduler(FlowMatchScheduler, SchedulerMixin, ConfigMixin):
+    """
+    在 FlowMatchScheduler 的基础上，提供便捷的配对接口：
+    - 默认返回形状为 [num_timesteps, 2] 的张量，每一行为 (t, t)
+    - 允许通过 set_pair_postprocess(fn) 设置一个后处理函数以修改配对行为
+    """
+
+    @register_to_config
+    def __init__(
+        self,
+        num_inference_steps=100,
+        num_train_timesteps=1000,
+        shift=3.0,
+        sigma_max=1.0,
+        sigma_min=0.003 / 1.002,
+        inverse_timesteps=False,
+        extra_one_step=False,
+        reverse_sigmas=False,
+        exponential_shift=False,
+        exponential_shift_mu=None,
+        shift_terminal=None,
+    ):
+        self._pair_postprocess_fn = None
+        self._pair_postprocess_requires_source = False
+        self.pair_timesteps: torch.Tensor | None = None
+        self.pair_sigmas: torch.Tensor | None = None
+        self.timesteps: torch.Tensor | None = None
+        self.sigmas: torch.Tensor | None = None
+        super().__init__(
+            num_inference_steps=num_inference_steps,
+            num_train_timesteps=num_train_timesteps,
+            shift=shift,
+            sigma_max=sigma_max,
+            sigma_min=sigma_min,
+            inverse_timesteps=inverse_timesteps,
+            extra_one_step=extra_one_step,
+            reverse_sigmas=reverse_sigmas,
+            exponential_shift=exponential_shift,
+            exponential_shift_mu=exponential_shift_mu,
+            shift_terminal=shift_terminal,
+        )
+
+    def set_pair_postprocess(self, fn):
+        """
+        设置一个后处理函数，用于在默认配对生成后进行自定义修改。
+        要求：
+        - fn(pairs: torch.Tensor) -> torch.Tensor
+        - 返回的张量必须与输入 pairs 形状一致，否则直接 raise。
+        """
+        if fn is not None and not callable(fn):
+            raise TypeError("pair_postprocess 必须是可调用对象或 None")
+        self._pair_postprocess_fn = fn
+        self._pair_postprocess_requires_source = (
+            False if fn is None else bool(getattr(fn, "_requires_source", False))
+        )
+        if self.timesteps is None or self.sigmas is None:
+            raise RuntimeError("调度器未初始化，请先调用 set_timesteps()")
+        self._refresh_pair_cache()
+
+    def set_pair_postprocess_by_name(self, name: str | None, **kwargs):
+        """
+        通过名称快速配置后处理函数。
+        支持：
+        - None/"none"/"off"/"false"/"no": 关闭
+        - "quadratic_perp_bulge_swap": x2=x+d, y2=x-d，其中 d=4*amp*s*(1-s), s=t/T
+        - "v2a_sequential": 假设原 pairs 为 (t,t)，用列0步长2采样一半序列，先让列0按该序列变化，再让列1按该序列变化
+        - "a2v_sequential": 同上，先列1后列0
+        - "dual_sigma_shift": 仅使用 timestep 数量，重新按照 FlowMatchScheduler 的 sigma 变换逻辑为两列独立构建调度，可配置 visual_shift/audio_shift
+
+        额外参数：
+        - amp: 浮点幅度，默认 150.0
+        """
+        if name is None or str(name).lower() in ("none", "off", "false", "no"):
+            self.set_pair_postprocess(None)
+            return
+        if name == "quadratic_perp_bulge_swap":
+            amp = float(kwargs.get("amp", 150.0))
+
+            def _quadratic_perp_bulge_swap(pairs: torch.Tensor):
+                if (
+                    not isinstance(pairs, torch.Tensor)
+                    or pairs.ndim != 2
+                    or pairs.shape[1] != 2
+                ):
+                    raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor")
+                x = pairs[:, 0]
+                T = float(self.num_train_timesteps)
+                s = x / T
+                d = 4.0 * amp * s * (1.0 - s)
+                x2 = x + d
+                y2 = x - d
+                return torch.stack([x2, y2], dim=1)
+
+            self.set_pair_postprocess(_quadratic_perp_bulge_swap)
+            return
+        if name == "v2a_sequential":
+
+            def _v2a(pairs: torch.Tensor):
+                if (
+                    not isinstance(pairs, torch.Tensor)
+                    or pairs.ndim != 2
+                    or pairs.shape[1] != 2
+                ):
+                    raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor")
+                N = pairs.shape[0]
+                base = pairs[:, 0]
+                seq_half = base[::2]
+                m = int(seq_half.shape[0])
+                col0 = torch.cat([seq_half, seq_half[-1:].repeat(m)], dim=0)[:N]
+                col1 = torch.cat([seq_half[0:1].repeat(m), seq_half], dim=0)[:N]
+                return torch.stack(
+                    [
+                        col0.to(dtype=pairs.dtype, device=pairs.device),
+                        col1.to(dtype=pairs.dtype, device=pairs.device),
+                    ],
+                    dim=1,
+                )
+
+            self.set_pair_postprocess(_v2a)
+            return
+        if name == "a2v_sequential":
+
+            def _a2v(pairs: torch.Tensor):
+                if (
+                    not isinstance(pairs, torch.Tensor)
+                    or pairs.ndim != 2
+                    or pairs.shape[1] != 2
+                ):
+                    raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor")
+                N = pairs.shape[0]
+                base = pairs[:, 0]
+                seq_half = base[::2]
+                m = int(seq_half.shape[0])
+                col0 = torch.cat([seq_half[0:1].repeat(m), seq_half], dim=0)[:N]
+                col1 = torch.cat([seq_half, seq_half[-1:].repeat(m)], dim=0)[:N]
+                return torch.stack(
+                    [
+                        col0.to(dtype=pairs.dtype, device=pairs.device),
+                        col1.to(dtype=pairs.dtype, device=pairs.device),
+                    ],
+                    dim=1,
+                )
+
+            self.set_pair_postprocess(_a2v)
+            return
+        if name == "v2a":
+
+            def _v2a_classic(pairs: torch.Tensor):
+                if (
+                    not isinstance(pairs, torch.Tensor)
+                    or pairs.ndim != 2
+                    or pairs.shape[1] != 2
+                ):
+                    raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor")
+                zeros = torch.zeros_like(pairs[:, 0])
+                return torch.stack([zeros, pairs[:, 1]], dim=1)
+
+            self.set_pair_postprocess(_v2a_classic)
+            return
+        if name == "a2v":
+
+            def _a2v_classic(pairs: torch.Tensor):
+                if (
+                    not isinstance(pairs, torch.Tensor)
+                    or pairs.ndim != 2
+                    or pairs.shape[1] != 2
+                ):
+                    raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor")
+                zeros = torch.zeros_like(pairs[:, 1])
+                return torch.stack([pairs[:, 0], zeros], dim=1)
+
+            self.set_pair_postprocess(_a2v_classic)
+            return
+        if name == "dual_sigma_shift":
+            visual_shift = float(kwargs.get("visual_shift", self.shift))
+            audio_shift = float(kwargs.get("audio_shift", self.shift))
+            visual_denoising_strength = float(
+                kwargs.get("visual_denoising_strength", 1.0)
+            )
+            audio_denoising_strength = float(
+                kwargs.get("audio_denoising_strength", 1.0)
+            )
+            visual_mu = kwargs.get(
+                "visual_exponential_shift_mu", self.exponential_shift_mu
+            )
+            audio_mu = kwargs.get(
+                "audio_exponential_shift_mu", self.exponential_shift_mu
+            )
+
+            def _dual_sigma_shift(pairs: torch.Tensor, *, source: str):
+                if not isinstance(pairs, torch.Tensor):
+                    raise TypeError("pairs 必须是 torch.Tensor")
+                if pairs.ndim != 2 or pairs.shape[1] != 2:
+                    raise ValueError("pairs 必须是形状 [N, 2] 的 torch.Tensor")
+                if pairs.shape[0] == 0:
+                    raise ValueError("pairs 的长度必须大于 0")
+                if source not in ("timesteps", "sigmas"):
+                    raise ValueError("source 仅支持 'timesteps' 或 'sigmas'")
+
+                num_steps = pairs.shape[0]
+                device = pairs.device
+                dtype = pairs.dtype
+
+                def _build_column(
+                    shift_value: float, denoising_strength: float, mu_override
+                ):
+                    if shift_value <= 0:
+                        raise ValueError("shift 必须为正数")
+                    if denoising_strength <= 0:
+                        raise ValueError("denoising_strength 必须为正数")
+
+                    sigma_start = (
+                        self.sigma_min
+                        + (self.sigma_max - self.sigma_min) * denoising_strength
+                    )
+                    if self.extra_one_step:
+                        base = torch.linspace(
+                            sigma_start,
+                            self.sigma_min,
+                            num_steps + 1,
+                            device=device,
+                            dtype=dtype,
+                        )[:-1]
+                    else:
+                        base = torch.linspace(
+                            sigma_start,
+                            self.sigma_min,
+                            num_steps,
+                            device=device,
+                            dtype=dtype,
+                        )
+
+                    if self.inverse_timesteps:
+                        base = torch.flip(base, dims=[0])
+
+                    if self.exponential_shift:
+                        mu_value = mu_override
+                        if mu_value is None:
+                            raise RuntimeError(
+                                "启用了 exponential_shift 但未提供 exponential_shift_mu"
+                            )
+                        exp_mu = math.exp(float(mu_value))
+                        base = exp_mu / (exp_mu + (1 / base - 1))
+                    else:
+                        base = shift_value * base / (1 + (shift_value - 1) * base)
+
+                    if self.shift_terminal is not None:
+                        one_minus_z = 1 - base
+                        scale_factor = one_minus_z[-1] / (1 - self.shift_terminal)
+                        base = 1 - (one_minus_z / scale_factor)
+
+                    if self.reverse_sigmas:
+                        base = 1 - base
+
+                    if source == "timesteps":
+                        return base * self.num_train_timesteps
+                    return base
+
+                col0 = _build_column(visual_shift, visual_denoising_strength, visual_mu)
+                col1 = _build_column(audio_shift, audio_denoising_strength, audio_mu)
+                return torch.stack([col0, col1], dim=1)
+
+            _dual_sigma_shift._requires_source = True
+            self.set_pair_postprocess(_dual_sigma_shift)
+            return
+        raise ValueError(f"Unknown pair_postprocess name: {name}")
+
+    def _make_pairs_from_vector(self, vec: torch.Tensor) -> torch.Tensor:
+        if vec.ndim != 1:
+            raise ValueError("vec must be 1D")
+        return torch.stack([vec, vec], dim=1)
+
+    def get_pairs(self, source: str = "timesteps") -> torch.Tensor:
+        if source == "timesteps":
+            if self.pair_timesteps is None:
+                self._refresh_pair_cache()
+            return self.pair_timesteps
+        if source == "sigmas":
+            if self.pair_sigmas is None:
+                self._refresh_pair_cache()
+            return self.pair_sigmas
+        raise ValueError("source must be 'timesteps' or 'sigmas'")
+
+    def timestep_to_sigma(self, timestep: torch.Tensor | float) -> torch.Tensor:
+        """根据给定的 timestep（标量）返回对应的 sigma（按最近邻在 self.timesteps 中查找）。"""
+        t_value = float(timestep)
+        t_cpu = torch.tensor(t_value)
+        idx = torch.argmin((self.train_timesteps - t_cpu).abs())
+        return self.train_sigmas[idx]
+
+    def step_from_to(
+        self,
+        model_output: torch.Tensor,
+        timestep_from: torch.Tensor,
+        timestep_to: torch.Tensor | None,
+        sample: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        使用显式给定的 (from, to) timestep 对，按照对应的 sigma 差进行一步更新：
+            x_{to} = x_{from} + model_output * (sigma(to) - sigma(from))
+        该方法可用于两个模态分别沿着它们各自列序列推进。
+        """
+        sigma_from = self.timestep_to_sigma(timestep_from)
+        if timestep_to is None:
+            sigma_to = torch.tensor(
+                1.0 if (self.inverse_timesteps or self.reverse_sigmas) else 0.0,
+                device=sigma_from.device,
+                dtype=sigma_from.dtype,
+            )
+        else:
+            sigma_to = self.timestep_to_sigma(timestep_to)
+        prev_sample = sample + model_output * (sigma_to - sigma_from)
+        return prev_sample
+
+    def _refresh_pair_cache(self) -> None:
+        if self.timesteps is None or self.sigmas is None:
+            raise RuntimeError("调度器未初始化，请先调用 set_timesteps()")
+
+        def _apply_postprocess(pairs: torch.Tensor, source: str) -> torch.Tensor:
+            if self._pair_postprocess_fn is None:
+                return pairs
+            if self._pair_postprocess_requires_source:
+                modified = self._pair_postprocess_fn(pairs, source=source)
+            else:
+                modified = self._pair_postprocess_fn(pairs)
+            if not isinstance(modified, torch.Tensor):
+                raise TypeError("pair_postprocess 返回值必须是 torch.Tensor")
+            if modified.shape != pairs.shape:
+                raise ValueError("pair_postprocess 返回的张量形状必须与输入一致")
+            return modified
+
+        base_pairs_timesteps = self._make_pairs_from_vector(self.timesteps)
+        base_pairs_sigmas = self._make_pairs_from_vector(self.sigmas)
+
+        self.pair_timesteps = _apply_postprocess(base_pairs_timesteps, "timesteps")
+        self.pair_sigmas = _apply_postprocess(base_pairs_sigmas, "sigmas")
+
+
+EntryClass = FlowMatchPairScheduler

python/sglang/multimodal_gen/runtime/models/vaes/common.py

@@ -611,7 +611,9 @@ class DiagonalGaussianDistribution:
         return x
 
     def kl(
-        self, other: Optional["DiagonalGaussianDistribution"] = None
+        self,
+        other: Optional["DiagonalGaussianDistribution"] = None,
+        dims: tuple[int, ...] = (1, 2, 3),
     ) -> torch.Tensor:
         if self.deterministic:
             return torch.Tensor([0.0])
@@ -619,7 +621,7 @@ class DiagonalGaussianDistribution:
             if other is None:
                 return 0.5 * torch.sum(
                     torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar,
-                    dim=[1, 2, 3],
+                    dim=dims,
                 )
             else:
                 return 0.5 * torch.sum(
@@ -628,7 +630,7 @@ class DiagonalGaussianDistribution:
                     - 1.0
                     - self.logvar
                     + other.logvar,
-                    dim=[1, 2, 3],
+                    dim=dims,
                 )
 
     def nll(

python/sglang/multimodal_gen/runtime/models/vaes/dac.py

@@ -0,0 +1,627 @@
+# Copied and adapted from: https://github.com/descriptinc/descript-audio-codec
+
+# SPDX-License-Identifier: MIT
+
+import math
+from bisect import bisect_right
+from typing import Union
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+
+from sglang.multimodal_gen.configs.models.vaes.dac import DacVAEConfig
+from sglang.multimodal_gen.runtime.models.vaes.common import (
+    DiagonalGaussianDistribution,
+)
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+class VectorQuantize(nn.Module):
+    """
+    Implementation of VQ similar to Karpathy's repo:
+    https://github.com/karpathy/deep-vector-quantization
+    Additionally uses following tricks from Improved VQGAN
+    (https://arxiv.org/pdf/2110.04627.pdf):
+        1. Factorized codes: Perform nearest neighbor lookup in low-dimensional space
+            for improved codebook usage
+        2. l2-normalized codes: Converts euclidean distance to cosine similarity which
+            improves training stability
+    """
+
+    def __init__(self, input_dim: int, codebook_size: int, codebook_dim: int):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+
+        self.in_proj = nn.Conv1d(input_dim, codebook_dim, kernel_size=1)
+        self.out_proj = nn.Conv1d(codebook_dim, input_dim, kernel_size=1)
+        self.codebook = nn.Embedding(codebook_size, codebook_dim)
+
+    def forward(self, z):
+        """Quantize the input tensor using a fixed codebook and return the corresponding codebook vectors.
+
+        Args:
+            z (torch.Tensor): Input tensor with shape ``[B, D, T]``.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``.
+                - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict
+                  vectors closer to codebook entries.
+                - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook.
+                - indices (torch.Tensor): Codebook indices (quantized discrete representation) with shape ``[B, T]``.
+                - z_e (torch.Tensor): Projected latents (continuous representation before quantization) with shape ``[B, D, T]``.
+        """
+
+        # Factorized codes (ViT-VQGAN) Project input into low-dimensional space
+        z_e = self.in_proj(z)  # z_e : (B x D x T)
+        z_q, indices = self.decode_latents(z_e)
+
+        commitment_loss = F.mse_loss(z_e, z_q.detach(), reduction="none").mean([1, 2])
+        codebook_loss = F.mse_loss(z_q, z_e.detach(), reduction="none").mean([1, 2])
+
+        z_q = (
+            z_e + (z_q - z_e).detach()
+        )  # noop in forward pass, straight-through gradient estimator in backward pass
+
+        z_q = self.out_proj(z_q)
+
+        return z_q, commitment_loss, codebook_loss, indices, z_e
+
+    def embed_code(self, embed_id):
+        return F.embedding(embed_id, self.codebook.weight)
+
+    def decode_code(self, embed_id):
+        return self.embed_code(embed_id).transpose(1, 2)
+
+    def decode_latents(self, latents):
+        encodings = rearrange(latents, "b d t -> (b t) d")
+        codebook = self.codebook.weight  # codebook: (N x D)
+
+        # L2 normalize encodings and codebook (ViT-VQGAN)
+        encodings = F.normalize(encodings)
+        codebook = F.normalize(codebook)
+
+        # Compute euclidean distance with codebook
+        dist = (
+            encodings.pow(2).sum(1, keepdim=True)
+            - 2 * encodings @ codebook.t()
+            + codebook.pow(2).sum(1, keepdim=True).t()
+        )
+        indices = rearrange((-dist).max(1)[1], "(b t) -> b t", b=latents.size(0))
+        z_q = self.decode_code(indices)
+        return z_q, indices
+
+
+class ResidualVectorQuantize(nn.Module):
+    """
+    Introduced in SoundStream: An end2end neural audio codec
+    https://arxiv.org/abs/2107.03312
+    """
+
+    def __init__(
+        self,
+        input_dim: int = 512,
+        n_codebooks: int = 9,
+        codebook_size: int = 1024,
+        codebook_dim: Union[int, list] = 8,
+        quantizer_dropout: float = 0.0,
+    ):
+        super().__init__()
+        if isinstance(codebook_dim, int):
+            codebook_dim = [codebook_dim for _ in range(n_codebooks)]
+
+        self.n_codebooks = n_codebooks
+        self.codebook_dim = codebook_dim
+        self.codebook_size = codebook_size
+        dim_offsets = [0]
+        for dim in self.codebook_dim:
+            dim_offsets.append(dim_offsets[-1] + dim)
+        self._codebook_dim_offsets = tuple(dim_offsets)
+
+        self.quantizers = nn.ModuleList(
+            [
+                VectorQuantize(input_dim, codebook_size, codebook_dim[i])
+                for i in range(n_codebooks)
+            ]
+        )
+        self.quantizer_dropout = quantizer_dropout
+
+    def forward(self, z, n_quantizers: int = None):
+        """Quantize the input tensor using a fixed set of codebooks and return the corresponding codebook vectors.
+
+        Args:
+            z (torch.Tensor): Input tensor with shape ``[B, D, T]``.
+            n_quantizers (int, optional): Number of quantizers to use. If ``None``,
+                all quantizers are used. When ``n_quantizers`` < ``self.n_codebooks``,
+                quantizer dropout is applied. Note: if ``self.quantizer_dropout`` > 0
+                and in training mode, this argument is ignored and a random number of
+                quantizers is used.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``.
+                - codes (torch.Tensor): Codebook indices for each codebook with shape ``[B, N, T]``
+                  (quantized discrete representation of input).
+                - latents (torch.Tensor): Projected latents with shape ``[B, N*D, T]``
+                  (continuous representation before quantization).
+                - commitment_loss (torch.Tensor): Commitment loss scalar to train encoder to predict
+                  vectors closer to codebook entries.
+                - codebook_loss (torch.Tensor): Codebook loss scalar to update the codebook.
+        """
+        z_q = 0
+        residual = z
+        commitment_loss = 0
+        codebook_loss = 0
+
+        codebook_indices = []
+        latents = []
+
+        if n_quantizers is None:
+            n_quantizers = self.n_codebooks
+        quantizers = self.quantizers
+        if self.training:
+            batch_size = z.shape[0]
+            device = z.device
+            n_quantizers = torch.full(
+                (batch_size,),
+                self.n_codebooks + 1,
+                device=device,
+                dtype=torch.long,
+            )
+            if self.quantizer_dropout > 0:
+                dropout = torch.randint(
+                    1,
+                    self.n_codebooks + 1,
+                    (batch_size,),
+                    device=device,
+                )
+                n_dropout = int(batch_size * self.quantizer_dropout)
+                if n_dropout > 0:
+                    n_quantizers[:n_dropout] = dropout[:n_dropout]
+
+            for i, quantizer in enumerate(quantizers):
+                z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                    residual
+                )
+
+                # Create mask to apply quantizer dropout
+                mask = i < n_quantizers
+                z_q = z_q + z_q_i * mask[:, None, None]
+                residual = residual - z_q_i
+
+                # Sum losses
+                commitment_loss += (commitment_loss_i * mask).mean()
+                codebook_loss += (codebook_loss_i * mask).mean()
+
+                codebook_indices.append(indices_i)
+                latents.append(z_e_i)
+        else:
+            for i, quantizer in enumerate(quantizers):
+                if i >= n_quantizers:
+                    break
+                z_q_i, commitment_loss_i, codebook_loss_i, indices_i, z_e_i = quantizer(
+                    residual
+                )
+                z_q = z_q + z_q_i
+                residual = residual - z_q_i
+
+                commitment_loss += commitment_loss_i.mean()
+                codebook_loss += codebook_loss_i.mean()
+
+                codebook_indices.append(indices_i)
+                latents.append(z_e_i)
+
+        codes = torch.stack(codebook_indices, dim=1)
+        latents = torch.cat(latents, dim=1)
+
+        return z_q, codes, latents, commitment_loss, codebook_loss
+
+    def from_codes(self, codes: torch.Tensor):
+        """Reconstruct the continuous representation from quantized codes.
+
+        Args:
+            codes (torch.Tensor): Quantized discrete representation with shape ``[B, N, T]``.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized continuous representation with shape ``[B, D, T]``.
+                - z_p (torch.Tensor): Concatenated latent space representation with shape ``[B, N*D, T]``.
+                - codes (torch.Tensor): Original input codebook indices with shape ``[B, N, T]``.
+        """
+        z_q = 0.0
+        z_p = []
+        n_codebooks = codes.shape[1]
+        for i in range(n_codebooks):
+            z_p_i = self.quantizers[i].decode_code(codes[:, i, :])
+            z_p.append(z_p_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+        return z_q, torch.cat(z_p, dim=1), codes
+
+    def from_latents(self, latents: torch.Tensor):
+        """Reconstruct the continuous representation from unquantized latents.
+
+        Args:
+            latents (torch.Tensor): Continuous representation after projection with shape ``[B, N*D, T]``.
+
+        Returns:
+            tuple: A tuple containing:
+                - z_q (torch.Tensor): Quantized representation of full-projected space with shape ``[B, D, T]``.
+                - z_p (torch.Tensor): Quantized representation of latent space with shape ``[B, N*D, T]``.
+                - codes (torch.Tensor): Codebook indices with shape ``[B, N, T]``.
+        """
+        z_q = 0
+        z_p = []
+        codes = []
+        dims = self._codebook_dim_offsets
+        n_codebooks = bisect_right(dims, latents.shape[1]) - 1
+        for i in range(n_codebooks):
+            j, k = dims[i], dims[i + 1]
+            z_p_i, codes_i = self.quantizers[i].decode_latents(latents[:, j:k, :])
+            z_p.append(z_p_i)
+            codes.append(codes_i)
+
+            z_q_i = self.quantizers[i].out_proj(z_p_i)
+            z_q = z_q + z_q_i
+
+        return z_q, torch.cat(z_p, dim=1), torch.stack(codes, dim=1)
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim: int = 16, dilation: int = 1):
+        super().__init__()
+        pad = ((7 - 1) * dilation) // 2
+        self.block = nn.Sequential(
+            Snake1d(dim),
+            nn.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
+            Snake1d(dim),
+            nn.Conv1d(dim, dim, kernel_size=1),
+        )
+
+    def forward(self, x):
+        y = self.block(x)
+        pad = (x.shape[-1] - y.shape[-1]) // 2
+        if pad > 0:
+            x = x[..., pad:-pad]
+        return x + y
+
+
+class EncoderBlock(nn.Module):
+    def __init__(self, dim: int = 16, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            ResidualUnit(dim // 2, dilation=1),
+            ResidualUnit(dim // 2, dilation=3),
+            ResidualUnit(dim // 2, dilation=9),
+            Snake1d(dim // 2),
+            nn.Conv1d(
+                dim // 2,
+                dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+            ),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        d_model: int = 64,
+        strides: list = [2, 4, 8, 8],
+        d_latent: int = 64,
+    ):
+        super().__init__()
+        # Create first convolution
+        self.block = [nn.Conv1d(1, d_model, kernel_size=7, padding=3)]
+
+        # Create EncoderBlocks that double channels as they downsample by `stride`
+        for stride in strides:
+            d_model *= 2
+            self.block += [EncoderBlock(d_model, stride=stride)]
+
+        # Create last convolution
+        self.block += [
+            Snake1d(d_model),
+            nn.Conv1d(d_model, d_latent, kernel_size=3, padding=1),
+        ]
+
+        # Wrap black into nn.Sequential
+        self.block = nn.Sequential(*self.block)
+        self.enc_dim = d_model
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim),
+            nn.ConvTranspose1d(
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                output_padding=stride % 2,
+            ),
+            ResidualUnit(output_dim, dilation=1),
+            ResidualUnit(output_dim, dilation=3),
+            ResidualUnit(output_dim, dilation=9),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        input_channel,
+        channels,
+        rates,
+        d_out: int = 1,
+    ):
+        super().__init__()
+
+        # Add first conv layer
+        layers = [nn.Conv1d(input_channel, channels, kernel_size=7, padding=3)]
+
+        # Add upsampling + MRF blocks
+        for i, stride in enumerate(rates):
+            input_dim = channels // 2**i
+            output_dim = channels // 2 ** (i + 1)
+            layers += [DecoderBlock(input_dim, output_dim, stride)]
+
+        # Add final conv layer
+        layers += [
+            Snake1d(output_dim),
+            nn.Conv1d(output_dim, d_out, kernel_size=7, padding=3),
+            nn.Tanh(),
+        ]
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class DAC(nn.Module):
+    def __init__(
+        self,
+        config: DacVAEConfig,
+    ):
+        super().__init__()
+
+        self.continuous = config.continuous
+        self.decoder_dim = config.decoder_dim
+        self.decoder_rates = config.decoder_rates
+        self.encoder_dim = config.encoder_dim
+        self.encoder_rates = config.encoder_rates
+        self.hop_length = math.prod(config.encoder_rates)
+        self.sample_rate = config.sample_rate
+
+        if config.latent_dim is None:
+            latent_dim = config.encoder_dim * (2 ** len(config.encoder_rates))
+        else:
+            latent_dim = config.latent_dim
+
+        self.latent_dim = latent_dim
+
+        if config.load_encoder:
+            self.encoder = Encoder(config.encoder_dim, config.encoder_rates, latent_dim)
+
+        if not config.continuous:
+            self.n_codebooks = config.n_codebooks
+            self.codebook_size = config.codebook_size
+            self.codebook_dim = config.codebook_dim
+            self.quantizer = ResidualVectorQuantize(
+                input_dim=latent_dim,
+                n_codebooks=config.n_codebooks,
+                codebook_size=config.codebook_size,
+                codebook_dim=config.codebook_dim,
+                quantizer_dropout=config.quantizer_dropout,
+            )
+        else:
+            self.quant_conv = torch.nn.Conv1d(latent_dim, 2 * latent_dim, 1)
+            self.post_quant_conv = torch.nn.Conv1d(latent_dim, latent_dim, 1)
+
+        if config.load_decoder:
+            self.decoder = Decoder(
+                latent_dim,
+                config.decoder_dim,
+                config.decoder_rates,
+            )
+
+        self.apply(self.init_weights)
+
+    @staticmethod
+    def init_weights(m):
+        if isinstance(m, nn.Conv1d):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            nn.init.constant_(m.bias, 0)
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def preprocess(self, audio_data, sample_rate):
+        if sample_rate is None:
+            sample_rate = self.sample_rate
+        assert sample_rate == self.sample_rate
+
+        length = audio_data.shape[-1]
+        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
+        audio_data = nn.functional.pad(audio_data, (0, right_pad))
+
+        return audio_data
+
+    def encode(
+        self,
+        audio_data: torch.Tensor,
+        n_quantizers: int = None,
+    ):
+        """Encode audio data into latent representations.
+
+        This method processes audio through the encoder network and optionally applies
+        vector quantization (in VQ mode) or projects to a Gaussian distribution (in
+        continuous mode) to produce latent representations.
+
+        Args:
+            audio_data (torch.Tensor): Audio data to encode, with shape ``[B, 1, T]``.
+            n_quantizers (int, optional): Number of quantizers to use. If ``None``,
+                all quantizers are used. Only applicable in VQ mode (``continuous=False``).
+
+        Returns:
+            tuple: A tuple containing:
+                - z (torch.Tensor): Encoded representation. In VQ mode, this is the
+                  quantized continuous representation with shape ``[B, D, T]``. In
+                  continuous mode, this is a ``DiagonalGaussianDistribution`` object.
+                - codes (torch.Tensor or None): Codebook indices with shape ``[B, N, T]``
+                  in VQ mode, ``None`` in continuous mode.
+                - latents (torch.Tensor or None): Projected latents with shape ``[B, N*D, T]``
+                  in VQ mode, ``None`` in continuous mode.
+                - commitment_loss (torch.Tensor): Commitment loss scalar.
+                - codebook_loss (torch.Tensor): Codebook loss scalar.
+
+        Note:
+            In continuous mode, the encoded representation is projected through a
+            quantization convolution layer and wrapped in a ``DiagonalGaussianDistribution``
+            for VAE training.
+        """
+        z = self.encoder(audio_data)  # [B x D x T]
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.quantizer(
+                z, n_quantizers
+            )
+        else:
+            z = self.quant_conv(z)  # [B x 2D x T]
+            z = DiagonalGaussianDistribution(z)
+            codes, latents, commitment_loss, codebook_loss = None, None, 0, 0
+
+        return z, codes, latents, commitment_loss, codebook_loss
+
+    def decode(self, z: torch.Tensor):
+        """Decode latent representations back to audio waveforms.
+
+        This method takes latent representations (either quantized from VQ mode or sampled
+        from the posterior in continuous mode) and reconstructs the corresponding audio
+        through the decoder network.
+
+        Args:
+            z (torch.Tensor): Latent representation to decode, with shape ``[B, D, T]``.
+                In VQ mode (``continuous=False``), this is the quantized continuous
+                representation. In continuous mode (``continuous=True``), this is sampled
+                from the posterior distribution.
+
+        Returns:
+            torch.Tensor: Decoded audio data with shape ``[B, 1, T']``. The output length
+            T' is determined by the decoder's upsampling rates and may differ from the
+            input temporal dimension T.
+
+        Note:
+            In continuous mode (``continuous=True``), the input is first passed through
+            a post-quantization convolution layer before being fed to the decoder.
+        """
+        if not self.continuous:
+            audio = self.decoder(z)
+        else:
+            z = self.post_quant_conv(z)
+            audio = self.decoder(z)
+
+        return audio
+
+    def forward(
+        self,
+        audio_data: torch.Tensor,
+        sample_rate: int = None,
+        n_quantizers: int = None,
+    ):
+        """Model forward pass.
+
+        Args:
+            audio_data (torch.Tensor): Audio to encode, shape [B, 1, T].
+            sample_rate (int, optional): Sample rate in Hz. Defaults to
+                ``self.sample_rate`` when ``None``.
+            n_quantizers (int, optional): Number of quantizers to use. When ``None``,
+                all quantizers are used. Only used in VQ mode (``continuous=False``).
+
+        Returns:
+            dict: A dictionary containing different keys depending on the mode:
+
+            **VQ Mode (``continuous=False``):**
+                - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length].
+                - "z" (torch.Tensor): Quantized continuous representation, shape [B, D, T].
+                - "codes" (torch.Tensor): Codebook indices, shape [B, N, T].
+                - "latents" (torch.Tensor): Projected latents, shape [B, N*D, T].
+                - "vq/commitment_loss" (torch.Tensor): Commitment loss.
+                - "vq/codebook_loss" (torch.Tensor): Codebook loss.
+
+            **Continuous Mode (``continuous=True``):**
+                - "audio" (torch.Tensor): Decoded audio, shape [B, 1, length].
+                - "z" (torch.Tensor): Latent representation, shape [B, D, T].
+                - "kl_loss" (torch.Tensor): KL divergence loss (for VAE training).
+        """
+        length = audio_data.shape[-1]
+        audio_data = self.preprocess(audio_data, sample_rate)
+        if not self.continuous:
+            z, codes, latents, commitment_loss, codebook_loss = self.encode(
+                audio_data, n_quantizers
+            )
+
+            x = self.decode(z)
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "codes": codes,
+                "latents": latents,
+                "vq/commitment_loss": commitment_loss,
+                "vq/codebook_loss": codebook_loss,
+            }
+        else:
+            posterior, _, _, _, _ = self.encode(audio_data, n_quantizers)
+            z = posterior.sample()
+            x = self.decode(z)
+
+            kl_loss = posterior.kl(dims=(1, 2))
+            kl_loss = kl_loss.mean()
+
+            return {
+                "audio": x[..., :length],
+                "z": z,
+                "kl_loss": kl_loss,
+            }
+
+
+EntryClass = DAC

python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py

@@ -0,0 +1,119 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+MOVA pipeline integration (native SGLang pipeline).
+"""
+
+from __future__ import annotations
+
+from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
+from sglang.multimodal_gen.configs.sample.mova import MOVASamplingParams
+from sglang.multimodal_gen.runtime.models.model_stages.mova import (
+    MOVADecodingStage,
+    MOVADenoisingStage,
+    MOVALatentPreparationStage,
+    MOVATimestepPreparationStage,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
+    ComposedPipelineBase,
+)
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    ConditioningStage,
+    ImageVAEEncodingStage,
+    InputValidationStage,
+    TextEncodingStage,
+)
+from sglang.multimodal_gen.runtime.server_args import ServerArgs
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class MOVAPipeline(ComposedPipelineBase):
+    """MOVA pipeline with SGLang stage orchestration."""
+
+    pipeline_name = "MOVA"
+    is_video_pipeline = True
+    _required_config_modules = [
+        "video_vae",
+        "audio_vae",
+        "text_encoder",
+        "tokenizer",
+        "scheduler",
+        "video_dit",
+        "video_dit_2",
+        "audio_dit",
+        "dual_tower_bridge",
+    ]
+    pipeline_config_cls = MOVAPipelineConfig
+    sampling_params_cls = MOVASamplingParams
+
+    def initialize_pipeline(self, server_args: ServerArgs) -> None:
+        """
+        Initialize the pipeline.
+
+        MOVA supports Context Parallel (sequence parallel) through USPAttention,
+        which uses Ulysses-style all-to-all communication for distributed attention.
+        """
+        if server_args.sp_degree > 1:
+            logger.info(
+                "MOVA Context Parallel enabled with sp_degree=%d. "
+                "Using USPAttention for distributed self-attention.",
+                server_args.sp_degree,
+            )
+
+    def create_pipeline_stages(self, server_args: ServerArgs) -> None:
+        self.add_stage(
+            stage_name="input_validation_stage", stage=InputValidationStage()
+        )
+        self.add_stage(
+            stage_name="prompt_encoding_stage",
+            stage=TextEncodingStage(
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
+            ),
+        )
+        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+        if getattr(self.get_module("video_dit"), "require_vae_embedding", True):
+            self.add_stage(
+                stage_name="image_latent_preparation_stage",
+                stage=ImageVAEEncodingStage(vae=self.get_module("video_vae")),
+            )
+        self.add_stage(
+            stage_name="mova_latent_preparation_stage",
+            stage=MOVALatentPreparationStage(
+                audio_vae=self.get_module("audio_vae"),
+                require_vae_embedding=getattr(
+                    self.get_module("video_dit"), "require_vae_embedding", True
+                ),
+            ),
+        )
+        self.add_stage(
+            stage_name="mova_timestep_preparation_stage",
+            stage=MOVATimestepPreparationStage(
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+        self.add_stage(
+            stage_name="mova_denoising_stage",
+            stage=MOVADenoisingStage(
+                video_dit=self.get_module("video_dit"),
+                video_dit_2=self.get_module("video_dit_2"),
+                audio_dit=self.get_module("audio_dit"),
+                dual_tower_bridge=self.get_module("dual_tower_bridge"),
+                scheduler=self.get_module("scheduler"),
+            ),
+        )
+        self.add_stage(
+            stage_name="mova_decoding_stage",
+            stage=MOVADecodingStage(
+                video_vae=self.get_module("video_vae"),
+                audio_vae=self.get_module("audio_vae"),
+            ),
+        )
+
+
+class MOVAPipelineAlias(MOVAPipeline):
+    pipeline_name = "MOVAPipeline"
+
+
+EntryClass = [MOVAPipeline, MOVAPipelineAlias]

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -182,12 +182,25 @@ class ComposedPipelineBase(ABC):
             "boundary_ratio" in model_index
             and model_index["boundary_ratio"] is not None
         ):
-            logger.info(
-                "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
+            has_transformer = (
+                "transformer" in model_index
+                or "transformer_2" in model_index
+                or "transformer" in self.required_config_modules
+                or "transformer_2" in self.required_config_modules
             )
-            self.required_config_modules.append("transformer_2")
+            if has_transformer:
+                logger.info(
+                    "MoE pipeline detected. Adding transformer_2 to self.required_config_modules..."
+                )
+                if "transformer_2" not in self.required_config_modules:
+                    self.required_config_modules.append("transformer_2")
+            else:
+                logger.info(
+                    "Boundary ratio found in model_index.json without transformers; "
+                    "using it for pipeline config only."
+                )
             logger.info(
-                "MoE pipeline detected. Setting boundary ratio to %s",
+                "Setting boundary ratio to %s",
                 model_index["boundary_ratio"],
             )
             server_args.pipeline_config.dit_config.boundary_ratio = model_index[

python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py

@@ -90,11 +90,13 @@ class Req:
 
     # Latent tensors
     latents: torch.Tensor | None = None
+    y: torch.Tensor | None = None
     # Flux-2
     latent_ids: torch.Tensor | None = None
 
-    # Audio Latents (LTX-2)
+    # Audio Latents
     audio_latents: torch.Tensor | None = None
+    audio_noise: torch.Tensor | None = None
     raw_audio_latent_shape: tuple[int, ...] | None = None
 
     # Audio Parameters
@@ -114,6 +116,7 @@ class Req:
 
     # Timesteps
     timesteps: torch.Tensor | None = None
+    paired_timesteps: torch.Tensor | None = None
     timestep: torch.Tensor | float | int | None = None
     step_index: int | None = None
 
@@ -154,6 +157,8 @@ class Req:
 
     # results
     output: torch.Tensor | None = None
+    audio: torch.Tensor | None = None
+    audio_sample_rate: int | None = None
 
     def __init__(self, **kwargs):
         # Initialize dataclass fields

python/sglang/multimodal_gen/runtime/pipelines_core/stages/input_validation.py

@@ -11,6 +11,7 @@ from PIL import Image
 
 from sglang.multimodal_gen.configs.pipeline_configs import WanI2V480PConfig
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
+from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
 from sglang.multimodal_gen.runtime.models.vision_utils import load_image, load_video
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
@@ -46,6 +47,25 @@ class InputValidationStage(PipelineStage):
         super().__init__()
         self.vae_image_processor = vae_image_processor
 
+    @staticmethod
+    def _calculate_dimensions_from_area(
+        max_area: float, aspect_ratio: float, mod_value: int
+    ) -> tuple[int, int]:
+        """
+        Calculate output dimensions based on maximum area and aspect ratio.
+
+        Args:
+            max_area: Maximum area constraint for the output
+            aspect_ratio: Target aspect ratio (height/width)
+            mod_value: Value to round dimensions to (typically vae_scale * patch_size)
+
+        Returns:
+            Tuple of (width, height) rounded to multiples of mod_value
+        """
+        height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+        width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+        return width, height
+
     def _generate_seeds(self, batch: Req, server_args: ServerArgs):
         """Generate seeds for the inference"""
         seed = batch.seed
@@ -168,13 +188,45 @@ class InputValidationStage(PipelineStage):
                 server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
                 * server_args.pipeline_config.dit_config.arch_config.patch_size[1]
             )
-            height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-            width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+            width, height = self._calculate_dimensions_from_area(
+                max_area, aspect_ratio, mod_value
+            )
 
             batch.condition_image = batch.condition_image.resize((width, height))
             batch.height = height
             batch.width = width
 
+        elif issubclass(type(server_args.pipeline_config), MOVAPipelineConfig):
+            # resize image only, MOVA
+            image = batch.condition_image
+            if isinstance(image, list):
+                image = image[0]  # not support multi image input yet.
+
+            max_area = server_args.pipeline_config.max_area
+            if hasattr(batch, "height") and hasattr(batch, "width"):
+                aspect_ratio = batch.height / batch.width
+            else:
+                aspect_ratio = (
+                    batch.sampling_params.height / batch.sampling_params.width
+                )
+            mod_value = (
+                server_args.pipeline_config.vae_config.arch_config.scale_factor_spatial
+                * server_args.pipeline_config.dit_config.arch_config.patch_size[1]
+            )
+            width, height = self._calculate_dimensions_from_area(
+                max_area, aspect_ratio, mod_value
+            )
+
+            config = server_args.pipeline_config
+            image, (final_w, final_h) = (
+                server_args.pipeline_config.preprocess_condition_image(
+                    image, width, height, self.vae_image_processor
+                )
+            )
+            batch.condition_image = image
+            batch.width = final_w
+            batch.height = final_h
+
     def forward(
         self,
         batch: Req,

python/sglang/multimodal_gen/runtime/server_args.py

@@ -342,6 +342,11 @@ class ServerArgs:
         default_factory=lambda: {
             "transformer": True,
             "vae": True,
+            "video_vae": True,
+            "audio_vae": True,
+            "video_dit": True,
+            "audio_dit": True,
+            "dual_tower_bridge": True,
         }
     )
 

python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py

@@ -412,20 +412,6 @@ def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
             "Only HuggingFace diffusers format is supported."
         )
 
-    # Check for transformer and vae directories
-    transformer_dir = os.path.join(model_path, "transformer")
-    vae_dir = os.path.join(model_path, "vae")
-
-    if not os.path.exists(transformer_dir):
-        raise ValueError(
-            f"Model directory {model_path} does not contain a transformer/ directory."
-        )
-
-    if not os.path.exists(vae_dir):
-        raise ValueError(
-            f"Model directory {model_path} does not contain a vae/ directory."
-        )
-
     # Load the config
     with open(config_path) as f:
         config = json.load(f)
@@ -435,6 +421,37 @@ def verify_model_config_and_directory(model_path: str) -> dict[str, Any]:
         raise ValueError("model_index.json does not contain _diffusers_version")
 
     logger.info("Diffusers version: %s", config["_diffusers_version"])
+
+    component_keys = [
+        key
+        for key, value in config.items()
+        if isinstance(value, (list, tuple))
+        and len(value) == 2
+        and all(isinstance(item, str) for item in value)
+    ]
+    if component_keys:
+        missing_components = [
+            component_key
+            for component_key in component_keys
+            if not os.path.exists(os.path.join(model_path, component_key))
+        ]
+        if missing_components:
+            missing_str = ", ".join(missing_components)
+            raise ValueError(
+                f"Model directory {model_path} is missing required component "
+                f"directories: {missing_str}."
+            )
+    else:
+        transformer_dir = os.path.join(model_path, "transformer")
+        vae_dir = os.path.join(model_path, "vae")
+        if not os.path.exists(transformer_dir):
+            raise ValueError(
+                f"Model directory {model_path} does not contain a transformer/ directory."
+            )
+        if not os.path.exists(vae_dir):
+            raise ValueError(
+                f"Model directory {model_path} does not contain a vae/ directory."
+            )
     return cast(dict[str, Any], config)
 
 
@@ -543,14 +560,35 @@ def maybe_download_model(
 
     def _verify_model_complete(path: str) -> bool:
         """Check if model directory has required subdirectories."""
+        config_path = os.path.join(path, "model_index.json")
+        if not os.path.exists(config_path):
+            return False
+
+        try:
+            with open(config_path) as config_file:
+                model_index = json.load(config_file)
+        except Exception as exc:
+            logger.warning(
+                "Failed to read model_index.json at %s: %s", config_path, exc
+            )
+            return False
+
+        component_keys = [
+            key
+            for key, value in model_index.items()
+            if isinstance(value, (list, tuple))
+            and len(value) == 2
+            and all(isinstance(item, str) for item in value)
+        ]
+        if component_keys:
+            return all(
+                os.path.exists(os.path.join(path, component_key))
+                for component_key in component_keys
+            )
+
         transformer_dir = os.path.join(path, "transformer")
         vae_dir = os.path.join(path, "vae")
-        config_path = os.path.join(path, "model_index.json")
-        return (
-            os.path.exists(config_path)
-            and os.path.exists(transformer_dir)
-            and os.path.exists(vae_dir)
-        )
+        return os.path.exists(transformer_dir) and os.path.exists(vae_dir)
 
     # 1. Local path check: if path exists locally, verify it's complete (skip for LoRA)
     if os.path.exists(model_name_or_path):
@@ -584,7 +622,7 @@ def maybe_download_model(
                 return model_name_or_path
         else:
             logger.warning(
-                "Local model at %s appears incomplete (missing transformer/ or vae/), "
+                "Local model at %s appears incomplete (missing required components), "
                 "will attempt re-download",
                 model_name_or_path,
             )